Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / dtensor / dtensor.py: 88%
398 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-04 05:18 +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.device_mesh import _mesh_resources
25from hyper_parallel.core.dtensor._collective_utils import mesh_broadcast, mesh_scatter
26from hyper_parallel.core.dtensor.layout import Layout, DeviceMesh, _get_slice_tensor_by_layout
27from hyper_parallel.core.dtensor.placement_types import Partial, Placement, Replicate, StridedShard
28from hyper_parallel.platform import get_platform
29from hyper_parallel.platform.platform import PlatformType
30from hyper_parallel.core.utils import compute_local_shape_and_global_offset
32platform = get_platform()
33DTensorBase = platform.DTensorBase
34Tensor = platform.Tensor
36logger = logging.getLogger(__name__)
39class SkipDTensorDispatch():
40 """Context manager that disables DTensor op dispatch for the enclosed block.
42 Args:
43 no_skip: Optional set of op callables or canonical op name strings that
44 should still be dispatched through DTensor even within this context.
45 All other ops bypass DTensor dispatch and operate on local tensors.
47 Example:
48 >>> import torch
49 >>> with SkipDTensorDispatch(no_skip={torch.zeros_like}):
50 ... # zeros_like still goes through DTensor dispatch;
51 ... # everything else uses the local tensor path.
52 ... result = torch.zeros_like(dtensor)
53 """
55 def __init__(self, no_skip: Optional[Set] = None):
56 self._no_skip_names: frozenset = frozenset()
57 if no_skip:
58 names = set()
59 for op in no_skip:
60 if isinstance(op, str):
61 names.add(op)
62 else:
63 names.add(platform.get_op_name(op))
64 self._no_skip_names = frozenset(names)
65 self._dispatch_token = None
66 self._ops_token = None
68 def __enter__(self):
69 # pylint: disable=C0415
70 from hyper_parallel.core.shard._op_dispatch import _dtensor_dispatch_disabled, _no_skip_ops
71 self._dispatch_token = _dtensor_dispatch_disabled.set(True)
72 if self._no_skip_names:
73 self._ops_token = _no_skip_ops.set(_no_skip_ops.get() | self._no_skip_names)
75 def __exit__(self, exc_type, exc_val, exc_tb):
76 # pylint: disable=C0415
77 from hyper_parallel.core.shard._op_dispatch import _dtensor_dispatch_disabled, _no_skip_ops
78 if self._ops_token is not None:
79 _no_skip_ops.reset(self._ops_token)
80 self._ops_token = None
81 _dtensor_dispatch_disabled.reset(self._dispatch_token)
82 self._dispatch_token = None
85# Cache for _build_layout to avoid redundant Layout computations
86# Key: (device_mesh.to_hash(), tuple(placements), tensor_dim)
87# Value: Layout
88_LAYOUT_CACHE = {}
91def _is_alias_placements(placements) -> bool:
92 """
93 Check if placements use alias strings rather than Placement objects.
95 Alias placements use mesh dimension names (strings) to specify
96 the sharding strategy, e.g., ("dp", "tp") or (("dp", "tp"), "None").
97 All elements must be strings or tuples of strings for the sequence
98 to be recognized as alias-style.
100 Args:
101 placements: A sequence of placement specifications.
103 Returns:
104 bool: True if all elements are alias strings or tuples of strings.
105 """
106 if len(placements) == 0:
107 return False
108 for p in placements:
109 if isinstance(p, str):
110 continue
111 if isinstance(p, tuple) and len(p) > 0 and all(isinstance(x, str) for x in p):
112 continue
113 return False
114 return True
117def _build_layout(
118 device_mesh: DeviceMesh,
119 placements: Union[Sequence[Placement], Sequence[Union[str, Tuple[str, ...]]]],
120 tensor_dim: int
121) -> Layout:
122 """
123 Build Layout from device_mesh and placements.
125 This function uses a cache to avoid redundant Layout computations
126 for the same (device_mesh, placements, tensor_dim) combination.
128 Args:
129 device_mesh: The device mesh describing the device topology.
130 placements: Supports two styles:
131 - Placement objects (Shard, Replicate, etc.)
132 - Alias strings ("dp", "None", ("dp", "tp"), etc.), length must
133 equal the number of tensor dimensions (``tensor_dim``).
134 tensor_dim: Number of dimensions in the tensor.
136 Returns:
137 Layout: The built layout object.
139 Raises:
140 ValueError: If alias placements length does not match tensor dimensions.
141 """
142 mesh_key = device_mesh.to_hash()
143 placements_key = tuple(placements)
144 cache_key = (mesh_key, placements_key, tensor_dim)
146 if cache_key in _LAYOUT_CACHE:
147 return _LAYOUT_CACHE[cache_key]
149 layout = Layout.from_device_mesh(device_mesh)
151 if _is_alias_placements(placements):
152 if len(placements) != tensor_dim:
153 raise ValueError(
154 f"Alias placements length ({len(placements)}) must equal "
155 f"tensor dimensions ({tensor_dim})."
156 )
157 result = layout(*placements)
158 else:
159 result = layout(placements)
160 result.placement_to_tensor_map(tensor_dim)
162 _LAYOUT_CACHE[cache_key] = result
164 return result
167def _is_broadcastable(src_shape: Sequence[int], dst_shape: Sequence[int]) -> bool:
168 """Return True iff ``src_shape`` is broadcastable to ``dst_shape``.
170 Standard NumPy / PyTorch right-aligned broadcast rule: ``src`` cannot
171 have more dimensions than ``dst``; each right-aligned dimension pair
172 must be equal, or ``src``'s dimension must be 1.
173 """
174 src_shape = tuple(src_shape)
175 dst_shape = tuple(dst_shape)
176 if len(src_shape) > len(dst_shape):
177 return False
178 for i in range(1, len(src_shape) + 1):
179 s, d = src_shape[-i], dst_shape[-i]
180 if s not in (d, 1):
181 return False
182 return True
185def _device_spec(device: Any) -> Tuple[str, Optional[int]]:
186 """Return normalised ``(device_type, device_index)``.
188 Handles device objects and strings such as ``"npu:0"``.
189 """
190 device_type = getattr(device, "type", None)
191 # Only read .index from objects that also have a .type attribute
192 # (i.e. torch.device). Avoids capturing str.index on plain strings.
193 device_index = (
194 getattr(device, "index", None) if device_type is not None else None
195 )
196 device_text = str(device).lower()
198 if device_type is None:
199 parts = device_text.split(":", maxsplit=1)
200 device_type = parts[0]
201 if len(parts) == 2 and parts[1].isdigit():
202 device_index = int(parts[1])
204 return str(device_type).lower(), device_index
207class DTensor(DTensorBase):
208 """
209 DTensor - Distributed Tensor
211 A DTensor represents a tensor that is distributed across multiple devices
212 according to a DeviceMesh and placement specifications.
214 Args:
215 local_tensor (Tensor): The local tensor shard on this device.
216 device_mesh (DeviceMesh): The device mesh describing the device topology.
217 placements: The placement strategy. Supports two styles:
218 - Placement objects (e.g., ``[Shard(0), Replicate()]``).
219 - Alias strings (e.g., ``("dp", "None")`` or
220 ``(("dp", "tp"), "None")``), length must equal the number of
221 tensor dimensions.
223 Example:
224 >>> mesh = init_device_mesh(device_type="npu", mesh_shape=(2, 2), mesh_dim_names=("dp", "tp"))
225 >>> local_tensor = Tensor(np.ones((4, 4)))
226 >>> # Placement style
227 >>> dtensor = DTensor.from_local(local_tensor, mesh, [Shard(0), Replicate()])
228 >>> # Alias style — length matches tensor dims
229 >>> dtensor = DTensor.from_local(local_tensor, mesh, ("dp", "None"))
230 """
231 _local_tensor: Tensor
232 _device_mesh: DeviceMesh
233 _placements: Sequence[Placement]
235 def __init_data__(
236 self,
237 local_tensor: Tensor,
238 device_mesh: DeviceMesh,
239 placements: Union[Sequence[Placement], Sequence[Union[str, Tuple[str, ...]]]],
240 layout: Optional[Layout] = None,
241 ):
242 self._local_tensor = local_tensor
243 self._device_mesh = device_mesh
244 # Fast path: when an already-built Layout is supplied (e.g. output layouts
245 # cached by infer_layout and passed straight through wrap_output), reuse it
246 # directly and skip _build_layout (which otherwise recomputes device_mesh.to_hash(),
247 # tuple(placements) and a cache lookup on every single output construction).
248 self._layout = layout if layout is not None else _build_layout(
249 device_mesh, placements, len(local_tensor.shape)
250 )
251 self._placements = tuple(self._layout.placements)
253 @property
254 def device_mesh(self) -> DeviceMesh:
255 """The device mesh of this DTensor."""
256 return self._device_mesh
258 @property
259 def placements(self) -> Sequence[Placement]:
260 """The placements of this DTensor."""
261 return self._placements
263 @property
264 def layout(self) -> Layout:
265 """Internal layout for redistribution (for backward compatibility)."""
266 if not hasattr(self, '_layout'):
267 return None
268 return self._layout
270 @staticmethod
271 def from_local(
272 local_tensor: Tensor,
273 device_mesh: DeviceMesh,
274 placements: Union[Sequence[Placement], Sequence[Union[str, Tuple[str, ...]]]],
275 *,
276 run_check: bool = False,
277 shape: Optional[Tuple[int, ...]] = None,
278 stride: Optional[Tuple[int, ...]] = None,
279 ) -> 'DTensor':
280 """
281 Create a DTensor from a local tensor with device mesh and placements.
283 Args:
284 local_tensor (Tensor): The local tensor shard on this device.
285 device_mesh (DeviceMesh): The device mesh describing the device topology.
286 placements: The placement strategy. Supports two styles:
287 - Placement objects (e.g., ``[Shard(0), Replicate()]``).
288 - Alias strings (e.g., ``("dp", "None")`` or
289 ``(("dp", "tp"), "None")``), length must equal the number
290 of tensor dimensions.
291 run_check (bool, optional): When ``True``, perform cross-rank metadata
292 checks and broadcast replicate placements from the mesh source rank.
293 Default: ``False``.
294 shape (tuple[int, ...], optional): Global DTensor shape hint for uneven
295 sharding when ``run_check=True``. Reserved for future use.
296 stride (tuple[int, ...], optional): Global stride hint. Reserved for
297 future use together with ``shape``.
299 Returns:
300 DTensor: A new DTensor instance.
302 Example:
303 >>> mesh = init_device_mesh(device_type="npu", mesh_shape=(2, 2), mesh_dim_names=("dp", "tp"))
304 >>> local_tensor = Tensor(np.ones((4, 4)))
305 >>> dtensor = DTensor.from_local(local_tensor, mesh, [Shard(0), Replicate()])
306 >>> dtensor = DTensor.from_local(local_tensor, mesh, ("dp", "None"))
307 """
308 if run_check:
309 # pylint: disable=C0415
310 from hyper_parallel.core.dtensor._from_local_utils import run_from_local_checks
311 layout = _build_layout(device_mesh, placements, len(local_tensor.shape))
312 run_from_local_checks(
313 local_tensor,
314 device_mesh,
315 layout.placements,
316 shape=shape,
317 stride=stride,
318 )
319 return DTensor(local_tensor, device_mesh, placements)
321 @staticmethod
322 def from_local_with_layout(local_tensor: Tensor, layout: Layout) -> 'DTensor':
323 """Fast DTensor construction from a local tensor and a pre-built Layout.
325 Unlike :meth:`from_local`, this does NOT rebuild the layout via
326 ``_build_layout`` — it hands the already-built ``layout`` straight to
327 ``__init_data__``. Intended for hot paths (e.g. ``wrap_output``) where the
328 output Layout was already inferred and cached by ``infer_layout``, so
329 recomputing ``device_mesh.to_hash()`` / ``tuple(placements)`` / the layout
330 cache lookup on every output is pure waste.
332 ``layout.placements`` (a plain attribute) is passed only to satisfy the
333 constructor's non-None check; ``__init_data__`` ignores it when ``layout``
334 is supplied.
335 """
336 return DTensor(local_tensor, layout.mesh, layout.placements, layout)
338 def _alias_placements(self) -> Sequence[Placement]:
339 """Return alias_placements from layout, falling back to _placements."""
340 if hasattr(self, '_layout') and self._layout:
341 return self._layout.alias_placements
342 return self._placements
344 def _from_converted_local(self, local_tensor: Tensor) -> 'DTensor':
345 """Rebuild converted DTensor data without preserving Parameter identity."""
346 cls = DTensor if isinstance(self, platform.Parameter) else self.__class__
347 return cls(local_tensor, device_mesh=self._device_mesh,
348 placements=self._alias_placements())
350 def to(self, *args, **kwargs):
351 """Move the DTensor to a different device or dtype.
353 Delegates to the underlying local tensor's ``to`` method and
354 reconstructs a DTensor preserving device_mesh and placements.
356 Args:
357 *args (tuple): Arguments passed to the underlying tensor's ``to``
358 method (e.g., device or dtype).
359 **kwargs (dict): Keyword arguments for the tensor conversion
360 (e.g., dtype, device, non_blocking).
362 Returns:
363 DTensor: A new DTensor with the converted local tensor.
364 """
365 new_local = self._local_tensor.to(*args, **kwargs)
366 return self._from_converted_local(new_local)
368 def float(self):
369 """Convert the DTensor to float dtype.
371 Returns:
372 DTensor: A new DTensor with float32 local tensor.
373 """
374 new_local = self._local_tensor.float()
375 return self._from_converted_local(new_local)
377 def type_as(self, other: Tensor) -> "DTensor":
378 """Cast this DTensor to the dtype of ``other``.
380 This is a **local** operation — no communication. Each shard
381 independently casts its elements to the target dtype.
383 Only the **dtype** of ``other`` is read; its shape, values, and
384 layout are ignored. The returned DTensor preserves the
385 device-mesh and placements of ``self`` unchanged.
387 Args:
388 other (Tensor): A tensor whose ``.dtype`` will be used as the
389 target type. May be a plain :class:`Tensor` or a
390 :class:`DTensor`. Must reside on the same device as
391 ``self``.
393 Returns:
394 DTensor: A new DTensor with the converted local tensor. When
395 ``self.dtype == other.dtype`` the method returns ``self``
396 unchanged (no-op).
398 Raises:
399 ValueError: If ``other`` is not a Tensor.
400 ValueError: If ``self`` has Partial placement (cast does not
401 commute with reduction).
402 ValueError: If ``self`` and ``other`` are on different devices.
404 Note:
405 This implementation intentionally covers **dtype-only**
406 conversion. PyTorch's native ``type_as`` may also handle
407 cross-device transfers, but a DTensor cannot silently change
408 its backend device while retaining the old ``DeviceMesh``.
409 Use :meth:`to` for explicit device + dtype conversion.
411 Example:
412 >>> # x is a DTensor of float16, y is a plain float32 Tensor
413 >>> # on the same device.
414 >>> z = x.type_as(y)
415 >>> z.dtype == y.dtype
416 True
417 """
418 if not isinstance(other, Tensor):
419 raise ValueError(
420 f"type_as() argument must be a Tensor, but got "
421 f"{type(other).__name__}."
422 )
423 if hasattr(self, '_layout') and self._layout is not None:
424 if self._layout.is_partial():
425 raise ValueError(
426 "DTensor.type_as does not support Partial input; "
427 "call reduce_partial() first."
428 )
430 other_local = other.to_local() if isinstance(other, DTensor) else other
431 if self._local_tensor.device != other_local.device:
432 raise ValueError(
433 "DTensor.type_as requires self and other to be on the "
434 "same device. Use to() for explicit device + dtype "
435 "conversion."
436 )
438 target_dtype = other.dtype
439 if self.dtype == target_dtype:
440 return self
441 new_local = self._local_tensor.to(dtype=target_dtype)
442 return self._from_converted_local(new_local)
444 def _validate_factory_device(self, device: Any) -> None:
445 """Raise :class:`ValueError` if ``device`` does not match the DTensor's device."""
446 requested_type, requested_index = _device_spec(device)
447 local_type, local_index = _device_spec(self._local_tensor.device)
448 if (
449 requested_type != local_type
450 or (
451 requested_index is not None
452 and requested_index != local_index
453 )
454 ):
455 raise ValueError(
456 f"DTensor requires device to match the input DTensor "
457 f"device {self._local_tensor.device}, but got {device}."
458 )
460 def _new_const_tensor_op(
461 self,
462 method_name: str,
463 size: Union[int, Sequence[int]],
464 *,
465 dtype: Optional[Any] = None,
466 device: Optional[Any] = None,
467 requires_grad: bool = False,
468 layout: Optional[Any] = None,
469 pin_memory: bool = False,
470 ) -> 'DTensor':
471 """Create an all-``Replicate`` constant DTensor.
473 Shared implementation for ``new_zeros`` and ``new_ones``.
475 ``self`` is only used as a dtype/device reference and mesh source;
476 its values are ignored. The output is always **all-Replicate**
477 because every device produces identical data independently.
479 Args:
480 method_name:
481 ``"new_zeros"`` or ``"new_ones"`` — the local tensor
482 factory method to call.
483 size:
484 Output shape — an int or a sequence of ints.
485 dtype:
486 Desired dtype. Defaults to ``self.dtype`` on Torch.
487 device:
488 Must match ``self``'s device (Torch only).
489 requires_grad:
490 Forwarded on Torch; rejected on MindSpore.
491 layout:
492 Forwarded on Torch; rejected on MindSpore.
493 pin_memory:
494 Forwarded on Torch; rejected on MindSpore.
496 Returns:
497 DTensor: A new DTensor with all-``Replicate`` placements on
498 ``self``'s ``DeviceMesh``.
500 Raises:
501 ValueError: If a Torch-only kwarg is used on MindSpore, or
502 ``device`` does not match the DTensor's device.
503 """
504 if isinstance(size, int):
505 size = (size,)
507 if platform.platform_type == PlatformType.MINDSPORE:
508 if device is not None or layout is not None or requires_grad or pin_memory:
509 raise ValueError(
510 f"DTensor.{method_name} only supports size and dtype "
511 "on MindSpore."
512 )
513 local_kwargs = {}
514 if dtype is not None:
515 local_kwargs["dtype"] = dtype
516 else:
517 local_kwargs = {}
518 if dtype is not None:
519 local_kwargs["dtype"] = dtype
520 if device is not None:
521 self._validate_factory_device(device)
522 # An unindexed device such as "cuda" resolves to the framework's
523 # current device, which may differ from this DTensor's local device.
524 local_kwargs["device"] = self._local_tensor.device
525 if requires_grad:
526 local_kwargs["requires_grad"] = True
527 if layout is not None:
528 local_kwargs["layout"] = layout
529 if pin_memory:
530 local_kwargs["pin_memory"] = True
532 factory = getattr(self._local_tensor, method_name)
533 local_result = factory(size, **local_kwargs)
535 replicated_placements = [Replicate()] * self._device_mesh.ndim
536 return DTensor.from_local(
537 local_result, self._device_mesh, replicated_placements,
538 )
540 def new_zeros(
541 self,
542 size: Union[int, Sequence[int]],
543 *,
544 dtype: Optional[Any] = None,
545 device: Optional[Any] = None,
546 requires_grad: bool = False,
547 layout: Optional[Any] = None,
548 pin_memory: bool = False,
549 ) -> 'DTensor':
550 """Create an all-Replicate DTensor filled with zeros.
552 The output is always **fully replicated** across every device in
553 ``self``'s ``DeviceMesh``, regardless of how ``self`` is sharded.
555 Args:
556 size:
557 Output shape — an int or a sequence of ints.
558 dtype:
559 Desired dtype. Defaults to ``self.dtype`` (Torch).
560 Not forwarded to MindSpore unless explicitly set.
561 device:
562 Must match ``self``'s device. Not supported on MindSpore.
563 requires_grad:
564 Forwarded on Torch; rejected on MindSpore.
565 layout:
566 Forwarded on Torch; rejected on MindSpore.
567 pin_memory:
568 Forwarded on Torch; rejected on MindSpore.
570 Returns:
571 DTensor: A new all-Replicate DTensor filled with zeros.
572 """
573 return self._new_const_tensor_op(
574 "new_zeros", size,
575 dtype=dtype,
576 device=device,
577 requires_grad=requires_grad,
578 layout=layout,
579 pin_memory=pin_memory,
580 )
582 def new_ones(
583 self,
584 size: Union[int, Sequence[int]],
585 *,
586 dtype: Optional[Any] = None,
587 device: Optional[Any] = None,
588 requires_grad: bool = False,
589 layout: Optional[Any] = None,
590 pin_memory: bool = False,
591 ) -> 'DTensor':
592 """Create an all-Replicate DTensor filled with ones.
594 The output is always **fully replicated** across every device in
595 ``self``'s ``DeviceMesh``, regardless of how ``self`` is sharded.
597 Args:
598 size:
599 Output shape — an int or a sequence of ints.
600 dtype:
601 Desired dtype. Defaults to ``self.dtype`` (Torch).
602 Not forwarded to MindSpore unless explicitly set.
603 device:
604 Must match ``self``'s device. Not supported on MindSpore.
605 requires_grad:
606 Forwarded on Torch; rejected on MindSpore.
607 layout:
608 Forwarded on Torch; rejected on MindSpore.
609 pin_memory:
610 Forwarded on Torch; rejected on MindSpore.
612 Returns:
613 DTensor: A new all-Replicate DTensor filled with ones.
614 """
615 return self._new_const_tensor_op(
616 "new_ones", size,
617 dtype=dtype,
618 device=device,
619 requires_grad=requires_grad,
620 layout=layout,
621 pin_memory=pin_memory,
622 )
624 def to_local(self) -> Tensor:
625 """
626 Convert DTensor to local tensor.
628 Returns:
629 Tensor: The local tensor shard on this device.
630 """
631 return self._local_tensor
633 def tolist(self):
634 """
635 Convert the DTensor to a nested Python list or number.
637 This operation gathers the complete tensor on every participating rank
638 before converting it to Python values. It is an **implicit collective**:
639 all ranks in the DeviceMesh must participate.
641 Returns:
642 Union[list, int, float, bool]: A nested Python list, or a Python
643 number for a scalar DTensor.
645 Note:
646 This triggers ``full_tensor()`` under the hood, which performs
647 all-gather communication. For large tensors, prefer slicing or
648 index-based access to avoid materialising the full tensor.
650 If you only need the **local shard** as a list, use
651 ``dtensor.to_local().tolist()`` instead — that path has zero
652 communication overhead.
654 Example:
655 >>> mesh = init_device_mesh("npu", (2,), ("dp",))
656 >>> x = distribute_tensor(torch.arange(8).reshape(4, 2), mesh, [Shard(0)])
657 >>> x.tolist() # full data: [[0,1],[2,3],[4,5],[6,7]]
658 >>> x.to_local().tolist() # local shard only (no comm)
659 """
660 return self.full_tensor().tolist()
662 def copy_(self, src: "DTensor", non_blocking: bool = False) -> "DTensor":
663 """In-place copy of ``src`` into this DTensor's local shard.
665 Delegates to ``Tensor.copy_`` on the underlying local tensors.
666 Follows standard ``Tensor.copy_`` semantics: version counter is
667 bumped and autograd edges are created when grad is enabled.
669 Constraints on ``src``:
670 * must be a ``DTensor`` on the same ``DeviceMesh`` as ``self``;
671 * its placements must equal ``self.placements``, OR
672 ``src._local_tensor.numel() == 1`` (single-element broadcast);
673 * its local shape must equal or be broadcastable to
674 ``self._local_tensor.shape``.
676 No redistribute / implicit slicing is performed; src dtype is cast
677 to self dtype in-place.
679 Args:
680 src (DTensor): Source DTensor satisfying the constraints above.
681 non_blocking (bool): Forwarded to the underlying ``copy_``.
683 Returns:
684 DTensor: ``self``.
686 Raises:
687 TypeError: if ``src`` is not a ``DTensor``.
688 ValueError: if mesh, placement, or shape constraint is violated.
689 """
690 if not isinstance(src, DTensor):
691 raise TypeError(
692 f"For DTensor.copy_, src should be a DTensor, but got {type(src).__name__}."
693 )
694 src_local = src.to_local()
695 if src.device_mesh is not self._device_mesh:
696 raise ValueError(
697 f"For DTensor.copy_, src and self should share the same DeviceMesh, "
698 f"but got src.device_mesh={src.device_mesh!r}, "
699 f"self._device_mesh={self._device_mesh!r}."
700 )
702 placement_eq = tuple(src.placements) == tuple(self._placements)
703 shape_eq = src_local.shape == self._local_tensor.shape
704 src_is_scalar = src_local.numel() == 1
706 if not placement_eq and not src_is_scalar:
707 raise ValueError(
708 f"For DTensor.copy_, src.placements should equal self.placements "
709 f"or src.numel() should be 1, but got "
710 f"src.placements={src.placements}, "
711 f"self.placements={self._placements}, "
712 f"src.numel()={src_local.numel()}."
713 )
714 if not shape_eq and not src_is_scalar and not _is_broadcastable(
715 src_local.shape, self._local_tensor.shape
716 ):
717 raise ValueError(
718 f"For DTensor.copy_, src local shape should be broadcastable to "
719 f"self local shape, but got "
720 f"src.shape={tuple(src_local.shape)}, "
721 f"self.shape={tuple(self._local_tensor.shape)}."
722 )
724 self._local_tensor.copy_(src_local, non_blocking=non_blocking)
725 return self
727 def zero_(self) -> "DTensor":
728 """In-place fill with zeros. Returns ``self``."""
729 self._local_tensor.zero_()
730 return self
732 def fill_(self, value) -> "DTensor":
733 """In-place fill with ``value``. Returns ``self``."""
734 self._local_tensor.fill_(value)
735 return self
737 @property
738 def shape(self) -> Tuple[int, ...]:
739 """
740 The global shape of this DTensor.
742 Returns:
743 Tuple[int, ...]: The global tensor shape.
744 """
745 return self._layout.get_global_shape(self._local_tensor.shape)
747 def size(self, dim=None):
748 """Return the global shape, consistent with .shape.
750 Without ``dim`` returns a tuple matching ``self.shape``.
751 With ``dim`` returns the size of that dimension.
752 """
753 global_shape = self.shape
754 if dim is not None:
755 return global_shape[dim]
756 return global_shape
758 def numel(self) -> int:
759 """Return the number of elements in this DTensor."""
760 return int(np.prod(self.shape))
762 @property
763 def local_shape(self) -> Tuple[int, ...]:
764 """
765 The local shape of this DTensor on this device.
767 Returns:
768 Tuple[int, ...]: The local tensor shape.
769 """
770 return self._local_tensor.shape
772 def redistribute(
773 self,
774 device_mesh: DeviceMesh,
775 placements: Union[Sequence[Placement], Sequence[Union[str, Tuple[str, ...]]]]
776 ) -> 'DTensor':
777 """
778 Redistribute this DTensor to a new device mesh and placements.
780 Args:
781 device_mesh (DeviceMesh): The target device mesh.
782 placements: The target placements. Supports Placement objects
783 or alias strings.
785 Returns:
786 DTensor: A new DTensor with the specified distribution.
788 Example:
789 >>> new_dtensor = dtensor.redistribute(mesh, [Replicate(), Shard(1)])
790 >>> new_dtensor = dtensor.redistribute(mesh, ("None", "tp"))
791 """
792 logger.debug(
793 "redistribute: shape=%s, src_placements=%s -> dst_placements=%s, "
794 "mesh_shape=%s, local_shape=%s",
795 tuple(self.shape),
796 tuple(self._placements),
797 tuple(placements),
798 tuple(device_mesh.shape),
799 tuple(self._local_tensor.shape),
800 )
802 # Build dst_layout from device_mesh and placements
803 dst_layout = _build_layout(
804 device_mesh, placements, len(self._local_tensor.shape)
805 )
807 # pylint: disable=C0415
808 from hyper_parallel.core.dtensor.tensor_redistribution import _tensor_redistribution
809 out = _tensor_redistribution.redistribution(self, dst_layout)
810 return out
812 def reduce_partial(self) -> 'DTensor':
813 """
814 Reduce partial sharding state for this DTensor.
816 Returns:
817 DTensor: A new DTensor with partial state reduced.
818 """
819 if not self._layout:
820 return self
821 to_layout = cp.deepcopy(self._layout)
822 to_layout.reset_partial()
823 # pylint: disable=C0415
824 from hyper_parallel.core.dtensor.tensor_redistribution import _tensor_redistribution
825 out = _tensor_redistribution.reduce_partial(self, to_layout)
826 return out
828 def full_tensor(self) -> Tensor:
829 """
830 Return the full tensor of this DTensor.
832 Returns:
833 Tensor: A Tensor object that represents the full tensor of this DTensor.
834 The returned tensor contains the complete data gathered from
835 all ranks.
837 Note:
838 This operation involves communication across all ranks in the DeviceMesh,
839 which may be expensive for large tensors. Use with caution in
840 performance-critical code paths.
842 Example:
843 >>> # Assume dtensor is sharded across multiple devices
844 >>> local_tensor = dtensor.to_local() # Returns only the local shard
845 >>> full_tensor = dtensor.full_tensor() # Returns the complete tensor
846 """
847 if not self._layout:
848 return self._local_tensor
850 # Create a fully replicated layout
851 replicated_layout = cp.deepcopy(self._layout)
853 # Set all placements to Replicate and convert to tensor_map
854 replicated_placements = [Replicate()] * len(replicated_layout.mesh_shape)
855 replicated_layout.set_placements(replicated_placements)
856 replicated_layout.placement_to_tensor_map(len(self._local_tensor.shape))
858 # Clear partial status from original layout since Replicate has no partial
859 replicated_layout.reset_partial()
861 # Redistribute to the replicated layout and return local tensor
862 # pylint: disable=C0415
863 from hyper_parallel.core.dtensor.tensor_redistribution import _tensor_redistribution
864 out = _tensor_redistribution.redistribution(self, replicated_layout)
865 return out.to_local()
868def _normalize_shard_dim(dim: int, ndim: int) -> int:
869 return dim + ndim if dim < 0 else dim
872def _distribute_tensor_with_communication(
873 tensor: Tensor,
874 device_mesh: DeviceMesh,
875 placements: Sequence[Placement],
876 src_data_rank: int,
877) -> Tensor:
878 """Scatter/broadcast a logical global tensor along mesh dimensions (PyTorch parity)."""
879 local = tensor
880 if len(placements) < device_mesh.ndim:
881 raise ValueError(
882 f"placements length ({len(placements)}) must be at least device_mesh.ndim "
883 f"({device_mesh.ndim}) when src_data_rank is set"
884 )
885 for mesh_dim in range(device_mesh.ndim):
886 placement = placements[mesh_dim]
887 if isinstance(placement, StridedShard):
888 raise NotImplementedError(
889 "distribute_tensor with src_data_rank does not support StridedShard yet; "
890 "pass src_data_rank=None for local-only sharding."
891 )
892 if placement.is_shard():
893 shard_dim = _normalize_shard_dim(placement.dim, local.ndim)
894 num_chunks = device_mesh.size(mesh_dim)
895 if num_chunks <= 0:
896 raise ValueError(f"invalid mesh dim size {num_chunks} on mesh_dim={mesh_dim}")
897 chunks = tuple(local.chunk(num_chunks, dim=shard_dim))
898 if not chunks:
899 raise ValueError(f"cannot shard dim {shard_dim} into {num_chunks} chunks")
900 output = platform.empty_like(chunks[0])
901 local = mesh_scatter(output, chunks, device_mesh, mesh_dim, group_src=src_data_rank)
902 elif placement.is_replicate() or placement.is_partial():
903 local = mesh_broadcast(local, device_mesh, mesh_dim, group_src=src_data_rank)
904 if isinstance(placement, Partial):
905 warnings.warn(
906 f"Partial placement {placement} during distribute_tensor: "
907 "broadcast only; partial partition is not applied yet.",
908 stacklevel=3,
909 )
910 else:
911 raise RuntimeError(
912 f"unsupported placement {placement} on device mesh dimension {mesh_dim}"
913 )
914 return local
917def distribute_tensor(
918 tensor: Tensor,
919 device_mesh: DeviceMesh,
920 placements: Union[Sequence[Placement], Sequence[Union[str, Tuple[str, ...]]]],
921 *,
922 src_data_rank: Optional[int] = None,
923) -> DTensor:
924 """
925 Distribute a global tensor to the device mesh according to the placements.
927 Args:
928 tensor (Tensor): The global tensor to be distributed. All ranks
929 should have the same tensor data.
930 device_mesh (DeviceMesh): The device mesh describing the device topology.
931 placements: The placement strategy. Supports two styles:
932 - Placement objects (e.g., ``[Shard(0), Replicate()]``).
933 - Alias strings (e.g., ``("dp", "None")`` or
934 ``(("dp", "tp"), "None")``), length must equal the number of
935 tensor dimensions.
937 Returns:
938 DTensor: A new DTensor with the local shard on each rank.
940 Note:
941 When ``src_data_rank`` is an ``int`` (e.g. ``0``), shard/replicate
942 placements use scatter/broadcast from the source rank on each mesh axis,
943 matching PyTorch ``distribute_tensor``. When ``src_data_rank=None``
944 (default), each rank slices its local tensor without communication
945 (legacy Hyper behavior; all ranks must hold the same global data).
947 Example:
948 >>> mesh = init_device_mesh(device_type="npu", mesh_shape=(2, 2), mesh_dim_names=("dp", "tp"))
949 >>> global_tensor = Tensor(np.arange(16).reshape(4, 4))
950 >>> dtensor = distribute_tensor(global_tensor, mesh, [Shard(0), Replicate()])
951 >>> dtensor = distribute_tensor(global_tensor, mesh, ("dp", "None"))
952 """
953 layout = _build_layout(device_mesh, placements, len(tensor.shape))
954 if src_data_rank is None:
955 local_tensor = _get_slice_tensor_by_layout(tensor, layout)
956 else:
957 local_tensor = _distribute_tensor_with_communication(
958 tensor, device_mesh, layout.placements, src_data_rank
959 )
960 return DTensor(local_tensor, device_mesh, layout.alias_placements)
963def _distribute_module_param_source(param: Any) -> Tensor:
964 """Tensor data used as the global tensor for :func:`distribute_tensor` (PyTorch uses ``param.data``)."""
965 if hasattr(param, "data"):
966 return param.data
967 return platform.get_param_local_data(param)
970def _distribute_module_new_parameter(key: str, dtensor: DTensor, requires_grad: bool) -> Any:
971 """Build a framework :class:`Parameter` holding *dtensor* (Torch vs MindSpore kwargs differ)."""
972 if platform.platform_type == PlatformType.MINDSPORE:
973 return platform.Parameter(dtensor, name=key, requires_grad=requires_grad)
974 return platform.Parameter(dtensor, requires_grad=requires_grad)
977def _distribute_module_set_param(module: Any, key: str, new_param: Any) -> None:
978 """Register or assign a parameter on *module* (``nn.Module`` or MindSpore ``Cell``)."""
979 if hasattr(module, "register_parameter"):
980 module.register_parameter(key, new_param)
981 return
982 if hasattr(module, "_params"):
983 module._params[key] = new_param
984 if hasattr(module, "_params_list"):
985 module._params_list[key] = new_param
986 if key in module.__dict__:
987 module.__dict__[key] = new_param
988 return
989 raise TypeError(
990 f"distribute_module expects nn.Module-like objects with register_parameter or _params; "
991 f"got {type(module)}."
992 )
995def _distribute_module_iter_params(module: Any) -> list:
996 """Return ``[(name, param), ...]`` for direct parameters (``_parameters`` or ``_params``)."""
997 if hasattr(module, "_parameters"):
998 return list(module._parameters.items())
999 if hasattr(module, "_params"):
1000 return list(module._params.items())
1001 return []
1004def _distribute_module_iter_buffers(module: Any) -> list:
1005 """Return ``[(name, buffer), ...]`` if the module has ``_buffers`` (PyTorch ``nn.Module``)."""
1006 if hasattr(module, "_buffers"):
1007 return list(module._buffers.items())
1008 return []
1011def _distribute_module_named_modules(module: Any):
1012 """``nn.Module.named_modules`` or MindSpore ``Cell.cells_and_names`` (submodule FQNs)."""
1013 if hasattr(module, "named_modules"):
1014 return module.named_modules()
1015 if hasattr(module, "cells_and_names"):
1016 return module.cells_and_names()
1017 raise TypeError(
1018 f"distribute_module expects module-like objects with named_modules or cells_and_names; "
1019 f"got {type(module)}."
1020 )
1023def _distribute_module_named_parameters(module: Any):
1024 """``nn.Module.named_parameters(recurse=False)`` or MindSpore ``Cell.parameters_and_names(expand=False)``."""
1025 if hasattr(module, "named_parameters"):
1026 return module.named_parameters(recurse=False)
1027 if hasattr(module, "parameters_and_names"):
1028 return module.parameters_and_names(expand=False)
1029 raise TypeError(
1030 f"distribute_module expects module-like objects with named_parameters or parameters_and_names; "
1031 f"got {type(module)}."
1032 )
1035def _replicate_submodule_params_buffers(
1036 sub_mod: Any,
1037 device_mesh: DeviceMesh,
1038 *,
1039 module_prefix: str = "",
1040) -> None:
1041 """Convert plain params/buffers on *sub_mod* to fully replicated :class:`DTensor`."""
1042 full_replicate = [Replicate()] * device_mesh.ndim
1043 for key, param in _distribute_module_iter_params(sub_mod):
1044 if param is None or isinstance(param, DTensorBase):
1045 continue
1046 src = _distribute_module_param_source(param)
1047 requires_grad = bool(getattr(param, "requires_grad", True))
1048 dt = distribute_tensor(src, device_mesh, full_replicate)
1049 param_name = f"{module_prefix}.{key}" if module_prefix else key
1050 new_param = _distribute_module_new_parameter(param_name, dt, requires_grad)
1051 _distribute_module_set_param(sub_mod, key, new_param)
1052 for key, buffer in _distribute_module_iter_buffers(sub_mod):
1053 if buffer is None or isinstance(buffer, DTensorBase):
1054 continue
1055 sub_mod._buffers[key] = distribute_tensor(buffer, device_mesh, full_replicate)
1058def _distribute_module_run_partition_and_replicate(
1059 module: Any,
1060 device_mesh: DeviceMesh,
1061 partition_fn: Optional[Callable[[str, Any, DeviceMesh], None]],
1062) -> None:
1063 """Call optional ``partition_fn`` per ``named_modules`` and replicate remaining tensors."""
1064 if partition_fn is None:
1065 for mod_name, submod in _distribute_module_named_modules(module):
1066 _replicate_submodule_params_buffers(submod, device_mesh, module_prefix=mod_name)
1067 return
1068 for mod_name, submod in _distribute_module_named_modules(module):
1069 partition_fn(mod_name, submod, device_mesh)
1070 _replicate_submodule_params_buffers(submod, device_mesh, module_prefix=mod_name)
1073def _distribute_module_register_input_fn(
1074 module: Any,
1075 device_mesh: DeviceMesh,
1076 input_fn: Callable[..., Any],
1077) -> None:
1078 """Register *input_fn* as a forward pre-hook on *module* (2- or 3-arg, PyTorch-compatible)."""
1079 num_args = len(inspect.signature(input_fn).parameters)
1080 if num_args == 2:
1081 warnings.warn(
1082 "Deprecating input_fn that takes two arguments (inputs, device_mesh), "
1083 "please use input_fn that takes in (module, inputs, device_mesh) instead!",
1084 FutureWarning,
1085 stacklevel=3,
1086 )
1087 module.register_forward_pre_hook(
1088 lambda _, inputs: input_fn(inputs, device_mesh)
1089 )
1090 elif num_args == 3:
1091 module.register_forward_pre_hook(
1092 lambda mod, inputs: input_fn(mod, inputs, device_mesh)
1093 )
1094 else:
1095 raise ValueError(
1096 f"input_fn should take in 2 or 3 arguments, but got {num_args} arguments!"
1097 )
1100def _distribute_module_register_output_fn(
1101 module: Any,
1102 device_mesh: DeviceMesh,
1103 output_fn: Callable[..., Any],
1104) -> None:
1105 """Register *output_fn* as a forward hook on *module* (2- or 3-arg, PyTorch-compatible)."""
1106 num_args = len(inspect.signature(output_fn).parameters)
1107 if num_args == 2:
1108 warnings.warn(
1109 "Deprecating output_fn that takes two arguments (outputs, device_mesh), "
1110 "please use output_fn that takes in (module, outputs, device_mesh) instead!",
1111 FutureWarning,
1112 stacklevel=3,
1113 )
1114 module.register_forward_hook(
1115 lambda mod, inputs, outputs: output_fn(outputs, device_mesh)
1116 )
1117 elif num_args == 3:
1118 module.register_forward_hook(
1119 lambda mod, inputs, outputs: output_fn(mod, outputs, device_mesh)
1120 )
1121 else:
1122 raise ValueError(
1123 f"output_fn should take in 2 or 3 arguments, but got {num_args} arguments!"
1124 )
1127def distribute_module(
1128 module: Any,
1129 device_mesh: Optional[DeviceMesh] = None,
1130 partition_fn: Optional[Callable[[str, Any, DeviceMesh], None]] = None,
1131 input_fn: Optional[Callable[..., Any]] = None,
1132 output_fn: Optional[Callable[..., Any]] = None,
1133) -> Any:
1134 """PyTorch ``distribute_module`` parity: shard/replicate params and optional I/O hooks.
1136 Unsharded parameters and buffers become fully replicated :class:`DTensor` after
1137 ``partition_fn``. ``input_fn`` / ``output_fn`` attach only to the root *module*.
1139 Args:
1140 module: Root ``nn.Module`` or MindSpore ``Cell`` with compatible APIs.
1141 device_mesh: Placement mesh; if ``None``, uses ``_mesh_resources.get_current_mesh()``.
1142 partition_fn: Per ``named_modules`` callback before replicate pass; ``None`` replicates all.
1143 input_fn: ``(module, inputs, mesh)`` or deprecated ``(inputs, mesh)`` pre-hook.
1144 output_fn: ``(module, outputs, mesh)`` or deprecated ``(outputs, mesh)`` forward hook.
1146 Returns:
1147 *module* in place, with distributed tensors where applied.
1149 Raises:
1150 RuntimeError: If called twice on the same *module*.
1151 ValueError: If ``input_fn`` / ``output_fn`` arity is not 2 or 3.
1153 Note:
1154 XLA / ``torch_xla`` is not supported; strided device :class:`DTensor` only.
1155 """
1156 if getattr(module, "_distribute_module_applied", False):
1157 raise RuntimeError(
1158 "distribute_module should only be called once on a module, "
1159 "but it has already been called on this module!"
1160 )
1161 device_mesh = device_mesh or _mesh_resources.get_current_mesh()
1162 _distribute_module_run_partition_and_replicate(module, device_mesh, partition_fn)
1163 if input_fn is not None:
1164 _distribute_module_register_input_fn(module, device_mesh, input_fn)
1165 if output_fn is not None:
1166 _distribute_module_register_output_fn(module, device_mesh, output_fn)
1167 module._distribute_module_applied = True
1168 return module
1171def _dtensor_init_helper(
1172 init_op,
1173 size,
1174 device_mesh,
1175 placements,
1176 *,
1177 rng_tracked: bool = False,
1178 **kwargs,
1179) -> DTensor:
1180 """
1181 Helper function to create and initialize a distributed tensor.
1183 Args:
1184 size: Shape of the tensor.
1185 dtype: Data type of the tensor.
1186 device: Target device for the tensor.
1187 requires_grad: Whether the tensor requires gradient.
1188 rng_tracked: When ``True``, initialize via :class:`OffsetBasedRNGTracker`
1189 so shard/replicate random semantics match PyTorch DTensor factories.
1191 Returns:
1192 DTensor: The initialized distributed tensor.
1193 """
1194 # get local tensor shape
1195 local_shape = compute_local_shape_and_global_offset(
1196 size, device_mesh, placements
1197 )
1199 # initialize the local tensor
1200 if init_op is platform.full:
1201 fill_value = kwargs.pop("fill_value", 0)
1202 local_tensor = init_op(local_shape, fill_value, **kwargs)
1203 elif rng_tracked:
1204 # pylint: disable=C0415
1205 from hyper_parallel.core.dtensor.random import is_rng_supported_mesh, OffsetBasedRNGTracker
1206 from hyper_parallel.core.shard._op_dispatch import _OP_DISPATCHER
1208 layout = _build_layout(device_mesh, placements, len(local_shape))
1209 if is_rng_supported_mesh(device_mesh):
1210 if _OP_DISPATCHER._rng_tracker is None:
1211 _OP_DISPATCHER._rng_tracker = OffsetBasedRNGTracker(run_state_sync=False)
1212 with _OP_DISPATCHER._rng_tracker._distribute_region(
1213 device_mesh,
1214 layout.placements,
1215 size,
1216 ):
1217 local_tensor = init_op(local_shape, **kwargs)
1218 else:
1219 local_tensor = init_op(local_shape, **kwargs)
1220 else:
1221 local_tensor = init_op(local_shape, **kwargs)
1223 return DTensor.from_local(
1224 local_tensor,
1225 device_mesh,
1226 placements,
1227 )
1230def ones(
1231 size,
1232 device_mesh,
1233 placements,
1234) -> DTensor:
1235 """
1236 Returns a :class:`DTensor` filled with the scalar value 1, with the shape defined
1237 by the variable argument ``size``.
1239 Args:
1240 size (Union[tuple[int], list[int], int, Tensor]): The specified shape of output tensor. Only positive integer or
1241 tuple or Tensor containing positive integers are allowed. If it is a Tensor,
1242 it must be a 0-D or 1-D Tensor with int32 or int64 dtypes.
1244 Keyword args:
1245 device_mesh: :class:`DeviceMesh` type, contains the mesh info of ranks
1246 placements: a sequence of :class:`Placement` type: ``Shard``, ``Replicate``
1248 Returns:
1249 A :class:`DTensor` object on each rank
1250 """
1251 ones_ = platform.ones
1252 return _dtensor_init_helper(
1253 ones_,
1254 size,
1255 device_mesh=device_mesh,
1256 placements=placements,
1257 )
1260def empty(
1261 size,
1262 device_mesh,
1263 placements,
1264) -> DTensor:
1265 """
1266 Returns a :class:`DTensor` filled with uninitialized data. The shape of the :class:`DTensor`
1267 is defined by the variable argument ``size``.
1269 Args:
1270 size (Union[tuple[int], list[int], int]): The specified shape of output tensor. Can be variable numbers of
1271 positive integers or tuple or list containing positive integers.
1273 Keyword args:
1274 device_mesh: :class:`DeviceMesh` type, contains the mesh info of ranks
1275 placements: a sequence of :class:`Placement` type: ``Shard``, ``Replicate``
1277 Returns:
1278 A :class:`DTensor` object on each rank
1279 """
1280 empty_ = platform.empty
1281 return _dtensor_init_helper(
1282 empty_,
1283 size,
1284 device_mesh=device_mesh,
1285 placements=placements,
1286 )
1289def full(
1290 size,
1291 fill_value,
1292 *,
1293 device_mesh,
1294 placements,
1295) -> DTensor:
1296 """
1297 Returns a :class:`DTensor` filled with ``fill_value`` according to ``device_mesh`` and
1298 ``placements``, with the shape defined by the argument ``size``.
1300 Args:
1301 size (Union[tuple[int], list[int]]): The specified shape of output tensor.
1302 fill_value (Union[numbers.Number, Tensor]): Value to fill the returned tensor. It can be a scalar number, a 0-D
1303 Tensor, or a 1-D Tensor with only one element.
1305 Keyword args:
1306 device_mesh: :class:`DeviceMesh` type, contains the mesh info of ranks.
1307 placements: a sequence of :class:`Placement` type: ``Shard``, ``Replicate``
1309 Returns:
1310 A :class:`DTensor` object on each rank
1311 """
1312 full_ = platform.full
1313 return _dtensor_init_helper(
1314 full_,
1315 size,
1316 fill_value=fill_value,
1317 device_mesh=device_mesh,
1318 placements=placements,
1319 )
1322def zeros(
1323 size,
1324 device_mesh,
1325 placements,
1326) -> DTensor:
1327 """
1328 Returns a :class:`DTensor` filled with the scalar value 0.
1330 Args:
1331 size (Union[tuple[int], list[int], int, Tensor]): The specified shape of output tensor. Only positive integer or
1332 tuple or Tensor containing positive integers are allowed. If it is a Tensor,
1333 it must be a 0-D or 1-D Tensor with int32 or int64 dtypes.
1334 Keyword args:
1335 device_mesh: :class:`DeviceMesh` type, contains the mesh info of ranks
1336 placements: a sequence of :class:`Placement` type: ``Shard``, ``Replicate``
1338 Returns:
1339 A :class:`DTensor` object on each rank
1340 """
1341 zeros_ = platform.zeros
1342 return _dtensor_init_helper(
1343 zeros_,
1344 size,
1345 device_mesh=device_mesh,
1346 placements=placements,
1347 )
1350def rand(
1351 size,
1352 device_mesh,
1353 placements,
1354 **kwargs,
1355) -> DTensor:
1356 """
1357 Returns a :class:`DTensor` filled with random numbers from a uniform
1358 distribution on ``[0, 1)``.
1360 Args:
1361 size: Global output shape.
1362 device_mesh: :class:`DeviceMesh` for the distributed layout.
1363 placements: Per-mesh-dimension :class:`Placement` values.
1364 **kwargs: Forwarded to the platform ``rand`` call (for example ``dtype``).
1366 Returns:
1367 A :class:`DTensor` object on each rank.
1368 """
1369 return _dtensor_init_helper(
1370 platform.rand,
1371 size,
1372 device_mesh=device_mesh,
1373 placements=placements,
1374 rng_tracked=True,
1375 **kwargs,
1376 )
1379def randn(
1380 size,
1381 device_mesh,
1382 placements,
1383 **kwargs,
1384) -> DTensor:
1385 """
1386 Returns a :class:`DTensor` filled with random numbers from a normal
1387 distribution with mean ``0`` and variance ``1``.
1389 Args:
1390 size: Global output shape.
1391 device_mesh: :class:`DeviceMesh` for the distributed layout.
1392 placements: Per-mesh-dimension :class:`Placement` values.
1393 **kwargs: Forwarded to the platform ``randn`` call (for example ``dtype``).
1395 Returns:
1396 A :class:`DTensor` object on each rank.
1397 """
1398 return _dtensor_init_helper(
1399 platform.randn,
1400 size,
1401 device_mesh=device_mesh,
1402 placements=placements,
1403 rng_tracked=True,
1404 **kwargs,
1405 )