Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / fully_shard / utils.py: 66%

1# Copyright 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 helpers, policies and mesh metadata for fully_shard. 

16 

17The fully_shard core used to reach process-group management, module traversal and 

18tensor plumbing through the platform abstraction layer. Only the torch backend is 

19supported now, so those implementations live here and every caller calls them 

20directly. 

21""" 

22# pylint: disable=C9006,C9007 

23from __future__ import annotations 

24 

25from collections import OrderedDict 

26from collections.abc import Mapping 

27from dataclasses import dataclass, fields, replace 

28from typing import Any, Callable, Dict, Optional 

29 

30import torch 

31from torch import Tensor, nn 

32from torch.nn.utils.rnn import PackedSequence 

33 

34from hyper_parallel.core.dtensor.device_mesh import DeviceMesh 

35from hyper_parallel.core.dtensor.placement_types import Placement 

36from hyper_parallel.core.utils.communication import get_device_handle 

37from hyper_parallel.core.utils.communication import get_group_local_rank 

38 

39DType = torch.dtype 

40 

41 

42# --------------------------------------------------------------------------- 

43# Module traversal 

44# --------------------------------------------------------------------------- 

45 

46def get_modules(module: nn.Module): 

47 """Return every sub-module contained in ``module``. 

48 

49 Args: 

50 module: The root module to traverse. 

51 

52 Returns: 

53 An iterator over the module tree. 

54 """ 

55 return module.modules() 

56 

57 

58def parameters_dict(cell: nn.Module): 

59 """Return every named parameter registered by the module tree. 

60 

61 Args: 

62 cell: The root module to traverse. 

63 

64 Returns: 

65 An iterator of ``(name, parameter)`` pairs. 

66 """ 

67 return cell.named_parameters() 

68 

69 

70def buffers_dict(cell: nn.Module): 

71 """Return every named buffer registered by the module tree. 

72 

73 Args: 

74 cell: The root module to traverse. 

75 

76 Returns: 

77 An iterator of ``(name, buffer)`` pairs. 

78 """ 

79 return cell.named_buffers() 

80 

81 

82# --------------------------------------------------------------------------- 

83# Tensor plumbing 

84# --------------------------------------------------------------------------- 

85 

86def load_into_param(param: Tensor, data: Tensor) -> None: 

87 """Write ``data`` into ``param``, materialising a meta local tensor first. 

88 

89 Args: 

90 param: The destination parameter, possibly a DTensor. 

91 data: The data to write into it. 

92 """ 

93 from hyper_parallel.core.dtensor.dtensor import DTensor # pylint: disable=C0415 

94 

95 if isinstance(param, DTensor): 

96 local = param._local_tensor # pylint: disable=protected-access 

97 if local.is_meta: 

98 orig_requires_grad = param.requires_grad 

99 param._local_tensor = data # pylint: disable=protected-access 

100 if data.requires_grad != orig_requires_grad: 

101 param.requires_grad_(orig_requires_grad) 

102 else: 

103 local.copy_(data) 

104 else: 

105 param.copy_(data) 

106 

107 

108def cast_fp_tensor(dtype: DType, x: Any) -> Any: 

109 """Cast a floating-point tensor to ``dtype``, leaving anything else alone. 

110 

111 Args: 

112 dtype: The target dtype. 

113 x: The candidate tensor. 

114 

115 Returns: 

116 Any: The cast tensor, or ``x`` unchanged when it is not a floating-point tensor. 

117 """ 

118 if not isinstance(x, Tensor) or not torch.is_floating_point(x) or x.dtype == dtype: 

119 return x 

120 return x.to(dtype) 

121 

122 

123def apply_to_tensors(fn: Callable[[Tensor], Any], container: Any) -> Any: 

124 """Recursively apply ``fn`` to every tensor inside ``container``. 

125 

126 Handles tensors, dataclasses, ordered and plain dicts, named tuples, 

127 lists, tuples, sets and packed sequences. 

128 

129 Args: 

130 fn: The callable applied to each tensor. 

131 container: The structure to walk. 

132 

133 Returns: 

134 Any: A structure of the same shape with every tensor replaced by ``fn``'s result. 

135 """ 

136 

137 def apply(x): 

138 if isinstance(x, Tensor): 

139 return fn(x) 

140 if hasattr(x, "__dataclass_fields__"): 

141 dc = replace(x) 

142 changes = {f.name: apply(getattr(dc, f.name)) for f in fields(dc)} 

143 return replace(dc, **changes) 

144 if isinstance(x, OrderedDict): 

145 od = x.__class__() 

146 for key, value in x.items(): 

147 od[key] = apply(value) 

148 return od 

149 if isinstance(x, PackedSequence): 

150 return x._replace(data=apply(x.data)) 

151 if isinstance(x, dict): 

152 return {key: apply(value) for key, value in x.items()} 

153 if isinstance(x, tuple) and hasattr(x, "_asdict") and hasattr(x, "_fields"): 

154 res = (apply(el) for el in x) 

155 return type(x)(*res) 

156 if isinstance(x, (list, tuple, set)): 

157 return type(x)(apply(el) for el in x) 

158 return x 

159 

160 return apply(container) 

161 

162 

163def profiler_record(name: str): 

164 """Return the torch profiler annotation context for ``name``. 

