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