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

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 dtensor base""" 

16from typing import Tuple, Dict, Any, Optional 

17import torch 

18from torch import Tensor 

19 

20 

21class DTensorBase(Tensor): 

22 """torch dtensor base""" 

23 

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

25 """ 

26 Create a new DTensorBase instance. 

27 

28 Args: 

29 local_tensor: The local tensor shard or another DTensorBase instance. 

30 device_mesh: The device mesh describing the device topology. 

31 placements: The placement strategy for each mesh dimension. 

32 layout: Optional pre-built Layout reused directly by ``__init_data__`` 

33 (skips ``_build_layout``; see ``DTensor.from_local_with_layout``). 

34 shape: Optional logical global tensor shape. 

35 """ 

36 if isinstance(local_tensor, DTensorBase): 

37 # Copy from existing DTensorBase — use alias_placements to preserve multi-axis ordering 

38 t = Tensor._make_subclass(cls, local_tensor._local_tensor, local_tensor._local_tensor.requires_grad) 

39 copy_placements = local_tensor.layout.alias_placements if local_tensor.layout else local_tensor.placements 

40 t.__init_data__( 

41 local_tensor._local_tensor, 

42 local_tensor.device_mesh, 

43 copy_placements, 

44 shape=getattr(local_tensor, "_global_shape", None), 

45 ) 

46 return t 

47 

48 if device_mesh is None: 

49 raise ValueError("device_mesh is None, must provide a DeviceMesh instance") 

50 if placements is None: 

51 raise ValueError("placements is None, must provide placements") 

52 

53 # Create Tensor subclass instance, sharing local_tensor's underlying storage 

54 t = Tensor._make_subclass(cls, local_tensor, local_tensor.requires_grad) 

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

56 return t 

57 

58 # pylint: disable=W0613, G.NAM.05 

59 @classmethod 

60 def __torch_function__( 

61 cls, 

62 func: torch._C._FunctionBase, 

63 types: Tuple[type, ...], 

64 args: Tuple[Any, ...] = (), 

65 kwargs: Optional[Dict[str, Any]] = None 

66 ) -> Any: 

67 """ 

68 Override PyTorch's __torch_function__ to intercept tensor operations. 

69 

70 This method dispatches operations through the distributed operator dispatcher 

71 to handle DTensor-specific layout inference and redistribution. 

72 

73 Args: 

74 func (torch._C._FunctionBase): The PyTorch function being called. 

75 types (Tuple[type, ...]): The types of tensors involved in the operation. 

76 args (Tuple[Any, ...]): Positional arguments passed to the function. 

77 kwargs (Optional[Dict[str, Any]]): Keyword arguments passed to the function. 

78 

79 Returns: 

80 Any: The result of the dispatched operation, typically a DTensor or tuple of DTensors. 

81 """ 

82 kwargs = kwargs or {} 

83 # pylint: disable=C0415 

84 from hyper_parallel.core.shard._op_dispatch import _OP_DISPATCHER 

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

86 return out 

87 

88 @property 

89 def grad(self) -> Optional[Tensor]: 

90 """ 

91 Get the gradient tensor of the local tensor. 

92 

93 Returns: 

94 Optional[Tensor]: The gradient tensor, or None if no gradient is set. 

95 """ 

96 return self._local_tensor.grad 

97 

98 @grad.setter 

99 def grad(self, value: Optional[Tensor]) -> None: 

100 """ 

101 Set the gradient tensor for the local tensor. 

102 

103 Args: 

104 value (Optional[Tensor]): The gradient tensor to set, or None to clear. 

105 """ 

106 self._local_tensor.grad = value 

107 

108 @property 

109 def requires_grad(self) -> bool: 

110 """ 

111 Check if gradient computation is enabled for this tensor. 

112 

113 Returns: 

114 bool: True if gradients should be computed for this tensor. 

115 """ 

116 return self._local_tensor.requires_grad 

117 

118 @requires_grad.setter 

119 def requires_grad(self, value: bool) -> None: 

120 """ 

121 Enable or disable gradient computation for this tensor. 

122 

123 Args: 

124 value (bool): True to enable gradient computation, False to disable. 

