已合并
[1/2] Add tp cases #18964
dilililiwhy创建于 2025年3月14日
[1/2] Add tp cases #18964
已合并
从refs/pull/18964/head合入到master
共 2 个文件变更+794-0
| @@ -0,0 +1,354 @@ | |||
| 1 | +# Owner(s): ["oncall: distributed"] | ||
| 2 | +from collections import OrderedDict | ||
| 3 | +from copy import deepcopy | ||
| 4 | + | ||
| 5 | +import torch | ||
| 6 | +from torch.distributed._tensor import DeviceMesh, DTensor, Replicate, Shard | ||
| 7 | +from torch.distributed.tensor.debug import CommDebugMode | ||
| 8 | +from torch.distributed.tensor.parallel.api import parallelize_module | ||
| 9 | +from torch.distributed.tensor.parallel.style import ( | ||
| 10 | + ColwiseParallel, | ||
| 11 | + PrepareModuleInput, | ||
| 12 | + PrepareModuleOutput, | ||
| 13 | + RowwiseParallel, | ||
| 14 | +) | ||
| 15 | +from torch.testing._internal.common_utils import run_tests | ||
| 16 | +from torch.testing._internal.distributed._tensor.common_dtensor import ( | ||
| 17 | + DTensorTestBase, | ||
| 18 | + MLPModule, | ||
| 19 | + MLPStacked, | ||
| 20 | +) | ||
| 21 | + | ||
| 22 | +import torch_npu | ||
| 23 | +from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU | ||
| 24 | + | ||
| 25 | + | ||
| 26 | +class DummyModule(torch.nn.Module): | ||
| 27 | + def __init__(self) -> None: | ||
| 28 | + super().__init__() | ||
| 29 | + | ||
| 30 | + def forward(self, x): | ||
| 31 | + return x | ||
| 32 | + | ||
| 33 | + | ||
| 34 | +class TensorParallelAPITests(DTensorTestBase): | ||
| 35 | + | ||
| 36 | + def world_size(self): | ||
| 37 | + return 2 | ||
| 38 | + | ||
| 39 | + def _compare_params( | ||
| 40 | + self, | ||
| 41 | + local_module, | ||
| 42 | + dist_module, | ||
| 43 | + rank0_only, | ||
| 44 | + skip_rowwise_bias=False, | ||
| 45 | + compare_grad=False, | ||
| 46 | + ): | ||
| 47 | + replicate = [Replicate()] | ||
| 48 | + for name, param in local_module.named_parameters(): | ||
| 49 | + dist_param = dist_module.get_parameter(name) | ||
| 50 | + param = param.grad if compare_grad else param | ||
| 51 | + dist_param = dist_param.grad if compare_grad else dist_param | ||
| 52 | + if ( | ||
| 53 | + (not rank0_only) | ||
| 54 | + or (self.rank == 0) | ||
| 55 | + or ( | ||
| 56 | + name not in ["net2.bias"] | ||
| 57 | + and not skip_rowwise_bias | ||
| 58 | + or name not in ["bias", "net2.bias"] | ||
| 59 | + ) | ||
| 60 | + ): | ||
| 61 | + self.assertEqual( | ||
| 62 | + param, | ||
| 63 | + dist_param.redistribute( | ||
| 64 | + device_mesh=dist_param.device_mesh, placements=replicate | ||
| 65 | + ).to_local(), | ||
| 66 | + f"{name} not equal between dist and non-dist", | ||
| 67 | + ) | ||
| 68 | + | ||
| 69 | + def _compare_module( | ||
| 70 | + self, local_module, dist_module, inp_size, rank0_only=True, rowwise=False | ||
| 71 | + ): | ||
| 72 | + LR = 0.25 # the learning rate we use for testing | ||
| 73 | + local_optim = torch.optim.SGD(local_module.parameters(), lr=LR) | ||
| 74 | + dist_optim = torch.optim.SGD(dist_module.parameters(), lr=LR) | ||
| 75 | + torch.manual_seed(0) | ||
| 76 | + inp = torch.rand(*inp_size, device=self.device_type) | ||
| 77 | + self._compare_params(local_module, dist_module, rank0_only) | ||
| 78 | + | ||
| 79 | + # check forward correctness | ||
| 80 | + local_output = local_module(inp) | ||
| 81 | + inp = inp.chunk(self.world_size, dim=-1)[self.rank] if rowwise else inp | ||
| 82 | + dist_output = dist_module(inp) | ||
| 83 | + dist_output = ( | ||
| 84 | + dist_output.redistribute(dist_output.device_mesh, [Replicate()]).to_local() | ||
| 85 | + if isinstance(dist_output, DTensor) | ||
| 86 | + else dist_output | ||
| 87 | + ) | ||
| 88 | + self.assertEqual(local_output, dist_output) | ||
| 89 | + | ||
| 90 | + local_output.sum().backward() | ||
| 91 | + dist_output.sum().backward() | ||
| 92 | + | ||
| 93 | + # check backward and ensure gradients are same | ||
| 94 | + self._compare_params(local_module, dist_module, rank0_only, rowwise, True) | ||
| 95 | + | ||
| 96 | + local_optim.step() | ||
| 97 | + dist_optim.step() | ||
| 98 | + self._compare_params(local_module, dist_module, rank0_only, rowwise) | ||
| 99 | + | ||
| 100 | + | ||
| 101 | + | ||
| 102 | + def test_parallelize_mlp_with_module_api(self): | ||
| 103 | + inp_size = [12, 10] | ||
| 104 | + model = MLPModule(self.device_type) | ||
| 105 | + model_tp = deepcopy(model) | ||
| 106 | + | ||
| 107 | + # Parallelize module. | ||
| 108 | + device_mesh = DeviceMesh(self.device_type, torch.arange(self.world_size)) | ||
| 109 | + model_tp = parallelize_module( | ||
| 110 | + model_tp, | ||
| 111 | + device_mesh, | ||
| 112 | + { | ||
| 113 | + "net1": ColwiseParallel(output_layouts=Replicate()), | ||
| 114 | + "net2": ColwiseParallel(output_layouts=Replicate()), | ||
| 115 | + }, | ||
| 116 | + ) | ||
| 117 | + self._compare_module(model, model_tp, inp_size, rank0_only=False) | ||
| 118 | + | ||
| 119 | + | ||
| 120 | + | ||
| 121 | + def test_parallelize_mlp_with_module_api_nested(self): | ||
| 122 | + inp_size = [12, 10] | ||
| 123 | + model = torch.nn.Sequential( | ||
| 124 | + OrderedDict([("dummy_encoder", MLPModule(self.device_type))]) | ||
| 125 | + ) | ||
| 126 | + model_tp = deepcopy(model) | ||
| 127 | + | ||
| 128 | + # Parallelize module. | ||
| 129 | + device_mesh = DeviceMesh(self.device_type, torch.arange(self.world_size)) | ||
| 130 | + model_tp = parallelize_module( | ||
| 131 | + model_tp, | ||
| 132 | + device_mesh, | ||
| 133 | + { | ||
| 134 | + "dummy_encoder.net1": ColwiseParallel(output_layouts=Replicate()), | ||
| 135 | + "dummy_encoder.net2": ColwiseParallel(output_layouts=Replicate()), | ||
| 136 | + }, | ||
| 137 | + ) | ||
| 138 | + self._compare_module(model, model_tp, inp_size, rank0_only=False) | ||
| 139 | + | ||
| 140 | + | ||
| 141 | + | ||
| 142 | + def test_linear_row_wise_parallel(self): | ||
| 143 | + # test RowwiseParallel | ||
| 144 | + inp_size = [9, 16] | ||
| 145 | + rowwise = RowwiseParallel() | ||
| 146 | + | ||
| 147 | + torch.manual_seed(5) | ||
| 148 | + model = torch.nn.Linear(16, 10, device=self.device_type) | ||
| 149 | + model_tp = deepcopy(model) | ||
| 150 | + | ||
| 151 | + # parallelize model_tp | ||
| 152 | + device_mesh = DeviceMesh(self.device_type, list(range(self.world_size))) | ||
| 153 | + model_tp = parallelize_module(model_tp, device_mesh, rowwise) | ||
| 154 | + | ||
| 155 | + # let each rank generate unique local input | ||
| 156 | + torch.manual_seed(self.rank) | ||
| 157 | + self._compare_module(model, model_tp, inp_size, rowwise=True) | ||
| 158 | + | ||
| 159 | + | ||
| 160 | + | ||
| 161 | + def test_linear_col_wise_parallel(self): | ||
| 162 | + # test ColwiseParallel | ||
| 163 | + inp_size = [8, 10] | ||
| 164 | + colwise = ColwiseParallel(output_layouts=Replicate()) | ||
| 165 | + | ||
| 166 | + torch.manual_seed(5) | ||
| 167 | + model = torch.nn.Linear(10, 16, device=self.device_type) | ||
| 168 | + model_tp = deepcopy(model) | ||
| 169 | + | ||
| 170 | + # parallelize model_tp | ||
| 171 | + device_mesh = DeviceMesh(self.device_type, list(range(self.world_size))) | ||
| 172 | + model_tp = parallelize_module(model_tp, device_mesh, colwise) | ||
| 173 | + | ||
| 174 | + self._compare_module(model, model_tp, inp_size) | ||
| 175 | + | ||
| 176 | + | ||
| 177 | + | ||
| 178 | + def test_prepare_module_input(self): | ||
| 179 | + module = DummyModule() | ||
| 180 | + device_mesh = DeviceMesh(self.device_type, list(range(self.world_size))) | ||
| 181 | + parallelize_module( | ||
| 182 | + module, | ||
| 183 | + device_mesh, | ||
| 184 | + PrepareModuleInput( | ||
| 185 | + input_layouts=Shard(0), desired_input_layouts=Replicate() | ||
| 186 | + ), | ||
| 187 | + ) | ||
| 188 | + inp = torch.rand(5, 7, device=self.device_type) | ||
| 189 | + output = module(inp).redistribute(device_mesh, [Shard(0)]).to_local() | ||
| 190 | + self.assertEqual(inp, output) | ||
| 191 | + | ||
| 192 | + | ||
| 193 | + | ||
| 194 | + def test_prepare_module_output(self): | ||
| 195 | + module = DummyModule() | ||
| 196 | + device_mesh = DeviceMesh(self.device_type, list(range(self.world_size))) | ||
| 197 | + parallelize_module( | ||
| 198 | + module, | ||
| 199 | + device_mesh, | ||
| 200 | + PrepareModuleOutput( | ||
| 201 | + output_layouts=Replicate(), desired_output_layouts=Shard(0) | ||
| 202 | + ), | ||
| 203 | + ) | ||
| 204 | + torch.manual_seed(15) | ||
| 205 | + inp = torch.rand(16, 7, device=self.device_type) | ||
| 206 | + dtensor = DTensor.from_local(inp, device_mesh, [Replicate()], run_check=False) | ||
| 207 | + output = module(dtensor) | ||
| 208 | + inp = dtensor.redistribute(device_mesh, [Shard(0)]).to_local() | ||
| 209 | + self.assertEqual(inp, output) | ||
| 210 | + | ||
| 211 | + | ||
| 212 | + | ||
| 213 | + def test_parallelize_module_with_star(self): | ||
| 214 | + inp_size = [12, 10] | ||
| 215 | + model = MLPModule(self.device_type) | ||
| 216 | + device_mesh = DeviceMesh(self.device_type, torch.arange(self.world_size)) | ||
| 217 | + | ||
| 218 | + model_tp = deepcopy(model) | ||
| 219 | + model_tp = parallelize_module( | ||
| 220 | + model_tp, | ||
| 221 | + device_mesh, | ||
| 222 | + { | ||
| 223 | + "net*": ColwiseParallel(output_layouts=Replicate()), | ||
| 224 | + }, | ||
| 225 | + ) | ||
| 226 | + self._compare_module(model, model_tp, inp_size, rank0_only=False) | ||
| 227 | + | ||
| 228 | + | ||
| 229 | + | ||
| 230 | + def test_parallelize_module_src_data_rank(self): | ||
| 231 | + # set seed different for each rank | ||
| 232 | + torch.manual_seed(self.rank) | ||
| 233 | + model = MLPModule(self.device_type) | ||
| 234 | + device_mesh = DeviceMesh(self.device_type, torch.arange(self.world_size)) | ||
| 235 | + | ||
| 236 | + comm_mode = CommDebugMode() | ||
| 237 | + | ||
| 238 | + # test src_data_rank == 1 | ||
| 239 | + with comm_mode: | ||
| 240 | + model_tp = deepcopy(model) | ||
| 241 | + model_tp = parallelize_module( | ||
| 242 | + model_tp, | ||
| 243 | + device_mesh, | ||
| 244 | + { | ||
| 245 | + "net*": ColwiseParallel(output_layouts=Replicate()), | ||
| 246 | + }, | ||
| 247 | + src_data_rank=1, | ||
| 248 | + ) | ||
| 249 | + | ||
| 250 | + self.assertTrue(comm_mode.get_total_counts() > 0) | ||
| 251 | + tp_full_params = [param.full_tensor() for param in model_tp.parameters()] | ||
| 252 | + if self.rank == 1: | ||
| 253 | + orig_model_params = list(model.parameters()) | ||
| 254 | + for idx, param in enumerate(tp_full_params): | ||
| 255 | + self.assertEqual(param, orig_model_params[idx]) | ||
| 256 | + | ||
| 257 | + # test src_data_rank == None | ||
| 258 | + model_tp_no_comm = deepcopy(model) | ||
| 259 | + with comm_mode: | ||
| 260 | + parallelize_module( | ||
| 261 | + model_tp_no_comm, | ||
| 262 | + device_mesh, | ||
| 263 | + { | ||
| 264 | + "net1": ColwiseParallel(), | ||
| 265 | + "net2": RowwiseParallel(), | ||
| 266 | + }, | ||
| 267 | + src_data_rank=None, | ||
| 268 | + ) | ||
| 269 | + self.assertEqual(comm_mode.get_total_counts(), 0) | ||
| 270 | + | ||
| 271 | + | ||
| 272 | + | ||
| 273 | + def test_parallelize_module_with_question(self): | ||
| 274 | + inp_size = [12, 10] | ||
| 275 | + model = MLPModule(self.device_type) | ||
| 276 | + device_mesh = DeviceMesh(self.device_type, torch.arange(self.world_size)) | ||
| 277 | + | ||
| 278 | + model_tp = deepcopy(model) | ||
| 279 | + model_tp = parallelize_module( | ||
| 280 | + model_tp, | ||
| 281 | + device_mesh, | ||
| 282 | + { | ||
| 283 | + "net?": ColwiseParallel(output_layouts=Replicate()), | ||
| 284 | + }, | ||
| 285 | + ) | ||
| 286 | + self._compare_module(model, model_tp, inp_size, rank0_only=False) | ||
| 287 | + | ||
| 288 | + | ||
| 289 | + | ||
| 290 | + def test_parallelize_module_with_digit(self): | ||
| 291 | + inp_size = [12, 10] | ||
| 292 | + model = MLPModule(self.device_type) | ||
| 293 | + device_mesh = DeviceMesh(self.device_type, torch.arange(self.world_size)) | ||
| 294 | + | ||
| 295 | + model_tp = deepcopy(model) | ||
| 296 | + model_tp = parallelize_module( | ||
| 297 | + model_tp, | ||
| 298 | + device_mesh, | ||
| 299 | + { | ||
| 300 | + "net[1-2]": ColwiseParallel(output_layouts=Replicate()), | ||
| 301 | + }, | ||
| 302 | + ) | ||
| 303 | + self._compare_module(model, model_tp, inp_size, rank0_only=False) | ||
| 304 | + | ||
| 305 | + | ||
| 306 | + | ||
| 307 | + def test_parallelize_module_multi_wildcard(self): | ||
| 308 | + inp_size = [12, 10] | ||
| 309 | + model = MLPStacked(self.device_type, n_layers=2) | ||
| 310 | + device_mesh = DeviceMesh(self.device_type, torch.arange(self.world_size)) | ||
| 311 | + | ||
| 312 | + model_tp = deepcopy(model) | ||
| 313 | + model_tp = parallelize_module( | ||
| 314 | + model_tp, | ||
| 315 | + device_mesh, | ||
| 316 | + { | ||
| 317 | + "layers.*.net[1]": ColwiseParallel(), | ||
| 318 | + "layers.*.net[2]": RowwiseParallel(), | ||
| 319 | + }, | ||
| 320 | + ) | ||
| 321 | + self._compare_module(model, model_tp, inp_size, rank0_only=False) | ||
| 322 | + | ||
| 323 | + | ||
| 324 | + | ||
| 325 | + def test_under_devicemesh_context(self): | ||
| 326 | + # test ColwiseParallel | ||
| 327 | + inp_size = [8, 10] | ||
| 328 | + colwise = ColwiseParallel(output_layouts=Replicate()) | ||
| 329 | + | ||
| 330 | + torch.manual_seed(5) | ||
| 331 | + model = torch.nn.Linear(10, 16, device=self.device_type) | ||
| 332 | + model_tp = deepcopy(model) | ||
| 333 | + | ||
| 334 | + # Call parallelize_module under DeviceMesh context. | ||
| 335 | + device_mesh = DeviceMesh(self.device_type, list(range(self.world_size))) | ||
| 336 | + with device_mesh: | ||
| 337 | + model_tp = parallelize_module(model_tp, parallelize_plan=colwise) | ||
| 338 | + | ||
| 339 | + self._compare_module(model, model_tp, inp_size) | ||
| 340 | + | ||
| 341 | + | ||
| 342 | + | ||
| 343 | + def test_empty_plan(self): | ||
| 344 | + torch.manual_seed(5) | ||
| 345 | + model = torch.nn.Linear(10, 16, device=self.device_type) | ||
| 346 | + | ||
| 347 | + # Call parallelize_module with empty plan. | ||
| 348 | + # Goal is not to crash. | ||
| 349 | + device_mesh = DeviceMesh(self.device_type, list(range(self.world_size))) | ||
| 350 | + parallelize_module(model, device_mesh) | ||
| 351 | + | ||
| 352 | + | ||
| 353 | +if __name__ == "__main__": | ||
| 354 | + run_tests() | ||
| @@ -0,0 +1,440 @@ | |||
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates | ||
| 2 | +# Owner(s): ["oncall: distributed"] | ||
| 3 | + | ||
| 4 | +from copy import deepcopy | ||
| 5 | + | ||
| 6 | +import torch | ||
| 7 | + | ||
| 8 | +import torch.nn as nn | ||
| 9 | +from torch.distributed._tensor import ( | ||
| 10 | + distribute_tensor, | ||
| 11 | + DTensor, | ||
| 12 | + init_device_mesh, | ||
| 13 | + Replicate, | ||
| 14 | + Shard, | ||
| 15 | +) | ||
| 16 | +from torch.distributed.tensor.debug import CommDebugMode | ||
| 17 | +from torch.distributed.tensor.parallel import parallelize_module | ||
| 18 | +from torch.distributed.tensor.parallel.style import ( | ||
| 19 | + ColwiseParallel, | ||
| 20 | + PrepareModuleInput, | ||
| 21 | + PrepareModuleOutput, | ||
| 22 | + RowwiseParallel, | ||
| 23 | + SequenceParallel, | ||
| 24 | +) | ||
| 25 | +from torch.distributed.tensor.placement_types import _Partial | ||
| 26 | +from torch.testing._internal.common_utils import run_tests | ||
| 27 | +from torch.testing._internal.distributed._tensor.common_dtensor import ( | ||
| 28 | + DTensorTestBase, | ||
| 29 | + RMSNormPython, | ||
| 30 | +) | ||
| 31 | + | ||
| 32 | +import torch_npu | ||
| 33 | +from torch_npu.testing.common_utils import SupportedDevices | ||
| 34 | +from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU | ||
| 35 | + | ||
| 36 | + | ||
| 37 | +c10d_functional = torch.ops.c10d_functional | ||
| 38 | + | ||
| 39 | + | ||
| 40 | +class TensorParallelStyleTest(DTensorTestBase): | ||
| 41 | + | ||
| 42 | + def world_size(self): | ||
| 43 | + return 2 | ||
| 44 | + | ||
| 45 | + | ||
| 46 | + | ||
| 47 | + def test_colwise_parallel_style(self): | ||
| 48 | + mesh = init_device_mesh(self.device_type, (self.world_size,)) | ||
| 49 | + | ||
| 50 | + comm_mode = CommDebugMode() | ||
| 51 | + tensor = torch.rand(8, 16, device=self.device_type, requires_grad=True) | ||
| 52 | + model = nn.Linear(16, 16, device=self.device_type) | ||
| 53 | + | ||
| 54 | + default_col_parallel = ColwiseParallel() | ||
| 55 | + colwise_mod = parallelize_module(deepcopy(model), mesh, default_col_parallel) | ||
| 56 | + with comm_mode: | ||
| 57 | + out = colwise_mod(tensor) | ||
| 58 | + # ensure output shard on the last dim | ||
| 59 | + self.assertEqual(out.shape, (8, 16 // self.world_size)) | ||
| 60 | + # ensure no communication happened in fwd | ||
| 61 | + self.assertEqual(comm_mode.get_total_counts(), 0) | ||
| 62 | + | ||
| 63 | + out.sum().backward() | ||
| 64 | + # allreduce in bwd | ||
| 65 | + self.assertEqual(comm_mode.get_comm_counts()[c10d_functional.all_reduce], 1) | ||
| 66 | + self.assertEqual(comm_mode.get_total_counts(), 1) | ||
| 67 | + | ||
| 68 | + sharded_col_parallel = ColwiseParallel(input_layouts=Shard(0)) | ||
| 69 | + colwise_mod = parallelize_module(deepcopy(model), mesh, sharded_col_parallel) | ||
| 70 | + with comm_mode: | ||
| 71 | + out = colwise_mod(tensor) | ||
| 72 | + # ensure output shard on the last dim | ||
| 73 | + self.assertEqual(out.shape, (8 * self.world_size, 16 // self.world_size)) | ||
| 74 | + # allgather in fwd | ||
| 75 | + self.assertEqual( | ||
| 76 | + comm_mode.get_comm_counts()[c10d_functional.all_gather_into_tensor], 1 | ||
| 77 | + ) | ||
| 78 | + self.assertEqual(comm_mode.get_total_counts(), 1) | ||
| 79 | + | ||
| 80 | + out.sum().backward() | ||
| 81 | + # reduce_scatter in bwd | ||
| 82 | + self.assertEqual( | ||
| 83 | + comm_mode.get_comm_counts()[c10d_functional.reduce_scatter_tensor], 1 | ||
| 84 | + ) | ||
| 85 | + self.assertEqual(comm_mode.get_total_counts(), 2) | ||
| 86 | + | ||
| 87 | + | ||
| 88 | + | ||
| 89 | + def test_colwise_parallel_embedding(self): | ||
| 90 | + mesh = init_device_mesh(self.device_type, (self.world_size,)) | ||
| 91 | + | ||
| 92 | + comm_mode = CommDebugMode() | ||
| 93 | + tensor = torch.arange(8, device=self.device_type).reshape(4, 2) | ||
| 94 | + model = nn.Embedding(16, 16, device=self.device_type) | ||
| 95 | + | ||
| 96 | + default_col_parallel = ColwiseParallel() | ||
| 97 | + colwise_mod = parallelize_module(deepcopy(model), mesh, default_col_parallel) | ||
| 98 | + with comm_mode: | ||
| 99 | + out = colwise_mod(tensor) | ||
| 100 | + # ensure output shard on the last dim | ||
| 101 | + self.assertEqual(out.shape, (4, 2, 16 // self.world_size)) | ||
| 102 | + # ensure no communication happened in fwd | ||
| 103 | + self.assertEqual(comm_mode.get_total_counts(), 0) | ||
| 104 | + | ||
| 105 | + out.sum().backward() | ||
| 106 | + # no comm in bwd | ||
| 107 | + self.assertEqual(comm_mode.get_total_counts(), 0) | ||
| 108 | + | ||
| 109 | + | ||
| 110 | + | ||
| 111 | + def test_rowwise_parallel_style(self): | ||
| 112 | + mesh = init_device_mesh(self.device_type, (self.world_size,)) | ||
| 113 | + | ||
| 114 | + comm_mode = CommDebugMode() | ||
| 115 | + tensor = torch.rand( | ||
| 116 | + 8, 16 // self.world_size, device=self.device_type, requires_grad=True | ||
| 117 | + ) | ||
| 118 | + model = nn.Linear(16, 16, device=self.device_type) | ||
| 119 | + | ||
| 120 | + default_row_parallel = RowwiseParallel() | ||
| 121 | + rowwise_mod = parallelize_module(deepcopy(model), mesh, default_row_parallel) | ||
| 122 | + with comm_mode: | ||
| 123 | + out = rowwise_mod(tensor) | ||
| 124 | + # ensure output replicated | ||
| 125 | + self.assertEqual(out.shape, (8, 16)) | ||
| 126 | + # allreduce in fwd | ||
| 127 | + self.assertEqual(comm_mode.get_comm_counts()[c10d_functional.all_reduce], 1) | ||
| 128 | + self.assertEqual(comm_mode.get_total_counts(), 1) | ||
| 129 | + | ||
| 130 | + out.sum().backward() | ||
| 131 | + # no op in bwd | ||
| 132 | + self.assertEqual(comm_mode.get_total_counts(), 1) | ||
| 133 | + | ||
| 134 | + sharded_row_parallel = RowwiseParallel(output_layouts=Shard(0)) | ||
| 135 | + rowwise_mod = parallelize_module(deepcopy(model), mesh, sharded_row_parallel) | ||
| 136 | + with comm_mode: | ||
| 137 | + out = rowwise_mod(tensor) | ||
| 138 | + # ensure output replicated | ||
| 139 | + self.assertEqual(out.shape, (8 // self.world_size, 16)) | ||
| 140 | + # reduce_scatter in fwd | ||
| 141 | + self.assertEqual( | ||
| 142 | + comm_mode.get_comm_counts()[c10d_functional.reduce_scatter_tensor], 1 | ||
| 143 | + ) | ||
| 144 | + self.assertEqual(comm_mode.get_total_counts(), 1) | ||
| 145 | + | ||
| 146 | + out.sum().backward() | ||
| 147 | + # allgather in bwd | ||
| 148 | + self.assertEqual( | ||
| 149 | + comm_mode.get_comm_counts()[c10d_functional.all_gather_into_tensor], 1 | ||
| 150 | + ) | ||
| 151 | + self.assertEqual(comm_mode.get_total_counts(), 2) | ||
| 152 | + | ||
| 153 | + | ||
| 154 | + | ||
| 155 | + def test_rowwise_parallel_embedding(self): | ||
| 156 | + mesh = init_device_mesh(self.device_type, (self.world_size,)) | ||
| 157 | + | ||
| 158 | + comm_mode = CommDebugMode() | ||
| 159 | + tensor = torch.arange(8, device=self.device_type).reshape(4, 2) | ||
| 160 | + model = nn.Embedding(16, 16, device=self.device_type) | ||
| 161 | + | ||
| 162 | + rowwise_mod = parallelize_module( | ||
| 163 | + deepcopy(model), mesh, RowwiseParallel(input_layouts=Replicate()) | ||
| 164 | + ) | ||
| 165 | + with comm_mode: | ||
| 166 | + out = rowwise_mod(tensor) | ||
| 167 | + # ensure output shard on the last dim | ||
| 168 | + self.assertEqual(out.shape, (4, 2, 16)) | ||
| 169 | + # ensure allreduce communication happened in fwd | ||
| 170 | + self.assertEqual(comm_mode.get_total_counts(), 1) | ||
| 171 | + self.assertEqual(comm_mode.get_comm_counts()[c10d_functional.all_reduce], 1) | ||
| 172 | + | ||
| 173 | + out.sum().backward() | ||
| 174 | + # no comm in bwd | ||
| 175 | + self.assertEqual(comm_mode.get_total_counts(), 1) | ||
| 176 | + | ||
| 177 | + sharded_row_parallel = RowwiseParallel( | ||
| 178 | + input_layouts=Replicate(), output_layouts=Shard(1) | ||
| 179 | + ) | ||
| 180 | + | ||
| 181 | + rowwise_mod = parallelize_module(deepcopy(model), mesh, sharded_row_parallel) | ||
| 182 | + | ||
| 183 | + inp_indices = torch.arange(8, device=self.device_type) | ||
| 184 | + with comm_mode: | ||
| 185 | + out = rowwise_mod(inp_indices) | ||
| 186 | + # ensure output shard on the last dim | ||
| 187 | + self.assertEqual(out.shape, (8, 16 // self.world_size)) | ||
| 188 | + # reduce scatter in fwd | ||
| 189 | + self.assertEqual(comm_mode.get_total_counts(), 1) | ||
| 190 | + self.assertEqual( | ||
| 191 | + comm_mode.get_comm_counts()[c10d_functional.reduce_scatter_tensor], 1 | ||
| 192 | + ) | ||
| 193 | + out.sum().backward() | ||
| 194 | + # allgather comm in bwd | ||
| 195 | + self.assertEqual(comm_mode.get_total_counts(), 2) | ||
| 196 | + self.assertEqual( | ||
| 197 | + comm_mode.get_comm_counts()[c10d_functional.all_gather_into_tensor], 1 | ||
| 198 | + ) | ||
| 199 | + | ||
| 200 | + | ||
| 201 | + | ||
| 202 | + def test_prepare_module_input(self): | ||
| 203 | + mesh = init_device_mesh(self.device_type, (self.world_size,)) | ||
| 204 | + | ||
| 205 | + tensor = torch.ones(2, 16, device=self.device_type) | ||
| 206 | + expected_tensor = torch.ones(2 * self.world_size, 16, device=self.device_type) | ||
| 207 | + prepare_inp_style = PrepareModuleInput( | ||
| 208 | + input_layouts=Shard(0), desired_input_layouts=Replicate() | ||
| 209 | + ) | ||
| 210 | + | ||
| 211 | + model = nn.Identity() | ||
| 212 | + allgather_mod = parallelize_module(model, mesh, prepare_inp_style) | ||
| 213 | + output = allgather_mod(tensor).full_tensor() | ||
| 214 | + self.assertEqual(output, expected_tensor) | ||
| 215 | + | ||
| 216 | + | ||
| 217 | + | ||
| 218 | + def test_prepare_module_input_multiple_inputs(self): | ||
| 219 | + mesh = init_device_mesh(self.device_type, (self.world_size,)) | ||
| 220 | + | ||
| 221 | + class TestModule(torch.nn.Module): | ||
| 222 | + def __init__(self) -> None: | ||
| 223 | + super().__init__() | ||
| 224 | + self.linear = torch.nn.Linear(8, 8) | ||
| 225 | + | ||
| 226 | + def forward(self, x, y): | ||
| 227 | + return self.linear(x) + y | ||
| 228 | + | ||
| 229 | + # Raise assertion error if input_layouts and desired_input_layouts do not have same length. | ||
| 230 | + test_mod = TestModule().to(self.device_type) | ||
| 231 | + with self.assertRaisesRegex( | ||
| 232 | + AssertionError, | ||
| 233 | + "input_layouts and desired_input_layouts should have same length!", | ||
| 234 | + ): | ||
| 235 | + PrepareModuleInput( | ||
| 236 | + input_layouts=Shard(0), desired_input_layouts=(Replicate(), None) | ||
| 237 | + ) | ||
| 238 | + # Raise assertion error if module inputs and input_layouts do not have same length. | ||
| 239 | + prepare_inps_short_dimension = PrepareModuleInput( | ||
| 240 | + input_layouts=Shard(0), desired_input_layouts=Replicate() | ||
| 241 | + ) | ||
| 242 | + parallelize_module(test_mod.linear, mesh, ColwiseParallel()) | ||
| 243 | + parallelize_module(test_mod, mesh, prepare_inps_short_dimension) | ||
| 244 | + with self.assertRaisesRegex( | ||
| 245 | + ValueError, "module inputs and input_layouts should have same length!" | ||
| 246 | + ): | ||
| 247 | + output = test_mod( | ||
| 248 | + torch.randn(2, 8, device=self.device_type), | ||
| 249 | + torch.ones( | ||
| 250 | + self.world_size * 2, 8 // self.world_size, device=self.device_type | ||
| 251 | + ), | ||
| 252 | + ) | ||
| 253 | + | ||
| 254 | + test_mod = TestModule().to(self.device_type) | ||
| 255 | + prepare_inps = PrepareModuleInput( | ||
| 256 | + input_layouts=(Shard(0), None), desired_input_layouts=(Replicate(), None) | ||
| 257 | + ) | ||
| 258 | + | ||
| 259 | + parallelize_module(test_mod.linear, mesh, ColwiseParallel()) | ||
| 260 | + parallelize_module(test_mod, mesh, prepare_inps) | ||
| 261 | + output = test_mod( | ||
| 262 | + torch.randn(2, 8, device=self.device_type), | ||
| 263 | + torch.ones( | ||
| 264 | + self.world_size * 2, 8 // self.world_size, device=self.device_type | ||
| 265 | + ), | ||
| 266 | + ) | ||
| 267 | + self.assertEqual(output.shape, (self.world_size * 2, 8 // self.world_size)) | ||
| 268 | + | ||
| 269 | + | ||
| 270 | + | ||
| 271 | + def test_prepare_module_kwargs_input(self): | ||
| 272 | + mesh = init_device_mesh(self.device_type, (self.world_size,)) | ||
| 273 | + | ||
| 274 | + class TestKwargModule(torch.nn.Module): | ||
| 275 | + def __init__(self) -> None: | ||
| 276 | + super().__init__() | ||
| 277 | + self.linear = torch.nn.Linear(8, 8) | ||
| 278 | + | ||
| 279 | + def forward(self, x, *, y, z=2): | ||
| 280 | + return self.linear(x) + y + z | ||
| 281 | + | ||
| 282 | + test_mod = TestKwargModule().to(self.device_type) | ||
| 283 | + prepare_inps_simple = PrepareModuleInput( | ||
| 284 | + input_kwarg_layouts={"y": Shard(0)}, | ||
| 285 | + desired_input_kwarg_layouts={"y": Replicate()}, | ||
| 286 | + ) | ||
| 287 | + parallelize_module( | ||
| 288 | + test_mod.linear, mesh, ColwiseParallel(use_local_output=False) | ||
| 289 | + ) | ||
| 290 | + parallelize_module(test_mod, mesh, prepare_inps_simple) | ||
| 291 | + | ||
| 292 | + comm_mode = CommDebugMode() | ||
| 293 | + with comm_mode: | ||
| 294 | + output = test_mod( | ||
| 295 | + torch.randn(1 * self.world_size, 8, device=self.device_type), | ||
| 296 | + y=torch.ones(1, 8, device=self.device_type), | ||
| 297 | + ) | ||
| 298 | + | ||
| 299 | + self.assertEqual(comm_mode.get_total_counts(), 1) | ||
| 300 | + self.assertEqual(output.shape, (1 * self.world_size, 8)) | ||
| 301 | + | ||
| 302 | + class TestKwargOnlyModule(torch.nn.Module): | ||
| 303 | + def __init__(self) -> None: | ||
| 304 | + super().__init__() | ||
| 305 | + self.linear = torch.nn.Linear(8, 8) | ||
| 306 | + | ||
| 307 | + def forward(self, *, x, y=2, z=None): | ||
| 308 | + return self.linear(x) + y + z | ||
| 309 | + | ||
| 310 | + test_kwonly_mod = TestKwargOnlyModule().to(self.device_type) | ||
| 311 | + prepare_inps_simple = PrepareModuleInput( | ||
| 312 | + input_kwarg_layouts={"x": Shard(0), "z": Shard(0)}, | ||
| 313 | + desired_input_kwarg_layouts={"x": Replicate(), "z": Replicate()}, | ||
| 314 | + ) | ||
| 315 | + parallelize_module( | ||
| 316 | + test_kwonly_mod.linear, mesh, ColwiseParallel(use_local_output=False) | ||
| 317 | + ) | ||
| 318 | + parallelize_module(test_kwonly_mod, mesh, prepare_inps_simple) | ||
| 319 | + | ||
| 320 | + with comm_mode: | ||
| 321 | + output = test_kwonly_mod( | ||
| 322 | + x=torch.randn(1, 8, device=self.device_type), | ||
| 323 | + z=torch.ones(1, 8, device=self.device_type), | ||
| 324 | + ) | ||
| 325 | + | ||
| 326 | + self.assertEqual(comm_mode.get_total_counts(), 2) | ||
| 327 | + self.assertEqual(output.shape, (1 * self.world_size, 8)) | ||
| 328 | + | ||
| 329 | + # test the case where x is a DTensor | ||
| 330 | + x_dt = DTensor.from_local( | ||
| 331 | + torch.randn(1, 8, device=self.device_type), mesh, [Shard(0)] | ||
| 332 | + ) | ||
| 333 | + with comm_mode: | ||
| 334 | + output = test_kwonly_mod( | ||
| 335 | + x=x_dt, z=torch.ones(1, 8, device=self.device_type) | ||
| 336 | + ) | ||
| 337 | + | ||
| 338 | + self.assertEqual(comm_mode.get_total_counts(), 2) | ||
| 339 | + self.assertEqual(output.shape, (1 * self.world_size, 8)) | ||
| 340 | + | ||
| 341 | + | ||
| 342 | + | ||
| 343 | + def test_prepare_module_output(self): | ||
| 344 | + mesh = init_device_mesh(self.device_type, (self.world_size,)) | ||
| 345 | + | ||
| 346 | + tensor = torch.ones(8, 16, device=self.device_type) | ||
| 347 | + expected_tensor = torch.ones(8 // self.world_size, 16, device=self.device_type) | ||
| 348 | + prepare_out_style = PrepareModuleOutput( | ||
| 349 | + output_layouts=Replicate(), desired_output_layouts=Shard(0) | ||
| 350 | + ) | ||
| 351 | + | ||
| 352 | + model = nn.Identity() | ||
| 353 | + chunk_mod = parallelize_module(model, mesh, prepare_out_style) | ||
| 354 | + output = chunk_mod(tensor) | ||
| 355 | + self.assertEqual(output, expected_tensor) | ||
| 356 | + | ||
| 357 | + | ||
| 358 | + | ||
| 359 | + | ||
| 360 | + def test_sequence_parallel_style(self): | ||
| 361 | + mesh = init_device_mesh(self.device_type, (self.world_size,)) | ||
| 362 | + | ||
| 363 | + comm_mode = CommDebugMode() | ||
| 364 | + batch, N, embedding_dim = 20, 8, 12 | ||
| 365 | + | ||
| 366 | + global_input = torch.rand( | ||
| 367 | + batch, | ||
| 368 | + N * self.world_size, | ||
| 369 | + embedding_dim, | ||
| 370 | + device=self.device_type, | ||
| 371 | + requires_grad=True, | ||
| 372 | + ) | ||
| 373 | + sharded_input = distribute_tensor(global_input, mesh, [Shard(1)]) | ||
| 374 | + | ||
| 375 | + # test LayerNorm | ||
| 376 | + for elementwise_affine in [True, False]: | ||
| 377 | + norm = nn.LayerNorm( | ||
| 378 | + embedding_dim, | ||
| 379 | + elementwise_affine=elementwise_affine, | ||
| 380 | + device=self.device_type, | ||
| 381 | + ) | ||
| 382 | + sp_norm = parallelize_module(deepcopy(norm), mesh, SequenceParallel()) | ||
| 383 | + | ||
| 384 | + output = norm(global_input) | ||
| 385 | + output.sum().backward() | ||
| 386 | + | ||
| 387 | + with comm_mode: | ||
| 388 | + sharded_out = sp_norm(sharded_input) | ||
| 389 | + grad_out = torch.ones_like(sharded_out) | ||
| 390 | + sharded_out.backward(grad_out) | ||
| 391 | + self.assertIsInstance(sharded_out, DTensor) | ||
| 392 | + self.assertEqual(sharded_out.placements, (Shard(1),)) | ||
| 393 | + self.assertEqual(comm_mode.get_total_counts(), 0) | ||
| 394 | + self.assertEqual( | ||
| 395 | + comm_mode.get_comm_counts()[c10d_functional.all_reduce], 0 | ||
| 396 | + ) | ||
| 397 | + if elementwise_affine: | ||
| 398 | + self.assertEqual(sp_norm.weight.grad.placements, (_Partial(),)) | ||
| 399 | + self.assertEqual(sp_norm.bias.grad.placements, (_Partial(),)) | ||
| 400 | + | ||
| 401 | + self.assertEqual(sharded_out.full_tensor(), output) | ||
| 402 | + | ||
| 403 | + # test RMSNorm | ||
| 404 | + rmsnorm = RMSNormPython(embedding_dim).to(self.device_type) | ||
| 405 | + sp_rmsnorm = parallelize_module(deepcopy(rmsnorm), mesh, SequenceParallel()) | ||
| 406 | + | ||
| 407 | + output = rmsnorm(global_input) | ||
| 408 | + output.sum().backward() | ||
| 409 | + | ||
| 410 | + with comm_mode: | ||
| 411 | + sharded_out = sp_rmsnorm(sharded_input) | ||
| 412 | + grad_out = torch.ones_like(sharded_out) | ||
| 413 | + sharded_out.backward(grad_out) | ||
| 414 | + self.assertIsInstance(sharded_out, DTensor) | ||
| 415 | + self.assertEqual(sharded_out.placements, (Shard(1),)) | ||
| 416 | + self.assertEqual(sp_rmsnorm.weight.grad.placements, (_Partial(),)) | ||
| 417 | + self.assertEqual(comm_mode.get_total_counts(), 0) | ||
| 418 | + self.assertEqual(comm_mode.get_comm_counts()[c10d_functional.all_reduce], 0) | ||
| 419 | + | ||
| 420 | + self.assertEqual(sharded_out.full_tensor(), output) | ||
| 421 | + | ||
| 422 | + # test sharded on non-sequence dim input | ||
| 423 | + sharded_batch_input = distribute_tensor(global_input, mesh, [Shard(0)]) | ||
| 424 | + rmsnorm = RMSNormPython(embedding_dim).to(self.device_type) | ||
| 425 | + sp_rmsnorm = parallelize_module(deepcopy(rmsnorm), mesh, SequenceParallel()) | ||
| 426 | + | ||
| 427 | + with comm_mode: | ||
| 428 | + sharded_out = sp_rmsnorm(sharded_batch_input) | ||
| 429 | + grad_out = torch.ones_like(sharded_out) | ||
| 430 | + sharded_out.backward(grad_out) | ||
| 431 | + self.assertIsInstance(sharded_out, DTensor) | ||
| 432 | + # output still sharded on sequence dimension | ||
| 433 | + self.assertEqual(sharded_out.placements, (Shard(1),)) | ||
| 434 | + self.assertEqual(sp_rmsnorm.weight.grad.placements, (_Partial(),)) | ||
| 435 | + # communication happens in both fwd/bwd to redistribute input | ||
| 436 | + self.assertEqual(comm_mode.get_total_counts(), 2) | ||
| 437 | + | ||
| 438 | + | ||
| 439 | +if __name__ == "__main__": | ||
| 440 | + run_tests() | ||