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

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

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

33 if isinstance(local_tensor, DTensorBase): 

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

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

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

37 t.__init_data__(local_tensor._local_tensor, local_tensor.device_mesh, copy_placements) 

38 return t 

39 

40 if device_mesh is None: 

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

42 if placements is None: 

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

44 

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

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

47 t.__init_data__(local_tensor, device_mesh, placements) 

48 return t 

49 

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

51 @classmethod 

52 def __torch_function__( 

53 cls, 

54 func: torch._C._FunctionBase, 

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

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

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

58 ) -> Any: 

59 """ 

60 Override PyTorch's __torch_function__ to intercept tensor operations. 

61 

62 This method dispatches operations through the distributed operator dispatcher 

63 to handle DTensor-specific layout inference and redistribution. 

64 

65 Args: 

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

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

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

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

70 

71 Returns: 

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

73 """ 

74 kwargs = kwargs or {} 

75 # pylint: disable=C0415 

76 from hyper_parallel.core.shard._op_dispatch import _OP_DISPATCHER 

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

78 return out 

79 

80 @property 

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

82 """ 

83 Get the gradient tensor of the local tensor. 

84 

85 Returns: 

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

87 """ 

88 return self._local_tensor.grad 

89 

90 @grad.setter 

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

92 """ 

93 Set the gradient tensor for the local tensor. 

94 

95 Args: 

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

97 """ 

98 self._local_tensor.grad = value 

99 

100 @property 

101 def requires_grad(self) -> bool: 

102 """ 

103 Check if gradient computation is enabled for this tensor. 

104 

105 Returns: 

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

107 """ 

108 return self._local_tensor.requires_grad 

109 

110 @requires_grad.setter 

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

112 """ 

113 Enable or disable gradient computation for this tensor. 

114 

115 Args: 

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

117 """ 

118 self._local_tensor.requires_grad_(value) 

119 # Sync DTensor wrapper's requires_grad 

120 super().requires_grad_(value) 

121 

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

123 """ 

124 Enable or disable gradient computation in-place. 

125 

126 Args: 

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

128 

129 Returns: 

130 DTensorBase: Self for method chaining. 

131 """ 

132 self._local_tensor.requires_grad_(requires_grad) 

133 super().requires_grad_(requires_grad) 

134 return self 

135 

136 @property 

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

138 """ 

139 Get the gradient function that created this tensor. 

140 

141 Returns: 

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

143 """ 

144 return self._local_tensor.grad_fn 

145 

146 def grad_zero_(self): 

147 """ 

148 Zero out the gradient tensor in-place. 

149 

150 Returns: 

151 DTensorBase: Self for method chaining. 

152 """ 

153 if self._local_tensor.grad is not None: 

154 self._local_tensor.grad.zero_() 

155 return self 

156 

157 def detach(self): 

158 """ 

159 Create a detached DTensor that does not require gradient. 

160 

161 Returns: 

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

163 """ 

164 detached_local = self._local_tensor.detach() 

165 return self.__class__(detached_local, device_mesh=self._device_mesh, placements=self._alias_placements()) 

166 

167 def detach_(self): 

168 """ 

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

170 

171 Returns: 

172 DTensorBase: Self for method chaining. 

173 """ 

174 self._local_tensor.detach_() 

175 super().detach_() 

176 return self 

177 

178 # ====================== Computation graph related overrides ====================== 

179 @property 

180 def is_leaf(self) -> bool: 

181 """ 

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

183 

184 Returns: 

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

186 """ 

187 return self._local_tensor.is_leaf 

188 

189 @property 

190 def retains_grad(self) -> bool: 

191 """ 

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

193 

194 Returns: 

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

196 """ 

197 return self._local_tensor.retains_grad 

198 

199 @retains_grad.setter 

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

201 """ 

202 Enable or disable gradient retention for this tensor. 

203 

204 Args: 

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

206 """ 

207 self._local_tensor.retains_grad_(value) 

208 

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

210 """ 

211 Compute the gradients for this tensor. 

212 

213 Args: 

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

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

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

217 """ 

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

219 

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

221 @property 

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

223 """ 

224 Get the device on which this tensor is stored. 

225 

226 Returns: 

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

228 """ 

229 return self._local_tensor.device 

230 

231 @property 

232 # pylint: disable=C2801 

233 def data(self): 

234 """Return the underlying Tensor's data view, bypassing DTensor wrappers.""" 

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

236 

237 @data.setter 

238 # pylint: disable=C2801 

239 def data(self, value): 

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

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

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

243 # and only rebinds _local_tensor through DTensor dispatch. 

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

245 Tensor.data.__set__(self, local_value) 

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

247 

248 @property 

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

250 """ 

251 Get the data type of this tensor. 

252 

253 Returns: 

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

255 """ 

256 return self._local_tensor.dtype 

257 

258 @property 

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

260 """ 

261 Get the shape of this tensor. 

262 

263 Returns: 

264 torch.Size: The shape of the tensor. 

265 """ 

266 return self._local_tensor.shape 

267 

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

269 """ 

270 Convert this tensor to the specified dtype. 

271 

272 Args: 

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

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

275 

276 Returns: 

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

278 """ 

279 if dtype is None: 

280 return self._local_tensor.type() 

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

282 return self.__class__(new_local, device_mesh=self._device_mesh, placements=self._alias_placements()) 

283 

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

285 """ 

286 Get the size of this tensor. 

287 

288 Args: 

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

290 

291 Returns: 

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

293 """ 

294 return self._local_tensor.size(dim) 

295 

296 @property 

297 def ndim(self) -> int: 

298 """ 

299 Get the number of dimensions of this tensor. 

300 

301 Returns: 

302 int: The number of dimensions. 

303 """ 

304 return self._local_tensor.ndim 

305 

306 def data_ptr(self) -> int: 

307 """ 

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

309 

310 Returns: 

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

312 """ 

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

314 return self._local_tensor.data_ptr() 

315 

316 def numel(self) -> int: 

317 """ 

318 Get the total number of elements in this tensor. 

319 

320 Returns: 

321 int: The total number of elements. 

322 """ 

323 return self._local_tensor.numel() 

324 

325 # ====================== Auxiliary print ====================== 

326 def _alias_placements(self): 

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

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

329 return self._layout.alias_placements 

330 return self._placements 

331 

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

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

334 

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

336 a DTensor with device_mesh and placements preserved. Uses _make_subclass 

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

338 

339 Args: 

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

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

342 

343 Returns: 

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

345 """ 

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

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

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

349 return new_dt 

350 

351 def __repr__(self) -> str: 

352 return ( 

353 f"DTensor(\n" 

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

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

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

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

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

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

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

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

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

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

364 f")" 

365 )