Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / mindspore / dtensor.py: 45%
118 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-04 05:18 +0800
1# Copyright 2025-2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""mindspore dtensor base"""
16from mindspore._c_expression import NoFallbackGuard, _DisableMsDispatchMode
17from mindspore.common.tensor import Tensor
18from mindspore.common.initializer import initializer
21class DTensorBase(Tensor):
22 """
23 DTensorBase - Base class for distributed tensors in MindSpore.
25 This class extends Tensor to support distributed tensor operations with
26 device mesh and placement specifications.
27 """
29 def __new__(cls, local_tensor, device_mesh=None, placements=None, layout=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. When supplied, ``__init_data__``
38 reuses it directly and skips ``_build_layout`` (hot-path fast
39 construction, see ``DTensor.from_local_with_layout``).
40 device: The device type (default: "Ascend").
41 """
42 # Fast path: a pre-built layout is only ever supplied by the internal
43 # wrap_output / from_local_with_layout hot path, where local_tensor is a
44 # freshly produced plain op-output Tensor (never a DTensorBase) already on
45 # the compute device, and device_mesh/placements are known-valid. Skip the
46 # ABCMeta isinstance(local_tensor, DTensorBase) check (MindSpore Tensor's
47 # metaclass is ABCMeta, so that isinstance is ~10x a normal one), the three
48 # None guards, and the device-placement guard — all pure per-output overhead.
49 if layout is not None:
50 t = Tensor._make_subclass(cls, local_tensor)
51 t.__init_data__(local_tensor, device_mesh, placements, layout)
52 return t
54 npu_device = "Ascend"
55 if isinstance(local_tensor, DTensorBase):
56 src = local_tensor
57 local_tensor = src.to_local()
58 device_mesh = src.device_mesh
59 placements = src._alias_placements()
60 else:
61 if local_tensor is None:
62 raise ValueError(
63 "DTensorBase: local_tensor must not be None when constructing from a raw tensor."
64 )
65 if device_mesh is None:
66 raise ValueError(
67 "DTensorBase: device_mesh must be a DeviceMesh instance, got None."
68 )
69 if placements is None:
70 raise ValueError(
71 "DTensorBase: placements must be a sequence of Placement objects, got None."
72 )
74 if local_tensor.has_init:
75 local_tensor.init_device = npu_device
76 else:
77 dev = local_tensor.device
78 if dev != "meta" and not dev.startswith(npu_device):
79 local_tensor = local_tensor.to(npu_device)
81 t = Tensor._make_subclass(cls, local_tensor)
82 t.__init_data__(local_tensor, device_mesh, placements, layout)
83 return t
85 def asnumpy(self):
86 """
87 Numpy value of local tensor.
88 """
89 return self._local_tensor.asnumpy()
91 def __str__(self):
92 return str(self._local_tensor)
94 def __copy__(self):
95 """
96 Create a shallow copy of the DTensorBase instance.
98 This method ensures that device_mesh and placements are correctly
99 propagated when creating a copy (e.g., for optimizer states).
100 """
101 # Get device_mesh and placements from layout (prefer alias_placements to preserve multi-axis ordering)
102 device_mesh = getattr(self, '_device_mesh', None)
103 placements = None
105 if hasattr(self, '_layout') and self._layout is not None:
106 if device_mesh is None:
107 device_mesh = self._layout.mesh
108 placements = self._layout.alias_placements
110 if placements is None:
111 placements = getattr(self, '_placements', None)
113 if device_mesh is None or placements is None:
114 raise ValueError(
115 "DTensorBase.__copy__: cannot copy without device_mesh and placements; "
116 f"device_mesh={device_mesh!r}, placements={placements!r}. "
117 "Ensure the tensor was constructed with a valid layout."
118 )
120 if self._local_tensor.has_init:
121 obj = DTensorBase.__new__(
122 type(self),
123 initializer(self._local_tensor.init, self._local_tensor.shape, self._local_tensor.dtype),
124 device_mesh,
125 placements
126 )
127 else:
128 obj = DTensorBase.__new__(
129 type(self),
130 self._local_tensor.clone(),
131 device_mesh,
132 placements
133 )
134 filtered_dict = {k: v for k, v in self.__dict__.items() if k != '_local_tensor'}
135 obj.__dict__.update(filtered_dict)
136 return obj
138 # pylint: disable=W0211, W0102, C0415, G.NAM.05
139 def __fallback__(self, func, args={}, kwargs=None):
140 if kwargs is None:
141 kwargs = {}
142 from hyper_parallel.core.shard._op_dispatch import _OP_DISPATCHER
143 with NoFallbackGuard():
144 out = _OP_DISPATCHER.dispatch(func, args, kwargs)
145 return out
147 # pylint: disable=W0212
148 def _need_contiguous(self):
149 """_need_contiguous"""
150 return self._local_tensor._need_contiguous()
152 @property
153 def device(self):
154 """Device info for dtensor"""
155 device_info = self._local_tensor.device
156 return device_info.split(':', 1)[0]
158 @property
159 # pylint: disable=C2801
160 def data(self):
161 """Return the underlying tensor data, preserving the DTensorBase subclass."""
162 return Tensor.data.__get__(self, type(self))
164 @data.setter
165 # pylint: disable=C2801
166 def data(self, value):
167 """Set the underlying tensor data, extracting the local shard if a DTensor is given."""
168 local_value = value.to_local() if isinstance(value, DTensorBase) else value
169 with _DisableMsDispatchMode():
170 Tensor.data.__set__(self, local_value)
171 Tensor.data.__set__(self._local_tensor, local_value)
173 # pylint: disable=W0212
174 def set_data(self, data, slice_shape=False):
175 """
176 Set shape/dtype/storage for dtensor and local tensor.
178 Args:
179 data (Tensor): New tensor payload.
180 slice_shape (bool): Kept for MindSpore `Parameter.set_data` API
181 compatibility. Static-graph slicing semantics are not used by
182 hyper_parallel, so this flag is accepted but ignored.
183 """
184 _ = slice_shape
185 if not isinstance(data, Tensor):
186 raise ValueError(f"The data type {type(data)} is not Tensor")
187 if data.has_init:
188 data.init_data()
189 data = data.to(self.device)
190 if isinstance(data, DTensorBase):
191 self._local_tensor._update_data(data.to_local())
192 self._device_mesh = data.device_mesh
193 self._placements = data.placements
194 self._layout = data.layout
195 self._update_data(self._local_tensor)
196 return
198 self._local_tensor._update_data(data)
199 self._update_data(data)
201 @property
202 def has_init(self):
203 """
204 Property to check if the initialization state is set in the local tensor.
206 Returns:
207 bool: True if the local tensor has the 'has_init' attribute, False otherwise.
208 """
209 if not hasattr(self._local_tensor, "has_init"):
210 return False
211 return self._local_tensor.has_init
213 @property
214 def init(self):
215 """
216 Property to get the initialization value from the local tensor.
218 Returns:
219 Any: The initialization value stored in the local tensor if the 'init' attribute exists;
220 None if the 'init' attribute is not present in the local tensor.
221 """
222 if not hasattr(self._local_tensor, "init"):
223 return None
224 return self._local_tensor.init
226 @init.setter
227 def init(self, init_value):
228 """
229 Setter for the initialization value, which assigns the value to the local tensor's 'init' attribute.
231 Args:
232 init_value: The value to be set as the initialization value in the local tensor.
233 """
234 self._local_tensor.init = init_value
236 @property
237 def local_param_info(self):
238 """
239 Property to get the param_info value from the local tensor.
241 Returns:
242 Any: The param_info value stored in the local tensor if the 'param_info' attribute exists;
243 None if the 'param_info' attribute is not present in the local tensor.
244 """
245 if not hasattr(self._local_tensor, "param_info"):
246 return None
247 return self._local_tensor.param_info
249 @local_param_info.setter
250 def local_param_info(self, local_param_info_value):
251 """
252 Setter for local_param_info value, which assigns the value to the local tensor's 'param_info' attribute.
254 Args:
255 local_param_info_value: The value to be set as the param_info value in the local tensor.
256 """
257 self._local_tensor.param_info = local_param_info_value
259 def _alias_placements(self):
260 """Return alias_placements from layout, falling back to _placements."""
261 if hasattr(self, '_layout') and self._layout is not None:
262 return self._layout.alias_placements
263 return self._placements
265 def to(self, *args, **kwargs):
266 """Move the DTensor to a different device or dtype.
268 This method overrides the base Tensor.to() to properly reconstruct
269 a DTensor with device_mesh and placements preserved. Uses _make_subclass
270 to avoid issues with Parameter subclasses that don't accept extra kwargs.
272 Args:
273 *args: Arguments passed to the underlying tensor's to() method.
274 **kwargs: Keyword arguments for the tensor conversion.
276 Returns:
277 DTensorBase: A new DTensor with the converted local tensor.
278 """
279 new_local = self._local_tensor.to(*args, **kwargs)
280 new_dt = Tensor._make_subclass(type(self), new_local)
281 new_dt.__init_data__(new_local, self._device_mesh, self._alias_placements())
282 return new_dt