165 

166 Args: 

167 name: Label shown in profiler traces for the enclosed region. 

168 

169 Returns: 

170 A context manager that records the region when the profiler is active. 

171 """ 

172 return torch.profiler.record_function(name) 

173 

174 

175# --------------------------------------------------------------------------- 

176# Gradient-ready stream 

177# --------------------------------------------------------------------------- 

178# The handle and its stream are per-process singletons: at most one gradient 

179# reduction is in flight across the whole process, so the state belongs at 

180# module scope rather than on any scheduler or parameter object. 

181_GRAD_READY_STATE: Dict[str, Any] = { 

182 "handle": None, 

183 "post_process": None, 

184 "stream": None, 

185} 

186 

187 

188def _process_pending_grad_handle() -> None: 

189 """Wait for the in-flight gradient handle and run its post-process callback.""" 

190 handle = _GRAD_READY_STATE["handle"] 

191 if handle is None: 

192 return 

193 handle.wait() 

194 post_process = _GRAD_READY_STATE["post_process"] 

195 if post_process is not None: 

196 post_process() 

197 

198 

199def _grad_ready_stream(): 

200 """Return the stream used to order gradient-handle waits, creating it lazily.""" 

201 if _GRAD_READY_STATE["stream"] is None: 

202 _GRAD_READY_STATE["stream"] = get_device_handle().Stream() 

203 return _GRAD_READY_STATE["stream"] 

204 

205 

206def grad_ready_stream(): 

207 """Return the context manager that switches to the gradient-ready stream. 

208 

209 Returns: 

210 The device stream context manager guarding the synchronisation stream. 

211 """ 

212 return get_device_handle().stream(_grad_ready_stream()) 

213 

214 

215def set_grad_reduce_handle(handle: Any, post_process: Optional[Callable[[], None]] = None) -> None: 

216 """Record a new in-flight gradient reduction handle. 

217 

218 Any previously recorded handle is waited on first, so at most one handle is 

219 ever outstanding. 

220 

221 Args: 

222 handle: The async work handle returned by the reduction collective. 

223 post_process: Callback run once ``handle`` completes, or ``None``. 

224 """ 

225 with grad_ready_stream(): 

226 _process_pending_grad_handle() 

227 _GRAD_READY_STATE["handle"] = handle 

228 _GRAD_READY_STATE["post_process"] = post_process 

229 

230 

231def wait_grad_handle() -> None: 

232 """Block until the in-flight gradient reduction handle completes and clear it.""" 

233 handle = _GRAD_READY_STATE["handle"] 

234 if handle is None: 

235 return 

236 with grad_ready_stream(): 

237 _process_pending_grad_handle() 

238 sync_event = _grad_ready_stream().record_event() 

239 sync_event.wait() 

240 _GRAD_READY_STATE["handle"] = None 

241 _GRAD_READY_STATE["post_process"] = None 

242 

243 

244# --------------------------------------------------------------------------- 

245# Policies and mesh metadata 

246# --------------------------------------------------------------------------- 

247 

248@dataclass 

249class MixedPrecisionPolicy: 

250 """ 

251 Configures mixed precision training for HSDP. 

252 

253 This policy controls data type casting during forward/backward computation 

254 and gradient reduction, enabling memory savings and potential speedups. 

255 

256 Attributes: 

257 param_dtype: Data type for parameter computation. If None, uses original dtype. 

258 reduce_dtype: Data type for gradient reduction. If None, uses param_dtype. 

259 output_dtype: Data type for module outputs. If None, no casting applied. 

260 custom_params: Per-parameter overrides that take precedence over this policy's 

261 ``param_dtype`` and ``reduce_dtype`` defaults. 

262 """ 

263 param_dtype: Optional[DType] = None 

264 reduce_dtype: Optional[DType] = None 

265 output_dtype: Optional[DType] = None 

266 cast_forward_inputs: bool = True 

267 apply_grad_on_fp32_main_grad: bool = False 

268 custom_params: Optional[Mapping[nn.Parameter, Mapping[str, Optional[DType]]]] = None 

269 

270 def __post_init__(self) -> None: 

271 """Validate overrides and index dtype configurations by source Parameter ID.""" 

272 if self.custom_params is None: 

273 return 

274 self._check_custom_param() 

275 self.custom_params = { 

276 id(parameter): dict(custom_dtypes) 

277 for parameter, custom_dtypes in self.custom_params.items() 

278 } 

279 

280 def _check_custom_param(self) -> None: 

281 """Validate configured parameter keys and dtype values without changing the mapping.""" 

282 if not isinstance(self.custom_params, Mapping): 

283 raise ValueError("custom_params must map Parameter objects to dtype dictionaries") 

284 for parameter, custom_dtypes in self.custom_params.items(): 

285 if not isinstance(parameter, nn.Parameter): 

286 raise ValueError("custom_params keys must be Parameter objects") 

287 if not isinstance(custom_dtypes, Mapping): 

288 raise ValueError("custom_params values must be dtype dictionaries") 

289 if custom_dtypes.keys() - {"param_dtype", "reduce_dtype"}: 

290 raise ValueError("custom_params only supports param_dtype and reduce_dtype overrides") 

291 for dtype_name, dtype in custom_dtypes.items(): 

292 if dtype is not None and not isinstance(dtype, torch.dtype): 

293 raise ValueError(f"custom_params {dtype_name} must be torch.dtype or None, got {dtype!r}") 

294 

295 

296@dataclass 

297class OffloadPolicy: 

298 """ 

