Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / dtensor / dtensor.py: 87%
308 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-13 05:07 +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._collective_utils import mesh_broadcast, mesh_scatter
25from hyper_parallel.core.dtensor.layout import Layout, DeviceMesh, _get_slice_tensor_by_layout
26from hyper_parallel.core.dtensor.placement_types import Partial, Placement, Replicate, StridedShard
27from hyper_parallel.platform import get_platform
28from hyper_parallel.platform.platform import PlatformType
29from hyper_parallel.core.utils import compute_local_shape_and_global_offset
31platform = get_platform()
32DTensorBase = platform.DTensorBase
33Tensor = platform.Tensor
36class SkipDTensorDispatch():
37 """Context manager that disables DTensor op dispatch for the enclosed block.
39 Args:
40 no_skip: Optional set of op callables or canonical op name strings that
41 should still be dispatched through DTensor even within this context.
42 All other ops bypass DTensor dispatch and operate on local tensors.
44 Example:
45 >>> import torch
46 >>> with SkipDTensorDispatch(no_skip={torch.zeros_like}):
47 ... # zeros_like still goes through DTensor dispatch;
48 ... # everything else uses the local tensor path.
49 ... result = torch.zeros_like(dtensor)
50 """
52 def __init__(self, no_skip: Optional[Set] = None):
53 self._no_skip_names: frozenset = frozenset()
54 if no_skip:
55 names = set()
56 for op in no_skip:
57 if isinstance(op, str):
58 names.add(op)
59 else:
60 names.add(platform.get_op_name(op))
61 self._no_skip_names = frozenset(names)
62 self._dispatch_token = None
63 self._ops_token = None
65 def __enter__(self):
66 # pylint: disable=C0415
67 from hyper_parallel.core.shard._op_dispatch import _dtensor_dispatch_disabled, _no_skip_ops
68 self._dispatch_token = _dtensor_dispatch_disabled.set(True)
69 if self._no_skip_names:
70 self._ops_token = _no_skip_ops.set(_no_skip_ops.get() | self._no_skip_names)
72 def __exit__(self, exc_type, exc_val, exc_tb):
73 # pylint: disable=C0415
74 from hyper_parallel.core.shard._op_dispatch import _dtensor_dispatch_disabled, _no_skip_ops
75 if self._ops_token is not None:
76 _no_skip_ops.reset(self._ops_token)
77 self._ops_token = None
78 _dtensor_dispatch_disabled.reset(self._dispatch_token)
79 self._dispatch_token = None
82# Cache for _build_layout to avoid redundant Layout computations
83# Key: (device_mesh.to_hash(), tuple(placements), tensor_dim)
84# Value: Layout
85_LAYOUT_CACHE = {}
88def _is_alias_placements(placements) -> bool:
89 """
90 Check if placements use alias strings rather than Placement objects.
92 Alias placements use mesh dimension names (strings) to specify
93 the sharding strategy, e.g., ("dp", "tp") or (("dp", "tp"), "None").
94 All elements must be strings or tuples of strings for the sequence
95 to be recognized as alias-style.
97 Args:
98 placements: A sequence of placement specifications.
100 Returns:
101 bool: True if all elements are alias strings or tuples of strings.
102 """
103 if len(placements) == 0:
104 return False
105 for p in placements:
106 if isinstance(p, str):
107 continue
108 if isinstance(p, tuple) and len(p) > 0 and all(isinstance(x, str) for x in p):
109 continue
110 return False
111 return True
114def _build_layout(
115 device_mesh: DeviceMesh,
116 placements: Union[Sequence[Placement], Sequence[Union[str, Tuple[str, ...]]]],
117 tensor_dim: int
118) -> Layout:
119 """
120 Build Layout from device_mesh and placements.
122 This function uses a cache to avoid redundant Layout computations
123 for the same (device_mesh, placements, tensor_dim) combination.
125 Args:
126 device_mesh: The device mesh describing the device topology.
127 placements: Supports two styles:
128 - Placement objects (Shard, Replicate, etc.)
129 - Alias strings ("dp", "None", ("dp", "tp"), etc.), length must
130 equal the number of tensor dimensions (``tensor_dim``).
131 tensor_dim: Number of dimensions in the tensor.
133 Returns:
134 Layout: The built layout object.
136 Raises:
137 ValueError: If alias placements length does not match tensor dimensions.
138 """
139 mesh_key = device_mesh.to_hash()
140 placements_key = tuple(placements)
141 cache_key = (mesh_key, placements_key, tensor_dim)
143 if cache_key in _LAYOUT_CACHE:
144 return _LAYOUT_CACHE[cache_key]
146 layout = Layout.from_device_mesh(device_mesh)
148 if _is_alias_placements(placements):
149 if len(placements) != tensor_dim:
150 raise ValueError(
151 f"Alias placements length ({len(placements)}) must equal "
152 f"tensor dimensions ({tensor_dim})."
153 )
154 result = layout(*placements)
155 else:
156 result = layout(placements)
157 result.placement_to_tensor_map(tensor_dim)
159 _LAYOUT_CACHE[cache_key] = result
161 return result
164def _is_broadcastable(src_shape: Sequence[int], dst_shape: Sequence[int]) -> bool:
165 """Return True iff ``src_shape`` is broadcastable to ``dst_shape``.
167 Standard NumPy / PyTorch right-aligned broadcast rule: ``src`` cannot
168 have more dimensions than ``dst``; each right-aligned dimension pair
169 must be equal, or ``src``'s dimension must be 1.
170 """
171 src_shape = tuple(src_shape)
172 dst_shape = tuple(dst_shape)
173 if len(src_shape) > len(dst_shape):
174 return False
175 for i in range(1, len(src_shape) + 1):
176 s, d = src_shape[-i], dst_shape[-i]
177 if s not in (d, 1):
178 return False
179 return True
182class DTensor(DTensorBase):
183 """
184 DTensor - Distributed Tensor
186 A DTensor represents a tensor that is distributed across multiple devices
187 according to a DeviceMesh and placement specifications.
189 Args:
190 local_tensor (Tensor): The local tensor shard on this device.
191 device_mesh (DeviceMesh): The device mesh describing the device topology.
192 placements: The placement strategy. Supports two styles:
193 - Placement objects (e.g., ``[Shard(0), Replicate()]``).
194 - Alias strings (e.g., ``("dp", "None")`` or
195 ``(("dp", "tp"), "None")``), length must equal the number of
196 tensor dimensions.
198 Example:
199 >>> mesh = init_device_mesh(device_type="npu", mesh_shape=(2, 2), mesh_dim_names=("dp", "tp"))
200 >>> local_tensor = Tensor(np.ones((4, 4)))
201 >>> # Placement style
202 >>> dtensor = DTensor.from_local(local_tensor, mesh, [Shard(0), Replicate()])
203 >>> # Alias style — length matches tensor dims
204 >>> dtensor = DTensor.from_local(local_tensor, mesh, ("dp", "None"))
205 """
206 _local_tensor: Tensor
207 _device_mesh: DeviceMesh
208 _placements: Sequence[Placement]
210 def __init_data__(
211 self,
212 local_tensor: Tensor,
213 device_mesh: DeviceMesh,
214 placements: Union[Sequence[Placement], Sequence[Union[str, Tuple[str, ...]]]]
215 ):
216 self._local_tensor = local_tensor
217 self._device_mesh = device_mesh
218 self._layout = _build_layout(
219 device_mesh, placements, len(local_tensor.shape)
220 )
221 self._placements = tuple(self._layout.placements)
223 @property
224 def device_mesh(self) -> DeviceMesh:
225 """The device mesh of this DTensor."""
226 return self._device_mesh
228 @property
229 def placements(self) -> Sequence[Placement]:
230 """The placements of this DTensor."""
231 return self._placements
233 @property
234 def layout(self) -> Layout:
235 """Internal layout for redistribution (for backward compatibility)."""
236 if not hasattr(self, '_layout'):
237 return None
238 return self._layout
240 @staticmethod
241 def from_local(
242 local_tensor: Tensor,
243 device_mesh: DeviceMesh,
244 placements: Union[Sequence[Placement], Sequence[Union[str, Tuple[str, ...]]]]
245 ) -> 'DTensor':
246 """
247 Create a DTensor from a local tensor with device mesh and placements.
249 Args:
250 local_tensor (Tensor): The local tensor shard on this device.
251 device_mesh (DeviceMesh): The device mesh describing the device topology.
252 placements: The placement strategy. Supports two styles:
253 - Placement objects (e.g., ``[Shard(0), Replicate()]``).
254 - Alias strings (e.g., ``("dp", "None")`` or
255 ``(("dp", "tp"), "None")``), length must equal the number
256 of tensor dimensions.
258 Returns:
259 DTensor: A new DTensor instance.
261 Example:
262 >>> mesh = init_device_mesh(device_type="npu", mesh_shape=(2, 2), mesh_dim_names=("dp", "tp"))
263 >>> local_tensor = Tensor(np.ones((4, 4)))
264 >>> dtensor = DTensor.from_local(local_tensor, mesh, [Shard(0), Replicate()])
265 >>> dtensor = DTensor.from_local(local_tensor, mesh, ("dp", "None"))
266 """
267 return DTensor(local_tensor, device_mesh, placements)
269 def _alias_placements(self) -> Sequence[Placement]:
270 """Return alias_placements from layout, falling back to _placements."""
271 if hasattr(self, '_layout') and self._layout:
272 return self._layout.alias_placements
273 return self._placements
275 def _from_converted_local(self, local_tensor: Tensor) -> 'DTensor':
276 """Rebuild converted DTensor data without preserving Parameter identity."""
277 cls = DTensor if isinstance(self, platform.Parameter) else self.__class__
278 return cls(local_tensor, device_mesh=self._device_mesh,
279 placements=self._alias_placements())
281 def to(self, *args, **kwargs):
282 """Move the DTensor to a different device or dtype.
284 Delegates to the underlying local tensor's ``to`` method and
285 reconstructs a DTensor preserving device_mesh and placements.
287 Args:
288 *args (tuple): Arguments passed to the underlying tensor's ``to``
289 method (e.g., device or dtype).
290 **kwargs (dict): Keyword arguments for the tensor conversion
291 (e.g., dtype, device, non_blocking).
293 Returns:
294 DTensor: A new DTensor with the converted local tensor.
295 """
296 new_local = self._local_tensor.to(*args, **kwargs)
297 return self._from_converted_local(new_local)
299 def float(self):
300 """Convert the DTensor to float dtype.
302 Returns:
303 DTensor: A new DTensor with float32 local tensor.
304 """
305 new_local = self._local_tensor.float()
306 return self._from_converted_local(new_local)
308 def to_local(self) -> Tensor:
309 """
310 Convert DTensor to local tensor.
312 Returns:
313 Tensor: The local tensor shard on this device.
314 """
315 return self._local_tensor
317 def copy_(self, src: "DTensor", non_blocking: bool = False) -> "DTensor":
318 """In-place copy of ``src`` into this DTensor's local shard.
320 Delegates to ``Tensor.copy_`` on the underlying local tensors.
321 Follows standard ``Tensor.copy_`` semantics: version counter is
322 bumped and autograd edges are created when grad is enabled.
324 Constraints on ``src``:
325 * must be a ``DTensor`` on the same ``DeviceMesh`` as ``self``;
326 * its placements must equal ``self.placements``, OR
327 ``src._local_tensor.numel() == 1`` (single-element broadcast);
328 * its local shape must equal or be broadcastable to
329 ``self._local_tensor.shape``.
331 No redistribute / implicit slicing is performed; src dtype is cast
332 to self dtype in-place.
334 Args:
335 src (DTensor): Source DTensor satisfying the constraints above.
336 non_blocking (bool): Forwarded to the underlying ``copy_``.
338 Returns:
339 DTensor: ``self``.
341 Raises:
342 TypeError: if ``src`` is not a ``DTensor``.
343 ValueError: if mesh, placement, or shape constraint is violated.
344 """
345 if not isinstance(src, DTensor):
346 raise TypeError(
347 f"For DTensor.copy_, src should be a DTensor, but got {type(src).__name__}."
348 )
349 src_local = src.to_local()
350 if src.device_mesh is not self._device_mesh:
351 raise ValueError(
352 f"For DTensor.copy_, src and self should share the same DeviceMesh, "
353 f"but got src.device_mesh={src.device_mesh!r}, "
354 f"self._device_mesh={self._device_mesh!r}."
355 )
357 placement_eq = tuple(src.placements) == tuple(self._placements)
358 shape_eq = src_local.shape == self._local_tensor.shape
359 src_is_scalar = src_local.numel() == 1
361 if not placement_eq and not src_is_scalar:
362 raise ValueError(
363 f"For DTensor.copy_, src.placements should equal self.placements "
364 f"or src.numel() should be 1, but got "
365 f"src.placements={src.placements}, "
366 f"self.placements={self._placements}, "
367 f"src.numel()={src_local.numel()}."
368 )
369 if not shape_eq and not src_is_scalar and not _is_broadcastable(
370 src_local.shape, self._local_tensor.shape
371 ):
372 raise ValueError(
373 f"For DTensor.copy_, src local shape should be broadcastable to "
374 f"self local shape, but got "
375 f"src.shape={tuple(src_local.shape)}, "
376 f"self.shape={tuple(self._local_tensor.shape)}."
377 )
379 self._local_tensor.copy_(src_local, non_blocking=non_blocking)
380 return self
382 def zero_(self) -> "DTensor":
383 """In-place fill with zeros. Returns ``self``."""
384 self._local_tensor.zero_()
385 return self
387 def fill_(self, value) -> "DTensor":
388 """In-place fill with ``value``. Returns ``self``."""
389 self._local_tensor.fill_(value)
390 return self
392 @property
393 def shape(self) -> Tuple[int, ...]:
394 """
395 The global shape of this DTensor.
397 Returns:
398 Tuple[int, ...]: The global tensor shape.
399 """
400 return self._layout.get_global_shape(self._local_tensor.shape)
402 def size(self, dim=None):
403 """Return the global shape, consistent with .shape.
405 Without ``dim`` returns a tuple matching ``self.shape``.
406 With ``dim`` returns the size of that dimension.
407 """
408 global_shape = self.shape
409 if dim is not None:
410 return global_shape[dim]
411 return global_shape
413 def numel(self) -> int:
414 """Return the number of elements in this DTensor."""
415 return int(np.prod(self.shape))
417 @property
418 def local_shape(self) -> Tuple[int, ...]:
419 """
420 The local shape of this DTensor on this device.
422 Returns:
423 Tuple[int, ...]: The local tensor shape.
424 """
425 return self._local_tensor.shape
427 def redistribute(
428 self,
429 device_mesh: DeviceMesh,
430 placements: Union[Sequence[Placement], Sequence[Union[str, Tuple[str, ...]]]]
431 ) -> 'DTensor':
432 """
433 Redistribute this DTensor to a new device mesh and placements.
435 Args:
436 device_mesh (DeviceMesh): The target device mesh.
437 placements: The target placements. Supports Placement objects
438 or alias strings.
440 Returns:
441 DTensor: A new DTensor with the specified distribution.
443 Example:
444 >>> new_dtensor = dtensor.redistribute(mesh, [Replicate(), Shard(1)])
445 >>> new_dtensor = dtensor.redistribute(mesh, ("None", "tp"))
446 """
447 # Build dst_layout from device_mesh and placements
448 dst_layout = _build_layout(
449 device_mesh, placements, len(self._local_tensor.shape)
450 )
452 # pylint: disable=C0415
453 from hyper_parallel.core.dtensor.tensor_redistribution import _tensor_redistribution
454 out = _tensor_redistribution.redistribution(self, dst_layout)
455 return out
457 def reduce_partial(self) -> 'DTensor':
458 """
459 Reduce partial sharding state for this DTensor.
461 Returns:
462 DTensor: A new DTensor with partial state reduced.
463 """
464 if not self._layout:
465 return self
466 to_layout = cp.deepcopy(self._layout)
467 to_layout.reset_partial()
468 # pylint: disable=C0415
469 from hyper_parallel.core.dtensor.tensor_redistribution import _tensor_redistribution
470 out = _tensor_redistribution.reduce_partial(self, to_layout)
471 return out
473 def full_tensor(self) -> Tensor:
474 """
475 Return the full tensor of this DTensor.
477 Returns:
478 Tensor: A Tensor object that represents the full tensor of this DTensor.
479 The returned tensor contains the complete data gathered from
480 all ranks.
482 Note:
483 This operation involves communication across all ranks in the DeviceMesh,
484 which may be expensive for large tensors. Use with caution in
485 performance-critical code paths.
487 Example:
488 >>> # Assume dtensor is sharded across multiple devices
489 >>> local_tensor = dtensor.to_local() # Returns only the local shard
490 >>> full_tensor = dtensor.full_tensor() # Returns the complete tensor
491 """
492 if not self._layout:
493 return self._local_tensor
495 # Create a fully replicated layout
496 replicated_layout = cp.deepcopy(self._layout)
498 # Set all placements to Replicate and convert to tensor_map
499 replicated_placements = [Replicate()] * len(replicated_layout.mesh_shape)
500 replicated_layout.set_placements(replicated_placements)
501 replicated_layout.placement_to_tensor_map(len(self._local_tensor.shape))
503 # Clear partial status from original layout since Replicate has no partial
504 replicated_layout.reset_partial()
506 # Redistribute to the replicated layout and return local tensor
507 # pylint: disable=C0415
508 from hyper_parallel.core.dtensor.tensor_redistribution import _tensor_redistribution
509 out = _tensor_redistribution.redistribution(self, replicated_layout)
510 return out.to_local()
513def _normalize_shard_dim(dim: int, ndim: int) -> int:
514 return dim + ndim if dim < 0 else dim
517def _distribute_tensor_with_communication(
518 tensor: Tensor,
519 device_mesh: DeviceMesh,
520 placements: Sequence[Placement],
521 src_data_rank: int,
522) -> Tensor:
523 """Scatter/broadcast a logical global tensor along mesh dimensions (PyTorch parity)."""
524 local = tensor
525 if len(placements) < device_mesh.ndim:
526 raise ValueError(
527 f"placements length ({len(placements)}) must be at least device_mesh.ndim "
528 f"({device_mesh.ndim}) when src_data_rank is set"
529 )
530 for mesh_dim in range(device_mesh.ndim):
531 placement = placements[mesh_dim]
532 if isinstance(placement, StridedShard):
533 raise NotImplementedError(
534 "distribute_tensor with src_data_rank does not support StridedShard yet; "
535 "pass src_data_rank=None for local-only sharding."
536 )
537 if placement.is_shard():
538 shard_dim = _normalize_shard_dim(placement.dim, local.ndim)
539 num_chunks = device_mesh.size(mesh_dim)
540 if num_chunks <= 0:
541 raise ValueError(f"invalid mesh dim size {num_chunks} on mesh_dim={mesh_dim}")
542 chunks = tuple(local.chunk(num_chunks, dim=shard_dim))
543 if not chunks:
544 raise ValueError(f"cannot shard dim {shard_dim} into {num_chunks} chunks")
545 output = platform.empty_like(chunks[0])
546 local = mesh_scatter(output, chunks, device_mesh, mesh_dim, group_src=src_data_rank)
547 elif placement.is_replicate() or placement.is_partial():
548 local = mesh_broadcast(local, device_mesh, mesh_dim, group_src=src_data_rank)
549 if isinstance(placement, Partial):
550 warnings.warn(
551 f"Partial placement {placement} during distribute_tensor: "
552 "broadcast only; partial partition is not applied yet.",
553 stacklevel=3,
554 )
555 else:
556 raise RuntimeError(
557 f"unsupported placement {placement} on device mesh dimension {mesh_dim}"
558 )
559 return local
562def distribute_tensor(
563 tensor: Tensor,
564 device_mesh: DeviceMesh,
565 placements: Union[Sequence[Placement], Sequence[Union[str, Tuple[str, ...]]]],
566 *,
567 src_data_rank: Optional[int] = None,
568) -> DTensor:
569 """
570 Distribute a global tensor to the device mesh according to the placements.
572 Args:
573 tensor (Tensor): The global tensor to be distributed. All ranks
574 should have the same tensor data.
575 device_mesh (DeviceMesh): The device mesh describing the device topology.
576 placements: The placement strategy. Supports two styles:
577 - Placement objects (e.g., ``[Shard(0), Replicate()]``).
578 - Alias strings (e.g., ``("dp", "None")`` or
579 ``(("dp", "tp"), "None")``), length must equal the number of
580 tensor dimensions.
582 Returns:
583 DTensor: A new DTensor with the local shard on each rank.
585 Note:
586 When ``src_data_rank`` is an ``int`` (e.g. ``0``), shard/replicate
587 placements use scatter/broadcast from the source rank on each mesh axis,
588 matching PyTorch ``distribute_tensor``. When ``src_data_rank=None``
589 (default), each rank slices its local tensor without communication
590 (legacy Hyper behavior; all ranks must hold the same global data).
592 Example:
593 >>> mesh = init_device_mesh(device_type="npu", mesh_shape=(2, 2), mesh_dim_names=("dp", "tp"))
594 >>> global_tensor = Tensor(np.arange(16).reshape(4, 4))
595 >>> dtensor = distribute_tensor(global_tensor, mesh, [Shard(0), Replicate()])
596 >>> dtensor = distribute_tensor(global_tensor, mesh, ("dp", "None"))
597 """
598 layout = _build_layout(device_mesh, placements, len(tensor.shape))
599 if src_data_rank is None:
600 local_tensor = _get_slice_tensor_by_layout(tensor, layout)
601 else:
602 local_tensor = _distribute_tensor_with_communication(
603 tensor, device_mesh, layout.placements, src_data_rank
604 )
605 return DTensor(local_tensor, device_mesh, layout.alias_placements)
608def _distribute_module_param_source(param: Any) -> Tensor:
609 """Tensor data used as the global tensor for :func:`distribute_tensor` (PyTorch uses ``param.data``)."""
610 if hasattr(param, "data"):
611 return param.data
612 return platform.get_param_local_data(param)
615def _distribute_module_new_parameter(key: str, dtensor: DTensor, requires_grad: bool) -> Any:
616 """Build a framework :class:`Parameter` holding *dtensor* (Torch vs MindSpore kwargs differ)."""
617 if platform.platform_type == PlatformType.MINDSPORE:
618 return platform.Parameter(dtensor, name=key, requires_grad=requires_grad)
619 return platform.Parameter(dtensor, requires_grad=requires_grad)
622def _distribute_module_set_param(module: Any, key: str, new_param: Any) -> None:
623 """Register or assign a parameter on *module* (``nn.Module`` or MindSpore ``Cell``)."""
624 if hasattr(module, "register_parameter"):
625 module.register_parameter(key, new_param)
626 return
627 if hasattr(module, "_params"):
628 module._params[key] = new_param
629 if hasattr(module, "_params_list"):
630 module._params_list[key] = new_param
631 if key in module.__dict__:
632 module.__dict__[key] = new_param
633 return
634 raise TypeError(
635 f"distribute_module expects nn.Module-like objects with register_parameter or _params; "
636 f"got {type(module)}."
637 )
640def _distribute_module_iter_params(module: Any) -> list:
641 """Return ``[(name, param), ...]`` for direct parameters (``_parameters`` or ``_params``)."""
642 if hasattr(module, "_parameters"):
643 return list(module._parameters.items())
644 if hasattr(module, "_params"):
645 return list(module._params.items())
646 return []
649def _distribute_module_iter_buffers(module: Any) -> list:
650 """Return ``[(name, buffer), ...]`` if the module has ``_buffers`` (PyTorch ``nn.Module``)."""
651 if hasattr(module, "_buffers"):
652 return list(module._buffers.items())
653 return []
656def _distribute_module_named_modules(module: Any):
657 """``nn.Module.named_modules`` or MindSpore ``Cell.cells_and_names`` (submodule FQNs)."""
658 if hasattr(module, "named_modules"):
659 return module.named_modules()
660 if hasattr(module, "cells_and_names"):
661 return module.cells_and_names()
662 raise TypeError(
663 f"distribute_module expects module-like objects with named_modules or cells_and_names; "
664 f"got {type(module)}."
665 )
668def _replicate_submodule_params_buffers(
669 sub_mod: Any,
670 device_mesh: DeviceMesh,
671 *,
672 module_prefix: str = "",
673) -> None:
674 """Convert plain params/buffers on *sub_mod* to fully replicated :class:`DTensor`."""
675 full_replicate = [Replicate()] * device_mesh.ndim
676 for key, param in _distribute_module_iter_params(sub_mod):
677 if param is None or isinstance(param, DTensorBase):
678 continue
679 src = _distribute_module_param_source(param)
680 requires_grad = bool(getattr(param, "requires_grad", True))
681 dt = distribute_tensor(src, device_mesh, full_replicate)
682 param_name = f"{module_prefix}.{key}" if module_prefix else key
683 new_param = _distribute_module_new_parameter(param_name, dt, requires_grad)
684 _distribute_module_set_param(sub_mod, key, new_param)
685 for key, buffer in _distribute_module_iter_buffers(sub_mod):
686 if buffer is None or isinstance(buffer, DTensorBase):
687 continue
688 sub_mod._buffers[key] = distribute_tensor(buffer, device_mesh, full_replicate)
691def _distribute_module_run_partition_and_replicate(
692 module: Any,
693 device_mesh: DeviceMesh,
694 partition_fn: Optional[Callable[[str, Any, DeviceMesh], None]],
695) -> None:
696 """Call optional ``partition_fn`` per ``named_modules`` and replicate remaining tensors."""
697 if partition_fn is None:
698 for mod_name, submod in _distribute_module_named_modules(module):
699 _replicate_submodule_params_buffers(submod, device_mesh, module_prefix=mod_name)
700 return
701 for mod_name, submod in _distribute_module_named_modules(module):
702 partition_fn(mod_name, submod, device_mesh)
703 _replicate_submodule_params_buffers(submod, device_mesh, module_prefix=mod_name)
706def _distribute_module_register_input_fn(
707 module: Any,
708 device_mesh: DeviceMesh,
709 input_fn: Callable[..., Any],
710) -> None:
711 """Register *input_fn* as a forward pre-hook on *module* (2- or 3-arg, PyTorch-compatible)."""
712 num_args = len(inspect.signature(input_fn).parameters)
713 if num_args == 2:
714 warnings.warn(
715 "Deprecating input_fn that takes two arguments (inputs, device_mesh), "
716 "please use input_fn that takes in (module, inputs, device_mesh) instead!",
717 FutureWarning,
718 stacklevel=3,
719 )
720 module.register_forward_pre_hook(
721 lambda _, inputs: input_fn(inputs, device_mesh)
722 )
723 elif num_args == 3:
724 module.register_forward_pre_hook(
725 lambda mod, inputs: input_fn(mod, inputs, device_mesh)
726 )
727 else:
728 raise ValueError(
729 f"input_fn should take in 2 or 3 arguments, but got {num_args} arguments!"
730 )
733def _distribute_module_register_output_fn(
734 module: Any,
735 device_mesh: DeviceMesh,
736 output_fn: Callable[..., Any],
737) -> None:
738 """Register *output_fn* as a forward hook on *module* (2- or 3-arg, PyTorch-compatible)."""
739 num_args = len(inspect.signature(output_fn).parameters)
740 if num_args == 2:
741 warnings.warn(
742 "Deprecating output_fn that takes two arguments (outputs, device_mesh), "
743 "please use output_fn that takes in (module, outputs, device_mesh) instead!",
744 FutureWarning,
745 stacklevel=3,
746 )
747 module.register_forward_hook(
748 lambda mod, inputs, outputs: output_fn(outputs, device_mesh)
749 )
750 elif num_args == 3:
751 module.register_forward_hook(
752 lambda mod, inputs, outputs: output_fn(mod, outputs, device_mesh)
753 )
754 else:
755 raise ValueError(
756 f"output_fn should take in 2 or 3 arguments, but got {num_args} arguments!"
757 )
760def distribute_module(
761 module: Any,
762 device_mesh: Optional[DeviceMesh] = None,
763 partition_fn: Optional[Callable[[str, Any, DeviceMesh], None]] = None,
764 input_fn: Optional[Callable[..., Any]] = None,
765 output_fn: Optional[Callable[..., Any]] = None,
766) -> Any:
767 """PyTorch ``distribute_module`` parity: shard/replicate params and optional I/O hooks.
769 Unsharded parameters and buffers become fully replicated :class:`DTensor` after
770 ``partition_fn``. ``input_fn`` / ``output_fn`` attach only to the root *module*.
772 Args:
773 module: Root ``nn.Module`` or MindSpore ``Cell`` with compatible APIs.
774 device_mesh: Placement mesh; if ``None``, uses ``_mesh_resources.get_current_mesh()``.
775 partition_fn: Per ``named_modules`` callback before replicate pass; ``None`` replicates all.
776 input_fn: ``(module, inputs, mesh)`` or deprecated ``(inputs, mesh)`` pre-hook.
777 output_fn: ``(module, outputs, mesh)`` or deprecated ``(outputs, mesh)`` forward hook.
779 Returns:
780 *module* in place, with distributed tensors where applied.
782 Raises:
783 RuntimeError: If called twice on the same *module*.
784 ValueError: If ``input_fn`` / ``output_fn`` arity is not 2 or 3.
786 Note:
787 XLA / ``torch_xla`` is not supported; strided device :class:`DTensor` only.
788 """
789 if getattr(module, "_distribute_module_applied", False):
790 raise RuntimeError(
791 "distribute_module should only be called once on a module, "
792 "but it has already been called on this module!"
793 )
794 device_mesh = device_mesh or _mesh_resources.get_current_mesh()
795 _distribute_module_run_partition_and_replicate(module, device_mesh, partition_fn)
796 if input_fn is not None:
797 _distribute_module_register_input_fn(module, device_mesh, input_fn)
798 if output_fn is not None:
799 _distribute_module_register_output_fn(module, device_mesh, output_fn)
800 module._distribute_module_applied = True
801 return module
804def _dtensor_init_helper(
805 init_op,
806 size,
807 device_mesh,
808 placements,
809 **kwargs,
810) -> DTensor:
811 """
812 Helper function to create and initialize a distributed tensor.
814 Args:
815 size: Shape of the tensor.
816 dtype: Data type of the tensor.
817 device: Target device for the tensor.
818 requires_grad: Whether the tensor requires gradient.
820 Returns:
821 DTensor: The initialized distributed tensor.
822 """
823 # get local tensor shape
824 local_shape = compute_local_shape_and_global_offset(
825 size, device_mesh, placements
826 )
828 # initialize the local tensor
829 if init_op is platform.full:
830 fill_value = kwargs.pop("fill_value", 0)
831 local_tensor = init_op(local_shape, fill_value, **kwargs)
832 else:
833 local_tensor = init_op(local_shape, **kwargs)
835 return DTensor.from_local(
836 local_tensor,
837 device_mesh,
838 placements,
839 )
842def ones(
843 size,
844 device_mesh,
845 placements,
846) -> DTensor:
847 """
848 Returns a :class:`DTensor` filled with the scalar value 1, with the shape defined
849 by the variable argument ``size``.
851 Args:
852 size (Union[tuple[int], list[int], int, Tensor]): The specified shape of output tensor. Only positive integer or
853 tuple or Tensor containing positive integers are allowed. If it is a Tensor,
854 it must be a 0-D or 1-D Tensor with int32 or int64 dtypes.
856 Keyword args:
857 device_mesh: :class:`DeviceMesh` type, contains the mesh info of ranks
858 placements: a sequence of :class:`Placement` type: ``Shard``, ``Replicate``
860 Returns:
861 A :class:`DTensor` object on each rank
862 """
863 ones_ = platform.ones
864 return _dtensor_init_helper(
865 ones_,
866 size,
867 device_mesh=device_mesh,
868 placements=placements,
869 )
872def empty(
873 size,
874 device_mesh,
875 placements,
876) -> DTensor:
877 """
878 Returns a :class:`DTensor` filled with uninitialized data. The shape of the :class:`DTensor`
879 is defined by the variable argument ``size``.
881 Args:
882 size (Union[tuple[int], list[int], int]): The specified shape of output tensor. Can be variable numbers of
883 positive integers or tuple or list containing positive integers.
885 Keyword args:
886 device_mesh: :class:`DeviceMesh` type, contains the mesh info of ranks
887 placements: a sequence of :class:`Placement` type: ``Shard``, ``Replicate``
889 Returns:
890 A :class:`DTensor` object on each rank
891 """
892 empty_ = platform.empty
893 return _dtensor_init_helper(
894 empty_,
895 size,
896 device_mesh=device_mesh,
897 placements=placements,
898 )
901def full(
902 size,
903 fill_value,
904 *,
905 device_mesh,
906 placements,
907) -> DTensor:
908 """
909 Returns a :class:`DTensor` filled with ``fill_value`` according to ``device_mesh`` and
910 ``placements``, with the shape defined by the argument ``size``.
912 Args:
913 size (Union[tuple[int], list[int]]): The specified shape of output tensor.
914 fill_value (Union[numbers.Number, Tensor]): Value to fill the returned tensor. It can be a scalar number, a 0-D
915 Tensor, or a 1-D Tensor with only one element.
917 Keyword args:
918 device_mesh: :class:`DeviceMesh` type, contains the mesh info of ranks.
919 placements: a sequence of :class:`Placement` type: ``Shard``, ``Replicate``
921 Returns:
922 A :class:`DTensor` object on each rank
923 """
924 full_ = platform.full
925 return _dtensor_init_helper(
926 full_,
927 size,
928 fill_value=fill_value,
929 device_mesh=device_mesh,
930 placements=placements,
931 )
934def zeros(
935 size,
936 device_mesh,
937 placements,
938) -> DTensor:
939 """
940 Returns a :class:`DTensor` filled with the scalar value 0.
942 Args:
943 size (Union[tuple[int], list[int], int, Tensor]): The specified shape of output tensor. Only positive integer or
944 tuple or Tensor containing positive integers are allowed. If it is a Tensor,
945 it must be a 0-D or 1-D Tensor with int32 or int64 dtypes.
946 Keyword args:
947 device_mesh: :class:`DeviceMesh` type, contains the mesh info of ranks
948 placements: a sequence of :class:`Placement` type: ``Shard``, ``Replicate``
950 Returns:
951 A :class:`DTensor` object on each rank
952 """
953 zeros_ = platform.zeros
954 return _dtensor_init_helper(
955 zeros_,
956 size,
957 device_mesh=device_mesh,
958 placements=placements,
959 )