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, layout=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 """ 

35 if isinstance(local_tensor, DTensorBase): 

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

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

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

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

40 return t 

41 

42 if device_mesh is None: 

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

44 if placements is None: 

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

46 

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

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

49 t.__init_data__(local_tensor, device_mesh, placements, layout) 

50 return t 

51 

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

53 @classmethod 

54 def __torch_function__( 

55 cls, 

56 func: torch._C._FunctionBase, 

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

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

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

60 ) -> Any: 

61 """ 

62 Override PyTorch's __torch_function__ to intercept tensor operations. 

63 

64 This method dispatches operations through the distributed operator dispatcher 

65 to handle DTensor-specific layout inference and redistribution. 

66 

67 Args: 

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

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

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

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

72 

73 Returns: 

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

75 """ 

76 kwargs = kwargs or {} 

77 # pylint: disable=C0415 

78 from hyper_parallel.core.shard._op_dispatch import _OP_DISPATCHER 

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

80 return out 

81 

82 @property 

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

84 """ 

85 Get the gradient tensor of the local tensor. 

86 

87 Returns: 

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

89 """ 

90 return self._local_tensor.grad 

91 

92 @grad.setter 

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

94 """ 

95 Set the gradient tensor for the local tensor. 

96 

97 Args: 

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

99 """ 

100 self._local_tensor.grad = value 

101 

102 @property 

103 def requires_grad(self) -> bool: 

104 """ 

105 Check if gradient computation is enabled for this tensor. 

106 

107 Returns: 

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

109 """ 

110 return self._local_tensor.requires_grad 

111 

112 @requires_grad.setter 

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

114 """ 

115 Enable or disable gradient computation for this tensor. 

116 

117 Args: 

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

119 """ 

120 self._local_tensor.requires_grad_(value) 

121 # Sync DTensor wrapper's requires_grad 

122 super().requires_grad_(value) 

123 

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

125 """ 

126 Enable or disable gradient computation in-place. 

127 

128 Args: 

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

130 

131 Returns: 

132 DTensorBase: Self for method chaining. 

133 """ 

134 self._local_tensor.requires_grad_(requires_grad) 

135 super().requires_grad_(requires_grad) 

136 return self 

137 

138 @property 

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

140 """ 

141 Get the gradient function that created this tensor. 

142 

143 Returns: 

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

145 """ 

146 return self._local_tensor.grad_fn 

147 

148 def grad_zero_(self): 

149 """ 

150 Zero out the gradient tensor in-place. 

151 

152 Returns: 

153 DTensorBase: Self for method chaining. 

154 """ 

155 if self._local_tensor.grad is not None: 

156 self._local_tensor.grad.zero_() 

157 return self 

158 

159 def detach(self): 

160 """ 

161 Create a detached DTensor that does not require gradient. 

162 

163 Returns: 

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

165 """ 

166 detached_local = self._local_tensor.detach() 

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

168 

169 def detach_(self): 

170 """ 

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

172 

173 Returns: 

174 DTensorBase: Self for method chaining. 

175 """ 

176 self._local_tensor.detach_() 

177 super().detach_() 

178 return self 

179 

180 # ====================== Computation graph related overrides ====================== 

181 @property 

182 def is_leaf(self) -> bool: 

183 """ 

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

185 

186 Returns: 

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

188 """ 

189 return self._local_tensor.is_leaf 

190 

191 @property 

192 def retains_grad(self) -> bool: 

193 """ 

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

195 

196 Returns: 

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

198 """ 

199 return self._local_tensor.retains_grad 

200 

201 @retains_grad.setter 

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

203 """ 

204 Enable or disable gradient retention for this tensor. 

205 

206 Args: 

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

208 """ 

209 self._local_tensor.retains_grad_(value) 

210 

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

212 """ 

213 Compute the gradients for this tensor. 

214 

215 Args: 

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

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

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

219 """ 

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

221 

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

223 @property 

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

225 """ 

226 Get the device on which this tensor is stored. 

227 

228 Returns: 

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

230 """ 

231 return self._local_tensor.device 

232 

233 @property 

234 # pylint: disable=C2801 

235 def data(self): 

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

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

238 

239 @data.setter 

240 # pylint: disable=C2801 

241 def data(self, value): 

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

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

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

245 # and only rebinds _local_tensor through DTensor dispatch. 

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

247 Tensor.data.__set__(self, local_value) 

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

249 

250 @property 

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

252 """ 

253 Get the data type of this tensor. 

254 

255 Returns: 

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

257 """ 

258 return self._local_tensor.dtype 

259 

260 @property 

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

262 """ 

263 Get the shape of this tensor. 

264 

265 Returns: 

266 torch.Size: The shape of the tensor. 

267 """ 

268 return self._local_tensor.shape 

269 

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

271 """ 

272 Convert this tensor to the specified dtype. 

273 

274 Args: 

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

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

277 

278 Returns: 

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

280 """ 

281 if dtype is None: 

282 return self._local_tensor.type() 

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

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

285 

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

287 """ 

288 Get the size of this tensor. 

289 

290 Args: 

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

292 

293 Returns: 

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

295 """ 

296 return self._local_tensor.size(dim) 

297 

298 @property 

299 def ndim(self) -> int: 

300 """ 

301 Get the number of dimensions of this tensor. 

302 

303 Returns: 

304 int: The number of dimensions. 

305 """ 

306 return self._local_tensor.ndim 

307 

308 def data_ptr(self) -> int: 

309 """ 

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

311 

312 Returns: 

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

314 """ 

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

316 return self._local_tensor.data_ptr() 

317 

318 def numel(self) -> int: 

319 """ 

320 Get the total number of elements in this tensor. 

321 

322 Returns: 

323 int: The total number of elements. 

324 """ 

325 return self._local_tensor.numel() 

326 

327 # ====================== Auxiliary print ====================== 

328 def _alias_placements(self): 

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

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

331 return self._layout.alias_placements 

332 return self._placements 

333 

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

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

336 

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

338 a DTensor with device_mesh and placements preserved. Uses _make_subclass 

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

340 

341 Args: 

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

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

344 

345 Returns: 

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

347 """ 

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

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

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

351 return new_dt 

352 

353 def __repr__(self) -> str: 

354 return ( 

355 f"DTensor(\n" 

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

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

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

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

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

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

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

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

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

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

366 f")" 

367 )