299 Base class for offload policies. 

300 

301 This represents no offloading and serves as the default policy. 

302 Subclass this to implement custom offload strategies. 

303 """ 

304 

305 

306@dataclass 

307class CPUOffloadPolicy(OffloadPolicy): 

308 """ 

309 Offloads sharded parameters and gradients to CPU memory. 

310 

311 When enabled, sharded parameters are kept on CPU and copied to device 

312 before all-gather. Gradients are copied back to CPU after backward. 

313 This reduces NPU memory usage at the cost of additional data transfers. 

314 

315 Attributes: 

316 pin_memory: If True, pins CPU memory for faster H2D/D2H transfers 

317 and enables overlap with computation. Disable if CPU memory 

318 is constrained. (Default: True) 

319 """ 

320 pin_memory: bool = True 

321 

322 

323@dataclass 

324class CommFusionPolicy(): 

325 enable_comm_fusion: bool = False 

326 comm_fusion_zero_copy: bool = False 

327 

328 

329@dataclass 

330class DataParallelMeshInfo: 

331 mesh: DeviceMesh 

332 shard_mesh_dim: Optional[int] = None 

333 replicate_mesh_dim: Optional[int] = None 

334 

335 def __post_init__(self): 

336 if self.shard_mesh_dim is None and self.replicate_mesh_dim is None: 

337 raise AssertionError( 

338 "At least one of shard_mesh_dim and replicate_mesh_dim must not be None" 

339 ) 

340 

341 

342@dataclass 

343class FSDPMeshInfo(DataParallelMeshInfo): 

344 def __post_init__(self): 

345 super().__post_init__() 

346 if self.shard_mesh_dim is None: 

347 raise AssertionError("Expects non-None shard_mesh_dim") 

348 self.shard_mesh_size: int = self.mesh.mesh_shape[self.shard_mesh_dim] 

349 self.shard_process_group = self.mesh.get_group(self.shard_mesh_dim) 

350 self.shard_mesh_rank: int = get_group_local_rank(self.shard_process_group) 

351 

352 

353@dataclass 

354class DDPMeshInfo(DataParallelMeshInfo): 

355 def __post_init__(self): 

356 super().__post_init__() 

357 if self.replicate_mesh_dim is None: 

358 raise AssertionError("Expects non-None replicate_mesh_dim") 

359 self.replicate_mesh_size: int = self.mesh.mesh_shape[self.replicate_mesh_dim] 

360 self.replicate_process_group = self.mesh.get_group(self.replicate_mesh_dim) 

361 self.replicate_mesh_rank: int = get_group_local_rank(self.replicate_process_group) 

362 

363 

364@dataclass 

365class HSDPMeshInfo(FSDPMeshInfo, DDPMeshInfo): 

366 # pylint: disable=W0246 

367 def __post_init__(self): 

368 # Calls `FSDPMeshInfo` -> `DDPMeshInfo` -> `DataParallelMeshInfo` 

369 super().__post_init__() 

370 

371 

372@dataclass(frozen=True) 

373class SourceShardMetaInfo: 

374 """Describe a parameter's source TP/EP layout before fully_shard.""" 

375 

376 mesh: DeviceMesh 

377 placements: tuple[Placement, ...] 

378 origin_is_dtensor: bool = False 

379 

380 

381__all__ = [ 

382 # Module traversal and tensor plumbing. 

383 "DType", 

384 "get_modules", 

385 "parameters_dict", 

386 "buffers_dict", 

387 "load_into_param", 

388 "cast_fp_tensor", 

389 "apply_to_tensors", 

390 "profiler_record", 

391 # Grad-reduce handle choreography. 

392 "grad_ready_stream", 

393 "set_grad_reduce_handle", 

394 "wait_grad_handle", 

395 # Precision/offload policies and mesh metadata. 

396 "MixedPrecisionPolicy", 

397 "OffloadPolicy", 

398 "CPUOffloadPolicy", 

399 "CommFusionPolicy", 

400 "DataParallelMeshInfo", 

401 "FSDPMeshInfo", 

402 "DDPMeshInfo", 

403 "HSDPMeshInfo", 

404 "SourceShardMetaInfo", 

405 # Re-exported from :mod:`hyper_parallel.core.utils.communication` because 

406 # callers here and in the tests reach them through this module. 

407 "get_device_handle", 

408 "get_group_local_rank", 

409]