Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / dtensor / dtensor.py: 86%
279 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
1# Copyright 2025-2026 Huawei Technologies Co., Ltd
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 at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# 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 and
13# limitations under the License.
14# ============================================================================
15"""dtensor"""
16import copy as cp
17import inspect
18import warnings
19from typing import Any, Callable, Optional, Sequence, Set, Tuple, Union
21import numpy as np
23from hyper_parallel.core.dtensor.device_mesh import _mesh_resources
24from hyper_parallel.core.dtensor.layout import Layout, DeviceMesh, _get_slice_tensor_by_layout
25from hyper_parallel.core.dtensor.placement_types import Placement, Replicate
26from hyper_parallel.platform import get_platform
27from hyper_parallel.platform.platform import PlatformType
28from hyper_parallel.core.utils import compute_local_shape_and_global_offset
30platform = get_platform()
31DTensorBase = platform.DTensorBase
32Tensor = platform.Tensor
35class SkipDTensorDispatch():
36 """Context manager that disables DTensor op dispatch for the enclosed block.
38 Args:
39 no_skip: Optional set of op callables or canonical op name strings that
40 should still be dispatched through DTensor even within this context.
41 All other ops bypass DTensor dispatch and operate on local tensors.
43 Example:
44 >>> import torch
45 >>> with SkipDTensorDispatch(no_skip={torch.zeros_like}):
46 ... # zeros_like still goes through DTensor dispatch;
47 ... # everything else uses the local tensor path.
48 ... result = torch.zeros_like(dtensor)
49 """
51 def __init__(self, no_skip: Optional[Set] = None):
52 self._no_skip_names: frozenset = frozenset()
53 if no_skip:
54 names = set()
55 for op in no_skip:
56 if isinstance(op, str):
57 names.add(op)
58 else:
59 names.add(platform.get_op_name(op))
60 self._no_skip_names = frozenset(names)
61 self._dispatch_token = None
62 self._ops_token = None
64 def __enter__(self):
65 # pylint: disable=C0415
66 from hyper_parallel.core.shard._op_dispatch import _dtensor_dispatch_disabled, _no_skip_ops
67 self._dispatch_token = _dtensor_dispatch_disabled.set(True)
68 if self._no_skip_names:
69 self._ops_token = _no_skip_ops.set(_no_skip_ops.get() | self._no_skip_names)
71 def __exit__(self, exc_type, exc_val, exc_tb):
72 # pylint: disable=C0415
73 from hyper_parallel.core.shard._op_dispatch import _dtensor_dispatch_disabled, _no_skip_ops
74 if self._ops_token is not None:
75 _no_skip_ops.reset(self._ops_token)
76 self._ops_token = None
77 _dtensor_dispatch_disabled.reset(self._dispatch_token)
78 self._dispatch_token = None
81# Cache for _build_layout to avoid redundant Layout computations
82# Key: (device_mesh.to_hash(), tuple(placements), tensor_dim)
83# Value: Layout
84_LAYOUT_CACHE = {}
87def _is_alias_placements(placements) -> bool:
88 """
89 Check if placements use alias strings rather than Placement objects.
91 Alias placements use mesh dimension names (strings) to specify
92 the sharding strategy, e.g., ("dp", "tp") or (("dp", "tp"), "None").
93 All elements must be strings or tuples of strings for the sequence
94 to be recognized as alias-style.
96 Args:
97 placements: A sequence of placement specifications.
99 Returns:
100 bool: True if all elements are alias strings or tuples of strings.
101 """
102 if len(placements) == 0:
103 return False
104 for p in placements:
105 if isinstance(p, str):
106 continue
107 if isinstance(p, tuple) and len(p) > 0 and all(isinstance(x, str) for x in p):
108 continue
109 return False
110 return True
113def _build_layout(
114 device_mesh: DeviceMesh,
115 placements: Union[Sequence[Placement], Sequence[Union[str, Tuple[str, ...]]]],
116 tensor_dim: int
117) -> Layout:
118 """
119 Build Layout from device_mesh and placements.
121 This function uses a cache to avoid redundant Layout computations
122 for the same (device_mesh, placements, tensor_dim) combination.
124 Args:
125 device_mesh: The device mesh describing the device topology.
126 placements: Supports two styles:
127 - Placement objects (Shard, Replicate, etc.)
128 - Alias strings ("dp", "None", ("dp", "tp"), etc.), length must
129 equal the number of tensor dimensions (``tensor_dim``).
130 tensor_dim: Number of dimensions in the tensor.
132 Returns:
133 Layout: The built layout object.
135 Raises:
136 ValueError: If alias placements length does not match tensor dimensions.
137 """
138 mesh_key = device_mesh.to_hash()
139 placements_key = tuple(placements)
140 cache_key = (mesh_key, placements_key, tensor_dim)
142 if cache_key in _LAYOUT_CACHE:
143 return _LAYOUT_CACHE[cache_key]
145 layout = Layout.from_device_mesh(device_mesh)
147 if _is_alias_placements(placements):
148 if len(placements) != tensor_dim:
149 raise ValueError(
150 f"Alias placements length ({len(placements)}) must equal "
151 f"tensor dimensions ({tensor_dim})."
152 )
153 result = layout(*placements)
154 else:
155 result = layout(placements)
156 result.placement_to_tensor_map(tensor_dim)
158 _LAYOUT_CACHE[cache_key] = result
160 return result
163def _is_broadcastable(src_shape: Sequence[int], dst_shape: Sequence[int]) -> bool:
164 """Return True iff ``src_shape`` is broadcastable to ``dst_shape``.
166 Standard NumPy / PyTorch right-aligned broadcast rule: ``src`` cannot
167 have more dimensions than ``dst``; each right-aligned dimension pair
168 must be equal, or ``src``'s dimension must be 1.
169 """
170 src_shape = tuple(src_shape)
171 dst_shape = tuple(dst_shape)
172 if len(src_shape) > len(dst_shape):
173 return False
174 for i in range(1, len(src_shape) + 1):
175 s, d = src_shape[-i], dst_shape[-i]
176 if s not in (d, 1):
177 return False
178 return True
181class DTensor(DTensorBase):
182 """
183 DTensor - Distributed Tensor
185 A DTensor represents a tensor that is distributed across multiple devices
186 according to a DeviceMesh and placement specifications.
188 Args:
189 local_tensor (Tensor): The local tensor shard on this device.
190 device_mesh (DeviceMesh): The device mesh describing the device topology.
191 placements: The placement strategy. Supports two styles:
192 - Placement objects (e.g., ``[Shard(0), Replicate()]``).
193 - Alias strings (e.g., ``("dp", "None")`` or
194 ``(("dp", "tp"), "None")``), length must equal the number of
195 tensor dimensions.
197 Example:
198 >>> mesh = init_device_mesh(device_type="npu", mesh_shape=(2, 2), mesh_dim_names=("dp", "tp"))
199 >>> local_tensor = Tensor(np.ones((4, 4)))
200 >>> # Placement style
201 >>> dtensor = DTensor.from_local(local_tensor, mesh, [Shard(0), Replicate()])
202 >>> # Alias style — length matches tensor dims
203 >>> dtensor = DTensor.from_local(local_tensor, mesh, ("dp", "None"))
204 """
205 _local_tensor: Tensor
206 _device_mesh: DeviceMesh
207 _placements: Sequence[Placement]
209 def __init_data__(
210 self,
211 local_tensor: Tensor,
212 device_mesh: DeviceMesh,
213 placements: Union[Sequence[Placement], Sequence[Union[str, Tuple[str, ...]]]]
214 ):
215 self._local_tensor = local_tensor
216 self._device_mesh = device_mesh
217 self._layout = _build_layout(
218 device_mesh, placements, len(local_tensor.shape)
219 )
220 self._placements = tuple(self._layout.placements)
222 @property
223 def device_mesh(self) -> DeviceMesh:
224 """The device mesh of this DTensor."""
225 return self._device_mesh
227 @property
228 def placements(self) -> Sequence[Placement]:
229 """The placements of this DTensor."""
230 return self._placements
232 @property
233 def layout(self) -> Layout:
234 """Internal layout for redistribution (for backward compatibility)."""
235 if not hasattr(self, '_layout'):
236 return None
237 return self._layout
239 @staticmethod
240 def from_local(
241 local_tensor: Tensor,
242 device_mesh: DeviceMesh,
243 placements: Union[Sequence[Placement], Sequence[Union[str, Tuple[str, ...]]]]
244 ) -> 'DTensor':
245 """
246 Create a DTensor from a local tensor with device mesh and placements.
248 Args:
249 local_tensor (Tensor): The local tensor shard on this device.
250 device_mesh (DeviceMesh): The device mesh describing the device topology.
251 placements: The placement strategy. Supports two styles:
252 - Placement objects (e.g., ``[Shard(0), Replicate()]``).
253 - Alias strings (e.g., ``("dp", "None")`` or
254 ``(("dp", "tp"), "None")``), length must equal the number
255 of tensor dimensions.
257 Returns:
258 DTensor: A new DTensor instance.
260 Example:
261 >>> mesh = init_device_mesh(device_type="npu", mesh_shape=(2, 2), mesh_dim_names=("dp", "tp"))
262 >>> local_tensor = Tensor(np.ones((4, 4)))
263 >>> dtensor = DTensor.from_local(local_tensor, mesh, [Shard(0), Replicate()])
264 >>> dtensor = DTensor.from_local(local_tensor, mesh, ("dp", "None"))
265 """
266 return DTensor(local_tensor, device_mesh, placements)
268 def _alias_placements(self) -> Sequence[Placement]:
269 """Return alias_placements from layout, falling back to _placements."""
270 if hasattr(self, '_layout') and self._layout:
271 return self._layout.alias_placements
272 return self._placements
274 def _from_converted_local(self, local_tensor: Tensor) -> 'DTensor':
275 """Rebuild converted DTensor data without preserving Parameter identity."""
276 cls = DTensor if isinstance(self, platform.Parameter) else self.__class__
277 return cls(local_tensor, device_mesh=self._device_mesh,
278 placements=self._alias_placements())
280 def to(self, *args, **kwargs):
281 """Move the DTensor to a different device or dtype.
283 Delegates to the underlying local tensor's ``to`` method and
284 reconstructs a DTensor preserving device_mesh and placements.
286 Args:
287 *args (tuple): Arguments passed to the underlying tensor's ``to``
288 method (e.g., device or dtype).
289 **kwargs (dict): Keyword arguments for the tensor conversion
290 (e.g., dtype, device, non_blocking).
292 Returns:
293 DTensor: A new DTensor with the converted local tensor.
294 """
295 new_local = self._local_tensor.to(*args, **kwargs)
296 return self._from_converted_local(new_local)
298 def float(self):
299 """Convert the DTensor to float dtype.
301 Returns:
302 DTensor: A new DTensor with float32 local tensor.
303 """
304 new_local = self._local_tensor.float()
305 return self._from_converted_local(new_local)
307 def to_local(self) -> Tensor:
308 """
309 Convert DTensor to local tensor.
311 Returns:
312 Tensor: The local tensor shard on this device.
313 """
314 return self._local_tensor
316 def copy_(self, src: "DTensor", non_blocking: bool = False) -> "DTensor":
317 """In-place copy of ``src`` into this DTensor's local shard.
319 Delegates to ``Tensor.copy_`` on the underlying local tensors.
320 Follows standard ``Tensor.copy_`` semantics: version counter is
321 bumped and autograd edges are created when grad is enabled.
323 Constraints on ``src``:
324 * must be a ``DTensor`` on the same ``DeviceMesh`` as ``self``;
325 * its placements must equal ``self.placements``, OR
326 ``src._local_tensor.numel() == 1`` (single-element broadcast);
327 * its local shape must equal or be broadcastable to
328 ``self._local_tensor.shape``.
330 No redistribute / implicit slicing is performed; src dtype is cast
331 to self dtype in-place.
333 Args:
334 src (DTensor): Source DTensor satisfying the constraints above.
335 non_blocking (bool): Forwarded to the underlying ``copy_``.
337 Returns:
338 DTensor: ``self``.
340 Raises:
341 TypeError: if ``src`` is not a ``DTensor``.
342 ValueError: if mesh, placement, or shape constraint is violated.
343 """
344 if not isinstance(src, DTensor):
345 raise TypeError(
346 f"For DTensor.copy_, src should be a DTensor, but got {type(src).__name__}."
347 )
348 src_local = src.to_local()
349 if src.device_mesh is not self._device_mesh:
350 raise ValueError(
351 f"For DTensor.copy_, src and self should share the same DeviceMesh, "
352 f"but got src.device_mesh={src.device_mesh!r}, "
353 f"self._device_mesh={self._device_mesh!r}."
354 )
356 placement_eq = tuple(src.placements) == tuple(self._placements)
357 shape_eq = src_local.shape == self._local_tensor.shape
358 src_is_scalar = src_local.numel() == 1
360 if not placement_eq and not src_is_scalar:
361 raise ValueError(
362 f"For DTensor.copy_, src.placements should equal self.placements "
363 f"or src.numel() should be 1, but got "
364 f"src.placements={src.placements}, "
365 f"self.placements={self._placements}, "
366 f"src.numel()={src_local.numel()}."
367 )
368 if not shape_eq and not src_is_scalar and not _is_broadcastable(
369 src_local.shape, self._local_tensor.shape
370 ):
371 raise ValueError(
372 f"For DTensor.copy_, src local shape should be broadcastable to "
373 f"self local shape, but got "
374 f"src.shape={tuple(src_local.shape)}, "
375 f"self.shape={tuple(self._local_tensor.shape)}."
376 )
378 self._local_tensor.copy_(src_local, non_blocking=non_blocking)
379 return self
381 def zero_(self) -> "DTensor":
382 """In-place fill with zeros. Returns ``self``."""
383 self._local_tensor.zero_()
384 return self
386 def fill_(self, value) -> "DTensor":
387 """In-place fill with ``value``. Returns ``self``."""
388 self._local_tensor.fill_(value)
389 return self
391 @property
392 def shape(self) -> Tuple[int, ...]:
393 """
394 The global shape of this DTensor.
396 Returns:
397 Tuple[int, ...]: The global tensor shape.
398 """
399 return self._layout.get_global_shape(self._local_tensor.shape)
401 def size(self, dim=None):
402 """Return the global shape, consistent with .shape.
404 Without ``dim`` returns a tuple matching ``self.shape``.
405 With ``dim`` returns the size of that dimension.
406 """
407 global_shape = self.shape
408 if dim is not None:
409 return global_shape[dim]
410 return global_shape
412 def numel(self) -> int:
413 """Return the number of elements in this DTensor."""
414 return int(np.prod(self.shape))
416 @property
417 def local_shape(self) -> Tuple[int, ...]:
418 """
419 The local shape of this DTensor on this device.
421 Returns:
422 Tuple[int, ...]: The local tensor shape.
423 """
424 return self._local_tensor.shape
426 def redistribute(
427 self,
428 device_mesh: DeviceMesh,
429 placements: Union[Sequence[Placement], Sequence[Union[str, Tuple[str, ...]]]]
430 ) -> 'DTensor':
431 """
432 Redistribute this DTensor to a new device mesh and placements.
434 Args:
435 device_mesh (DeviceMesh): The target device mesh.
436 placements: The target placements. Supports Placement objects
437 or alias strings.
439 Returns:
440 DTensor: A new DTensor with the specified distribution.
442 Example:
443 >>> new_dtensor = dtensor.redistribute(mesh, [Replicate(), Shard(1)])
444 >>> new_dtensor = dtensor.redistribute(mesh, ("None", "tp"))
445 """
446 # Build dst_layout from device_mesh and placements
447 dst_layout = _build_layout(
448 device_mesh, placements, len(self._local_tensor.shape)
449 )
451 # pylint: disable=C0415
452 from hyper_parallel.core.dtensor.tensor_redistribution import _tensor_redistribution
453 out = _tensor_redistribution.redistribution(self, dst_layout)
454 return out
456 def reduce_partial(self) -> 'DTensor':
457 """
458 Reduce partial sharding state for this DTensor.
460 Returns:
461 DTensor: A new DTensor with partial state reduced.
462 """
463 if not self._layout:
464 return self
465 to_layout = cp.deepcopy(self._layout)
466 to_layout.reset_partial()
467 # pylint: disable=C0415
468 from hyper_parallel.core.dtensor.tensor_redistribution import _tensor_redistribution
469 out = _tensor_redistribution.reduce_partial(self, to_layout)
470 return out
472 def full_tensor(self) -> Tensor:
473 """
474 Return the full tensor of this DTensor.
476 Returns:
477 Tensor: A Tensor object that represents the full tensor of this DTensor.
478 The returned tensor contains the complete data gathered from
479 all ranks.
481 Note:
482 This operation involves communication across all ranks in the DeviceMesh,
483 which may be expensive for large tensors. Use with caution in
484 performance-critical code paths.
486 Example:
487 >>> # Assume dtensor is sharded across multiple devices
488 >>> local_tensor = dtensor.to_local() # Returns only the local shard
489 >>> full_tensor = dtensor.full_tensor() # Returns the complete tensor
490 """
491 if not self._layout:
492 return self._local_tensor
494 # Create a fully replicated layout
495 replicated_layout = cp.deepcopy(self._layout)
497 # Set all placements to Replicate and convert to tensor_map
498 replicated_placements = [Replicate()] * len(replicated_layout.mesh_shape)
499 replicated_layout.set_placements(replicated_placements)
500 replicated_layout.placement_to_tensor_map(len(self._local_tensor.shape))
502 # Clear partial status from original layout since Replicate has no partial
503 replicated_layout.reset_partial()
505 # Redistribute to the replicated layout and return local tensor
506 # pylint: disable=C0415
507 from hyper_parallel.core.dtensor.tensor_redistribution import _tensor_redistribution
508 out = _tensor_redistribution.redistribution(self, replicated_layout)
509 return out.to_local()
512def distribute_tensor(
513 tensor: Tensor,
514 device_mesh: DeviceMesh,
515 placements: Union[Sequence[Placement], Sequence[Union[str, Tuple[str, ...]]]]
516) -> DTensor:
517 """
518 Distribute a global tensor to the device mesh according to the placements.
520 Args:
521 tensor (Tensor): The global tensor to be distributed. All ranks
522 should have the same tensor data.
523 device_mesh (DeviceMesh): The device mesh describing the device topology.
524 placements: The placement strategy. Supports two styles:
525 - Placement objects (e.g., ``[Shard(0), Replicate()]``).
526 - Alias strings (e.g., ``("dp", "None")`` or
527 ``(("dp", "tp"), "None")``), length must equal the number of
528 tensor dimensions.
530 Returns:
531 DTensor: A new DTensor with the local shard on each rank.
533 Note:
534 This method assumes all ranks have the same global tensor. It slices
535 the tensor locally without communication. If ranks have different
536 data, use `from_local` instead.
538 Example:
539 >>> mesh = init_device_mesh(device_type="npu", mesh_shape=(2, 2), mesh_dim_names=("dp", "tp"))
540 >>> global_tensor = Tensor(np.arange(16).reshape(4, 4))
541 >>> dtensor = distribute_tensor(global_tensor, mesh, [Shard(0), Replicate()])
542 >>> dtensor = distribute_tensor(global_tensor, mesh, ("dp", "None"))
543 """
544 layout = _build_layout(device_mesh, placements, len(tensor.shape))
545 local_tensor = _get_slice_tensor_by_layout(tensor, layout)
546 return DTensor(local_tensor, device_mesh, layout.alias_placements)
549def _distribute_module_param_source(param: Any) -> Tensor:
550 """Tensor data used as the global tensor for :func:`distribute_tensor` (PyTorch uses ``param.data``)."""
551 if hasattr(param, "data"):
552 return param.data
553 return platform.get_param_local_data(param)
556def _distribute_module_new_parameter(key: str, dtensor: DTensor, requires_grad: bool) -> Any:
557 """Build a framework :class:`Parameter` holding *dtensor* (Torch vs MindSpore kwargs differ)."""
558 if platform.platform_type == PlatformType.MINDSPORE:
559 return platform.Parameter(dtensor, name=key, requires_grad=requires_grad)
560 return platform.Parameter(dtensor, requires_grad=requires_grad)
563def _distribute_module_set_param(module: Any, key: str, new_param: Any) -> None:
564 """Register or assign a parameter on *module* (``nn.Module`` or MindSpore ``Cell``)."""
565 if hasattr(module, "register_parameter"):
566 module.register_parameter(key, new_param)
567 return
568 if hasattr(module, "_params"):
569 module._params[key] = new_param
570 if hasattr(module, "_params_list"):
571 module._params_list[key] = new_param
572 if key in module.__dict__:
573 module.__dict__[key] = new_param
574 return
575 raise TypeError(
576 f"distribute_module expects nn.Module-like objects with register_parameter or _params; "
577 f"got {type(module)}."
578 )
581def _distribute_module_iter_params(module: Any) -> list:
582 """Return ``[(name, param), ...]`` for direct parameters (``_parameters`` or ``_params``)."""
583 if hasattr(module, "_parameters"):
584 return list(module._parameters.items())
585 if hasattr(module, "_params"):
586 return list(module._params.items())
587 return []
590def _distribute_module_iter_buffers(module: Any) -> list:
591 """Return ``[(name, buffer), ...]`` if the module has ``_buffers`` (PyTorch ``nn.Module``)."""
592 if hasattr(module, "_buffers"):
593 return list(module._buffers.items())
594 return []
597def _distribute_module_named_modules(module: Any):
598 """``nn.Module.named_modules`` or MindSpore ``Cell.cells_and_names`` (submodule FQNs)."""
599 if hasattr(module, "named_modules"):
600 return module.named_modules()
601 if hasattr(module, "cells_and_names"):
602 return module.cells_and_names()
603 raise TypeError(
604 f"distribute_module expects module-like objects with named_modules or cells_and_names; "
605 f"got {type(module)}."
606 )
609def _replicate_submodule_params_buffers(
610 sub_mod: Any,
611 device_mesh: DeviceMesh,
612 *,
613 module_prefix: str = "",
614) -> None:
615 """Convert plain params/buffers on *sub_mod* to fully replicated :class:`DTensor`."""
616 full_replicate = [Replicate()] * device_mesh.ndim
617 for key, param in _distribute_module_iter_params(sub_mod):
618 if param is None or isinstance(param, DTensorBase):
619 continue
620 src = _distribute_module_param_source(param)
621 requires_grad = bool(getattr(param, "requires_grad", True))
622 dt = distribute_tensor(src, device_mesh, full_replicate)
623 param_name = f"{module_prefix}.{key}" if module_prefix else key
624 new_param = _distribute_module_new_parameter(param_name, dt, requires_grad)
625 _distribute_module_set_param(sub_mod, key, new_param)
626 for key, buffer in _distribute_module_iter_buffers(sub_mod):
627 if buffer is None or isinstance(buffer, DTensorBase):
628 continue
629 sub_mod._buffers[key] = distribute_tensor(buffer, device_mesh, full_replicate)
632def _distribute_module_run_partition_and_replicate(
633 module: Any,
634 device_mesh: DeviceMesh,
635 partition_fn: Optional[Callable[[str, Any, DeviceMesh], None]],
636) -> None:
637 """Call optional ``partition_fn`` per ``named_modules`` and replicate remaining tensors."""
638 if partition_fn is None:
639 for mod_name, submod in _distribute_module_named_modules(module):
640 _replicate_submodule_params_buffers(submod, device_mesh, module_prefix=mod_name)
641 return
642 for mod_name, submod in _distribute_module_named_modules(module):
643 partition_fn(mod_name, submod, device_mesh)
644 _replicate_submodule_params_buffers(submod, device_mesh, module_prefix=mod_name)
647def _distribute_module_register_input_fn(
648 module: Any,
649 device_mesh: DeviceMesh,
650 input_fn: Callable[..., Any],
651) -> None:
652 """Register *input_fn* as a forward pre-hook on *module* (2- or 3-arg, PyTorch-compatible)."""
653 num_args = len(inspect.signature(input_fn).parameters)
654 if num_args == 2:
655 warnings.warn(
656 "Deprecating input_fn that takes two arguments (inputs, device_mesh), "
657 "please use input_fn that takes in (module, inputs, device_mesh) instead!",
658 FutureWarning,
659 stacklevel=3,
660 )
661 module.register_forward_pre_hook(
662 lambda _, inputs: input_fn(inputs, device_mesh)
663 )
664 elif num_args == 3:
665 module.register_forward_pre_hook(
666 lambda mod, inputs: input_fn(mod, inputs, device_mesh)
667 )
668 else:
669 raise ValueError(
670 f"input_fn should take in 2 or 3 arguments, but got {num_args} arguments!"
671 )
674def _distribute_module_register_output_fn(
675 module: Any,
676 device_mesh: DeviceMesh,
677 output_fn: Callable[..., Any],
678) -> None:
679 """Register *output_fn* as a forward hook on *module* (2- or 3-arg, PyTorch-compatible)."""
680 num_args = len(inspect.signature(output_fn).parameters)
681 if num_args == 2:
682 warnings.warn(
683 "Deprecating output_fn that takes two arguments (outputs, device_mesh), "
684 "please use output_fn that takes in (module, outputs, device_mesh) instead!",
685 FutureWarning,
686 stacklevel=3,
687 )
688 module.register_forward_hook(
689 lambda mod, inputs, outputs: output_fn(outputs, device_mesh)
690 )
691 elif num_args == 3:
692 module.register_forward_hook(
693 lambda mod, inputs, outputs: output_fn(mod, outputs, device_mesh)
694 )
695 else:
696 raise ValueError(
697 f"output_fn should take in 2 or 3 arguments, but got {num_args} arguments!"
698 )
701def distribute_module(
702 module: Any,
703 device_mesh: Optional[DeviceMesh] = None,
704 partition_fn: Optional[Callable[[str, Any, DeviceMesh], None]] = None,
705 input_fn: Optional[Callable[..., Any]] = None,
706 output_fn: Optional[Callable[..., Any]] = None,
707) -> Any:
708 """PyTorch ``distribute_module`` parity: shard/replicate params and optional I/O hooks.
710 Unsharded parameters and buffers become fully replicated :class:`DTensor` after
711 ``partition_fn``. ``input_fn`` / ``output_fn`` attach only to the root *module*.
713 Args:
714 module: Root ``nn.Module`` or MindSpore ``Cell`` with compatible APIs.
715 device_mesh: Placement mesh; if ``None``, uses ``_mesh_resources.get_current_mesh()``.
716 partition_fn: Per ``named_modules`` callback before replicate pass; ``None`` replicates all.
717 input_fn: ``(module, inputs, mesh)`` or deprecated ``(inputs, mesh)`` pre-hook.
718 output_fn: ``(module, outputs, mesh)`` or deprecated ``(outputs, mesh)`` forward hook.
720 Returns:
721 *module* in place, with distributed tensors where applied.
723 Raises:
724 RuntimeError: If called twice on the same *module*.
725 ValueError: If ``input_fn`` / ``output_fn`` arity is not 2 or 3.
727 Note:
728 XLA / ``torch_xla`` is not supported; strided device :class:`DTensor` only.
729 """
730 if getattr(module, "_distribute_module_applied", False):
731 raise RuntimeError(
732 "distribute_module should only be called once on a module, "
733 "but it has already been called on this module!"
734 )
735 device_mesh = device_mesh or _mesh_resources.get_current_mesh()
736 _distribute_module_run_partition_and_replicate(module, device_mesh, partition_fn)
737 if input_fn is not None:
738 _distribute_module_register_input_fn(module, device_mesh, input_fn)
739 if output_fn is not None:
740 _distribute_module_register_output_fn(module, device_mesh, output_fn)
741 module._distribute_module_applied = True
742 return module
745def _dtensor_init_helper(
746 init_op,
747 size,
748 device_mesh,
749 placements,
750 **kwargs,
751) -> DTensor:
752 """
753 Helper function to create and initialize a distributed tensor.
755 Args:
756 size: Shape of the tensor.
757 dtype: Data type of the tensor.
758 device: Target device for the tensor.
759 requires_grad: Whether the tensor requires gradient.
761 Returns:
762 DTensor: The initialized distributed tensor.
763 """
764 # get local tensor shape
765 local_shape = compute_local_shape_and_global_offset(
766 size, device_mesh, placements
767 )
769 # initialize the local tensor
770 if init_op is platform.full:
771 fill_value = kwargs.pop("fill_value", 0)
772 local_tensor = init_op(local_shape, fill_value, **kwargs)
773 else:
774 local_tensor = init_op(local_shape, **kwargs)
776 return DTensor.from_local(
777 local_tensor,
778 device_mesh,
779 placements,
780 )
783def ones(
784 size,
785 device_mesh,
786 placements,
787) -> DTensor:
788 """
789 Returns a :class:`DTensor` filled with the scalar value 1, with the shape defined
790 by the variable argument ``size``.
792 Args:
793 size (Union[tuple[int], list[int], int, Tensor]): The specified shape of output tensor. Only positive integer or
794 tuple or Tensor containing positive integers are allowed. If it is a Tensor,
795 it must be a 0-D or 1-D Tensor with int32 or int64 dtypes.
797 Keyword args:
798 device_mesh: :class:`DeviceMesh` type, contains the mesh info of ranks
799 placements: a sequence of :class:`Placement` type: ``Shard``, ``Replicate``
801 Returns:
802 A :class:`DTensor` object on each rank
803 """
804 ones_ = platform.ones
805 return _dtensor_init_helper(
806 ones_,
807 size,
808 device_mesh=device_mesh,
809 placements=placements,
810 )
813def empty(
814 size,
815 device_mesh,
816 placements,
817) -> DTensor:
818 """
819 Returns a :class:`DTensor` filled with uninitialized data. The shape of the :class:`DTensor`
820 is defined by the variable argument ``size``.
822 Args:
823 size (Union[tuple[int], list[int], int]): The specified shape of output tensor. Can be variable numbers of
824 positive integers or tuple or list containing positive integers.
826 Keyword args:
827 device_mesh: :class:`DeviceMesh` type, contains the mesh info of ranks
828 placements: a sequence of :class:`Placement` type: ``Shard``, ``Replicate``
830 Returns:
831 A :class:`DTensor` object on each rank
832 """
833 empty_ = platform.empty
834 return _dtensor_init_helper(
835 empty_,
836 size,
837 device_mesh=device_mesh,
838 placements=placements,
839 )
842def full(
843 size,
844 fill_value,
845 *,
846 device_mesh,
847 placements,
848) -> DTensor:
849 """
850 Returns a :class:`DTensor` filled with ``fill_value`` according to ``device_mesh`` and
851 ``placements``, with the shape defined by the argument ``size``.
853 Args:
854 size (Union[tuple[int], list[int]]): The specified shape of output tensor.
855 fill_value (Union[numbers.Number, Tensor]): Value to fill the returned tensor. It can be a scalar number, a 0-D
856 Tensor, or a 1-D Tensor with only one element.
858 Keyword args:
859 device_mesh: :class:`DeviceMesh` type, contains the mesh info of ranks.
860 placements: a sequence of :class:`Placement` type: ``Shard``, ``Replicate``
862 Returns:
863 A :class:`DTensor` object on each rank
864 """
865 full_ = platform.full
866 return _dtensor_init_helper(
867 full_,
868 size,
869 fill_value=fill_value,
870 device_mesh=device_mesh,
871 placements=placements,
872 )
875def zeros(
876 size,
877 device_mesh,
878 placements,
879) -> DTensor:
880 """
881 Returns a :class:`DTensor` filled with the scalar value 0.
883 Args:
884 size (Union[tuple[int], list[int], int, Tensor]): The specified shape of output tensor. Only positive integer or
885 tuple or Tensor containing positive integers are allowed. If it is a Tensor,
886 it must be a 0-D or 1-D Tensor with int32 or int64 dtypes.
887 Keyword args:
888 device_mesh: :class:`DeviceMesh` type, contains the mesh info of ranks
889 placements: a sequence of :class:`Placement` type: ``Shard``, ``Replicate``
891 Returns:
892 A :class:`DTensor` object on each rank
893 """
894 zeros_ = platform.zeros
895 return _dtensor_init_helper(
896 zeros_,
897 size,
898 device_mesh=device_mesh,
899 placements=placements,
900 )