Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / mindspore / dtensor.py: 45%

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 

19 

20 

21class DTensorBase(Tensor): 

22 """ 

23 DTensorBase - Base class for distributed tensors in MindSpore. 

24 

25 This class extends Tensor to support distributed tensor operations with 

26 device mesh and placement specifications. 

27 """ 

28 

29 def __new__(cls, local_tensor, device_mesh=None, placements=None, layout=None, shape=None): 

30 """ 

31 Create a new DTensorBase instance. 

32 

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 shape: Optional logical global tensor shape. 

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, shape) 

52 return t 

53 

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 shape = getattr(src, "_global_shape", None) 

61 else: 

62 if local_tensor is None: 

63 raise ValueError( 

64 "DTensorBase: local_tensor must not be None when constructing from a raw tensor." 

65 ) 

66 if device_mesh is None: 

67 raise ValueError( 

68 "DTensorBase: device_mesh must be a DeviceMesh instance, got None." 

69 ) 

70 if placements is None: 

71 raise ValueError( 

72 "DTensorBase: placements must be a sequence of Placement objects, got None." 

73 ) 

74 

75 if local_tensor.has_init: 

76 local_tensor.init_device = npu_device 

77 else: 

78 dev = local_tensor.device 

79 if dev != "meta" and not dev.startswith(npu_device): 

80 local_tensor = local_tensor.to(npu_device) 

81 

82 t = Tensor._make_subclass(cls, local_tensor) 

83 t.__init_data__(local_tensor, device_mesh, placements, layout, shape) 

84 return t 

85 

86 def asnumpy(self): 

87 """ 

88 Numpy value of local tensor. 

89 """ 

90 return self._local_tensor.asnumpy() 

91 

92 def __str__(self): 

93 return str(self._local_tensor) 

94 

95 def __copy__(self): 

96 """ 

97 Create a shallow copy of the DTensorBase instance. 

98 

99 This method ensures that device_mesh and placements are correctly 

100 propagated when creating a copy (e.g., for optimizer states). 

101 """ 

102 # Get device_mesh and placements from layout (prefer alias_placements to preserve multi-axis ordering) 

103 device_mesh = getattr(self, '_device_mesh', None) 

104 placements = None 

105 

106 if hasattr(self, '_layout') and self._layout is not None: 

107 if device_mesh is None: 

108 device_mesh = self._layout.mesh 

109 placements = self._layout.alias_placements 

110 

111 if placements is None: 

112 placements = getattr(self, '_placements', None) 

113 

114 if device_mesh is None or placements is None: 

115 raise ValueError( 

116 "DTensorBase.__copy__: cannot copy without device_mesh and placements; " 

117 f"device_mesh={device_mesh!r}, placements={placements!r}. " 

118 "Ensure the tensor was constructed with a valid layout." 

119 ) 

120 

121 if self._local_tensor.has_init: 

122 obj = DTensorBase.__new__( 

123 type(self), 

124 initializer(self._local_tensor.init, self._local_tensor.shape, self._local_tensor.dtype), 

125 device_mesh, 

126 placements, 

127 shape=getattr(self, "_global_shape", None), 

128 ) 

129 else: 

130 obj = DTensorBase.__new__( 

131 type(self), 

132 self._local_tensor.clone(), 

133 device_mesh, 

134 placements, 

135 shape=getattr(self, "_global_shape", None), 

136 ) 

137 filtered_dict = {k: v for k, v in self.__dict__.items() if k != '_local_tensor'} 

138 obj.__dict__.update(filtered_dict) 

139 return obj 

140 

141 # pylint: disable=W0211, W0102, C0415, G.NAM.05 

142 def __fallback__(self, func, args={}, kwargs=None): 

143 if kwargs is None: 

144 kwargs = {} 

145 from hyper_parallel.core.shard._op_dispatch import _OP_DISPATCHER 

146 with NoFallbackGuard(): 

147 out = _OP_DISPATCHER.dispatch(func, args, kwargs) 

148 return out 

149 

150 # pylint: disable=W0212 

151 def _need_contiguous(self): 

152 """_need_contiguous""" 

153 return self._local_tensor._need_contiguous() 

154 

155 @property 

156 def device(self): 

157 """Device info for dtensor""" 

158 device_info = self._local_tensor.device 

159 return device_info.split(':', 1)[0] 

160 

161 @property 

162 # pylint: disable=C2801 

163 def data(self): 

164 """Return the underlying tensor data, preserving the DTensorBase subclass.""" 

165 return Tensor.data.__get__(self, type(self)) 

166 

167 @data.setter 

168 # pylint: disable=C2801 

169 def data(self, value): 

170 """Set the underlying tensor data, extracting the local shard if a DTensor is given.""" 

171 local_value = value.to_local() if isinstance(value, DTensorBase) else value 

172 with _DisableMsDispatchMode(): 

173 Tensor.data.__set__(self, local_value) 

174 Tensor.data.__set__(self._local_tensor, local_value) 

175 

176 # pylint: disable=W0212 

177 def set_data(self, data, slice_shape=False): 

178 """ 

