Aluode/PerceptionLabPortable
0
1# Copyright 2024 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14from __future__ import annotations15 16import math17import operator18import os19import re20from functools import partial, reduce21 22import torch23import torch.distributed as dist24from torch import nn25 26from ..distributed import DistributedConfig27from ..utils import is_torch_greater_or_equal, logging28from ..utils.generic import GeneralInterface29 30 31logger = logging.get_logger(__name__)32 33# Cache this result has it's a C FFI call which can be pretty time-consuming34_torch_distributed_available = torch.distributed.is_available()35 36 37if is_torch_greater_or_equal("2.5") and _torch_distributed_available:38 from torch.distributed.tensor import DTensor, Placement, Replicate, Shard39 40 41def initialize_tensor_parallelism(tp_plan, tp_size=None):42 r"""43 Sets up the device mesh and initialized the backend for tensor parallelism.44 This function is called when the model is loaded and the TP plan is set to 'auto'.45 """46 if tp_plan is None:47 return None, None, None48 49 if not is_torch_greater_or_equal("2.5"):50 raise OSError("Tensor parallel is only supported for `torch>=2.5`.")51 52 # Detect the accelerator on the machine. If no accelerator is available, it returns CPU.53 device_type = torch._C._get_accelerator().type54 current_device = getattr(torch, device_type)55 if not torch.distributed.is_initialized():56 try:57 rank = int(os.environ["RANK"])58 local_rank = int(os.environ["LOCAL_RANK"])59 world_size = int(os.environ["WORLD_SIZE"])60 61 backend_map = {"cuda": "nccl", "cpu": "gloo", "xpu": "xccl", "hpu": "hccl"}62 backend = backend_map.get(device_type)63 if device_type == "cpu" and int(os.environ.get("CCL_WORKER_COUNT", "0")):64 backend = "ccl"65 if device_type == "xpu" and not is_torch_greater_or_equal("2.8", accept_dev=True):66 backend = "ccl"67 68 torch.distributed.init_process_group(backend=backend, rank=rank, world_size=world_size)69 current_device = getattr(torch, device_type)70 if device_type != "cpu":71 current_device.set_device(local_rank)72 73 except Exception as e:74 raise OSError(75 "We tried to initialize torch.distributed for you, but it failed. Make "76 "sure you init torch distributed in your script to use `tp_plan='auto'`."77 ) from e78 79 if device_type != "cpu":80 current_device.set_device(int(os.environ["LOCAL_RANK"]))81 index = current_device.current_device() if device_type != "cpu" else None82 tp_device = torch.device(device_type, index)83 84 # Silence output for non-primary ranks85 if index is not None and index > 0:86 import sys87 88 sys.stdout = open(os.devnull, "w")89 sys.stderr = open(os.devnull, "w")90 91 device_map = tp_device92 tp_size = tp_size if tp_size is not None else torch.distributed.get_world_size()93 device_mesh = torch.distributed.init_device_mesh(tp_device.type, (tp_size,))94 return tp_device, device_map, device_mesh, tp_size95 96 97def _blocks_to_block_sizes(total_size: int, blocks: int | list[int]) -> list[int]:98 """99 Convert block count or proportions to block sizes.100 101 This function accepts102 103 - The number of blocks (int), in which case the block size is104 total_size//blocks; or105 - A list of block sizes (list[int]).106 107 In the second case, if sum(blocks) < total_size, the ratios between108 the block sizes will be preserved. For instance, if blocks is109 [2, 1, 1] and total_size is 1024, the returned block sizes are110 [512, 256, 256].111 """112 if isinstance(blocks, list):113 total_blocks = sum(blocks)114 assert total_size % total_blocks == 0, f"Cannot split {total_size} in proportional blocks: {blocks}"115 part_size = total_size // total_blocks116 return [part_size * block for block in blocks]117 else:118 assert total_size % blocks == 0, f"Prepacked is not divisible by {blocks}"119 single_size = total_size // blocks120 return [single_size] * blocks121 122 123def _get_parameter_tp_plan(parameter_name: str, tp_plan: dict[str, str], is_weight=True) -> str | None:124 """125 Get the TP style for a parameter from the TP plan.126 127 The TP plan is a dictionary that maps parameter names to TP styles.128 The parameter name can be a generic name with wildcards (e.g. "*.weight") or a specific name (e.g. "layer_1.weight").129 130 The `is_weight` is important because for weights, we want to support `.weights` and `.bias` cases seamlessly! but131 not parent classes for `post_init` calls132 """133 generic_param_name = re.sub(r"\d+", "*", parameter_name)134 if generic_param_name in tp_plan:135 return tp_plan[generic_param_name]136 elif "." in generic_param_name and generic_param_name.rsplit(".", 1)[0] in tp_plan and is_weight:137 return tp_plan[generic_param_name.rsplit(".", 1)[0]]138 return None139 140 141str_to_dtype = {142 "BOOL": torch.bool,143 "U8": torch.uint8,144 "I8": torch.int8,145 "I16": torch.int16,146 "F16": torch.float16,147 "BF16": torch.bfloat16,148 "I32": torch.int32,149 "F32": torch.float32,150 "F64": torch.float64,151 "I64": torch.int64,152 "F8_E4M3": torch.float8_e4m3fn,153}154 155 156def get_packed_weights(param, empty_param, device_mesh, rank, dim):157 """158 When weights are packed (gate_up_proj), we need to make sure each shard gets its correct share.159 So if you have: gate_proj ( 16, 5120, 8190)160 and up_proj ( 16, 5120, 8190)161 packed as gate_up_proj ( 16, 5120, 2 * 8190)162 And you shard along the last dimension, you need to interleave the gate and up values:163 164 Now, if we shard along the last dimension across TP_size (Tensor Parallelism size), we must interleave the values from gate and up projections correctly.165 166 Let's take TP_size = 4 for an example:167 168 Packed tensor `gate_up_proj`169 ---------------------------------------------------------------170 [ G0 G1 G2 G3 | G4 G5 G6 G7 | ... | U0 U1 U2 U3 | U4 U5 U6 U7 | ... ]171 ↑─────────────↑ ↑─────────────↑ ↑─────────────↑ ↑─────────────↑172 Gate Slice 0 Gate Slice 1 Up Slice 0 Up Slice 1173 174 Explanation:175 - The first half of the tensor (left of the center) holds the gate_proj values.176 - The second half (right of the center) holds the up_proj values.177 - For TP=4, we divide each half into 4 slices. In this example, we show two slices for brevity.178 - Each shard receives one slice from the gate part and the corresponding slice from the up part.179 180 For instance:181 • Shard 0 gets: [ Gate Slice 0, Up Slice 0 ] = [ G0, G1, G2, G3, U0, U1, U2, U3 ]182 • Shard 1 gets: [ Gate Slice 1, Up Slice 1 ] = [ G4, G5, G6, G7, U4, U5, U6, U7 ]183 • … and so on.184 185 This ensures that each shard receives an equal portion of both gate and up projections, maintaining consistency across tensor parallelism.186 """187 slice_ = param188 total_size = empty_param.shape[dim]189 world_size = device_mesh.size()190 block_sizes = _blocks_to_block_sizes(total_size=total_size, blocks=2)191 192 tensors_slices = []193 block_offset = 0194 for block_size in block_sizes:195 shard_block_size = block_size // world_size196 start = rank * shard_block_size197 stop = (rank + 1) * shard_block_size198 tensors_slices += range(block_offset + start, block_offset + stop)199 block_offset += block_size200 201 slice_dtype = slice_.get_dtype()202 # Handle F8_E4M3 dtype by converting to float16 before slicing203 # Without upcasting, the slicing causes : RuntimeError: "index_cpu" not implemented for 'Float8_e4m3fn'204 casted = False205 if slice_dtype == "F8_E4M3" or slice_dtype == "F8_E5M2":206 slice_ = slice_[...].to(torch.float16)207 casted = True208 209 if dim == 0:210 tensor = slice_[tensors_slices, ...]211 elif dim == 1 or dim == -2:212 tensor = slice_[:, tensors_slices, ...]213 elif dim == 2 or dim == -1:214 tensor = slice_[..., tensors_slices]215 else:216 raise ValueError(f"Unsupported dim {dim}, only dim 0, 1 or 2 are supported")217 218 if casted:219 return tensor220 else:221 return tensor.to(str_to_dtype[slice_dtype])222 223 224def repack_weights(225 packed_parameter: torch.Tensor,226 sharded_dim: int, # The dimension index in the global tensor that was sharded227 world_size: int,228 num_blocks: int = 2,229) -> torch.Tensor:230 """231 Reorders a tensor that was reconstructed from sharded packed weights into its canonical packed format.232 233 For example, if a weight was packed (e.g., gate_proj and up_proj) and then sharded,234 DTensor.full_tensor() might produce an interleaved layout like [G0, U0, G1, U1, ...]235 along the sharded dimension. This function reorders it to [G0, G1, ..., U0, U1, ...].236 This is an inverse operation to get_packed_weights.237 238 Args:239 reconstructed_tensor: The tensor reconstructed from DTensor (e.g., via .full_tensor().contiguous()).240 sharded_dim: The dimension index in the reconstructed_tensor that was originally sharded.241 world_size: The tensor parallel world size.242 num_packed_projs: The number of projections that were packed together (e.g., 2 for gate_up_proj).243 244 Returns:245 The reordered tensor in canonical packed format.246 """247 248 if num_blocks != 2:249 raise ValueError(250 "Num blocks different from 2 is not supported yet. This is most likely a bug in your implementation as we only pack gate and up projections together."251 )252 253 actual_sharded_dim = sharded_dim if sharded_dim >= 0 else sharded_dim + packed_parameter.ndim254 total_size_on_sharded_dim = packed_parameter.shape[actual_sharded_dim]255 original_block_size_on_dim = total_size_on_sharded_dim // num_blocks256 shard_chunk_size = original_block_size_on_dim // world_size257 258 prefix_shape = packed_parameter.shape[:actual_sharded_dim]259 suffix_shape = packed_parameter.shape[actual_sharded_dim + 1 :]260 261 tensor_view = packed_parameter.view(262 *prefix_shape,263 world_size,264 num_blocks,265 shard_chunk_size,266 *suffix_shape,267 )268 269 # Permute to bring num_packed_projs first, then world_size, then shard_chunk_size270 # This groups all chunks of G together, then all chunks of U together.271 # Target order of these middle dimensions: (num_packed_projs, world_size, shard_chunk_size)272 # Current order of view's middle dimensions: (world_size, num_packed_projs, shard_chunk_size)273 # Absolute indices of the dimensions to be permuted (world_size, num_packed_projs)274 axis_ws_abs = len(prefix_shape)275 axis_npp_abs = len(prefix_shape) + 1276 277 permute_order = list(range(tensor_view.ndim))278 permute_order[axis_ws_abs], permute_order[axis_npp_abs] = permute_order[axis_npp_abs], permute_order[axis_ws_abs]279 280 tensor_permuted = tensor_view.permute(*permute_order)281 282 # Reshape back to the original tensor's ndim, with the sharded dimension now correctly ordered as [G_all, U_all].283 # The final shape should be the same as reconstructed_tensor.284 final_ordered_tensor = tensor_permuted.reshape_as(packed_parameter)285 286 return final_ordered_tensor287 288 289def get_tensor_shard(param, empty_param, device_mesh, rank, dim):290 """291 Generalized tensor sharding across a multi-dimensional device mesh.292 Extract only the fraction of the parameter owned by the given `rank` when the parameter would have gone sharding at provided `dim`.293 Extraction follows the pytorch `Shard` placement so that sharding and materializing back to full tensor follows `Shard` semantics.294 `Shard` follows torch.chunk style sharding of the tensor. We demonstrate some cases below on how sharding happens including some edge cases295 such as some ranks having an empty tensor as shard. Below implementation is robut to all these cases.296 297 Case (1)298 empty_param (16, 5120, 8190)299 dim 0300 device_mesh.size() 4301 rank 0 gets (4, 5120, 8190) (0 ... 4, 5120, 8190)302 rank 1 gets (4, 5120, 8190) (4 ... 8, 5120, 8190)303 rank 2 gets (4, 5120, 8190) (8 ... 12, 5120, 8190)304 rank 3 gets (4, 5120, 8190) (12 ... 16, 5120, 8190)305 306 Case (2)307 empty_param (16, 5120, 8190)308 dim 0309 device_mesh.size() 14310 rank 0 gets (2, 5120, 8190) (0 ... 2, 5120, 8190)311 rank 1 gets (2, 5120, 8190) (2 ... 4, 5120, 8190)312 rank 2 gets (2, 5120, 8190) (4 ... 6, 5120, 8190)313 rank 3 gets (2, 5120, 8190) (6 ... 8, 5120, 8190)314 rank 4 gets (2, 5120, 8190) (8 ... 10, 5120, 8190)315 rank 5 gets (2, 5120, 8190) (10 ... 12, 5120, 8190)316 rank 6 gets (2, 5120, 8190) (12 ... 14, 5120, 8190)317 rank 7 gets (2, 5120, 8190) (14 ... 16, 5120, 8190)318 rank 8 gets (0, 5120, 8190)319 rank 9 gets (0, 5120, 8190)320 rank 10 gets (0, 5120, 8190)321 rank 11 gets (0, 5120, 8190)322 rank 12 gets (0, 5120, 8190)323 rank 13 gets (0, 5120, 8190)324 325 Case (3)326 empty_param (16, 5120, 8190)327 dim 0328 device_mesh.size() 3329 rank 0 gets (6, 5120, 8190) (0 ... 6, 5120, 8190)330 rank 1 gets (6, 5120, 8190) (6 ... 12, 5120, 8190)331 rank 2 gets (4, 5120, 8190) (12 ... 16, 5120, 8190)332 333 In case (2), empty shards are returned with appropriate dimension to allow for operations to work smoothly.334 Args:335 param (torch.Tensor): The tensor to shard.336 empty_param (torch.Tensor): A tensor used for shape reference.337 device_mesh (torch.Tensor): Shape [d_0, ..., d_n] representing the mesh.338 rank (int): Global rank of the current process/device.339 dim (int): Dimension along which to shard the tensor.340 """341 param_dim = empty_param.dim()342 343 if dim < 0:344 dim = param_dim + dim345 if dim >= param_dim:346 raise ValueError(f"dim {dim} is out of bounds for tensor of dimension {param_dim}")347 348 # Flatten the mesh to get the total number of devices349 mesh_shape = device_mesh.shape350 world_size = reduce(operator.mul, mesh_shape)351 352 if rank >= world_size:353 raise ValueError(f"Rank {rank} is out of bounds for mesh size {world_size}")354 355 shard_size = math.ceil(empty_param.shape[dim] / world_size)356 start = rank * shard_size357 358 # Construct slicing index dynamically359 end = min(start + shard_size, empty_param.shape[dim])360 slice_indices = [slice(None)] * param_dim361 if start < empty_param.shape[dim]:362 slice_indices[dim] = slice(start, end)363 return param[tuple(slice_indices)]364 dimensions = list(param.shape)365 dimensions[dim] = 0366 return torch.empty(tuple(dimensions), dtype=torch.int64)367 368 369def distribute_module(370 module: nn.Module,371 device_mesh=None,372 input_fn=None,373 output_fn=None,374) -> nn.Module:375 """376 Copy pasted from torch's function but we remove the communications (partitioning)377 as well as buffer registering that is similarly not efficient.378 """379 if len(module._forward_pre_hooks) == 0:380 if input_fn is not None:381 module.register_forward_pre_hook(lambda mod, inputs: input_fn(mod, inputs, device_mesh))382 if output_fn is not None:383 module.register_forward_hook(lambda mod, inputs, outputs: output_fn(mod, outputs, device_mesh))384 return module385 386 387class TensorParallelLayer:388 """389 General tensor parallel layer for transformers.390 """391 392 use_dtensor = True393 394 @staticmethod395 def _prepare_input_fn(input_layouts, desired_input_layouts, mod, inputs, device_mesh): ...396 397 @staticmethod398 def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh): ...399 400 def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh):401 raise NotImplementedError402 403 def prepare_module_tp(self, module: nn.Module, device_mesh) -> nn.Module:404 if self.use_dtensor:405 distribute_module(406 module,407 device_mesh,408 partial(self._prepare_input_fn, self.input_layouts, self.desired_input_layouts),409 partial(self._prepare_output_fn, self.output_layouts, self.use_local_output),410 )411 412 413# use_dtensor needs to be set to false for nn.Parameter when you want to view, chunk, slice414# you name it. Whatever you want to do that is a bit unconventional, you need local tensors415class GatherParallel(TensorParallelLayer):416 """417 Simple class used to define the hooks to add to a layer when we just want to gather the outputs418 """419 420 def __init__(421 self,422 *,423 input_layouts: Placement | None = None,424 output_layouts: Placement | None = None,425 use_local_output: bool = True,426 ):427 super().__init__()428 self.input_layouts = (input_layouts or Replicate(),)429 self.output_layouts = output_layouts430 self.desired_input_layouts = (Replicate(),)431 self.use_local_output = use_local_output432 433 @staticmethod434 def _prepare_input_fn(input_layouts, desired_input_layouts, mod, inputs, device_mesh):435 mod.expert_parallel_group = device_mesh.get_group()436 if inputs and isinstance(inputs[0], DTensor):437 inputs = inputs[0].to_local()438 return inputs439 440 @staticmethod441 def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh):442 if isinstance(outputs, torch.Tensor):443 dist.all_reduce(outputs, op=dist.ReduceOp.SUM, async_op=False)444 else:445 dist.all_reduce(outputs[0], op=dist.ReduceOp.SUM, async_op=False)446 return outputs447 448 def prepare_module_tp(self, module: nn.Module, device_mesh) -> nn.Module:449 distribute_module(450 module,451 device_mesh,452 partial(self._prepare_input_fn, None, None),453 partial(self._prepare_output_fn, None, None),454 )455 456 457class IsolatedParallel(TensorParallelLayer):458 """459 This class is used to isolate computation in a TP layer from the rest of the world.460 Parameters need to be LOCAL, so not dtensors461 """462 463 @staticmethod464 def _prepare_input_fn(input_layouts, desired_input_layouts, mod, inputs, device_mesh=None):465 # annotate module input placements/sharding with input_layouts466 input_tensor = inputs[0]467 if isinstance(input_tensor, DTensor):468 input_tensor = input_tensor.to_local()469 return input_tensor470 471 @staticmethod472 def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh=None):473 # TODO: figure out dynamo support for instance method and switch this to instance method474 return outputs475 476 def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh):477 param = param[...].to(param_casting_dtype)478 if to_contiguous:479 param = param.contiguous()480 param = param / device_mesh.size() # TODO should be optionable481 # TODO: assumes parent module will allreduce the output afterwards (e.g rowlinear bias is IsolatedParallel and parent module is GatherParallel)482 return param483 484 def prepare_module_tp(self, module: nn.Module, device_mesh) -> nn.Module:485 distribute_module(486 module,487 device_mesh,488 partial(self._prepare_input_fn, None, None),489 partial(self._prepare_output_fn, None, None),490 )491 492 493class ReplicateParallel(TensorParallelLayer):494 """495 This class is used to replicate computation in a TP layer (used in SP regions when we don't use sequence parallelism for example)496 """497 498 def __init__(self, *, use_dtensor=True, use_local_output=True):499 super().__init__()500 self.input_layouts = (Replicate(),)501 self.output_layouts = (Replicate(),)502 self.desired_input_layouts = (Replicate(),)503 self.use_local_output = use_local_output504 self.use_dtensor = use_dtensor505 506 @staticmethod507 def _prepare_input_fn(input_layouts, desired_input_layouts, mod, inputs, device_mesh):508 # TODO: figure out dynamo support for instance method and switch this to instance method509 # annotate module input placements/sharding with input_layouts510 input_tensor = inputs[0]511 if not isinstance(input_tensor, DTensor):512 input_tensor = DTensor.from_local(input_tensor, device_mesh, input_layouts, run_check=False)513 514 return input_tensor515 516 @staticmethod517 def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh):518 return outputs.to_local() if use_local_output and isinstance(outputs, DTensor) else outputs519 520 def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh):521 param = param[...].to(param_casting_dtype)522 if to_contiguous:523 param = param.contiguous()524 param = DTensor.from_local(param, device_mesh, [Replicate()], run_check=False)525 return param526 527 528class ColwiseParallel(TensorParallelLayer):529 """530 General tensor parallel layer for transformers.531 """532 533 def __init__(534 self,535 *,536 input_layouts: Placement | None = None,537 output_layouts: Placement | None = None,538 use_local_output: bool = True,539 use_dtensor=True,540 ):541 super().__init__()542 self.input_layouts = (input_layouts or Replicate(),)543 self.output_layouts = (output_layouts or Shard(-1),)544 self.desired_input_layouts = (Replicate(),)545 self.use_local_output = use_local_output546 self.use_dtensor = use_dtensor547 548 @staticmethod549 def _prepare_input_fn(input_layouts, desired_input_layouts, mod, inputs, device_mesh):550 # TODO: figure out dynamo support for instance method and switch this to instance method551 # annotate module input placements/sharding with input_layouts552 input_tensor = inputs[0]553 if not isinstance(input_tensor, DTensor):554 input_tensor = DTensor.from_local(input_tensor, device_mesh, input_layouts, run_check=False)555 556 # transform the input layouts to the desired layouts of ColwiseParallel557 if input_layouts != desired_input_layouts:558 input_tensor = input_tensor.redistribute(placements=desired_input_layouts, async_op=False)559 return input_tensor560 561 def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh):562 # colwise shard weight/bias to Shard(0), weight be Shard(-2) (0 if you have 1 dim only)563 # means Colwise as Linear is input * weight^T + bias, where564 # weight would become Shard(1)565 if param_type == "bias":566 parameter = get_tensor_shard(param, empty_param, device_mesh, rank, -1)567 shard = [Shard(-1)]568 else:569 shard = [Shard(-2)]570 parameter = get_tensor_shard(param, empty_param, device_mesh, rank, -2)571 572 parameter = parameter.to(param_casting_dtype)573 if to_contiguous:574 parameter = parameter.contiguous()575 if self.use_dtensor:576 parameter = DTensor.from_local(577 parameter, device_mesh, shard, run_check=False, shape=empty_param.size(), stride=empty_param.stride()578 )579 return nn.Parameter(parameter, requires_grad=parameter.is_floating_point())580 581 @staticmethod582 def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh):583 # outputs is a shard on last dimension DTensor, i.e. Shard(-1)584 if outputs.placements != output_layouts:585 outputs = outputs.redistribute(placements=output_layouts, async_op=False)586 # back to local tensor587 return outputs.to_local() if use_local_output and isinstance(outputs, DTensor) else outputs588 589 590class PackedColwiseParallel(ColwiseParallel):591 def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh):592 # colwise shard weight/bias to Shard(0), weight be Shard(-2) (0 if you have 1 dim only)593 # means Colwise as Linear is input * weight^T + bias, where594 # weight would become Shard(1)595 parameter = get_packed_weights(param, empty_param, device_mesh, rank, -2)596 parameter = parameter.to(param_casting_dtype)597 if to_contiguous:598 parameter = parameter.contiguous()599 if self.use_dtensor:600 parameter = DTensor.from_local(parameter, device_mesh, [Shard(-2)], run_check=False)601 return nn.Parameter(parameter, requires_grad=parameter.is_floating_point())602 603 604class RowwiseParallel(TensorParallelLayer):605 """606 Partition a compatible nn.Module in a row-wise fashion. Currently supports nn.Linear and nn.Embedding.607 Users can compose it with ColwiseParallel to achieve the sharding of more complicated modules.608 (i.e. MLP, Attention)609 610 Keyword Args:611 input_layouts (Placement, optional):612 The DTensor layout of input tensor for the nn.Module, this is used to annotate the input tensor to613 become a DTensor. If not specified, we assume the input tensor to be sharded on the last dimension.614 output_layouts (Placement, optional):615 The DTensor layout of the output for the nn.Module, this is used to ensure the output of the nn.Module616 with the user desired layout. If not specified, the output tensor is replicated.617 use_local_output (bool, optional):618 Whether to use local :class:`torch.Tensor` instead of :class:`DTensor` for the module output, default: True.619 Returns:620 A :class:`ParallelStyle` object that represents Rowwise sharding of the nn.Module.621 """622 623 def __init__(624 self,625 *,626 input_layouts: Placement | None = None,627 output_layouts: Placement | None = None,628 use_local_output: bool = True,629 use_dtensor=True,630 ):631 super().__init__()632 self.input_layouts = (input_layouts or Shard(-1),)633 self.output_layouts = (output_layouts or Replicate(),)634 self.use_local_output = use_local_output635 self.use_dtensor = use_dtensor636 637 def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh):638 # Rowwise shard weight to Shard(1), bias to Replicate(), weight be Shard(1)639 # means Rowwise as nn.Linear is input * weight^T + bias, where640 # weight would become Shard(0)641 if param_type != "bias":642 parameter = get_tensor_shard(param, empty_param, device_mesh, rank, -1)643 shard = [Shard(-1)]644 else:645 shard = [Replicate()]646 parameter = param[:]647 648 parameter = parameter.to(param_casting_dtype)649 if to_contiguous:650 parameter = parameter.contiguous()651 if self.use_dtensor:652 parameter = DTensor.from_local(653 parameter, device_mesh, shard, run_check=False, shape=empty_param.size(), stride=empty_param.stride()654 )655 return nn.Parameter(parameter, requires_grad=parameter.is_floating_point())656 657 @staticmethod658 def _prepare_input_fn(input_layouts, desired_input_layouts, mod, inputs, device_mesh):659 if hasattr(mod, "bias") and mod.bias is not None:660 mod._bias = mod.bias.to_local()661 mod.bias = None662 663 input_tensor = inputs[0]664 if not isinstance(input_tensor, DTensor):665 input_tensor = DTensor.from_local(input_tensor, device_mesh, input_layouts, run_check=False)666 667 if input_layouts != desired_input_layouts:668 input_tensor = input_tensor.redistribute(placements=desired_input_layouts, async_op=True)669 return input_tensor670 671 @staticmethod672 def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh):673 # Rowwise sharding produces partial output, depending on output layouts:674 # 1. to replicate -> allreduce675 # 2. to shard -> reduce_scatter676 if outputs.placements != output_layouts:677 outputs = outputs.redistribute(placements=output_layouts, async_op=True)678 outputs = outputs.to_local() # otherwise the `+=` op will gather679 if hasattr(mod, "_bias"):680 outputs = outputs + mod._bias681 # back to local tensor if use_local_output is True682 return outputs683 684 def prepare_module_tp(self, module: nn.Module, device_mesh) -> nn.Module:685 module._distribute_module_applied = True686 if self.use_dtensor:687 if isinstance(module, nn.Linear):688 # rowwise linear runtime sharding requires input tensor shard on last dim689 self.desired_input_layouts: tuple[Placement, ...] = (Shard(-1),)690 elif isinstance(module, nn.Embedding):691 # rowwise embedding runtime sharding requires input tensor replicated692 self.desired_input_layouts = (Replicate(),)693 elif isinstance(module, nn.Parameter):694 # rowwise embedding runtime sharding requires input tensor replicated695 self.desired_input_layouts = (Shard(-1),)696 else:697 raise NotImplementedError("RowwiseParallel currently only support nn.Linear and nn.Embedding!")698 699 distribute_module(700 module,701 device_mesh,702 partial(self._prepare_input_fn, self.input_layouts, self.desired_input_layouts),703 partial(self._prepare_output_fn, self.output_layouts, self.use_local_output),704 )705 706 707class PackedRowwiseParallel(RowwiseParallel):708 def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh):709 # colwise shard weight/bias to Shard(0), weight be Shard(-2) (0 if you have 1 dim only)710 # means Colwise as Linear is input * weight^T + bias, where711 # weight would become Shard(1)712 parameter = get_packed_weights(param, empty_param, device_mesh, rank, -1)713 parameter = parameter.to(param_casting_dtype)714 if to_contiguous:715 parameter = parameter.contiguous()716 if self.use_dtensor:717 parameter = DTensor.from_local(parameter, device_mesh, [Shard(-1)], run_check=False)718 return nn.Parameter(parameter, requires_grad=parameter.is_floating_point())719 720 721class SequenceParallel(TensorParallelLayer):722 """723 SequenceParallel replicates a compatible ``nn.Module`` parameters and runs the sharded computation with724 input sharded on the sequence dimension. This currently supports ``nn.LayerNorm``, ``nn.Dropout``, and the725 `RMSNorm python implementation <https://github.com/facebookresearch/llama/blob/main/llama/model.py#L34>`__726 727 This style implements the operation that is described in the paper728 `Reducing Activation Recomputation in Large Transformer Models <https://huggingface.co/papers/2205.05198>`__729 730 If the input passed in to this ``nn.Module`` is a :class:`torch.Tensor`, it assumes that the input is already sharded731 on the sequence dimension and converts the input to a :class:`DTensor` sharded on the sequence dimension. If the input732 passed in to this ``nn.Module`` is already a :class:`DTensor` but is not sharded on the sequence dimension, it would733 redistribute the input to be sharded on the sequence dimension.734 735 The output of the ``nn.Module`` will be sharded on the sequence dimension.736 737 Keyword Args:738 sequence_dim (int, optional):739 The sequence dimension of the input tensor for the ``nn.Module``, this is used to annotate the input tensor to740 become a DTensor that is sharded on the sequence dimension, default: 1.741 use_local_output (bool, optional):742 Whether to use local :class:`torch.Tensor` instead of :class:`DTensor` for the module output, default: False.743 Returns:744 A :class:`ParallelStyle` object that represents Sequence Parallel of the ``nn.Module``.745 746 Example::747 >>> # xdoctest: +SKIP(failing)748 >>> from torch.distributed.tensor.parallel import parallelize_module, SequenceParallel749 >>> from torch.distributed.device_mesh import init_device_mesh750 >>> ...751 >>> m = Model(...) # m is a nn.Module that contains a "norm" nn.LayerNorm submodule752 >>> tp_mesh = init_device_mesh("cuda", (8,))753 >>>754 >>> # By default, the input of the "norm" will be converted to DTensor that shards on the sequence dim755 >>> # and the output of "norm" will return a sharded on sequence dimension :class:`DTensor`.756 >>>757 >>> sharded_mod = parallelize_module(m, tp_mesh, {"norm": SequenceParallel()}),758 >>> ...759 760 .. note:: SequenceParallel style assumes ones initialization if there are weights in the nn.Module (i.e.761 ``nn.LayerNorm`` or ``RMSNorm``, and they by default have ones initialization). If you have custom762 inits for the weights on those modules, you need to broadcast the weights before/after parallelizing763 to ensure that they are replicated.764 """765 766 def __init__(self, *, sequence_dim: int = 1, use_local_output: bool = False, use_dtensor=False):767 super().__init__()768 self.input_layouts = (Replicate(),)769 self.desired_input_layouts = (Shard(1),)770 self.output_layouts = (Replicate(),)771 self.use_local_output = use_local_output772 self.use_dtensor = True773 self.sequence_sharding = (Shard(sequence_dim),)774 self.use_local_output = use_local_output775 776 @staticmethod777 def _prepare_input_fn(input_layouts, desired_input_layouts, mod, inputs, device_mesh):778 input_tensor = inputs[0]779 if not isinstance(input_tensor, DTensor):780 input_tensor = DTensor.from_local(input_tensor, device_mesh, input_layouts, run_check=False)781 if input_layouts != desired_input_layouts:782 input_tensor = input_tensor.redistribute(placements=desired_input_layouts, async_op=True)783 return input_tensor784 785 @staticmethod786 def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh):787 outputs = outputs.redistribute(788 placements=(Replicate(),), async_op=True789 ) # maybe we have to replicate ? because next layer is not sharded790 return outputs.to_local() # if use_local_output else outputs791 792 def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh):793 # colwise shard weight/bias to Shard(0), weight be Shard(-2) (0 if you have 1 dim only)794 # means Colwise as Linear is input * weight^T + bias, where795 # weight would become Shard(1)796 parameter = param[...]797 parameter = parameter.to(param_casting_dtype)798 if to_contiguous:799 parameter = parameter.contiguous()800 if self.use_dtensor:801 parameter = DTensor.from_local(parameter, device_mesh, [Replicate()], run_check=False)802 return nn.Parameter(parameter, requires_grad=parameter.is_floating_point())803 804 805class GroupedGemmParallel(TensorParallelLayer):806 """807 Applies Expert Parallelism to MoE experts by loading the correct experts on each device.808 """809 810 def __init__(self):811 super().__init__()812 self.use_dtensor = False813 814 def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh):815 ep_rank = rank816 global_num_experts = empty_param.shape[0]817 if global_num_experts % device_mesh.size() != 0:818 raise ValueError(819 f"Global number of experts must be divisible by number of devices: {global_num_experts} % {device_mesh.size()} != 0"820 )821 local_num_experts = global_num_experts // device_mesh.size()822 param = param[ep_rank * local_num_experts : (ep_rank + 1) * local_num_experts].to(param_casting_dtype)823 if to_contiguous:824 param = param.contiguous()825 return param826 827 828class RouterParallel(TensorParallelLayer):829 """830 Allows to reshape the router scores to support running expert parallel.831 """832 833 def __init__(self, *args, **kwargs):834 self.args = args835 self.kwargs = kwargs836 self.use_dtensor = False837 838 @staticmethod839 def _prepare_input_fn(input_layouts, desired_input_layouts, mod, inputs, device_mesh):840 input_tensor = inputs[0]841 if isinstance(input_tensor, DTensor):842 raise NotImplementedError("RouterParallel does not support DTensor input for now")843 return input_tensor844 845 @staticmethod846 def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh):847 """848 Imagine if you had 4 tokens, top_k = 4, and 128experts.849 With EP = 8. The num_local_expert should be 128/8 = 16850 Imagine router_indices being:851 [ 52, 42, 119, 67],852 [102, 89, 61, 40],853 [ 82, 103, 4, 34],854 [ 93, 23, 109, 11],855 856 then you can map which rank should be getting which values857 858 [3, 2, 7, 4],859 [6, 5, 3, 2],860 [5, 6, 0, 2],861 [5, 1, 6, 0],862 863 Thus for say rank 0, you fill with 16 (num_local_expert) the index tensor864 865 [ 16, 16, 16, 16],866 [ 16, 16, 16, 16],867 [ 16, 16, 4, 16],868 [ 16, 16, 16, 11],869 870 This works well. For another rank you need to make sure you round to num_local_expert871 because the next operation will one hot encode the router index vector.872 873 This allows us to know directly which local expert is hit.874 Similarly the scores are indexed with something created form875 router_indices.876 877 The kinda naive training loop that we use for device_map "auto" uses a similar logic.878 Here we are just making each rank believe that he is alone, and he computes his part of the hiddenstates.879 Mask invalid indices with num_local_expert for one-hot encoding, so the computes will skip the masking index.880 """881 ep_rank, ep_size = device_mesh.get_local_rank(), device_mesh.size()882 if mod.num_experts % ep_size != 0:883 raise ValueError(884 f"The number of experts must be divisible by number of ep_size: {mod.num_experts} % {ep_size} != 0"885 )886 num_local_experts = mod.num_experts // ep_size887 router_scores, router_indices = outputs888 router_scores = router_scores[:, ep_rank * num_local_experts : (ep_rank + 1) * num_local_experts]889 router_indices = router_indices.masked_fill((router_indices // num_local_experts) != ep_rank, -1)890 # As -1 % 1 is 0, we can only use mask fill when num_local_experts is 1891 if num_local_experts > 1:892 router_indices = torch.fmod(router_indices, num_local_experts)893 else:894 router_indices = router_indices.masked_fill(router_indices > 0, 0).masked_fill(router_indices < 0, -1)895 router_indices = router_indices.masked_fill(896 router_indices == -1, num_local_experts897 ) # masking class for one hot898 return router_scores, router_indices899 900 def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh):901 # TODO: i'd like for this to be the default902 param = param[...].to(param_casting_dtype)903 if to_contiguous:904 param = param.contiguous()905 return param906 907 def prepare_module_tp(self, module: nn.Module, device_mesh) -> nn.Module:908 # TODO: need an abstract Parallel class that is different from TensorParallelLayer909 distribute_module(910 module,911 device_mesh,912 partial(self._prepare_input_fn, None, None),913 partial(self._prepare_output_fn, None, None),914 )915 916 917class ParallelInterface(GeneralInterface):918 # Class instance object, so that a call to `register` can be reflected into all other files correctly, even if919 # a new instance is created (in order to locally override a given entry)920 _global_mapping = (921 {922 "colwise": ColwiseParallel(),923 "rowwise": RowwiseParallel(),924 "colwise_rep": ColwiseParallel(output_layouts=Replicate()),925 "rowwise_rep": RowwiseParallel(input_layouts=Replicate()),926 "local_colwise": ColwiseParallel(use_dtensor=False),927 "local_rowwise": RowwiseParallel(use_dtensor=False),928 "local": IsolatedParallel(),929 "gather": GatherParallel(),930 "local_packed_rowwise": PackedRowwiseParallel(use_dtensor=False),931 "sequence_parallel": SequenceParallel(),932 "replicate": ReplicateParallel(),933 "grouped_gemm": GroupedGemmParallel(),934 "ep_router": RouterParallel(),935 }936 if is_torch_greater_or_equal("2.5") and _torch_distributed_available937 else {}938 )939 940 941ALL_PARALLEL_STYLES: ParallelInterface = ParallelInterface()942 943 944def convert_local_tensor_to_dtensor(945 parameter: torch.Tensor, parameter_name: str, device_mesh, tp_plan: dict[str, str]946) -> DTensor:947 """948 Converts a local variant of weights to a DTensor with corresponding placements. Shouldn't be done ever except of before saving the model.949 """950 _, param_type = parameter_name.rsplit(".", 1) if "." in parameter_name else parameter_name951 tp_style = _get_parameter_tp_plan(parameter_name, tp_plan)952 if not tp_style:953 return parameter954 955 if tp_style not in ["local_packed_rowwise", "local_rowwise", "local_colwise"]:956 return parameter957 # TODO: this logic should be wrapped in a function, this is copied from corresponding tp classes.958 if tp_style == "local_packed_rowwise":959 placements = [Shard(-1)]960 elif tp_style == "local_rowwise":961 if param_type == "bias":962 placements = [Replicate()]963 else:964 placements = [Shard(-1)]965 elif tp_style == "local_colwise":966 if param_type == "bias":967 placements = [Shard(-1)]968 else:969 placements = [Shard(-2)]970 return DTensor.from_local(parameter, device_mesh, placements, run_check=False)971 972 973def replace_state_dict_local_with_dtensor(974 state_dict: dict[str, torch.Tensor],975 tp_plan: dict[str, str],976 device_mesh,977) -> dict[str, torch.Tensor]:978 """979 Replaces all tensors that were sharded with `local_*` strategy with DTensor to make determining their proper size possible.980 """981 for key, value in state_dict.items():982 if isinstance(value, torch.Tensor) and not isinstance(value, DTensor):983 state_dict[key] = convert_local_tensor_to_dtensor(value, key, device_mesh, tp_plan)984 return state_dict985 986 987def add_tensor_parallel_hooks_to_module(988 model, module, tp_plan, layer_name, current_module_plan, device_mesh, parameter_name=None989):990 r"""991 This function is called in `PretrainedModel.post_init()`. It is responsible of adding hooks992 to the modules of the `model`, based on the `PretrainedModel._tp_plan`.993 994 This is the place where we add the `pre_forward` and `post_forwards` hooks. These are defined995 for each `TensorParallelLayer` as `_prepare_input_fn` and `_prepare_output_fn`.996 997 """998 if current_module_plan is not None:999 tp_layer = ALL_PARALLEL_STYLES[current_module_plan]1000 try:1001 tp_layer.prepare_module_tp(module, device_mesh)1002 except NotImplementedError as e:1003 print(1004 f"Trying to prepare {layer_name}, but it's not supported. Corresponding module: {module} Fix it's TP plan: {e}"1005 )1006 1007 module._hf_tp_plan = current_module_plan1008 module.__repr__ = lambda: f"{module.__repr__()}\nTP Plan: {current_module_plan}"1009 1010 1011def shard_and_distribute_module(1012 model, param, empty_param, parameter_name, param_casting_dtype, is_contiguous, rank, device_mesh1013): # TODO: rename to shard_and_distribute_param1014 r"""1015 This function is called in `from_pretrained` when loading a model's checkpoints.1016 It receives the pointer to the parameter (or the parameter itself) and takes care of "sharding".1017 All process run this function, so they just load the partition of the tensor that they require.1018 1019 Main uses cases:1020 - column / rowise parallelism, you just shard all the weights of the layer (weight and bias)1021 - packed layers: you slice the weights, then shard like above1022 - custom operation:1023 - you want to add an all-gather at the end of a local layer.1024 - you want to have a layer that is isolated from the rest of the world (because torch.DTensor does not work well with `.view` for instance)1025 1026 """1027 param_name, param_type = parameter_name.rsplit(".", 1) if "." in parameter_name else parameter_name1028 tp_plan = model.tp_plan or {}1029 module_to_tp = model.get_submodule(param_name) # TODO: can i loop over modules?1030 rank = int(rank)1031 current_shard_plan = _get_parameter_tp_plan(parameter_name, tp_plan)1032 1033 if dist.get_rank() == 0:1034 if current_shard_plan is None:1035 logger.info(f"Tensor sharding plan for {param_name} not found, using default 'replicate' plan.")1036 else:1037 logger.info(f"Tensor sharding plan for {param_name}: {current_shard_plan}")1038 1039 if current_shard_plan is not None:1040 try:1041 tp_layer = ALL_PARALLEL_STYLES[current_shard_plan]1042 param = tp_layer.partition_tensor(1043 param, empty_param, param_type, param_casting_dtype, is_contiguous, rank, device_mesh1044 )1045 except NotImplementedError as e:1046 print(1047 f"Trying to prepare {parameter_name}, but it's not supported. Corresponding module: {module_to_tp} Fix it's TP plan, current layer: {tp_layer} : {e}"1048 )1049 else:1050 param = param[:].to(param_casting_dtype)1051 1052 # SUPER IMPORTANT we have to use setattr1053 # otherwise loading is crazy slow1054 if not isinstance(param, torch.nn.Parameter):1055 param = torch.nn.Parameter(param, requires_grad=empty_param.is_floating_point())1056 setattr(module_to_tp, param_type, param)1057 # module_to_tp.load_state_dict({param_type: param}, strict=False, assign=True)1058 return param1059 1060 1061def verify_tp_plan(expected_keys: list[str], tp_plan: dict[str, str] | None):1062 """1063 Verify the TP plan of the model, log a warning if the layers that were not sharded and the rules that were not applied.1064 """1065 1066 if tp_plan is None:1067 return1068 1069 generic_keys = {re.sub(r"\d+", "*", key) for key in expected_keys}1070 unsharded_layers = set(generic_keys)1071 unused_rules = tp_plan1072 1073 for key in generic_keys:1074 param_name = key.rsplit(".", 1)[0] if "." in key else key1075 generic_param_name = re.sub(r"\d+", "*", param_name)1076 1077 if generic_param_name in tp_plan:1078 unused_rules.pop(generic_param_name)1079 unsharded_layers.discard(key)1080 elif "." in generic_param_name and (parent_param_name := generic_param_name.rsplit(".", 1)[0]) in tp_plan:1081 unused_rules.pop(parent_param_name)1082 unsharded_layers.discard(key)1083 else:1084 pass # we couldn't find the rule for this parameter, so it's not sharded1085 1086 if len(unused_rules) > 0:1087 logger.warning(f"The following TP rules were not applied on any of the layers: {unused_rules}")1088 if len(unsharded_layers) > 0:1089 logger.warning(f"The following layers were not sharded: {', '.join(unsharded_layers)}")1090 1091 1092def distribute_model(model, distributed_config, device_mesh, tp_size):1093 model._tp_size = tp_size1094 model._device_mesh = device_mesh1095 if distributed_config is not None:1096 if isinstance(distributed_config, dict):1097 distributed_config = DistributedConfig.from_dict(distributed_config)1098 model.config.distributed_config = distributed_config1099 model_plan = model.tp_plan1100 if model_plan is not None and is_torch_greater_or_equal("2.5") and _torch_distributed_available:1101 for v in model_plan.values():1102 if v not in ALL_PARALLEL_STYLES:1103 raise ValueError(f"Unsupported tensor parallel style {v}. Supported styles are {ALL_PARALLEL_STYLES}")1104 for name, module in model.named_modules():1105 if not getattr(module, "_is_hooked", False):1106 plan = _get_parameter_tp_plan(parameter_name=name, tp_plan=model_plan, is_weight=False)1107 add_tensor_parallel_hooks_to_module(1108 model=model,1109 module=module,1110 tp_plan=model_plan,1111 layer_name="",1112 current_module_plan=plan,1113 device_mesh=device_mesh,1114 )1115 module._is_hooked = True1116 return model1117