125 """ 

126 self._local_tensor.requires_grad_(value) 

127 # Sync DTensor wrapper's requires_grad 

128 super().requires_grad_(value) 

129 

130 def requires_grad_(self, requires_grad: bool = True): 

131 """ 

132 Enable or disable gradient computation in-place. 

133 

134 Args: 

135 requires_grad (bool): True to enable gradient computation. Default: True. 

136 

137 Returns: 

138 DTensorBase: Self for method chaining. 

139 """ 

140 self._local_tensor.requires_grad_(requires_grad) 

141 super().requires_grad_(requires_grad) 

142 return self 

143 

144 @property 

145 def grad_fn(self) -> Optional[torch.autograd.Function]: 

146 """ 

147 Get the gradient function that created this tensor. 

148 

149 Returns: 

150 Optional[torch.autograd.Function]: The gradient function, or None if not applicable. 

151 """ 

152 return self._local_tensor.grad_fn 

153 

154 def grad_zero_(self): 

155 """ 

156 Zero out the gradient tensor in-place. 

157 

158 Returns: 

159 DTensorBase: Self for method chaining. 

160 """ 

161 if self._local_tensor.grad is not None: 

162 self._local_tensor.grad.zero_() 

163 return self 

164 

165 def detach(self): 

166 """ 

167 Create a detached DTensor that does not require gradient. 

168 

169 Returns: 

170 DTensorBase: A new DTensor with the same data but detached from the computation graph. 

171 """ 

172 detached_local = self._local_tensor.detach() 

173 return self.__class__( 

174 detached_local, 

175 device_mesh=self._device_mesh, 

176 placements=self._alias_placements(), 

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

178 ) 

179 

180 def detach_(self): 

181 """ 

182 Detach this tensor from the computation graph in-place. 

183 

184 Returns: 

185 DTensorBase: Self for method chaining. 

186 """ 

187 self._local_tensor.detach_() 

188 super().detach_() 

189 return self 

190 

191 # ====================== Computation graph related overrides ====================== 

192 @property 

193 def is_leaf(self) -> bool: 

194 """ 

195 Check if this tensor is a leaf node in the computation graph. 

196 

197 Returns: 

198 bool: True if this is a leaf tensor (created by user, not by any operation). 

199 """ 

200 return self._local_tensor.is_leaf 

201 

202 @property 

203 def retains_grad(self) -> bool: 

204 """ 

205 Check if this tensor retains its gradient during backward pass. 

206 

207 Returns: 

208 bool: True if gradients are retained for non-leaf tensors. 

209 """ 

210 return self._local_tensor.retains_grad 

211 

212 @retains_grad.setter 

213 def retains_grad(self, value: bool) -> None: 

214 """ 

215 Enable or disable gradient retention for this tensor. 

216 

217 Args: 

218 value (bool): True to enable gradient retention. 

219 """ 

220 self._local_tensor.retains_grad_(value) 

221 

222 def backward(self, gradient=None, retain_graph=None, create_graph=False) -> None: 

223 """ 

224 Compute the gradients for this tensor. 

225 

226 Args: 

227 gradient (Optional[Tensor]): The gradient of the loss w.r.t. this tensor. 

228 retain_graph (Optional[bool]): Whether to retain the computation graph. 

229 create_graph (bool): Whether to create a graph of the gradient computation. 

230 """ 

231 self._local_tensor.backward(gradient, retain_graph, create_graph) 

232 

233 # ====================== Metadata related overrides (sync with local_tensor) ====================== 

234 @property 

235 def device(self) -> torch.device: 

236 """ 

237 Get the device on which this tensor is stored. 

238 

239 Returns: 

240 torch.device: The device object (e.g., 'cuda:0', 'cpu'). 

241 """ 

242 return self._local_tensor.device 

243 

244 @property 

245 # pylint: disable=C2801 

246 def data(self): 

247 """ 

248 Directory get Tensor.data relative storage. 

249 After DTensor object created, there are two reference on underlying storage. (.data and ._local_tensor) 

250 If not using DisableTorchFunctionSubclass, Tensor.__get__ will goto '__torch_function__', 

251 and finally return '._local_tensor' storage. 

