Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / dtensor / dtensor_base.py: 69%
109 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2025-2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Torch ``DTensorBase``: the platform-agnostic base of :class:`DTensor`."""
16# pylint: disable=C9006,C9007
17from typing import Tuple, Dict, Any, Optional
18import torch
19from torch import Tensor
22class DTensorBase(Tensor):
23 """Base class of :class:`DTensor`: a ``torch.Tensor`` subclass that owns a local shard.
25 Holds the local shard plus the ``DeviceMesh``/``Placements`` describing how that shard
26 maps onto the global tensor, and routes every op through the distributed dispatcher.
27 """
29 def __new__(cls, local_tensor, device_mesh=None, placements=None, layout=None, shape=None):
30 """
31 Create a new DTensorBase instance.
33 Args:
34 local_tensor: The local tensor shard or another DTensorBase instance.
35 device_mesh: The device mesh describing the device topology.
36 placements: The placement strategy for each mesh dimension.
37 layout: Optional pre-built Layout reused directly by ``__init_data__``
38 (skips ``_build_layout``; see ``DTensor.from_local_with_layout``).
39 shape: Optional logical global tensor shape.
40 """
41 if isinstance(local_tensor, DTensorBase):
42 # Copy from existing DTensorBase — use alias_placements to preserve multi-axis ordering
43 t = Tensor._make_subclass(cls, local_tensor._local_tensor, local_tensor._local_tensor.requires_grad)
44 copy_placements = local_tensor.layout.alias_placements if local_tensor.layout else local_tensor.placements
45 t.__init_data__(
46 local_tensor._local_tensor,
47 local_tensor.device_mesh,
48 copy_placements,
49 shape=getattr(local_tensor, "_global_shape", None),
50 )
51 return t
53 if device_mesh is None:
54 raise ValueError("device_mesh is None, must provide a DeviceMesh instance")
55 if placements is None:
56 raise ValueError("placements is None, must provide placements")
58 # Create Tensor subclass instance, sharing local_tensor's underlying storage
59 t = Tensor._make_subclass(cls, local_tensor, local_tensor.requires_grad)
60 t.__init_data__(local_tensor, device_mesh, placements, layout, shape)
61 return t
63 # pylint: disable=W0613, G.NAM.05
64 @classmethod
65 def __torch_function__(
66 cls,
67 func: torch._C._FunctionBase,
68 types: Tuple[type, ...],
69 args: Tuple[Any, ...] = (),
70 kwargs: Optional[Dict[str, Any]] = None
71 ) -> Any:
72 """
73 Override PyTorch's __torch_function__ to intercept tensor operations.
75 This method dispatches operations through the distributed operator dispatcher
76 to handle DTensor-specific layout inference and redistribution.
78 Args:
79 func (torch._C._FunctionBase): The PyTorch function being called.
80 types (Tuple[type, ...]): The types of tensors involved in the operation.
81 args (Tuple[Any, ...]): Positional arguments passed to the function.
82 kwargs (Optional[Dict[str, Any]]): Keyword arguments passed to the function.
84 Returns:
85 Any: The result of the dispatched operation, typically a DTensor or tuple of DTensors.
86 """
87 kwargs = kwargs or {}
88 # pylint: disable=C0415
89 from hyper_parallel.core.shard._op_dispatch import _OP_DISPATCHER
90 out = _OP_DISPATCHER.dispatch(func, args, kwargs)
91 return out
93 @property
94 def grad(self) -> Optional[Tensor]:
95 """
96 Get the gradient tensor of the local tensor.
98 Returns:
99 Optional[Tensor]: The gradient tensor, or None if no gradient is set.
100 """
101 return self._local_tensor.grad
103 @grad.setter
104 def grad(self, value: Optional[Tensor]) -> None:
105 """
106 Set the gradient tensor for the local tensor.
108 Args:
109 value (Optional[Tensor]): The gradient tensor to set, or None to clear.
110 """
111 self._local_tensor.grad = value
113 @property
114 def requires_grad(self) -> bool:
115 """
116 Check if gradient computation is enabled for this tensor.
118 Returns:
119 bool: True if gradients should be computed for this tensor.
120 """
121 return self._local_tensor.requires_grad
123 @requires_grad.setter
124 def requires_grad(self, value: bool) -> None:
125 """
126 Enable or disable gradient computation for this tensor.
128 Args:
129 value (bool): True to enable gradient computation, False to disable.
130 """
131 self._local_tensor.requires_grad_(value)
132 # Sync DTensor wrapper's requires_grad
133 super().requires_grad_(value)
135 def requires_grad_(self, requires_grad: bool = True):
136 """
137 Enable or disable gradient computation in-place.
139 Args:
140 requires_grad (bool): True to enable gradient computation. Default: True.
142 Returns:
143 DTensorBase: Self for method chaining.
144 """
145 self._local_tensor.requires_grad_(requires_grad)
146 super().requires_grad_(requires_grad)
147 return self
149 @property
150 def grad_fn(self) -> Optional[torch.autograd.Function]:
151 """
152 Get the gradient function that created this tensor.
154 Returns:
155 Optional[torch.autograd.Function]: The gradient function, or None if not applicable.
156 """
157 return self._local_tensor.grad_fn
159 def grad_zero_(self):
160 """
161 Zero out the gradient tensor in-place.
163 Returns:
164 DTensorBase: Self for method chaining.
165 """
166 if self._local_tensor.grad is not None:
167 self._local_tensor.grad.zero_()
168 return self
170 def detach(self):
171 """
172 Create a detached DTensor that does not require gradient.
174 Returns:
175 DTensorBase: A new DTensor with the same data but detached from the computation graph.
176 """
177 detached_local = self._local_tensor.detach()
178 return self.__class__(
179 detached_local,
180 device_mesh=self._device_mesh,
181 placements=self._alias_placements(),
182 shape=getattr(self, "_global_shape", None),
183 )
185 def detach_(self):
186 """
187 Detach this tensor from the computation graph in-place.
189 Returns:
190 DTensorBase: Self for method chaining.
191 """
192 self._local_tensor.detach_()
193 super().detach_()
194 return self
196 # ====================== Computation graph related overrides ======================
197 @property
198 def is_leaf(self) -> bool:
199 """
200 Check if this tensor is a leaf node in the computation graph.
202 Returns:
203 bool: True if this is a leaf tensor (created by user, not by any operation).
204 """
205 return self._local_tensor.is_leaf
207 @property
208 def retains_grad(self) -> bool:
209 """
210 Check if this tensor retains its gradient during backward pass.
212 Returns:
213 bool: True if gradients are retained for non-leaf tensors.
214 """
215 return self._local_tensor.retains_grad
217 @retains_grad.setter
218 def retains_grad(self, value: bool) -> None:
219 """
220 Enable or disable gradient retention for this tensor.
222 Args:
223 value (bool): True to enable gradient retention.
224 """
225 self._local_tensor.retains_grad_(value)
227 def backward(self, gradient=None, retain_graph=None, create_graph=False) -> None:
228 """
229 Compute the gradients for this tensor.
231 Args:
232 gradient (Optional[Tensor]): The gradient of the loss w.r.t. this tensor.
233 retain_graph (Optional[bool]): Whether to retain the computation graph.
234 create_graph (bool): Whether to create a graph of the gradient computation.
235 """
236 self._local_tensor.backward(gradient, retain_graph, create_graph)
238 # ====================== Metadata related overrides (sync with local_tensor) ======================
239 @property
240 def device(self) -> torch.device:
241 """
242 Get the device on which this tensor is stored.
244 Returns:
245 torch.device: The device object (e.g., 'cuda:0', 'cpu').
246 """
247 return self._local_tensor.device
249 @property
250 # pylint: disable=C2801
251 def data(self):
252 """
253 Directory get Tensor.data relative storage.
254 After DTensor object created, there are two reference on underlying storage. (.data and ._local_tensor)
255 If not using DisableTorchFunctionSubclass, Tensor.__get__ will goto '__torch_function__',
256 and finally return '._local_tensor' storage.
257 """
258 with getattr(torch, "_C").DisableTorchFunctionSubclass():
259 # Directory get Tensor.data relative storage.
260 # If not using DisableTorchFunctionSubclass, Tensor.__get__ will goto __torch_funtin
261 return Tensor.data.__get__(self, type(self))
263 @data.setter
264 # pylint: disable=C2801
265 def data(self, value):
266 """Set the underlying tensor data, extracting the local shard if a DTensor is given."""
267 local_value = value.to_local() if isinstance(value, DTensorBase) else value
268 # Tensor.data.__set__ on a Tensor subclass otherwise enters __torch_function__
269 # and only rebinds _local_tensor through DTensor dispatch.
270 with getattr(torch, "_C").DisableTorchFunctionSubclass():
271 Tensor.data.__set__(self, local_value)
272 Tensor.data.__set__(self._local_tensor, local_value)
274 @property
275 def dtype(self) -> torch.dtype:
276 """
277 Get the data type of this tensor.
279 Returns:
280 torch.dtype: The data type (e.g., torch.float32, torch.int64).
281 """
282 return self._local_tensor.dtype
284 @property
285 def shape(self) -> torch.Size:
286 """
287 Get the shape of this tensor.
289 Returns:
290 torch.Size: The shape of the tensor.
291 """
292 return self._local_tensor.shape
294 def type(self, dtype=None, non_blocking=False):
295 """
296 Convert this tensor to the specified dtype.
298 Args:
299 dtype (Optional[torch.dtype]): The target dtype. If None, returns the current type string.
300 non_blocking (bool): Whether to perform the operation asynchronously. Default: False.
302 Returns:
303 Union[str, DTensorBase]: The type string if dtype is None, otherwise a new DTensor.
304 """
305 if dtype is None:
306 return self._local_tensor.type()
307 new_local = self._local_tensor.to(dtype=dtype, non_blocking=non_blocking)
308 return self.__class__(
309 new_local,
310 device_mesh=self._device_mesh,
311 placements=self._alias_placements(),
312 shape=getattr(self, "_global_shape", None),
313 )
315 def size(self, dim: Optional[int] = None):
316 """
317 Get the size of this tensor.
319 Args:
320 dim (Optional[int]): The dimension to query. If None, returns the full shape.
322 Returns:
323 Union[torch.Size, int]: The shape or size along a specific dimension.
324 """
325 return self._local_tensor.size(dim)
327 @property
328 def ndim(self) -> int:
329 """
330 Get the number of dimensions of this tensor.
332 Returns:
333 int: The number of dimensions.
334 """
335 return self._local_tensor.ndim
337 def data_ptr(self) -> int:
338 """
339 Get the pointer to the data storage of the local tensor.
341 Returns:
342 int: The memory address of the tensor's data.
343 """
344 # Force return local_tensor's data pointer (ensure address consistency)
345 return self._local_tensor.data_ptr()
347 def numel(self) -> int:
348 """
349 Get the total number of elements in this tensor.
351 Returns:
352 int: The total number of elements.
353 """
354 return self._local_tensor.numel()
356 # ====================== Auxiliary print ======================
357 def _alias_placements(self):
358 """Return alias_placements from layout, falling back to _placements."""
359 if hasattr(self, '_layout') and self._layout is not None:
360 return self._layout.alias_placements
361 return self._placements
363 def to(self, *args, **kwargs):
364 """Move the DTensor to a different device or dtype.
366 This method overrides the base Tensor.to() to properly reconstruct
367 a DTensor with device_mesh and placements preserved. Uses _make_subclass
368 to avoid issues with Parameter subclasses that don't accept extra kwargs.
370 Args:
371 *args: Arguments passed to the underlying tensor's to() method.
372 **kwargs: Keyword arguments for the tensor conversion.
374 Returns:
375 DTensorBase: A new DTensor with the converted local tensor.
376 """
377 new_local = self._local_tensor.to(*args, **kwargs)
378 new_dt = Tensor._make_subclass(type(self), new_local, new_local.requires_grad)
379 new_dt.__init_data__(
380 new_local,
381 self._device_mesh,
382 self._alias_placements(),
383 shape=getattr(self, "_global_shape", None),
384 )
385 return new_dt
387 def __repr__(self) -> str:
388 return (
389 f"DTensor(\n"
390 f" local_tensor={self._local_tensor},\n"
391 f" device_mesh={self._device_mesh},\n"
392 f" placements={self._placements},\n"
393 f" layout={getattr(self, '_layout', None)},\n"
394 f" device={self.device},\n"
395 f" dtype={self.dtype},\n"
396 f" requires_grad={self.requires_grad},\n"
397 f" grad={self.grad},\n"
398 f" is_leaf={self.is_leaf},\n"
399 f" data_ptr={self.data_ptr()}\n"
400 f")"
401 )