Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / dtensor / dtensor.py: 88%
456 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-25 04:27 +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 logging
19import warnings
20from typing import Any, Callable, Optional, Sequence, Set, Tuple, Union
22import numpy as np
24from hyper_parallel.core.dtensor._collective_utils import mesh_broadcast, mesh_scatter
25from hyper_parallel.core.dtensor._ragged_utils import (
26 _compute_ragged_slice,
27 _layout_has_ragged_shard,
28 _normalize_global_shape,
29 _scatter_ragged_tensor,
30 _slice_ragged_tensor,
31)
32from hyper_parallel.core.dtensor.device_mesh import _mesh_resources
33from hyper_parallel.core.dtensor.layout import (
34 DeviceMesh,
35 Layout,
36 _get_slice_tensor_by_layout,
37)
38from hyper_parallel.core.dtensor.placement_types import Partial, Placement, Replicate, StridedShard
39from hyper_parallel.platform import get_platform
40from hyper_parallel.platform.platform import PlatformType
41from hyper_parallel.core.utils import compute_local_shape_and_global_offset
43platform = get_platform()
44DTensorBase = platform.DTensorBase
45Tensor = platform.Tensor
47logger = logging.getLogger(__name__)
50def _device_meshes_are_compatible(lhs: Any, rhs: Any) -> bool:
51 """Return whether two mesh objects describe the same device topology."""
52 if lhs is rhs:
53 return True
54 if not isinstance(lhs, DeviceMesh) or not isinstance(rhs, DeviceMesh):
55 return False
56 return lhs.device_type == rhs.device_type and lhs.to_hash() == rhs.to_hash()
59class SkipDTensorDispatch():
60 """Context manager that disables DTensor op dispatch for the enclosed block.
62 Args:
63 no_skip: Optional set of op callables or canonical op name strings that
64 should still be dispatched through DTensor even within this context.
65 All other ops bypass DTensor dispatch and operate on local tensors.
67 Example:
68 >>> import torch
69 >>> with SkipDTensorDispatch(no_skip={torch.zeros_like}):
70 ... # zeros_like still goes through DTensor dispatch;
71 ... # everything else uses the local tensor path.
72 ... result = torch.zeros_like(dtensor)
73 """
75 def __init__(self, no_skip: Optional[Set] = None):
76 self._no_skip_names: frozenset = frozenset()
77 if no_skip:
78 names = set()
79 for op in no_skip:
80 if isinstance(op, str):
81 names.add(op)
82 else:
83 names.add(platform.get_op_name(op))
84 self._no_skip_names = frozenset(names)
85 self._dispatch_token = None
86 self._ops_token = None
88 def __enter__(self):
89 # pylint: disable=C0415
90 from hyper_parallel.core.shard._op_dispatch import _dtensor_dispatch_disabled, _no_skip_ops
91 self._dispatch_token = _dtensor_dispatch_disabled.set(True)
92 if self._no_skip_names:
93 self._ops_token = _no_skip_ops.set(_no_skip_ops.get() | self._no_skip_names)
95 def __exit__(self, exc_type, exc_val, exc_tb):
96 # pylint: disable=C0415
97 from hyper_parallel.core.shard._op_dispatch import _dtensor_dispatch_disabled, _no_skip_ops
98 if self._ops_token is not None:
99 _no_skip_ops.reset(self._ops_token)
100 self._ops_token = None
101 _dtensor_dispatch_disabled.reset(self._dispatch_token)
102 self._dispatch_token = None
105# Cache for _build_layout to avoid redundant Layout computations
106# Key: (device_mesh.to_hash(), tuple(placements), tensor_dim)
107# Value: Layout
108_LAYOUT_CACHE = {}
111def _is_alias_placements(placements) -> bool:
112 """
113 Check if placements use alias strings rather than Placement objects.
115 Alias placements use mesh dimension names (strings) to specify
116 the sharding strategy, e.g., ("dp", "tp") or (("dp", "tp"), "None").
117 All elements must be strings or tuples of strings for the sequence
118 to be recognized as alias-style.
120 Args:
121 placements: A sequence of placement specifications.
123 Returns:
124 bool: True if all elements are alias strings or tuples of strings.
125 """
126 if len(placements) == 0:
127 return False
128 for p in placements:
129 if isinstance(p, str):
130 continue
131 if isinstance(p, tuple) and len(p) > 0 and all(isinstance(x, str) for x in p):
132 continue
133 return False
134 return True
137def _build_layout(
138 device_mesh: DeviceMesh,
139 placements: Union[Sequence[Placement], Sequence[Union[str, Tuple[str, ...]]]],
140 tensor_dim: int
141) -> Layout:
142 """
143 Build Layout from device_mesh and placements.
145 This function uses a cache to avoid redundant Layout computations
146 for the same (device_mesh, placements, tensor_dim) combination.
148 Args:
149 device_mesh: The device mesh describing the device topology.
150 placements: Supports two styles:
151 - Placement objects (Shard, Replicate, etc.)
152 - Alias strings ("dp", "None", ("dp", "tp"), etc.), length must
153 equal the number of tensor dimensions (``tensor_dim``).
154 tensor_dim: Number of dimensions in the tensor.
156 Returns:
157 Layout: The built layout object.
159 Raises:
160 ValueError: If alias placements length does not match tensor dimensions.
161 """
162 mesh_key = device_mesh.to_hash()
163 placements_key = tuple(placements)
164 cache_key = (mesh_key, placements_key, tensor_dim)
166 if cache_key in _LAYOUT_CACHE:
167 return _LAYOUT_CACHE[cache_key]
169 layout = Layout.from_device_mesh(device_mesh)
171 if _is_alias_placements(placements):
172 if len(placements) != tensor_dim:
173 raise ValueError(
174 f"Alias placements length ({len(placements)}) must equal "
175 f"tensor dimensions ({tensor_dim})."
176 )
177 result = layout(*placements)
178 else:
179 result = layout(placements)
180 result.placement_to_tensor_map(tensor_dim)
182 _LAYOUT_CACHE[cache_key] = result
184 return result
187def _is_broadcastable(src_shape: Sequence[int], dst_shape: Sequence[int]) -> bool:
188 """Return True iff ``src_shape`` is broadcastable to ``dst_shape``.
190 Standard NumPy / PyTorch right-aligned broadcast rule: ``src`` cannot
191 have more dimensions than ``dst``; each right-aligned dimension pair
192 must be equal, or ``src``'s dimension must be 1.
193 """
194 src_shape = tuple(src_shape)
195 dst_shape = tuple(dst_shape)
196 if len(src_shape) > len(dst_shape):
197 return False
198 for i in range(1, len(src_shape) + 1):
199 s, d = src_shape[-i], dst_shape[-i]
200 if s not in (d, 1):
201 return False
202 return True
205def _device_spec(device: Any) -> Tuple[str, Optional[int]]:
206 """Return normalised ``(device_type, device_index)``.
208 Handles device objects and strings such as ``"npu:0"``.
209 """
210 device_type = getattr(device, "type", None)
211 # Only read .index from objects that also have a .type attribute
212 # (i.e. torch.device). Avoids capturing str.index on plain strings.
213 device_index = (
214 getattr(device, "index", None) if device_type is not None else None
215 )
216 device_text = str(device).lower()
218 if device_type is None:
219 parts = device_text.split(":", maxsplit=1)
220 device_type = parts[0]
221 if len(parts) == 2 and parts[1].isdigit():
222 device_index = int(parts[1])
224 return str(device_type).lower(), device_index
227class DTensor(DTensorBase):
228 """
229 DTensor - Distributed Tensor
231 A DTensor represents a tensor that is distributed across multiple devices
232 according to a DeviceMesh and placement specifications.
234 Args:
235 local_tensor (Tensor): The local tensor shard on this device.
236 device_mesh (DeviceMesh): The device mesh describing the device topology.
237 placements: The placement strategy. Supports two styles:
238 - Placement objects (e.g., ``[Shard(0), Replicate()]``).
239 - Alias strings (e.g., ``("dp", "None")`` or
240 ``(("dp", "tp"), "None")``), length must equal the number of
241 tensor dimensions.
243 Example:
244 >>> mesh = init_device_mesh(device_type="npu", mesh_shape=(2, 2), mesh_dim_names=("dp", "tp"))
245 >>> local_tensor = Tensor(np.ones((4, 4)))
246 >>> # Placement style
247 >>> dtensor = DTensor.from_local(local_tensor, mesh, [Shard(0), Replicate()])
248 >>> # Alias style — length matches tensor dims
249 >>> dtensor = DTensor.from_local(local_tensor, mesh, ("dp", "None"))
250 """
251 _local_tensor: Tensor
252 _device_mesh: DeviceMesh
253 _placements: Sequence[Placement]
255 def __init_data__(
256 self,
257 local_tensor: Tensor,
258 device_mesh: DeviceMesh,
259 placements: Union[Sequence[Placement], Sequence[Union[str, Tuple[str, ...]]]],
260 layout: Optional[Layout] = None,
261 shape: Optional[Tuple[int, ...]] = None,
262 ):
263 self._local_tensor = local_tensor
264 self._device_mesh = device_mesh
265 tensor_dim = len(shape) if shape is not None else len(local_tensor.shape)
266 # Fast path: when an already-built Layout is supplied (e.g. output layouts
267 # cached by infer_layout and passed straight through wrap_output), reuse it
268 # directly and skip _build_layout (which otherwise recomputes device_mesh.to_hash(),
269 # tuple(placements) and a cache lookup on every single output construction).
270 self._layout = layout if layout is not None else _build_layout(
271 device_mesh, placements, tensor_dim
272 )
273 self._placements = tuple(self._layout.placements)
274 is_ragged = _layout_has_ragged_shard(self._layout)
275 if is_ragged and shape is None:
276 raise ValueError(
277 "DTensor.from_local with RaggedShard requires an explicit global shape"
278 )
279 if shape is not None:
280 self._global_shape = _normalize_global_shape(shape)
281 else:
282 self._global_shape = tuple(self._layout.get_global_shape(local_tensor.shape))
283 if (
284 shape is not None
285 and (
286 self._layout.tensor_map is None
287 or len(self._layout.tensor_map) != len(self._global_shape)
288 )
289 ):
290 raise ValueError(
291 "DTensor global shape rank must match layout tensor_map rank, "
292 f"got global_shape={self._global_shape!r}, tensor_map={self._layout.tensor_map!r}"
293 )
294 if is_ragged:
295 if hasattr(local_tensor, "is_contiguous") and not local_tensor.is_contiguous():
296 raise ValueError("RaggedShard local tensor must be contiguous")
297 if len(local_tensor.shape) != 1:
298 raise ValueError(
299 "RaggedShard local tensor must use one-dimensional flat storage, "
300 f"got local_shape={tuple(local_tensor.shape)!r}"
301 )
302 expected = _compute_ragged_slice(self._global_shape, self._layout)
303 if local_tensor.numel() != expected.local_numel:
304 raise ValueError(
305 "RaggedShard local tensor numel does not match its allocation, "
306 f"got actual={local_tensor.numel()}, expected={expected.local_numel}, "
307 f"global_shape={self._global_shape!r}, placement={self._layout.ragged_shard.placement!r}"
308 )
310 @property
311 def device_mesh(self) -> DeviceMesh:
312 """The device mesh of this DTensor."""
313 return self._device_mesh
315 @property
316 def placements(self) -> Sequence[Placement]:
317 """The placements of this DTensor."""
318 return self._placements
320 @property
321 def layout(self) -> Layout:
322 """Internal layout for redistribution (for backward compatibility)."""
323 if not hasattr(self, '_layout'):
324 return None
325 return self._layout
327 @staticmethod
328 def from_local(
329 local_tensor: Tensor,
330 device_mesh: DeviceMesh,
331 placements: Union[Sequence[Placement], Sequence[Union[str, Tuple[str, ...]]]],
332 *,
333 run_check: bool = False,
334 shape: Optional[Tuple[int, ...]] = None,
335 stride: Optional[Tuple[int, ...]] = None,
336 ) -> 'DTensor':
337 """
338 Create a DTensor from a local tensor with device mesh and placements.
340 Args:
341 local_tensor (Tensor): The local tensor shard on this device. For
342 ``RaggedShard``, the input may use its natural rank-local shape;
343 construction stores it as a one-dimensional view internally.
344 device_mesh (DeviceMesh): The device mesh describing the device topology.
345 placements: The placement strategy. Supports two styles:
346 - Placement objects (e.g., ``[Shard(0), Replicate()]``).
347 - Alias strings (e.g., ``("dp", "None")`` or
348 ``(("dp", "tp"), "None")``), length must equal the number
349 of tensor dimensions.
350 run_check (bool, optional): When ``True``, perform cross-rank metadata
351 checks and broadcast replicate placements from the mesh source rank.
352 Default: ``False``.
353 shape (tuple[int, ...], optional): Explicit logical global shape.
354 Required for RaggedShard.
355 stride (tuple[int, ...], optional): Explicit logical global stride.
356 Requires ``shape``. Normal layouts also retain compatibility
357 with shape-only construction.
359 Returns:
360 DTensor: A new DTensor instance.
362 Example:
363 >>> mesh = init_device_mesh(device_type="npu", mesh_shape=(2, 2), mesh_dim_names=("dp", "tp"))
364 >>> local_tensor = Tensor(np.ones((4, 4)))
365 >>> dtensor = DTensor.from_local(local_tensor, mesh, [Shard(0), Replicate()])
366 >>> dtensor = DTensor.from_local(local_tensor, mesh, ("dp", "None"))
367 """
368 tensor_dim = len(shape) if shape is not None else len(local_tensor.shape)
369 layout = _build_layout(device_mesh, placements, tensor_dim)
370 is_ragged = _layout_has_ragged_shard(layout)
371 if is_ragged:
372 if stride is not None and shape is None:
373 raise ValueError("stride requires an explicit shape")
374 if hasattr(local_tensor, "is_contiguous") and not local_tensor.is_contiguous():
375 raise ValueError("RaggedShard local tensor must be contiguous")
376 local_tensor = local_tensor.view(-1)
377 elif stride is not None and shape is None:
378 raise ValueError("stride requires an explicit shape")
379 if run_check:
380 # pylint: disable=C0415
381 from hyper_parallel.core.dtensor._from_local_utils import run_from_local_checks
382 run_from_local_checks(
383 local_tensor,
384 device_mesh,
385 layout.placements,
386 shape=shape,
387 )
388 if shape is not None and stride is not None:
389 layout = cp.deepcopy(layout)
390 layout.set_tensor_meta(shape, stride, local_tensor.dtype)
391 return DTensor(
392 local_tensor,
393 device_mesh,
394 layout.placements,
395 layout,
396 shape=shape,
397 )
399 @staticmethod
400 def from_local_with_layout(
401 local_tensor: Tensor,
402 layout: Layout,
403 *,
404 shape: Optional[Tuple[int, ...]] = None,
405 ) -> 'DTensor':
406 """Fast DTensor construction from a local tensor and a pre-built Layout.
408 Unlike :meth:`from_local`, this does NOT rebuild the layout via
409 ``_build_layout`` — it hands the already-built ``layout`` straight to
410 ``__init_data__``. Intended for hot paths (e.g. ``wrap_output``) where the
411 output Layout was already inferred and cached by ``infer_layout``, so
412 recomputing ``device_mesh.to_hash()`` / ``tuple(placements)`` / the layout
413 cache lookup on every output is pure waste.
415 ``layout.placements`` (a plain attribute) is passed only to satisfy the
416 constructor's non-None check; ``__init_data__`` ignores it when ``layout``
417 is supplied.
418 """
419 return DTensor(
420 local_tensor,
421 layout.mesh,
422 layout.placements,
423 layout,
424 shape=shape,
425 )
427 def _alias_placements(self) -> Sequence[Placement]:
428 """Return alias_placements from layout, falling back to _placements."""
429 if hasattr(self, '_layout') and self._layout:
430 return self._layout.alias_placements
431 return self._placements
433 def _from_converted_local(self, local_tensor: Tensor) -> 'DTensor':
434 """Rebuild converted DTensor data without preserving Parameter identity."""
435 cls = DTensor if isinstance(self, platform.Parameter) else self.__class__
436 if not isinstance(self._layout, Layout):
437 constructor_kwargs = {
438 "device_mesh": self._device_mesh,
439 "placements": self._alias_placements(),
440 }
441 if hasattr(self, "_global_shape"):
442 constructor_kwargs["shape"] = self._global_shape
443 return cls(local_tensor, **constructor_kwargs)
444 layout = cp.deepcopy(self._layout)
445 if layout.tensor_shape is not None:
446 layout.set_tensor_meta(
447 layout.tensor_shape,
448 layout.tensor_stride,
449 local_tensor.dtype,
450 )
451 return cls(
452 local_tensor,
453 device_mesh=self._device_mesh,
454 placements=layout.placements,
455 layout=layout,
456 shape=getattr(self, "_global_shape", None),
457 )
459 def to(self, *args, **kwargs):
460 """Move the DTensor to a different device or dtype.
462 Delegates to the underlying local tensor's ``to`` method and
463 reconstructs a DTensor preserving device_mesh and placements.
465 Args:
466 *args (tuple): Arguments passed to the underlying tensor's ``to``
467 method (e.g., device or dtype).
468 **kwargs (dict): Keyword arguments for the tensor conversion
469 (e.g., dtype, device, non_blocking).
471 Returns:
472 DTensor: A new DTensor with the converted local tensor.
473 """
474 new_local = self._local_tensor.to(*args, **kwargs)
475 return self._from_converted_local(new_local)
477 def float(self):
478 """Convert the DTensor to float dtype.
480 Returns:
481 DTensor: A new DTensor with float32 local tensor.
482 """
483 new_local = self._local_tensor.float()
484 return self._from_converted_local(new_local)
486 def type_as(self, other: Tensor) -> "DTensor":
487 """Cast this DTensor to the dtype of ``other``.
489 This is a **local** operation — no communication. Each shard
490 independently casts its elements to the target dtype.
492 Only the **dtype** of ``other`` is read; its shape, values, and
493 layout are ignored. The returned DTensor preserves the
494 device-mesh and placements of ``self`` unchanged.
496 Args:
497 other (Tensor): A tensor whose ``.dtype`` will be used as the
498 target type. May be a plain :class:`Tensor` or a
499 :class:`DTensor`. Must reside on the same device as
500 ``self``.
502 Returns:
503 DTensor: A new DTensor with the converted local tensor. When
504 ``self.dtype == other.dtype`` the method returns ``self``
505 unchanged (no-op).
507 Raises:
508 ValueError: If ``other`` is not a Tensor.
509 ValueError: If ``self`` has Partial placement (cast does not
510 commute with reduction).
511 ValueError: If ``self`` and ``other`` are on different devices.
513 Note:
514 This implementation intentionally covers **dtype-only**
515 conversion. PyTorch's native ``type_as`` may also handle
516 cross-device transfers, but a DTensor cannot silently change
517 its backend device while retaining the old ``DeviceMesh``.
518 Use :meth:`to` for explicit device + dtype conversion.
520 Example:
521 >>> # x is a DTensor of float16, y is a plain float32 Tensor
522 >>> # on the same device.
523 >>> z = x.type_as(y)
524 >>> z.dtype == y.dtype
525 True
526 """
527 if not isinstance(other, Tensor):
528 raise ValueError(
529 f"type_as() argument must be a Tensor, but got "
530 f"{type(other).__name__}."
531 )
532 if hasattr(self, '_layout') and self._layout is not None:
533 if self._layout.is_partial():
534 raise ValueError(
535 "DTensor.type_as does not support Partial input; "
536 "call reduce_partial() first."
537 )
539 other_local = other.to_local() if isinstance(other, DTensor) else other
540 if self._local_tensor.device != other_local.device:
541 raise ValueError(
542 "DTensor.type_as requires self and other to be on the "
543 "same device. Use to() for explicit device + dtype "
544 "conversion."
545 )
547 target_dtype = other.dtype
548 if self.dtype == target_dtype:
549 return self
550 new_local = self._local_tensor.to(dtype=target_dtype)
551 return self._from_converted_local(new_local)
553 def _validate_factory_device(self, device: Any) -> None:
554 """Raise :class:`ValueError` if ``device`` does not match the DTensor's device."""
555 requested_type, requested_index = _device_spec(device)
556 local_type, local_index = _device_spec(self._local_tensor.device)
557 if (
558 requested_type != local_type
559 or (
560 requested_index is not None
561 and requested_index != local_index
562 )
563 ):
564 raise ValueError(
565 f"DTensor requires device to match the input DTensor "
566 f"device {self._local_tensor.device}, but got {device}."
567 )
569 def _new_const_tensor_op(
570 self,
571 method_name: str,
572 size: Union[int, Sequence[int]],
573 *,
574 dtype: Optional[Any] = None,
575 device: Optional[Any] = None,
576 requires_grad: bool = False,
577 layout: Optional[Any] = None,
578 pin_memory: bool = False,
579 ) -> 'DTensor':
580 """Create an all-``Replicate`` constant DTensor.
582 Shared implementation for ``new_zeros`` and ``new_ones``.
584 ``self`` is only used as a dtype/device reference and mesh source;
585 its values are ignored. The output is always **all-Replicate**
586 because every device produces identical data independently.
588 Args:
589 method_name:
590 ``"new_zeros"`` or ``"new_ones"`` — the local tensor
591 factory method to call.
592 size:
593 Output shape — an int or a sequence of ints.
594 dtype:
595 Desired dtype. Defaults to ``self.dtype`` on Torch.
596 device:
597 Must match ``self``'s device (Torch only).
598 requires_grad:
599 Forwarded on Torch; rejected on MindSpore.
600 layout:
601 Forwarded on Torch; rejected on MindSpore.
602 pin_memory:
603 Forwarded on Torch; rejected on MindSpore.
605 Returns:
606 DTensor: A new DTensor with all-``Replicate`` placements on
607 ``self``'s ``DeviceMesh``.
609 Raises:
610 ValueError: If a Torch-only kwarg is used on MindSpore, or
611 ``device`` does not match the DTensor's device.
612 """
613 if isinstance(size, int):
614 size = (size,)
616 if platform.platform_type == PlatformType.MINDSPORE:
617 if device is not None or layout is not None or requires_grad or pin_memory:
618 raise ValueError(
619 f"DTensor.{method_name} only supports size and dtype "
620 "on MindSpore."
621 )
622 local_kwargs = {}
623 if dtype is not None:
624 local_kwargs["dtype"] = dtype
625 else:
626 local_kwargs = {}
627 if dtype is not None:
628 local_kwargs["dtype"] = dtype
629 if device is not None:
630 self._validate_factory_device(device)
631 # An unindexed device such as "cuda" resolves to the framework's
632 # current device, which may differ from this DTensor's local device.
633 local_kwargs["device"] = self._local_tensor.device
634 if requires_grad:
635 local_kwargs["requires_grad"] = True
636 if layout is not None:
637 local_kwargs["layout"] = layout
638 if pin_memory:
639 local_kwargs["pin_memory"] = True
641 factory = getattr(self._local_tensor, method_name)
642 local_result = factory(size, **local_kwargs)
644 replicated_placements = [Replicate()] * self._device_mesh.ndim
645 return DTensor.from_local(
646 local_result, self._device_mesh, replicated_placements,
647 )
649 def new_zeros(
650 self,
651 size: Union[int, Sequence[int]],
652 *,
653 dtype: Optional[Any] = None,
654 device: Optional[Any] = None,
655 requires_grad: bool = False,
656 layout: Optional[Any] = None,
657 pin_memory: bool = False,
658 ) -> 'DTensor':
659 """Create an all-Replicate DTensor filled with zeros.
661 The output is always **fully replicated** across every device in
662 ``self``'s ``DeviceMesh``, regardless of how ``self`` is sharded.
664 Args:
665 size:
666 Output shape — an int or a sequence of ints.
667 dtype:
668 Desired dtype. Defaults to ``self.dtype`` (Torch).
669 Not forwarded to MindSpore unless explicitly set.
670 device:
671 Must match ``self``'s device. Not supported on MindSpore.
672 requires_grad:
673 Forwarded on Torch; rejected on MindSpore.
674 layout:
675 Forwarded on Torch; rejected on MindSpore.
676 pin_memory:
677 Forwarded on Torch; rejected on MindSpore.
679 Returns:
680 DTensor: A new all-Replicate DTensor filled with zeros.
681 """
682 return self._new_const_tensor_op(
683 "new_zeros", size,
684 dtype=dtype,
685 device=device,
686 requires_grad=requires_grad,
687 layout=layout,
688 pin_memory=pin_memory,
689 )
691 def new_ones(
692 self,
693 size: Union[int, Sequence[int]],
694 *,
695 dtype: Optional[Any] = None,
696 device: Optional[Any] = None,
697 requires_grad: bool = False,
698 layout: Optional[Any] = None,
699 pin_memory: bool = False,
700 ) -> 'DTensor':
701 """Create an all-Replicate DTensor filled with ones.
703 The output is always **fully replicated** across every device in
704 ``self``'s ``DeviceMesh``, regardless of how ``self`` is sharded.
706 Args:
707 size:
708 Output shape — an int or a sequence of ints.
709 dtype:
710 Desired dtype. Defaults to ``self.dtype`` (Torch).
711 Not forwarded to MindSpore unless explicitly set.
712 device:
713 Must match ``self``'s device. Not supported on MindSpore.
714 requires_grad:
715 Forwarded on Torch; rejected on MindSpore.
716 layout:
717 Forwarded on Torch; rejected on MindSpore.
718 pin_memory:
719 Forwarded on Torch; rejected on MindSpore.
721 Returns:
722 DTensor: A new all-Replicate DTensor filled with ones.
723 """
724 return self._new_const_tensor_op(
725 "new_ones", size,
726 dtype=dtype,
727 device=device,
728 requires_grad=requires_grad,
729 layout=layout,
730 pin_memory=pin_memory,
731 )
733 def to_local(self) -> Tensor:
734 """
735 Convert DTensor to local tensor.
737 Returns:
738 Tensor: The local tensor shard on this device.
739 """
740 return self._local_tensor
742 def tolist(self):
743 """
744 Convert the DTensor to a nested Python list or number.
746 This operation gathers the complete tensor on every participating rank
747 before converting it to Python values. It is an **implicit collective**:
748 all ranks in the DeviceMesh must participate.
750 Returns:
751 Union[list, int, float, bool]: A nested Python list, or a Python
752 number for a scalar DTensor.
754 Note:
755 This triggers ``full_tensor()`` under the hood, which performs
756 all-gather communication. For large tensors, prefer slicing or
757 index-based access to avoid materialising the full tensor.
759 If you only need the **local shard** as a list, use
760 ``dtensor.to_local().tolist()`` instead — that path has zero
761 communication overhead.
763 Example:
764 >>> mesh = init_device_mesh("npu", (2,), ("dp",))
765 >>> x = distribute_tensor(torch.arange(8).reshape(4, 2), mesh, [Shard(0)])
766 >>> x.tolist() # full data: [[0,1],[2,3],[4,5],[6,7]]
767 >>> x.to_local().tolist() # local shard only (no comm)
768 """
769 return self.full_tensor().tolist()
771 def copy_(self, src: "DTensor", non_blocking: bool = False) -> "DTensor":
772 """In-place copy of ``src`` into this DTensor's local shard.
774 Delegates to ``Tensor.copy_`` on the underlying local tensors.
775 Follows standard ``Tensor.copy_`` semantics: version counter is
776 bumped and autograd edges are created when grad is enabled.
778 Constraints on ``src``:
779 * must be a ``DTensor`` on the same or an equivalent ``DeviceMesh`` topology as ``self``;
780 * its placements must equal ``self.placements``, OR
781 ``src._local_tensor.numel() == 1`` (single-element broadcast);
782 * its local shape must equal or be broadcastable to
783 ``self._local_tensor.shape``.
785 No redistribute / implicit slicing is performed; src dtype is cast
786 to self dtype in-place.
788 Args:
789 src (DTensor): Source DTensor satisfying the constraints above.
790 non_blocking (bool): Forwarded to the underlying ``copy_``.
792 Returns:
793 DTensor: ``self``.
795 Raises:
796 TypeError: if ``src`` is not a ``DTensor``.
797 ValueError: if mesh, placement, or shape constraint is violated.
798 """
799 if not isinstance(src, DTensor):
800 raise TypeError(
801 f"For DTensor.copy_, src should be a DTensor, but got {type(src).__name__}."
802 )
803 src_local = src.to_local()
804 if not _device_meshes_are_compatible(src.device_mesh, self._device_mesh):
805 raise ValueError(
806 f"For DTensor.copy_, src and self should share the same DeviceMesh, "
807 f"but got src.device_mesh={src.device_mesh!r}, "
808 f"self._device_mesh={self._device_mesh!r}."
809 )
811 placement_eq = tuple(src.placements) == tuple(self._placements)
812 shape_eq = src_local.shape == self._local_tensor.shape
813 src_is_scalar = src_local.numel() == 1
815 if not placement_eq and not src_is_scalar:
816 raise ValueError(
817 f"For DTensor.copy_, src.placements should equal self.placements "
818 f"or src.numel() should be 1, but got "
819 f"src.placements={src.placements}, "
820 f"self.placements={self._placements}, "
821 f"src.numel()={src_local.numel()}."
822 )
823 if not shape_eq and not src_is_scalar and not _is_broadcastable(
824 src_local.shape, self._local_tensor.shape
825 ):
826 raise ValueError(
827 f"For DTensor.copy_, src local shape should be broadcastable to "
828 f"self local shape, but got "
829 f"src.shape={tuple(src_local.shape)}, "
830 f"self.shape={tuple(self._local_tensor.shape)}."
831 )
833 self._local_tensor.copy_(src_local, non_blocking=non_blocking)
834 return self
836 def zero_(self) -> "DTensor":
837 """In-place fill with zeros. Returns ``self``."""
838 self._local_tensor.zero_()
839 return self
841 def fill_(self, value) -> "DTensor":
842 """In-place fill with ``value``. Returns ``self``."""
843 self._local_tensor.fill_(value)
844 return self
846 @property
847 def shape(self) -> Tuple[int, ...]:
848 """
849 The global shape of this DTensor.
851 Returns:
852 Tuple[int, ...]: The global tensor shape.
853 """
854 return self._global_shape
856 def size(self, dim=None):
857 """Return the global shape, consistent with .shape.
859 Without ``dim`` returns a tuple matching ``self.shape``.
860 With ``dim`` returns the size of that dimension.
861 """
862 global_shape = self.shape
863 if dim is not None:
864 return global_shape[dim]
865 return global_shape
867 def numel(self) -> int:
868 """Return the number of elements in this DTensor."""
869 return int(np.prod(self.shape))
871 @property
872 def ndim(self) -> int:
873 """Return the logical global tensor rank."""
874 return len(self._global_shape)
876 def dim(self) -> int:
877 """Return the logical global tensor rank."""
878 return len(self._global_shape)
880 @property
881 def local_shape(self) -> Tuple[int, ...]:
882 """
883 The local shape of this DTensor on this device.
885 Returns:
886 Tuple[int, ...]: The local tensor shape.
887 """
888 return self._local_tensor.shape
890 def redistribute(
891 self,
892 device_mesh: DeviceMesh,
893 placements: Union[Sequence[Placement], Sequence[Union[str, Tuple[str, ...]]]]
894 ) -> 'DTensor':
895 """
896 Redistribute this DTensor to a new device mesh and placements.
898 Args:
899 device_mesh (DeviceMesh): The target device mesh.
900 placements: The target placements. Supports Placement objects
901 or alias strings.
903 Returns:
904 DTensor: A new DTensor with the specified distribution.
906 Example:
907 >>> new_dtensor = dtensor.redistribute(mesh, [Replicate(), Shard(1)])
908 >>> new_dtensor = dtensor.redistribute(mesh, ("None", "tp"))
909 """
910 logger.debug(
911 "redistribute: shape=%s, src_placements=%s -> dst_placements=%s, "
912 "mesh_shape=%s, local_shape=%s",
913 tuple(self.shape),
914 tuple(self._placements),
915 tuple(placements),
916 tuple(device_mesh.shape),
917 tuple(self._local_tensor.shape),
918 )
920 # Build dst_layout from device_mesh and placements
921 dst_layout = _build_layout(
922 device_mesh, placements, len(self._global_shape)
923 )
925 # pylint: disable=C0415
926 from hyper_parallel.core.dtensor.tensor_redistribution import _tensor_redistribution
927 out = _tensor_redistribution.redistribution(self, dst_layout)
928 return out
930 def reduce_partial(self) -> 'DTensor':
931 """
932 Reduce partial sharding state for this DTensor.
934 Returns:
935 DTensor: A new DTensor with partial state reduced.
936 """
937 if not self._layout:
938 return self
939 to_layout = cp.deepcopy(self._layout)
940 to_layout.reset_partial()
941 # pylint: disable=C0415
942 from hyper_parallel.core.dtensor.tensor_redistribution import _tensor_redistribution
943 out = _tensor_redistribution.reduce_partial(self, to_layout)
944 return out
946 def full_tensor(self) -> Tensor:
947 """
948 Return the full tensor of this DTensor.
950 Returns:
951 Tensor: A Tensor object that represents the full tensor of this DTensor.
952 The returned tensor contains the complete data gathered from
953 all ranks.
955 Note:
956 This operation involves communication across all ranks in the DeviceMesh,
957 which may be expensive for large tensors. Use with caution in
958 performance-critical code paths.
960 Example:
961 >>> # Assume dtensor is sharded across multiple devices
962 >>> local_tensor = dtensor.to_local() # Returns only the local shard
963 >>> full_tensor = dtensor.full_tensor() # Returns the complete tensor
964 """
965 if not self._layout:
966 return self._local_tensor
968 # Create a fully replicated layout
969 replicated_layout = cp.deepcopy(self._layout)
971 # Set all placements to Replicate and convert to tensor_map
972 replicated_placements = [Replicate()] * len(replicated_layout.mesh_shape)
973 replicated_layout.set_placements(replicated_placements)
974 replicated_layout.placement_to_tensor_map(len(self._global_shape))
976 # Clear partial status from original layout since Replicate has no partial
977 replicated_layout.reset_partial()
979 # Redistribute to the replicated layout and return local tensor
980 # pylint: disable=C0415
981 from hyper_parallel.core.dtensor.tensor_redistribution import _tensor_redistribution
982 out = _tensor_redistribution.redistribution(self, replicated_layout)
983 return out.to_local()
986def _normalize_shard_dim(dim: int, ndim: int) -> int:
987 return dim + ndim if dim < 0 else dim
990def _distribute_tensor_with_communication(
991 tensor: Tensor,
992 device_mesh: DeviceMesh,
993 placements: Sequence[Placement],
994 src_data_rank: int,
995) -> Tensor:
996 """Scatter/broadcast a logical global tensor along mesh dimensions (PyTorch parity)."""
997 local = tensor
998 if len(placements) < device_mesh.ndim:
999 raise ValueError(
1000 f"placements length ({len(placements)}) must be at least device_mesh.ndim "
1001 f"({device_mesh.ndim}) when src_data_rank is set"
1002 )
1003 for mesh_dim in range(device_mesh.ndim):
1004 placement = placements[mesh_dim]
1005 if isinstance(placement, StridedShard):
1006 raise NotImplementedError(
1007 "distribute_tensor with src_data_rank does not support StridedShard yet; "
1008 "pass src_data_rank=None for local-only sharding."
1009 )
1010 if placement.is_shard():
1011 shard_dim = _normalize_shard_dim(placement.dim, local.ndim)
1012 num_chunks = device_mesh.size(mesh_dim)
1013 if num_chunks <= 0:
1014 raise ValueError(f"invalid mesh dim size {num_chunks} on mesh_dim={mesh_dim}")
1015 chunks = tuple(local.chunk(num_chunks, dim=shard_dim))
1016 if not chunks:
1017 raise ValueError(f"cannot shard dim {shard_dim} into {num_chunks} chunks")
1018 output = platform.empty_like(chunks[0])
1019 local = mesh_scatter(output, chunks, device_mesh, mesh_dim, group_src=src_data_rank)
1020 elif placement.is_replicate() or placement.is_partial():
1021 local = mesh_broadcast(local, device_mesh, mesh_dim, group_src=src_data_rank)
1022 if isinstance(placement, Partial):
1023 warnings.warn(
1024 f"Partial placement {placement} during distribute_tensor: "
1025 "broadcast only; partial partition is not applied yet.",
1026 stacklevel=3,
1027 )
1028 else:
1029 raise RuntimeError(
1030 f"unsupported placement {placement} on device mesh dimension {mesh_dim}"
1031 )
1032 return local
1035def distribute_tensor(
1036 tensor: Tensor,
1037 device_mesh: DeviceMesh,
1038 placements: Union[Sequence[Placement], Sequence[Union[str, Tuple[str, ...]]]],
1039 *,
1040 src_data_rank: Optional[int] = None,
1041) -> DTensor:
1042 """
1043 Distribute a global tensor to the device mesh according to the placements.
1045 Args:
1046 tensor (Tensor): The global tensor to be distributed. All ranks
1047 should have the same tensor data.
1048 device_mesh (DeviceMesh): The device mesh describing the device topology.
1049 placements: The placement strategy. Supports two styles:
1050 - Placement objects (e.g., ``[Shard(0), Replicate()]``).
1051 - Alias strings (e.g., ``("dp", "None")`` or
1052 ``(("dp", "tp"), "None")``), length must equal the number of
1053 tensor dimensions.
1055 Returns:
1056 DTensor: A new DTensor with the local shard on each rank.
1058 Note:
1059 When ``src_data_rank`` is an ``int`` (e.g. ``0``), shard/replicate
1060 placements use scatter/broadcast from the source rank on each mesh axis,
1061 matching PyTorch ``distribute_tensor``. When ``src_data_rank=None``
1062 (default), each rank slices its local tensor without communication
1063 (legacy Hyper behavior; all ranks must hold the same global data).
1065 Example:
1066 >>> mesh = init_device_mesh(device_type="npu", mesh_shape=(2, 2), mesh_dim_names=("dp", "tp"))
1067 >>> global_tensor = Tensor(np.arange(16).reshape(4, 4))
1068 >>> dtensor = distribute_tensor(global_tensor, mesh, [Shard(0), Replicate()])
1069 >>> dtensor = distribute_tensor(global_tensor, mesh, ("dp", "None"))
1070 """
1071 layout = _build_layout(device_mesh, placements, len(tensor.shape))
1072 if _layout_has_ragged_shard(layout):
1073 if src_data_rank is None:
1074 local_tensor = _slice_ragged_tensor(tensor, layout)
1075 else:
1076 local_tensor = _scatter_ragged_tensor(tensor, layout, src_data_rank)
1077 elif src_data_rank is None:
1078 local_tensor = _get_slice_tensor_by_layout(tensor, layout)
1079 else:
1080 local_tensor = _distribute_tensor_with_communication(
1081 tensor, device_mesh, layout.placements, src_data_rank
1082 )
1083 return DTensor.from_local_with_layout(
1084 local_tensor,
1085 layout,
1086 shape=tuple(tensor.shape),
1087 )
1090def _distribute_module_param_source(param: Any) -> Tensor:
1091 """Tensor data used as the global tensor for :func:`distribute_tensor` (PyTorch uses ``param.data``)."""
1092 if hasattr(param, "data"):
1093 return param.data
1094 return platform.get_param_local_data(param)
1097def _distribute_module_new_parameter(key: str, dtensor: DTensor, requires_grad: bool) -> Any:
1098 """Build a framework :class:`Parameter` holding *dtensor* (Torch vs MindSpore kwargs differ)."""
1099 if platform.platform_type == PlatformType.MINDSPORE:
1100 return platform.Parameter(dtensor, name=key, requires_grad=requires_grad)
1101 return platform.Parameter(dtensor, requires_grad=requires_grad)
1104def _distribute_module_set_param(module: Any, key: str, new_param: Any) -> None:
1105 """Register or assign a parameter on *module* (``nn.Module`` or MindSpore ``Cell``)."""
1106 if hasattr(module, "register_parameter"):
1107 module.register_parameter(key, new_param)
1108 return
1109 if hasattr(module, "_params"):
1110 module._params[key] = new_param
1111 if hasattr(module, "_params_list"):
1112 module._params_list[key] = new_param
1113 if key in module.__dict__:
1114 module.__dict__[key] = new_param
1115 return
1116 raise TypeError(
1117 f"distribute_module expects nn.Module-like objects with register_parameter or _params; "
1118 f"got {type(module)}."
1119 )
1122def _distribute_module_iter_params(module: Any) -> list:
1123 """Return ``[(name, param), ...]`` for direct parameters (``_parameters`` or ``_params``)."""
1124 if hasattr(module, "_parameters"):
1125 return list(module._parameters.items())
1126 if hasattr(module, "_params"):
1127 return list(module._params.items())
1128 return []
1131def _distribute_module_iter_buffers(module: Any) -> list:
1132 """Return ``[(name, buffer), ...]`` if the module has ``_buffers`` (PyTorch ``nn.Module``)."""
1133 if hasattr(module, "_buffers"):
1134 return list(module._buffers.items())
1135 return []
1138def _distribute_module_named_modules(module: Any):
1139 """``nn.Module.named_modules`` or MindSpore ``Cell.cells_and_names`` (submodule FQNs)."""
1140 if hasattr(module, "named_modules"):
1141 return module.named_modules()
1142 if hasattr(module, "cells_and_names"):
1143 return module.cells_and_names()
1144 raise TypeError(
1145 f"distribute_module expects module-like objects with named_modules or cells_and_names; "
1146 f"got {type(module)}."
1147 )
1150def _distribute_module_named_parameters(module: Any):
1151 """``nn.Module.named_parameters(recurse=False)`` or MindSpore ``Cell.parameters_and_names(expand=False)``."""
1152 if hasattr(module, "named_parameters"):
1153 return module.named_parameters(recurse=False)
1154 if hasattr(module, "parameters_and_names"):
1155 return module.parameters_and_names(expand=False)
1156 raise TypeError(
1157 f"distribute_module expects module-like objects with named_parameters or parameters_and_names; "
1158 f"got {type(module)}."
1159 )
1162def _replicate_submodule_params_buffers(
1163 sub_mod: Any,
1164 device_mesh: DeviceMesh,
1165 *,
1166 module_prefix: str = "",
1167) -> None:
1168 """Convert plain params/buffers on *sub_mod* to fully replicated :class:`DTensor`."""
1169 full_replicate = [Replicate()] * device_mesh.ndim
1170 for key, param in _distribute_module_iter_params(sub_mod):
1171 if param is None or isinstance(param, DTensorBase):
1172 continue
1173 src = _distribute_module_param_source(param)
1174 requires_grad = bool(getattr(param, "requires_grad", True))
1175 dt = distribute_tensor(src, device_mesh, full_replicate)
1176 param_name = f"{module_prefix}.{key}" if module_prefix else key
1177 new_param = _distribute_module_new_parameter(param_name, dt, requires_grad)
1178 _distribute_module_set_param(sub_mod, key, new_param)
1179 for key, buffer in _distribute_module_iter_buffers(sub_mod):
1180 if buffer is None or isinstance(buffer, DTensorBase):
1181 continue
1182 sub_mod._buffers[key] = distribute_tensor(buffer, device_mesh, full_replicate)
1185def _distribute_module_run_partition_and_replicate(
1186 module: Any,
1187 device_mesh: DeviceMesh,
1188 partition_fn: Optional[Callable[[str, Any, DeviceMesh], None]],
1189) -> None:
1190 """Call optional ``partition_fn`` per ``named_modules`` and replicate remaining tensors."""
1191 if partition_fn is None:
1192 for mod_name, submod in _distribute_module_named_modules(module):
1193 _replicate_submodule_params_buffers(submod, device_mesh, module_prefix=mod_name)
1194 return
1195 for mod_name, submod in _distribute_module_named_modules(module):
1196 partition_fn(mod_name, submod, device_mesh)
1197 _replicate_submodule_params_buffers(submod, device_mesh, module_prefix=mod_name)
1200def _distribute_module_register_input_fn(
1201 module: Any,
1202 device_mesh: DeviceMesh,
1203 input_fn: Callable[..., Any],
1204) -> None:
1205 """Register *input_fn* as a forward pre-hook on *module* (2- or 3-arg, PyTorch-compatible)."""
1206 num_args = len(inspect.signature(input_fn).parameters)
1207 if num_args == 2:
1208 warnings.warn(
1209 "Deprecating input_fn that takes two arguments (inputs, device_mesh), "
1210 "please use input_fn that takes in (module, inputs, device_mesh) instead!",
1211 FutureWarning,
1212 stacklevel=3,
1213 )
1214 module.register_forward_pre_hook(
1215 lambda _, inputs: input_fn(inputs, device_mesh)
1216 )
1217 elif num_args == 3:
1218 module.register_forward_pre_hook(
1219 lambda mod, inputs: input_fn(mod, inputs, device_mesh)
1220 )
1221 else:
1222 raise ValueError(
1223 f"input_fn should take in 2 or 3 arguments, but got {num_args} arguments!"
1224 )
1227def _distribute_module_register_output_fn(
1228 module: Any,
1229 device_mesh: DeviceMesh,
1230 output_fn: Callable[..., Any],
1231) -> None:
1232 """Register *output_fn* as a forward hook on *module* (2- or 3-arg, PyTorch-compatible)."""
1233 num_args = len(inspect.signature(output_fn).parameters)
1234 if num_args == 2:
1235 warnings.warn(
1236 "Deprecating output_fn that takes two arguments (outputs, device_mesh), "
1237 "please use output_fn that takes in (module, outputs, device_mesh) instead!",
1238 FutureWarning,
1239 stacklevel=3,
1240 )
1241 module.register_forward_hook(
1242 lambda mod, inputs, outputs: output_fn(outputs, device_mesh)
1243 )
1244 elif num_args == 3:
1245 module.register_forward_hook(
1246 lambda mod, inputs, outputs: output_fn(mod, outputs, device_mesh)
1247 )
1248 else:
1249 raise ValueError(
1250 f"output_fn should take in 2 or 3 arguments, but got {num_args} arguments!"
1251 )
1254def distribute_module(
1255 module: Any,
1256 device_mesh: Optional[DeviceMesh] = None,
1257 partition_fn: Optional[Callable[[str, Any, DeviceMesh], None]] = None,
1258 input_fn: Optional[Callable[..., Any]] = None,
1259 output_fn: Optional[Callable[..., Any]] = None,
1260) -> Any:
1261 """PyTorch ``distribute_module`` parity: shard/replicate params and optional I/O hooks.
1263 Unsharded parameters and buffers become fully replicated :class:`DTensor` after
1264 ``partition_fn``. ``input_fn`` / ``output_fn`` attach only to the root *module*.
1266 Args:
1267 module: Root ``nn.Module`` or MindSpore ``Cell`` with compatible APIs.
1268 device_mesh: Placement mesh; if ``None``, uses ``_mesh_resources.get_current_mesh()``.
1269 partition_fn: Per ``named_modules`` callback before replicate pass; ``None`` replicates all.
1270 input_fn: ``(module, inputs, mesh)`` or deprecated ``(inputs, mesh)`` pre-hook.
1271 output_fn: ``(module, outputs, mesh)`` or deprecated ``(outputs, mesh)`` forward hook.
1273 Returns:
1274 *module* in place, with distributed tensors where applied.
1276 Raises:
1277 RuntimeError: If called twice on the same *module*.
1278 ValueError: If ``input_fn`` / ``output_fn`` arity is not 2 or 3.
1280 Note:
1281 XLA / ``torch_xla`` is not supported; strided device :class:`DTensor` only.
1282 """
1283 if getattr(module, "_distribute_module_applied", False):
1284 raise RuntimeError(
1285 "distribute_module should only be called once on a module, "
1286 "but it has already been called on this module!"
1287 )
1288 device_mesh = device_mesh or _mesh_resources.get_current_mesh()
1289 _distribute_module_run_partition_and_replicate(module, device_mesh, partition_fn)
1290 if input_fn is not None:
1291 _distribute_module_register_input_fn(module, device_mesh, input_fn)
1292 if output_fn is not None:
1293 _distribute_module_register_output_fn(module, device_mesh, output_fn)
1294 module._distribute_module_applied = True
1295 return module
1298def _dtensor_init_helper(
1299 init_op,
1300 size,
1301 device_mesh,
1302 placements,
1303 *,
1304 rng_tracked: bool = False,
1305 **kwargs,
1306) -> DTensor:
1307 """
1308 Helper function to create and initialize a distributed tensor.
1310 Args:
1311 size: Shape of the tensor.
1312 dtype: Data type of the tensor.
1313 device: Target device for the tensor.
1314 requires_grad: Whether the tensor requires gradient.
1315 rng_tracked: When ``True``, initialize via :class:`OffsetBasedRNGTracker`
1316 so shard/replicate random semantics match PyTorch DTensor factories.
1318 Returns:
1319 DTensor: The initialized distributed tensor.
1320 """
1321 global_shape = (size,) if isinstance(size, int) else tuple(size)
1322 layout = _build_layout(device_mesh, placements, len(global_shape))
1323 if _layout_has_ragged_shard(layout):
1324 raise NotImplementedError(
1325 "RaggedShard tensor factories are not implemented in the DTensor metadata phase"
1326 )
1328 # get local tensor shape
1329 local_shape = compute_local_shape_and_global_offset(
1330 size, device_mesh, placements
1331 )
1333 # initialize the local tensor
1334 if init_op is platform.full:
1335 fill_value = kwargs.pop("fill_value", 0)
1336 local_tensor = init_op(local_shape, fill_value, **kwargs)
1337 elif rng_tracked:
1338 # pylint: disable=C0415
1339 from hyper_parallel.core.dtensor.random import is_rng_supported_mesh, OffsetBasedRNGTracker
1340 from hyper_parallel.core.shard._op_dispatch import _OP_DISPATCHER
1342 layout = _build_layout(device_mesh, placements, len(local_shape))
1343 if is_rng_supported_mesh(device_mesh):
1344 if _OP_DISPATCHER._rng_tracker is None:
1345 _OP_DISPATCHER._rng_tracker = OffsetBasedRNGTracker(run_state_sync=False)
1346 with _OP_DISPATCHER._rng_tracker._distribute_region(
1347 device_mesh,
1348 layout.placements,
1349 global_shape,
1350 ):
1351 local_tensor = init_op(local_shape, **kwargs)
1352 else:
1353 local_tensor = init_op(local_shape, **kwargs)
1354 else:
1355 local_tensor = init_op(local_shape, **kwargs)
1357 return DTensor.from_local(
1358 local_tensor,
1359 device_mesh,
1360 placements,
1361 )
1364def ones(
1365 size,
1366 device_mesh,
1367 placements,
1368) -> DTensor:
1369 """
1370 Returns a :class:`DTensor` filled with the scalar value 1, with the shape defined
1371 by the variable argument ``size``.
1373 Args:
1374 size (Union[tuple[int], list[int], int, Tensor]): The specified shape of output tensor. Only positive integer or
1375 tuple or Tensor containing positive integers are allowed. If it is a Tensor,
1376 it must be a 0-D or 1-D Tensor with int32 or int64 dtypes.
1378 Keyword args:
1379 device_mesh: :class:`DeviceMesh` type, contains the mesh info of ranks
1380 placements: a sequence of :class:`Placement` type: ``Shard``, ``Replicate``
1382 Returns:
1383 A :class:`DTensor` object on each rank
1384 """
1385 ones_ = platform.ones
1386 return _dtensor_init_helper(
1387 ones_,
1388 size,
1389 device_mesh=device_mesh,
1390 placements=placements,
1391 )
1394def empty(
1395 size,
1396 device_mesh,
1397 placements,
1398) -> DTensor:
1399 """
1400 Returns a :class:`DTensor` filled with uninitialized data. The shape of the :class:`DTensor`
1401 is defined by the variable argument ``size``.
1403 Args:
1404 size (Union[tuple[int], list[int], int]): The specified shape of output tensor. Can be variable numbers of
1405 positive integers or tuple or list containing positive integers.
1407 Keyword args:
1408 device_mesh: :class:`DeviceMesh` type, contains the mesh info of ranks
1409 placements: a sequence of :class:`Placement` type: ``Shard``, ``Replicate``
1411 Returns:
1412 A :class:`DTensor` object on each rank
1413 """
1414 empty_ = platform.empty
1415 return _dtensor_init_helper(
1416 empty_,
1417 size,
1418 device_mesh=device_mesh,
1419 placements=placements,
1420 )
1423def full(
1424 size,
1425 fill_value,
1426 *,
1427 device_mesh,
1428 placements,
1429) -> DTensor:
1430 """
1431 Returns a :class:`DTensor` filled with ``fill_value`` according to ``device_mesh`` and
1432 ``placements``, with the shape defined by the argument ``size``.
1434 Args:
1435 size (Union[tuple[int], list[int]]): The specified shape of output tensor.
1436 fill_value (Union[numbers.Number, Tensor]): Value to fill the returned tensor. It can be a scalar number, a 0-D
1437 Tensor, or a 1-D Tensor with only one element.
1439 Keyword args:
1440 device_mesh: :class:`DeviceMesh` type, contains the mesh info of ranks.
1441 placements: a sequence of :class:`Placement` type: ``Shard``, ``Replicate``
1443 Returns:
1444 A :class:`DTensor` object on each rank
1445 """
1446 full_ = platform.full
1447 return _dtensor_init_helper(
1448 full_,
1449 size,
1450 fill_value=fill_value,
1451 device_mesh=device_mesh,
1452 placements=placements,
1453 )
1456def zeros(
1457 size,
1458 device_mesh,
1459 placements,
1460) -> DTensor:
1461 """
1462 Returns a :class:`DTensor` filled with the scalar value 0.
1464 Args:
1465 size (Union[tuple[int], list[int], int, Tensor]): The specified shape of output tensor. Only positive integer or
1466 tuple or Tensor containing positive integers are allowed. If it is a Tensor,
1467 it must be a 0-D or 1-D Tensor with int32 or int64 dtypes.
1468 Keyword args:
1469 device_mesh: :class:`DeviceMesh` type, contains the mesh info of ranks
1470 placements: a sequence of :class:`Placement` type: ``Shard``, ``Replicate``
1472 Returns:
1473 A :class:`DTensor` object on each rank
1474 """
1475 zeros_ = platform.zeros
1476 return _dtensor_init_helper(
1477 zeros_,
1478 size,
1479 device_mesh=device_mesh,
1480 placements=placements,
1481 )
1484def rand(
1485 size,
1486 device_mesh,
1487 placements,
1488 **kwargs,
1489) -> DTensor:
1490 """
1491 Returns a :class:`DTensor` filled with random numbers from a uniform
1492 distribution on ``[0, 1)``.
1494 Args:
1495 size: Global output shape.
1496 device_mesh: :class:`DeviceMesh` for the distributed layout.
1497 placements: Per-mesh-dimension :class:`Placement` values.
1498 **kwargs: Forwarded to the platform ``rand`` call (for example ``dtype``).
1500 Returns:
1501 A :class:`DTensor` object on each rank.
1502 """
1503 return _dtensor_init_helper(
1504 platform.rand,
1505 size,
1506 device_mesh=device_mesh,
1507 placements=placements,
1508 rng_tracked=True,
1509 **kwargs,
1510 )
1513def randn(
1514 size,
1515 device_mesh,
1516 placements,
1517 **kwargs,
1518) -> DTensor:
1519 """
1520 Returns a :class:`DTensor` filled with random numbers from a normal
1521 distribution with mean ``0`` and variance ``1``.
1523 Args:
1524 size: Global output shape.
1525 device_mesh: :class:`DeviceMesh` for the distributed layout.
1526 placements: Per-mesh-dimension :class:`Placement` values.
1527 **kwargs: Forwarded to the platform ``randn`` call (for example ``dtype``).
1529 Returns:
1530 A :class:`DTensor` object on each rank.
1531 """
1532 return _dtensor_init_helper(
1533 platform.randn,
1534 size,
1535 device_mesh=device_mesh,
1536 placements=placements,
1537 rng_tracked=True,
1538 **kwargs,
1539 )