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