Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / dtensor / dtensor.py: 88%
439 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-21 04:29 +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 ) -> 'DTensor':
336 """
337 Create a DTensor from a local tensor with device mesh and placements.
339 Args:
340 local_tensor (Tensor): The local tensor shard on this device.
341 device_mesh (DeviceMesh): The device mesh describing the device topology.
342 placements: The placement strategy. Supports two styles:
343 - Placement objects (e.g., ``[Shard(0), Replicate()]``).
344 - Alias strings (e.g., ``("dp", "None")`` or
345 ``(("dp", "tp"), "None")``), length must equal the number
346 of tensor dimensions.
347 run_check (bool, optional): When ``True``, perform cross-rank metadata
348 checks and broadcast replicate placements from the mesh source rank.
349 Default: ``False``.
350 shape (tuple[int, ...], optional): Explicit logical global shape.
351 Required for RaggedShard.
353 Returns:
354 DTensor: A new DTensor instance.
356 Example:
357 >>> mesh = init_device_mesh(device_type="npu", mesh_shape=(2, 2), mesh_dim_names=("dp", "tp"))
358 >>> local_tensor = Tensor(np.ones((4, 4)))
359 >>> dtensor = DTensor.from_local(local_tensor, mesh, [Shard(0), Replicate()])
360 >>> dtensor = DTensor.from_local(local_tensor, mesh, ("dp", "None"))
361 """
362 if run_check:
363 # pylint: disable=C0415
364 from hyper_parallel.core.dtensor._from_local_utils import run_from_local_checks
365 tensor_dim = len(shape) if shape is not None else len(local_tensor.shape)
366 layout = _build_layout(device_mesh, placements, tensor_dim)
367 run_from_local_checks(
368 local_tensor,
369 device_mesh,
370 layout.placements,
371 shape=shape,
372 )
373 return DTensor(local_tensor, device_mesh, placements, shape=shape)
375 @staticmethod
376 def from_local_with_layout(
377 local_tensor: Tensor,
378 layout: Layout,
379 *,
380 shape: Optional[Tuple[int, ...]] = None,
381 ) -> 'DTensor':
382 """Fast DTensor construction from a local tensor and a pre-built Layout.
384 Unlike :meth:`from_local`, this does NOT rebuild the layout via
385 ``_build_layout`` — it hands the already-built ``layout`` straight to
386 ``__init_data__``. Intended for hot paths (e.g. ``wrap_output``) where the
387 output Layout was already inferred and cached by ``infer_layout``, so
388 recomputing ``device_mesh.to_hash()`` / ``tuple(placements)`` / the layout
389 cache lookup on every output is pure waste.
391 ``layout.placements`` (a plain attribute) is passed only to satisfy the
392 constructor's non-None check; ``__init_data__`` ignores it when ``layout``
393 is supplied.
394 """
395 return DTensor(
396 local_tensor,
397 layout.mesh,
398 layout.placements,
399 layout,
400 shape=shape,
401 )
403 def _alias_placements(self) -> Sequence[Placement]:
404 """Return alias_placements from layout, falling back to _placements."""
405 if hasattr(self, '_layout') and self._layout:
406 return self._layout.alias_placements
407 return self._placements
409 def _from_converted_local(self, local_tensor: Tensor) -> 'DTensor':
410 """Rebuild converted DTensor data without preserving Parameter identity."""
411 cls = DTensor if isinstance(self, platform.Parameter) else self.__class__
412 kwargs = {
413 "device_mesh": self._device_mesh,
414 "placements": self._alias_placements(),
415 }
416 if hasattr(self, "_global_shape"):
417 kwargs["shape"] = self._global_shape
418 return cls(local_tensor, **kwargs)
420 def to(self, *args, **kwargs):
421 """Move the DTensor to a different device or dtype.
423 Delegates to the underlying local tensor's ``to`` method and
424 reconstructs a DTensor preserving device_mesh and placements.
426 Args:
427 *args (tuple): Arguments passed to the underlying tensor's ``to``
428 method (e.g., device or dtype).
429 **kwargs (dict): Keyword arguments for the tensor conversion
430 (e.g., dtype, device, non_blocking).
432 Returns:
433 DTensor: A new DTensor with the converted local tensor.
434 """
435 new_local = self._local_tensor.to(*args, **kwargs)
436 return self._from_converted_local(new_local)
438 def float(self):
439 """Convert the DTensor to float dtype.
441 Returns:
442 DTensor: A new DTensor with float32 local tensor.
443 """
444 new_local = self._local_tensor.float()
445 return self._from_converted_local(new_local)
447 def type_as(self, other: Tensor) -> "DTensor":
448 """Cast this DTensor to the dtype of ``other``.
450 This is a **local** operation — no communication. Each shard
451 independently casts its elements to the target dtype.
453 Only the **dtype** of ``other`` is read; its shape, values, and
454 layout are ignored. The returned DTensor preserves the
455 device-mesh and placements of ``self`` unchanged.
457 Args:
458 other (Tensor): A tensor whose ``.dtype`` will be used as the
459 target type. May be a plain :class:`Tensor` or a
460 :class:`DTensor`. Must reside on the same device as
461 ``self``.
463 Returns:
464 DTensor: A new DTensor with the converted local tensor. When
465 ``self.dtype == other.dtype`` the method returns ``self``
466 unchanged (no-op).
468 Raises:
469 ValueError: If ``other`` is not a Tensor.
470 ValueError: If ``self`` has Partial placement (cast does not
471 commute with reduction).
472 ValueError: If ``self`` and ``other`` are on different devices.
474 Note:
475 This implementation intentionally covers **dtype-only**
476 conversion. PyTorch's native ``type_as`` may also handle
477 cross-device transfers, but a DTensor cannot silently change
478 its backend device while retaining the old ``DeviceMesh``.
479 Use :meth:`to` for explicit device + dtype conversion.
481 Example:
482 >>> # x is a DTensor of float16, y is a plain float32 Tensor
483 >>> # on the same device.
484 >>> z = x.type_as(y)
485 >>> z.dtype == y.dtype
486 True
487 """
488 if not isinstance(other, Tensor):
489 raise ValueError(
490 f"type_as() argument must be a Tensor, but got "
491 f"{type(other).__name__}."
492 )
493 if hasattr(self, '_layout') and self._layout is not None:
494 if self._layout.is_partial():
495 raise ValueError(
496 "DTensor.type_as does not support Partial input; "
497 "call reduce_partial() first."
498 )
500 other_local = other.to_local() if isinstance(other, DTensor) else other
501 if self._local_tensor.device != other_local.device:
502 raise ValueError(
503 "DTensor.type_as requires self and other to be on the "
504 "same device. Use to() for explicit device + dtype "
505 "conversion."
506 )
508 target_dtype = other.dtype
509 if self.dtype == target_dtype:
510 return self
511 new_local = self._local_tensor.to(dtype=target_dtype)
512 return self._from_converted_local(new_local)
514 def _validate_factory_device(self, device: Any) -> None:
515 """Raise :class:`ValueError` if ``device`` does not match the DTensor's device."""
516 requested_type, requested_index = _device_spec(device)
517 local_type, local_index = _device_spec(self._local_tensor.device)
518 if (
519 requested_type != local_type
520 or (
521 requested_index is not None
522 and requested_index != local_index
523 )
524 ):
525 raise ValueError(
526 f"DTensor requires device to match the input DTensor "
527 f"device {self._local_tensor.device}, but got {device}."
528 )
530 def _new_const_tensor_op(
531 self,
532 method_name: str,
533 size: Union[int, Sequence[int]],
534 *,
535 dtype: Optional[Any] = None,
536 device: Optional[Any] = None,
537 requires_grad: bool = False,
538 layout: Optional[Any] = None,
539 pin_memory: bool = False,
540 ) -> 'DTensor':
541 """Create an all-``Replicate`` constant DTensor.
543 Shared implementation for ``new_zeros`` and ``new_ones``.
545 ``self`` is only used as a dtype/device reference and mesh source;
546 its values are ignored. The output is always **all-Replicate**
547 because every device produces identical data independently.
549 Args:
550 method_name:
551 ``"new_zeros"`` or ``"new_ones"`` — the local tensor
552 factory method to call.
553 size:
554 Output shape — an int or a sequence of ints.
555 dtype:
556 Desired dtype. Defaults to ``self.dtype`` on Torch.
557 device:
558 Must match ``self``'s device (Torch only).
559 requires_grad:
560 Forwarded on Torch; rejected on MindSpore.
561 layout:
562 Forwarded on Torch; rejected on MindSpore.
563 pin_memory:
564 Forwarded on Torch; rejected on MindSpore.
566 Returns:
567 DTensor: A new DTensor with all-``Replicate`` placements on
568 ``self``'s ``DeviceMesh``.
570 Raises:
571 ValueError: If a Torch-only kwarg is used on MindSpore, or
572 ``device`` does not match the DTensor's device.
573 """
574 if isinstance(size, int):
575 size = (size,)
577 if platform.platform_type == PlatformType.MINDSPORE:
578 if device is not None or layout is not None or requires_grad or pin_memory:
579 raise ValueError(
580 f"DTensor.{method_name} only supports size and dtype "
581 "on MindSpore."
582 )
583 local_kwargs = {}
584 if dtype is not None:
585 local_kwargs["dtype"] = dtype
586 else:
587 local_kwargs = {}
588 if dtype is not None:
589 local_kwargs["dtype"] = dtype
590 if device is not None:
591 self._validate_factory_device(device)
592 # An unindexed device such as "cuda" resolves to the framework's
593 # current device, which may differ from this DTensor's local device.
594 local_kwargs["device"] = self._local_tensor.device
595 if requires_grad:
596 local_kwargs["requires_grad"] = True
597 if layout is not None:
598 local_kwargs["layout"] = layout
599 if pin_memory:
600 local_kwargs["pin_memory"] = True
602 factory = getattr(self._local_tensor, method_name)
603 local_result = factory(size, **local_kwargs)
605 replicated_placements = [Replicate()] * self._device_mesh.ndim
606 return DTensor.from_local(
607 local_result, self._device_mesh, replicated_placements,
608 )
610 def new_zeros(
611 self,
612 size: Union[int, Sequence[int]],
613 *,
614 dtype: Optional[Any] = None,
615 device: Optional[Any] = None,
616 requires_grad: bool = False,
617 layout: Optional[Any] = None,
618 pin_memory: bool = False,
619 ) -> 'DTensor':
620 """Create an all-Replicate DTensor filled with zeros.
622 The output is always **fully replicated** across every device in
623 ``self``'s ``DeviceMesh``, regardless of how ``self`` is sharded.
625 Args:
626 size:
627 Output shape — an int or a sequence of ints.
628 dtype:
629 Desired dtype. Defaults to ``self.dtype`` (Torch).
630 Not forwarded to MindSpore unless explicitly set.
631 device:
632 Must match ``self``'s device. Not supported on MindSpore.
633 requires_grad:
634 Forwarded on Torch; rejected on MindSpore.
635 layout:
636 Forwarded on Torch; rejected on MindSpore.
637 pin_memory:
638 Forwarded on Torch; rejected on MindSpore.
640 Returns:
641 DTensor: A new all-Replicate DTensor filled with zeros.
642 """
643 return self._new_const_tensor_op(
644 "new_zeros", size,
645 dtype=dtype,
646 device=device,
647 requires_grad=requires_grad,
648 layout=layout,
649 pin_memory=pin_memory,
650 )
652 def new_ones(
653 self,
654 size: Union[int, Sequence[int]],
655 *,
656 dtype: Optional[Any] = None,
657 device: Optional[Any] = None,
658 requires_grad: bool = False,
659 layout: Optional[Any] = None,
660 pin_memory: bool = False,
661 ) -> 'DTensor':
662 """Create an all-Replicate DTensor filled with ones.
664 The output is always **fully replicated** across every device in
665 ``self``'s ``DeviceMesh``, regardless of how ``self`` is sharded.
667 Args:
668 size:
669 Output shape — an int or a sequence of ints.
670 dtype:
671 Desired dtype. Defaults to ``self.dtype`` (Torch).
672 Not forwarded to MindSpore unless explicitly set.
673 device:
674 Must match ``self``'s device. Not supported on MindSpore.
675 requires_grad:
676 Forwarded on Torch; rejected on MindSpore.
677 layout:
678 Forwarded on Torch; rejected on MindSpore.
679 pin_memory:
680 Forwarded on Torch; rejected on MindSpore.
682 Returns:
683 DTensor: A new all-Replicate DTensor filled with ones.
684 """
685 return self._new_const_tensor_op(
686 "new_ones", size,
687 dtype=dtype,
688 device=device,
689 requires_grad=requires_grad,
690 layout=layout,
691 pin_memory=pin_memory,
692 )
694 def to_local(self) -> Tensor:
695 """
696 Convert DTensor to local tensor.
698 Returns:
699 Tensor: The local tensor shard on this device.
700 """
701 return self._local_tensor
703 def tolist(self):
704 """
705 Convert the DTensor to a nested Python list or number.
707 This operation gathers the complete tensor on every participating rank
708 before converting it to Python values. It is an **implicit collective**:
709 all ranks in the DeviceMesh must participate.
711 Returns:
712 Union[list, int, float, bool]: A nested Python list, or a Python
713 number for a scalar DTensor.
715 Note:
716 This triggers ``full_tensor()`` under the hood, which performs
717 all-gather communication. For large tensors, prefer slicing or
718 index-based access to avoid materialising the full tensor.
720 If you only need the **local shard** as a list, use
721 ``dtensor.to_local().tolist()`` instead — that path has zero
722 communication overhead.
724 Example:
725 >>> mesh = init_device_mesh("npu", (2,), ("dp",))
726 >>> x = distribute_tensor(torch.arange(8).reshape(4, 2), mesh, [Shard(0)])
727 >>> x.tolist() # full data: [[0,1],[2,3],[4,5],[6,7]]
728 >>> x.to_local().tolist() # local shard only (no comm)
729 """
730 return self.full_tensor().tolist()
732 def copy_(self, src: "DTensor", non_blocking: bool = False) -> "DTensor":
733 """In-place copy of ``src`` into this DTensor's local shard.
735 Delegates to ``Tensor.copy_`` on the underlying local tensors.
736 Follows standard ``Tensor.copy_`` semantics: version counter is
737 bumped and autograd edges are created when grad is enabled.
739 Constraints on ``src``:
740 * must be a ``DTensor`` on the same or an equivalent ``DeviceMesh`` topology as ``self``;
741 * its placements must equal ``self.placements``, OR
742 ``src._local_tensor.numel() == 1`` (single-element broadcast);
743 * its local shape must equal or be broadcastable to
744 ``self._local_tensor.shape``.
746 No redistribute / implicit slicing is performed; src dtype is cast
747 to self dtype in-place.
749 Args:
750 src (DTensor): Source DTensor satisfying the constraints above.
751 non_blocking (bool): Forwarded to the underlying ``copy_``.
753 Returns:
754 DTensor: ``self``.
756 Raises:
757 TypeError: if ``src`` is not a ``DTensor``.
758 ValueError: if mesh, placement, or shape constraint is violated.
759 """
760 if not isinstance(src, DTensor):
761 raise TypeError(
762 f"For DTensor.copy_, src should be a DTensor, but got {type(src).__name__}."
763 )
764 src_local = src.to_local()
765 if not _device_meshes_are_compatible(src.device_mesh, self._device_mesh):
766 raise ValueError(
767 f"For DTensor.copy_, src and self should share the same DeviceMesh, "
768 f"but got src.device_mesh={src.device_mesh!r}, "
769 f"self._device_mesh={self._device_mesh!r}."
770 )
772 placement_eq = tuple(src.placements) == tuple(self._placements)
773 shape_eq = src_local.shape == self._local_tensor.shape
774 src_is_scalar = src_local.numel() == 1
776 if not placement_eq and not src_is_scalar:
777 raise ValueError(
778 f"For DTensor.copy_, src.placements should equal self.placements "
779 f"or src.numel() should be 1, but got "
780 f"src.placements={src.placements}, "
781 f"self.placements={self._placements}, "
782 f"src.numel()={src_local.numel()}."
783 )
784 if not shape_eq and not src_is_scalar and not _is_broadcastable(
785 src_local.shape, self._local_tensor.shape
786 ):
787 raise ValueError(
788 f"For DTensor.copy_, src local shape should be broadcastable to "
789 f"self local shape, but got "
790 f"src.shape={tuple(src_local.shape)}, "
791 f"self.shape={tuple(self._local_tensor.shape)}."
792 )
794 self._local_tensor.copy_(src_local, non_blocking=non_blocking)
795 return self
797 def zero_(self) -> "DTensor":
798 """In-place fill with zeros. Returns ``self``."""
799 self._local_tensor.zero_()
800 return self
802 def fill_(self, value) -> "DTensor":
803 """In-place fill with ``value``. Returns ``self``."""
804 self._local_tensor.fill_(value)
805 return self
807 @property
808 def shape(self) -> Tuple[int, ...]:
809 """
810 The global shape of this DTensor.
812 Returns:
813 Tuple[int, ...]: The global tensor shape.
814 """
815 return self._global_shape
817 def size(self, dim=None):
818 """Return the global shape, consistent with .shape.
820 Without ``dim`` returns a tuple matching ``self.shape``.
821 With ``dim`` returns the size of that dimension.
822 """
823 global_shape = self.shape
824 if dim is not None:
825 return global_shape[dim]
826 return global_shape
828 def numel(self) -> int:
829 """Return the number of elements in this DTensor."""
830 return int(np.prod(self.shape))
832 @property
833 def ndim(self) -> int:
834 """Return the logical global tensor rank."""
835 return len(self._global_shape)
837 def dim(self) -> int:
838 """Return the logical global tensor rank."""
839 return len(self._global_shape)
841 @property
842 def local_shape(self) -> Tuple[int, ...]:
843 """
844 The local shape of this DTensor on this device.
846 Returns:
847 Tuple[int, ...]: The local tensor shape.
848 """
849 return self._local_tensor.shape
851 def redistribute(
852 self,
853 device_mesh: DeviceMesh,
854 placements: Union[Sequence[Placement], Sequence[Union[str, Tuple[str, ...]]]]
855 ) -> 'DTensor':
856 """
857 Redistribute this DTensor to a new device mesh and placements.
859 Args:
860 device_mesh (DeviceMesh): The target device mesh.
861 placements: The target placements. Supports Placement objects
862 or alias strings.
864 Returns:
865 DTensor: A new DTensor with the specified distribution.
867 Example:
868 >>> new_dtensor = dtensor.redistribute(mesh, [Replicate(), Shard(1)])
869 >>> new_dtensor = dtensor.redistribute(mesh, ("None", "tp"))
870 """
871 logger.debug(
872 "redistribute: shape=%s, src_placements=%s -> dst_placements=%s, "
873 "mesh_shape=%s, local_shape=%s",
874 tuple(self.shape),
875 tuple(self._placements),
876 tuple(placements),
877 tuple(device_mesh.shape),
878 tuple(self._local_tensor.shape),
879 )
881 # Build dst_layout from device_mesh and placements
882 dst_layout = _build_layout(
883 device_mesh, placements, len(self._global_shape)
884 )
886 # pylint: disable=C0415
887 from hyper_parallel.core.dtensor.tensor_redistribution import _tensor_redistribution
888 out = _tensor_redistribution.redistribution(self, dst_layout)
889 return out
891 def reduce_partial(self) -> 'DTensor':
892 """
893 Reduce partial sharding state for this DTensor.
895 Returns:
896 DTensor: A new DTensor with partial state reduced.
897 """
898 if not self._layout:
899 return self
900 to_layout = cp.deepcopy(self._layout)
901 to_layout.reset_partial()
902 # pylint: disable=C0415
903 from hyper_parallel.core.dtensor.tensor_redistribution import _tensor_redistribution
904 out = _tensor_redistribution.reduce_partial(self, to_layout)
905 return out
907 def full_tensor(self) -> Tensor:
908 """
909 Return the full tensor of this DTensor.
911 Returns:
912 Tensor: A Tensor object that represents the full tensor of this DTensor.
913 The returned tensor contains the complete data gathered from
914 all ranks.
916 Note:
917 This operation involves communication across all ranks in the DeviceMesh,
918 which may be expensive for large tensors. Use with caution in
919 performance-critical code paths.
921 Example:
922 >>> # Assume dtensor is sharded across multiple devices
923 >>> local_tensor = dtensor.to_local() # Returns only the local shard
924 >>> full_tensor = dtensor.full_tensor() # Returns the complete tensor
925 """
926 if not self._layout:
927 return self._local_tensor
929 # Create a fully replicated layout
930 replicated_layout = cp.deepcopy(self._layout)
932 # Set all placements to Replicate and convert to tensor_map
933 replicated_placements = [Replicate()] * len(replicated_layout.mesh_shape)
934 replicated_layout.set_placements(replicated_placements)
935 replicated_layout.placement_to_tensor_map(len(self._global_shape))
937 # Clear partial status from original layout since Replicate has no partial
938 replicated_layout.reset_partial()
940 # Redistribute to the replicated layout and return local tensor
941 # pylint: disable=C0415
942 from hyper_parallel.core.dtensor.tensor_redistribution import _tensor_redistribution
943 out = _tensor_redistribution.redistribution(self, replicated_layout)
944 return out.to_local()
947def _normalize_shard_dim(dim: int, ndim: int) -> int:
948 return dim + ndim if dim < 0 else dim
951def _distribute_tensor_with_communication(
952 tensor: Tensor,
953 device_mesh: DeviceMesh,
954 placements: Sequence[Placement],
955 src_data_rank: int,
956) -> Tensor:
957 """Scatter/broadcast a logical global tensor along mesh dimensions (PyTorch parity)."""
958 local = tensor
959 if len(placements) < device_mesh.ndim:
960 raise ValueError(
961 f"placements length ({len(placements)}) must be at least device_mesh.ndim "
962 f"({device_mesh.ndim}) when src_data_rank is set"
963 )
964 for mesh_dim in range(device_mesh.ndim):
965 placement = placements[mesh_dim]
966 if isinstance(placement, StridedShard):
967 raise NotImplementedError(
968 "distribute_tensor with src_data_rank does not support StridedShard yet; "
969 "pass src_data_rank=None for local-only sharding."
970 )
971 if placement.is_shard():
972 shard_dim = _normalize_shard_dim(placement.dim, local.ndim)
973 num_chunks = device_mesh.size(mesh_dim)
974 if num_chunks <= 0:
975 raise ValueError(f"invalid mesh dim size {num_chunks} on mesh_dim={mesh_dim}")
976 chunks = tuple(local.chunk(num_chunks, dim=shard_dim))
977 if not chunks:
978 raise ValueError(f"cannot shard dim {shard_dim} into {num_chunks} chunks")
979 output = platform.empty_like(chunks[0])
980 local = mesh_scatter(output, chunks, device_mesh, mesh_dim, group_src=src_data_rank)
981 elif placement.is_replicate() or placement.is_partial():
982 local = mesh_broadcast(local, device_mesh, mesh_dim, group_src=src_data_rank)
983 if isinstance(placement, Partial):
984 warnings.warn(
985 f"Partial placement {placement} during distribute_tensor: "
986 "broadcast only; partial partition is not applied yet.",
987 stacklevel=3,
988 )
989 else:
990 raise RuntimeError(
991 f"unsupported placement {placement} on device mesh dimension {mesh_dim}"
992 )
993 return local
996def distribute_tensor(
997 tensor: Tensor,
998 device_mesh: DeviceMesh,
999 placements: Union[Sequence[Placement], Sequence[Union[str, Tuple[str, ...]]]],
1000 *,
1001 src_data_rank: Optional[int] = None,
1002) -> DTensor:
1003 """
1004 Distribute a global tensor to the device mesh according to the placements.
1006 Args:
1007 tensor (Tensor): The global tensor to be distributed. All ranks
1008 should have the same tensor data.
1009 device_mesh (DeviceMesh): The device mesh describing the device topology.
1010 placements: The placement strategy. Supports two styles:
1011 - Placement objects (e.g., ``[Shard(0), Replicate()]``).
1012 - Alias strings (e.g., ``("dp", "None")`` or
1013 ``(("dp", "tp"), "None")``), length must equal the number of
1014 tensor dimensions.
1016 Returns:
1017 DTensor: A new DTensor with the local shard on each rank.
1019 Note:
1020 When ``src_data_rank`` is an ``int`` (e.g. ``0``), shard/replicate
1021 placements use scatter/broadcast from the source rank on each mesh axis,
1022 matching PyTorch ``distribute_tensor``. When ``src_data_rank=None``
1023 (default), each rank slices its local tensor without communication
1024 (legacy Hyper behavior; all ranks must hold the same global data).
1026 Example:
1027 >>> mesh = init_device_mesh(device_type="npu", mesh_shape=(2, 2), mesh_dim_names=("dp", "tp"))
1028 >>> global_tensor = Tensor(np.arange(16).reshape(4, 4))
1029 >>> dtensor = distribute_tensor(global_tensor, mesh, [Shard(0), Replicate()])
1030 >>> dtensor = distribute_tensor(global_tensor, mesh, ("dp", "None"))
1031 """
1032 layout = _build_layout(device_mesh, placements, len(tensor.shape))
1033 if _layout_has_ragged_shard(layout):
1034 if src_data_rank is None:
1035 local_tensor = _slice_ragged_tensor(tensor, layout)
1036 else:
1037 local_tensor = _scatter_ragged_tensor(tensor, layout, src_data_rank)
1038 elif src_data_rank is None:
1039 local_tensor = _get_slice_tensor_by_layout(tensor, layout)
1040 else:
1041 local_tensor = _distribute_tensor_with_communication(
1042 tensor, device_mesh, layout.placements, src_data_rank
1043 )
1044 return DTensor.from_local_with_layout(
1045 local_tensor,
1046 layout,
1047 shape=tuple(tensor.shape),
1048 )
1051def _distribute_module_param_source(param: Any) -> Tensor:
1052 """Tensor data used as the global tensor for :func:`distribute_tensor` (PyTorch uses ``param.data``)."""
1053 if hasattr(param, "data"):
1054 return param.data
1055 return platform.get_param_local_data(param)
1058def _distribute_module_new_parameter(key: str, dtensor: DTensor, requires_grad: bool) -> Any:
1059 """Build a framework :class:`Parameter` holding *dtensor* (Torch vs MindSpore kwargs differ)."""
1060 if platform.platform_type == PlatformType.MINDSPORE:
1061 return platform.Parameter(dtensor, name=key, requires_grad=requires_grad)
1062 return platform.Parameter(dtensor, requires_grad=requires_grad)
1065def _distribute_module_set_param(module: Any, key: str, new_param: Any) -> None:
1066 """Register or assign a parameter on *module* (``nn.Module`` or MindSpore ``Cell``)."""
1067 if hasattr(module, "register_parameter"):
1068 module.register_parameter(key, new_param)
1069 return
1070 if hasattr(module, "_params"):
1071 module._params[key] = new_param
1072 if hasattr(module, "_params_list"):
1073 module._params_list[key] = new_param
1074 if key in module.__dict__:
1075 module.__dict__[key] = new_param
1076 return
1077 raise TypeError(
1078 f"distribute_module expects nn.Module-like objects with register_parameter or _params; "
1079 f"got {type(module)}."
1080 )
1083def _distribute_module_iter_params(module: Any) -> list:
1084 """Return ``[(name, param), ...]`` for direct parameters (``_parameters`` or ``_params``)."""
1085 if hasattr(module, "_parameters"):
1086 return list(module._parameters.items())
1087 if hasattr(module, "_params"):
1088 return list(module._params.items())
1089 return []
1092def _distribute_module_iter_buffers(module: Any) -> list:
1093 """Return ``[(name, buffer), ...]`` if the module has ``_buffers`` (PyTorch ``nn.Module``)."""
1094 if hasattr(module, "_buffers"):
1095 return list(module._buffers.items())
1096 return []
1099def _distribute_module_named_modules(module: Any):
1100 """``nn.Module.named_modules`` or MindSpore ``Cell.cells_and_names`` (submodule FQNs)."""
1101 if hasattr(module, "named_modules"):
1102 return module.named_modules()
1103 if hasattr(module, "cells_and_names"):
1104 return module.cells_and_names()
1105 raise TypeError(
1106 f"distribute_module expects module-like objects with named_modules or cells_and_names; "
1107 f"got {type(module)}."
1108 )
1111def _distribute_module_named_parameters(module: Any):
1112 """``nn.Module.named_parameters(recurse=False)`` or MindSpore ``Cell.parameters_and_names(expand=False)``."""
1113 if hasattr(module, "named_parameters"):
1114 return module.named_parameters(recurse=False)
1115 if hasattr(module, "parameters_and_names"):
1116 return module.parameters_and_names(expand=False)
1117 raise TypeError(
1118 f"distribute_module expects module-like objects with named_parameters or parameters_and_names; "
1119 f"got {type(module)}."
1120 )
1123def _replicate_submodule_params_buffers(
1124 sub_mod: Any,
1125 device_mesh: DeviceMesh,
1126 *,
1127 module_prefix: str = "",
1128) -> None:
1129 """Convert plain params/buffers on *sub_mod* to fully replicated :class:`DTensor`."""
1130 full_replicate = [Replicate()] * device_mesh.ndim
1131 for key, param in _distribute_module_iter_params(sub_mod):
1132 if param is None or isinstance(param, DTensorBase):
1133 continue
1134 src = _distribute_module_param_source(param)
1135 requires_grad = bool(getattr(param, "requires_grad", True))
1136 dt = distribute_tensor(src, device_mesh, full_replicate)
1137 param_name = f"{module_prefix}.{key}" if module_prefix else key
1138 new_param = _distribute_module_new_parameter(param_name, dt, requires_grad)
1139 _distribute_module_set_param(sub_mod, key, new_param)
1140 for key, buffer in _distribute_module_iter_buffers(sub_mod):
1141 if buffer is None or isinstance(buffer, DTensorBase):
1142 continue
1143 sub_mod._buffers[key] = distribute_tensor(buffer, device_mesh, full_replicate)
1146def _distribute_module_run_partition_and_replicate(
1147 module: Any,
1148 device_mesh: DeviceMesh,
1149 partition_fn: Optional[Callable[[str, Any, DeviceMesh], None]],
1150) -> None:
1151 """Call optional ``partition_fn`` per ``named_modules`` and replicate remaining tensors."""
1152 if partition_fn is None:
1153 for mod_name, submod in _distribute_module_named_modules(module):
1154 _replicate_submodule_params_buffers(submod, device_mesh, module_prefix=mod_name)
1155 return
1156 for mod_name, submod in _distribute_module_named_modules(module):
1157 partition_fn(mod_name, submod, device_mesh)
1158 _replicate_submodule_params_buffers(submod, device_mesh, module_prefix=mod_name)
1161def _distribute_module_register_input_fn(
1162 module: Any,
1163 device_mesh: DeviceMesh,
1164 input_fn: Callable[..., Any],
1165) -> None:
1166 """Register *input_fn* as a forward pre-hook on *module* (2- or 3-arg, PyTorch-compatible)."""
1167 num_args = len(inspect.signature(input_fn).parameters)
1168 if num_args == 2:
1169 warnings.warn(
1170 "Deprecating input_fn that takes two arguments (inputs, device_mesh), "
1171 "please use input_fn that takes in (module, inputs, device_mesh) instead!",
1172 FutureWarning,
1173 stacklevel=3,
1174 )
1175 module.register_forward_pre_hook(
1176 lambda _, inputs: input_fn(inputs, device_mesh)
1177 )
1178 elif num_args == 3:
1179 module.register_forward_pre_hook(
1180 lambda mod, inputs: input_fn(mod, inputs, device_mesh)
1181 )
1182 else:
1183 raise ValueError(
1184 f"input_fn should take in 2 or 3 arguments, but got {num_args} arguments!"
1185 )
1188def _distribute_module_register_output_fn(
1189 module: Any,
1190 device_mesh: DeviceMesh,
1191 output_fn: Callable[..., Any],
1192) -> None:
1193 """Register *output_fn* as a forward hook on *module* (2- or 3-arg, PyTorch-compatible)."""
1194 num_args = len(inspect.signature(output_fn).parameters)
1195 if num_args == 2:
1196 warnings.warn(
1197 "Deprecating output_fn that takes two arguments (outputs, device_mesh), "
1198 "please use output_fn that takes in (module, outputs, device_mesh) instead!",
1199 FutureWarning,
1200 stacklevel=3,
1201 )
1202 module.register_forward_hook(
1203 lambda mod, inputs, outputs: output_fn(outputs, device_mesh)
1204 )
1205 elif num_args == 3:
1206 module.register_forward_hook(
1207 lambda mod, inputs, outputs: output_fn(mod, outputs, device_mesh)
1208 )
1209 else:
1210 raise ValueError(
1211 f"output_fn should take in 2 or 3 arguments, but got {num_args} arguments!"
1212 )
1215def distribute_module(
1216 module: Any,
1217 device_mesh: Optional[DeviceMesh] = None,
1218 partition_fn: Optional[Callable[[str, Any, DeviceMesh], None]] = None,
1219 input_fn: Optional[Callable[..., Any]] = None,
1220 output_fn: Optional[Callable[..., Any]] = None,
1221) -> Any:
1222 """PyTorch ``distribute_module`` parity: shard/replicate params and optional I/O hooks.
1224 Unsharded parameters and buffers become fully replicated :class:`DTensor` after
1225 ``partition_fn``. ``input_fn`` / ``output_fn`` attach only to the root *module*.
1227 Args:
1228 module: Root ``nn.Module`` or MindSpore ``Cell`` with compatible APIs.
1229 device_mesh: Placement mesh; if ``None``, uses ``_mesh_resources.get_current_mesh()``.
1230 partition_fn: Per ``named_modules`` callback before replicate pass; ``None`` replicates all.
1231 input_fn: ``(module, inputs, mesh)`` or deprecated ``(inputs, mesh)`` pre-hook.
1232 output_fn: ``(module, outputs, mesh)`` or deprecated ``(outputs, mesh)`` forward hook.
1234 Returns:
1235 *module* in place, with distributed tensors where applied.
1237 Raises:
1238 RuntimeError: If called twice on the same *module*.
1239 ValueError: If ``input_fn`` / ``output_fn`` arity is not 2 or 3.
1241 Note:
1242 XLA / ``torch_xla`` is not supported; strided device :class:`DTensor` only.
1243 """
1244 if getattr(module, "_distribute_module_applied", False):
1245 raise RuntimeError(
1246 "distribute_module should only be called once on a module, "
1247 "but it has already been called on this module!"
1248 )
1249 device_mesh = device_mesh or _mesh_resources.get_current_mesh()
1250 _distribute_module_run_partition_and_replicate(module, device_mesh, partition_fn)
1251 if input_fn is not None:
1252 _distribute_module_register_input_fn(module, device_mesh, input_fn)
1253 if output_fn is not None:
1254 _distribute_module_register_output_fn(module, device_mesh, output_fn)
1255 module._distribute_module_applied = True
1256 return module
1259def _dtensor_init_helper(
1260 init_op,
1261 size,
1262 device_mesh,
1263 placements,
1264 *,
1265 rng_tracked: bool = False,
1266 **kwargs,
1267) -> DTensor:
1268 """
1269 Helper function to create and initialize a distributed tensor.
1271 Args:
1272 size: Shape of the tensor.
1273 dtype: Data type of the tensor.
1274 device: Target device for the tensor.
1275 requires_grad: Whether the tensor requires gradient.
1276 rng_tracked: When ``True``, initialize via :class:`OffsetBasedRNGTracker`
1277 so shard/replicate random semantics match PyTorch DTensor factories.
1279 Returns:
1280 DTensor: The initialized distributed tensor.
1281 """
1282 global_shape = (size,) if isinstance(size, int) else tuple(size)
1283 layout = _build_layout(device_mesh, placements, len(global_shape))
1284 if _layout_has_ragged_shard(layout):
1285 raise NotImplementedError(
1286 "RaggedShard tensor factories are not implemented in the DTensor metadata phase"
1287 )
1289 # get local tensor shape
1290 local_shape = compute_local_shape_and_global_offset(
1291 size, device_mesh, placements
1292 )
1294 # initialize the local tensor
1295 if init_op is platform.full:
1296 fill_value = kwargs.pop("fill_value", 0)
1297 local_tensor = init_op(local_shape, fill_value, **kwargs)
1298 elif rng_tracked:
1299 # pylint: disable=C0415
1300 from hyper_parallel.core.dtensor.random import is_rng_supported_mesh, OffsetBasedRNGTracker
1301 from hyper_parallel.core.shard._op_dispatch import _OP_DISPATCHER
1303 layout = _build_layout(device_mesh, placements, len(local_shape))
1304 if is_rng_supported_mesh(device_mesh):
1305 if _OP_DISPATCHER._rng_tracker is None:
1306 _OP_DISPATCHER._rng_tracker = OffsetBasedRNGTracker(run_state_sync=False)
1307 with _OP_DISPATCHER._rng_tracker._distribute_region(
1308 device_mesh,
1309 layout.placements,
1310 global_shape,
1311 ):
1312 local_tensor = init_op(local_shape, **kwargs)
1313 else:
1314 local_tensor = init_op(local_shape, **kwargs)
1315 else:
1316 local_tensor = init_op(local_shape, **kwargs)
1318 return DTensor.from_local(
1319 local_tensor,
1320 device_mesh,
1321 placements,
1322 )
1325def ones(
1326 size,
1327 device_mesh,
1328 placements,
1329) -> DTensor:
1330 """
1331 Returns a :class:`DTensor` filled with the scalar value 1, with the shape defined
1332 by the variable argument ``size``.
1334 Args:
1335 size (Union[tuple[int], list[int], int, Tensor]): The specified shape of output tensor. Only positive integer or
1336 tuple or Tensor containing positive integers are allowed. If it is a Tensor,
1337 it must be a 0-D or 1-D Tensor with int32 or int64 dtypes.
1339 Keyword args:
1340 device_mesh: :class:`DeviceMesh` type, contains the mesh info of ranks
1341 placements: a sequence of :class:`Placement` type: ``Shard``, ``Replicate``
1343 Returns:
1344 A :class:`DTensor` object on each rank
1345 """
1346 ones_ = platform.ones
1347 return _dtensor_init_helper(
1348 ones_,
1349 size,
1350 device_mesh=device_mesh,
1351 placements=placements,
1352 )
1355def empty(
1356 size,
1357 device_mesh,
1358 placements,
1359) -> DTensor:
1360 """
1361 Returns a :class:`DTensor` filled with uninitialized data. The shape of the :class:`DTensor`
1362 is defined by the variable argument ``size``.
1364 Args:
1365 size (Union[tuple[int], list[int], int]): The specified shape of output tensor. Can be variable numbers of
1366 positive integers or tuple or list containing positive integers.
1368 Keyword args:
1369 device_mesh: :class:`DeviceMesh` type, contains the mesh info of ranks
1370 placements: a sequence of :class:`Placement` type: ``Shard``, ``Replicate``
1372 Returns:
1373 A :class:`DTensor` object on each rank
1374 """
1375 empty_ = platform.empty
1376 return _dtensor_init_helper(
1377 empty_,
1378 size,
1379 device_mesh=device_mesh,
1380 placements=placements,
1381 )
1384def full(
1385 size,
1386 fill_value,
1387 *,
1388 device_mesh,
1389 placements,
1390) -> DTensor:
1391 """
1392 Returns a :class:`DTensor` filled with ``fill_value`` according to ``device_mesh`` and
1393 ``placements``, with the shape defined by the argument ``size``.
1395 Args:
1396 size (Union[tuple[int], list[int]]): The specified shape of output tensor.
1397 fill_value (Union[numbers.Number, Tensor]): Value to fill the returned tensor. It can be a scalar number, a 0-D
1398 Tensor, or a 1-D Tensor with only one element.
1400 Keyword args:
1401 device_mesh: :class:`DeviceMesh` type, contains the mesh info of ranks.
1402 placements: a sequence of :class:`Placement` type: ``Shard``, ``Replicate``
1404 Returns:
1405 A :class:`DTensor` object on each rank
1406 """
1407 full_ = platform.full
1408 return _dtensor_init_helper(
1409 full_,
1410 size,
1411 fill_value=fill_value,
1412 device_mesh=device_mesh,
1413 placements=placements,
1414 )
1417def zeros(
1418 size,
1419 device_mesh,
1420 placements,
1421) -> DTensor:
1422 """
1423 Returns a :class:`DTensor` filled with the scalar value 0.
1425 Args:
1426 size (Union[tuple[int], list[int], int, Tensor]): The specified shape of output tensor. Only positive integer or
1427 tuple or Tensor containing positive integers are allowed. If it is a Tensor,
1428 it must be a 0-D or 1-D Tensor with int32 or int64 dtypes.
1429 Keyword args:
1430 device_mesh: :class:`DeviceMesh` type, contains the mesh info of ranks
1431 placements: a sequence of :class:`Placement` type: ``Shard``, ``Replicate``
1433 Returns:
1434 A :class:`DTensor` object on each rank
1435 """
1436 zeros_ = platform.zeros
1437 return _dtensor_init_helper(
1438 zeros_,
1439 size,
1440 device_mesh=device_mesh,
1441 placements=placements,
1442 )
1445def rand(
1446 size,
1447 device_mesh,
1448 placements,
1449 **kwargs,
1450) -> DTensor:
1451 """
1452 Returns a :class:`DTensor` filled with random numbers from a uniform
1453 distribution on ``[0, 1)``.
1455 Args:
1456 size: Global output shape.
1457 device_mesh: :class:`DeviceMesh` for the distributed layout.
1458 placements: Per-mesh-dimension :class:`Placement` values.
1459 **kwargs: Forwarded to the platform ``rand`` call (for example ``dtype``).
1461 Returns:
1462 A :class:`DTensor` object on each rank.
1463 """
1464 return _dtensor_init_helper(
1465 platform.rand,
1466 size,
1467 device_mesh=device_mesh,
1468 placements=placements,
1469 rng_tracked=True,
1470 **kwargs,
1471 )
1474def randn(
1475 size,
1476 device_mesh,
1477 placements,
1478 **kwargs,
1479) -> DTensor:
1480 """
1481 Returns a :class:`DTensor` filled with random numbers from a normal
1482 distribution with mean ``0`` and variance ``1``.
1484 Args:
1485 size: Global output shape.
1486 device_mesh: :class:`DeviceMesh` for the distributed layout.
1487 placements: Per-mesh-dimension :class:`Placement` values.
1488 **kwargs: Forwarded to the platform ``randn`` call (for example ``dtype``).
1490 Returns:
1491 A :class:`DTensor` object on each rank.
1492 """
1493 return _dtensor_init_helper(
1494 platform.randn,
1495 size,
1496 device_mesh=device_mesh,
1497 placements=placements,
1498 rng_tracked=True,
1499 **kwargs,
1500 )