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

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): 

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 device: The device type (default: "Ascend"). 

38 """ 

39 npu_device = "Ascend" 

40 if isinstance(local_tensor, DTensorBase): 

41 src = local_tensor 

42 local_tensor = src.to_local() 

43 device_mesh = src.device_mesh 

44 placements = src._alias_placements() 

45 else: 

46 if local_tensor is None: 

47 raise ValueError( 

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

49 ) 

50 if device_mesh is None: 

51 raise ValueError( 

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

53 ) 

54 if placements is None: 

55 raise ValueError( 

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

57 ) 

58 

59 if local_tensor.has_init: 

60 local_tensor.init_device = npu_device 

61 else: 

62 dev = local_tensor.device 

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

64 local_tensor = local_tensor.to(npu_device) 

65 

66 t = Tensor._make_subclass(cls, local_tensor) 

67 t.__init_data__(local_tensor, device_mesh, placements) 

68 return t 

69 

70 def asnumpy(self): 

71 """ 

72 Numpy value of local tensor. 

73 """ 

74 return self._local_tensor.asnumpy() 

75 

76 def __str__(self): 

77 return str(self._local_tensor) 

78 

79 def __copy__(self): 

80 """ 

81 Create a shallow copy of the DTensorBase instance. 

82 

83 This method ensures that device_mesh and placements are correctly 

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

85 """ 

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

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

88 placements = None 

89 

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

91 if device_mesh is None: 

92 device_mesh = self._layout.mesh 

93 placements = self._layout.alias_placements 

94 

95 if placements is None: 

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

97 

98 if device_mesh is None or placements is None: 

99 raise ValueError( 

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

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

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

103 ) 

104 

105 if self._local_tensor.has_init: 

106 obj = DTensorBase.__new__( 

107 type(self), 

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

109 device_mesh, 

110 placements 

111 ) 

112 else: 

113 obj = DTensorBase.__new__( 

114 type(self), 

115 self._local_tensor.clone(), 

116 device_mesh, 

117 placements 

118 ) 

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

120 obj.__dict__.update(filtered_dict) 

121 return obj 

122 

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

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

125 if kwargs is None: 

126 kwargs = {} 

127 from hyper_parallel.core.shard._op_dispatch import _OP_DISPATCHER 

128 with NoFallbackGuard(): 

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

130 return out 

131 

132 # pylint: disable=W0212 

133 def _need_contiguous(self): 

134 """_need_contiguous""" 

135 return self._local_tensor._need_contiguous() 

136 

137 @property 

138 def device(self): 

139 """Device info for dtensor""" 

140 device_info = self._local_tensor.device 

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

142 

143 @property 

144 # pylint: disable=C2801 

145 def data(self): 

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

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

148 

149 @data.setter 

150 # pylint: disable=C2801 

151 def data(self, value): 

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

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

154 with _DisableMsDispatchMode(): 

155 Tensor.data.__set__(self, local_value) 

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

157 

158 # pylint: disable=W0212 

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

160 """ 

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

162 

163 Args: 

164 data (Tensor): New tensor payload. 

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

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

167 hyper_parallel, so this flag is accepted but ignored. 

168 """ 

169 _ = slice_shape 

170 if not isinstance(data, Tensor): 

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

172 if data.has_init: 

173 data.init_data() 

174 data = data.to(self.device) 

175 if isinstance(data, DTensorBase): 

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

177 self._device_mesh = data.device_mesh 

178 self._placements = data.placements 

179 self._layout = data.layout 

180 self._update_data(self._local_tensor) 

181 return 

182 

183 self._local_tensor._update_data(data) 

184 self._update_data(data) 

185 

186 @property 

187 def has_init(self): 

188 """ 

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

190 

191 Returns: 

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

193 """ 

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

195 return False 

196 return self._local_tensor.has_init 

197 

198 @property 

199 def init(self): 

200 """ 

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

202 

203 Returns: 

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

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

206 """ 

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

208 return None 

209 return self._local_tensor.init 

210 

211 @init.setter 

212 def init(self, init_value): 

213 """ 

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

215 

216 Args: 

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

218 """ 

219 self._local_tensor.init = init_value 

220 

221 @property 

222 def local_param_info(self): 

223 """ 

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

225 

226 Returns: 

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

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

229 """ 

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

231 return None 

232 return self._local_tensor.param_info 

233 

234 @local_param_info.setter 

235 def local_param_info(self, local_param_info_value): 

236 """ 

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

238 

239 Args: 

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

241 """ 

242 self._local_tensor.param_info = local_param_info_value 

243 

244 def _alias_placements(self): 

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

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

247 return self._layout.alias_placements 

248 return self._placements 

249 

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

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

252 

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

254 a DTensor with device_mesh and placements preserved. Uses _make_subclass 

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

256 

257 Args: 

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

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

260 

261 Returns: 

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

263 """ 

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

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

266 new_dt.__init_data__(new_local, self._device_mesh, self._alias_placements()) 

267 return new_dt