252 """ 

253 with getattr(torch, "_C").DisableTorchFunctionSubclass(): 

254 # Directory get Tensor.data relative storage. 

255 # If not using DisableTorchFunctionSubclass, Tensor.__get__ will goto __torch_funtin 

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

257 

258 @data.setter 

259 # pylint: disable=C2801 

260 def data(self, value): 

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

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

263 # Tensor.data.__set__ on a Tensor subclass otherwise enters __torch_function__ 

264 # and only rebinds _local_tensor through DTensor dispatch. 

265 with getattr(torch, "_C").DisableTorchFunctionSubclass(): 

266 Tensor.data.__set__(self, local_value) 

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

268 

269 @property 

270 def dtype(self) -> torch.dtype: 

271 """ 

272 Get the data type of this tensor. 

273 

274 Returns: 

275 torch.dtype: The data type (e.g., torch.float32, torch.int64). 

276 """ 

277 return self._local_tensor.dtype 

278 

279 @property 

280 def shape(self) -> torch.Size: 

281 """ 

282 Get the shape of this tensor. 

283 

284 Returns: 

285 torch.Size: The shape of the tensor. 

286 """ 

287 return self._local_tensor.shape 

288 

289 def type(self, dtype=None, non_blocking=False): 

290 """ 

291 Convert this tensor to the specified dtype. 

292 

293 Args: 

294 dtype (Optional[torch.dtype]): The target dtype. If None, returns the current type string. 

295 non_blocking (bool): Whether to perform the operation asynchronously. Default: False. 

296 

297 Returns: 

298 Union[str, DTensorBase]: The type string if dtype is None, otherwise a new DTensor. 

299 """ 

300 if dtype is None: 

301 return self._local_tensor.type() 

302 new_local = self._local_tensor.to(dtype=dtype, non_blocking=non_blocking) 

303 return self.__class__( 

304 new_local, 

305 device_mesh=self._device_mesh, 

306 placements=self._alias_placements(), 

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

308 ) 

309 

310 def size(self, dim: Optional[int] = None): 

311 """ 

312 Get the size of this tensor. 

313 

314 Args: 

315 dim (Optional[int]): The dimension to query. If None, returns the full shape. 

316 

317 Returns: 

318 Union[torch.Size, int]: The shape or size along a specific dimension. 

319 """ 

320 return self._local_tensor.size(dim) 

321 

322 @property 

323 def ndim(self) -> int: 

324 """ 

325 Get the number of dimensions of this tensor. 

326 

327 Returns: 

328 int: The number of dimensions. 

329 """ 

330 return self._local_tensor.ndim 

331 

332 def data_ptr(self) -> int: 

333 """ 

334 Get the pointer to the data storage of the local tensor. 

335 

336 Returns: 

337 int: The memory address of the tensor's data. 

338 """ 

339 # Force return local_tensor's data pointer (ensure address consistency) 

340 return self._local_tensor.data_ptr() 

341 

342 def numel(self) -> int: 

343 """ 

344 Get the total number of elements in this tensor. 

345 

346 Returns: 

347 int: The total number of elements. 

348 """ 

349 return self._local_tensor.numel() 

350 

351 # ====================== Auxiliary print ====================== 

352 def _alias_placements(self): 

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

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

355 return self._layout.alias_placements 

356 return self._placements 

357 

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

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

360 

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

362 a DTensor with device_mesh and placements preserved. Uses _make_subclass 

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

364 

365 Args: 

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

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

368 

369 Returns: 

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

371 """ 

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

373 new_dt = Tensor._make_subclass(type(self), new_local, new_local.requires_grad) 

374 new_dt.__init_data__( 

375 new_local, 

376 self._device_mesh, 

377 self._alias_placements(), 

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

379 ) 

380 return new_dt 

381 

382 def __repr__(self) -> str: 

383 return ( 

384 f"DTensor(\n" 

385 f" local_tensor={self._local_tensor},\n" 

386 f" device_mesh={self._device_mesh},\n" 

387 f" placements={self._placements},\n" 

388 f" layout={getattr(self, '_layout', None)},\n" 

389 f" device={self.device},\n" 

390 f" dtype={self.dtype},\n" 

391 f" requires_grad={self.requires_grad},\n" 

392 f" grad={self.grad},\n" 

393 f" is_leaf={self.is_leaf},\n" 

394 f" data_ptr={self.data_ptr()}\n" 

395 f")" 

396 )