179 Set shape/dtype/storage for dtensor and local tensor. 

180 

181 Args: 

182 data (Tensor): New tensor payload. 

183 slice_shape (bool): Kept for MindSpore `Parameter.set_data` API 

184 compatibility. Static-graph slicing semantics are not used by 

185 hyper_parallel, so this flag is accepted but ignored. 

186 """ 

187 _ = slice_shape 

188 if not isinstance(data, Tensor): 

189 raise ValueError(f"The data type {type(data)} is not Tensor") 

190 if data.has_init: 

191 data.init_data() 

192 data = data.to(self.device) 

193 if isinstance(data, DTensorBase): 

194 self._local_tensor._update_data(data.to_local()) 

195 self._device_mesh = data.device_mesh 

196 self._placements = data.placements 

197 self._layout = data.layout 

198 self._global_shape = getattr(data, "_global_shape", tuple(data.shape)) 

199 self._update_data(self._local_tensor) 

200 return 

201 

202 self._local_tensor._update_data(data) 

203 self._update_data(data) 

204 

205 @property 

206 def has_init(self): 

207 """ 

208 Property to check if the initialization state is set in the local tensor. 

209 

210 Returns: 

211 bool: True if the local tensor has the 'has_init' attribute, False otherwise. 

212 """ 

213 if not hasattr(self._local_tensor, "has_init"): 

214 return False 

215 return self._local_tensor.has_init 

216 

217 @property 

218 def init(self): 

219 """ 

220 Property to get the initialization value from the local tensor. 

221 

222 Returns: 

223 Any: The initialization value stored in the local tensor if the 'init' attribute exists; 

224 None if the 'init' attribute is not present in the local tensor. 

225 """ 

226 if not hasattr(self._local_tensor, "init"): 

227 return None 

228 return self._local_tensor.init 

229 

230 @init.setter 

231 def init(self, init_value): 

232 """ 

233 Setter for the initialization value, which assigns the value to the local tensor's 'init' attribute. 

234 

235 Args: 

236 init_value: The value to be set as the initialization value in the local tensor. 

237 """ 

238 self._local_tensor.init = init_value 

239 

240 @property 

241 def local_param_info(self): 

242 """ 

243 Property to get the param_info value from the local tensor. 

244 

245 Returns: 

246 Any: The param_info value stored in the local tensor if the 'param_info' attribute exists; 

247 None if the 'param_info' attribute is not present in the local tensor. 

248 """ 

249 if not hasattr(self._local_tensor, "param_info"): 

250 return None 

251 return self._local_tensor.param_info 

252 

253 @local_param_info.setter 

254 def local_param_info(self, local_param_info_value): 

255 """ 

256 Setter for local_param_info value, which assigns the value to the local tensor's 'param_info' attribute. 

257 

258 Args: 

259 local_param_info_value: The value to be set as the param_info value in the local tensor. 

260 """ 

261 self._local_tensor.param_info = local_param_info_value 

262 

263 def _alias_placements(self): 

264 """Return alias_placements from layout, falling back to _placements.""" 

265 if hasattr(self, '_layout') and self._layout is not None: 

266 return self._layout.alias_placements 

267 return self._placements 

268 

269 def to(self, *args, **kwargs): 

270 """Move the DTensor to a different device or dtype. 

271 

272 This method overrides the base Tensor.to() to properly reconstruct 

273 a DTensor with device_mesh and placements preserved. Uses _make_subclass 

274 to avoid issues with Parameter subclasses that don't accept extra kwargs. 

275 

276 Args: 

277 *args: Arguments passed to the underlying tensor's to() method. 

278 **kwargs: Keyword arguments for the tensor conversion. 

279 

280 Returns: 

281 DTensorBase: A new DTensor with the converted local tensor. 

282 """ 

283 new_local = self._local_tensor.to(*args, **kwargs) 

284 new_dt = Tensor._make_subclass(type(self), new_local) 

285 new_dt.__init_data__( 

286 new_local, 

287 self._device_mesh, 

288 self._alias_placements(), 

289 shape=getattr(self, "_global_shape", None), 

290 ) 

291 return new_dt