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

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 """Return the underlying Tensor's data view, bypassing DTensor wrappers.""" 

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

249 

250 @data.setter 

251 # pylint: disable=C2801 

252 def data(self, value): 

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

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

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

256 # and only rebinds _local_tensor through DTensor dispatch. 

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

258 Tensor.data.__set__(self, local_value) 

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

260 

261 @property 

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

263 """ 

264 Get the data type of this tensor. 

265 

266 Returns: 

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

268 """ 

269 return self._local_tensor.dtype 

270 

271 @property 

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

273 """ 

274 Get the shape of this tensor. 

275 

276 Returns: 

277 torch.Size: The shape of the tensor. 

278 """ 

279 return self._local_tensor.shape 

280 

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

282 """ 

283 Convert this tensor to the specified dtype. 

284 

285 Args: 

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

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

288 

289 Returns: 

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

291 """ 

292 if dtype is None: 

293 return self._local_tensor.type() 

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

295 return self.__class__( 

296 new_local, 

297 device_mesh=self._device_mesh, 

298 placements=self._alias_placements(), 

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

300 ) 

301 

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

303 """ 

304 Get the size of this tensor. 

305 

306 Args: 

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

308 

309 Returns: 

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

311 """ 

312 return self._local_tensor.size(dim) 

313 

314 @property 

315 def ndim(self) -> int: 

316 """ 

317 Get the number of dimensions of this tensor. 

318 

319 Returns: 

320 int: The number of dimensions. 

321 """ 

322 return self._local_tensor.ndim 

323 

324 def data_ptr(self) -> int: 

325 """ 

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

327 

328 Returns: 

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

330 """ 

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

332 return self._local_tensor.data_ptr() 

333 

334 def numel(self) -> int: 

335 """ 

336 Get the total number of elements in this tensor. 

337 

338 Returns: 

339 int: The total number of elements. 

340 """ 

341 return self._local_tensor.numel() 

342 

343 # ====================== Auxiliary print ====================== 

344 def _alias_placements(self): 

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

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

347 return self._layout.alias_placements 

348 return self._placements 

349 

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

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

352 

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

354 a DTensor with device_mesh and placements preserved. Uses _make_subclass 

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

356 

357 Args: 

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

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

360 

361 Returns: 

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

363 """ 

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

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

366 new_dt.__init_data__( 

367 new_local, 

368 self._device_mesh, 

369 self._alias_placements(), 

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

371 ) 

372 return new_dt 

373 

374 def __repr__(self) -> str: 

375 return ( 

376 f"DTensor(\n" 

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

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

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

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

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

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

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

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

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

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

387 f")" 

388 )