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

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 dataclasses import dataclass, fields, replace 

27from typing import Any, Callable, Dict, Optional 

28 

29import torch 

30from torch import Tensor, nn 

31from torch.nn.utils.rnn import PackedSequence 

32 

33from hyper_parallel.core.dtensor.device_mesh import DeviceMesh 

34from hyper_parallel.core.dtensor.placement_types import Placement 

35from hyper_parallel.core.utils.communication import get_device_handle 

36from hyper_parallel.core.utils.communication import get_group_local_rank 

37 

38DType = torch.dtype 

39 

40 

41# --------------------------------------------------------------------------- 

42# Module traversal 

43# --------------------------------------------------------------------------- 

44 

45def get_modules(module: nn.Module): 

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

47 

48 Args: 

49 module: The root module to traverse. 

50 

51 Returns: 

52 An iterator over the module tree. 

53 """ 

54 return module.modules() 

55 

56 

57def parameters_dict(cell: nn.Module): 

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

59 

60 Args: 

61 cell: The root module to traverse. 

62 

63 Returns: 

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

65 """ 

66 return cell.named_parameters() 

67 

68 

69def buffers_dict(cell: nn.Module): 

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

71 

72 Args: 

73 cell: The root module to traverse. 

74 

75 Returns: 

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

77 """ 

78 return cell.named_buffers() 

79 

80 

81# --------------------------------------------------------------------------- 

82# Tensor plumbing 

83# --------------------------------------------------------------------------- 

84 

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

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

87 

88 Args: 

89 param: The destination parameter, possibly a DTensor. 

90 data: The data to write into it. 

91 """ 

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

93 

94 if isinstance(param, DTensor): 

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

96 if local.is_meta: 

97 orig_requires_grad = param.requires_grad 

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

99 if data.requires_grad != orig_requires_grad: 

100 param.requires_grad_(orig_requires_grad) 

101 else: 

102 local.copy_(data) 

103 else: 

104 param.copy_(data) 

105 

106 

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

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

109 

110 Args: 

111 dtype: The target dtype. 

112 x: The candidate tensor. 

113 

114 Returns: 

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

116 """ 

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

118 return x 

119 return x.to(dtype) 

120 

121 

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

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

124 

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

126 lists, tuples, sets and packed sequences. 

127 

128 Args: 

129 fn: The callable applied to each tensor. 

130 container: The structure to walk. 

131 

132 Returns: 

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

134 """ 

135 

136 def apply(x): 

137 if isinstance(x, Tensor): 

138 return fn(x) 

139 if hasattr(x, "__dataclass_fields__"): 

140 dc = replace(x) 

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

142 return replace(dc, **changes) 

143 if isinstance(x, OrderedDict): 

144 od = x.__class__() 

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

146 od[key] = apply(value) 

147 return od 

148 if isinstance(x, PackedSequence): 

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

150 if isinstance(x, dict): 

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

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

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

154 return type(x)(*res) 

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

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

157 return x 

158 

159 return apply(container) 

160 

161 

162def profiler_record(name: str): 

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

164 

165 Args: 

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

167 

168 Returns: 

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

170 """ 

171 return torch.profiler.record_function(name) 

172 

173 

174# --------------------------------------------------------------------------- 

175# Gradient-ready stream 

176# --------------------------------------------------------------------------- 

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

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

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

180_GRAD_READY_STATE: Dict[str, Any] = { 

181 "handle": None, 

182 "post_process": None, 

183 "stream": None, 

184} 

185 

186 

187def _process_pending_grad_handle() -> None: 

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

189 handle = _GRAD_READY_STATE["handle"] 

190 if handle is None: 

191 return 

192 handle.wait() 

193 post_process = _GRAD_READY_STATE["post_process"] 

194 if post_process is not None: 

195 post_process() 

196 

197 

198def _grad_ready_stream(): 

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

200 if _GRAD_READY_STATE["stream"] is None: 

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

202 return _GRAD_READY_STATE["stream"] 

203 

204 

205def grad_ready_stream(): 

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

207 

208 Returns: 

209 The device stream context manager guarding the synchronisation stream. 

210 """ 

211 return get_device_handle().stream(_grad_ready_stream()) 

212 

213 

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

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

216 

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

218 ever outstanding. 

219 

220 Args: 

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

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

223 """ 

224 with grad_ready_stream(): 

225 _process_pending_grad_handle() 

226 _GRAD_READY_STATE["handle"] = handle 

227 _GRAD_READY_STATE["post_process"] = post_process 

228 

229 

230def wait_grad_handle() -> None: 

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

232 handle = _GRAD_READY_STATE["handle"] 

233 if handle is None: 

234 return 

235 with grad_ready_stream(): 

236 _process_pending_grad_handle() 

237 sync_event = _grad_ready_stream().record_event() 

238 sync_event.wait() 

239 _GRAD_READY_STATE["handle"] = None 

240 _GRAD_READY_STATE["post_process"] = None 

241 

242 

243# --------------------------------------------------------------------------- 

244# Policies and mesh metadata 

245# --------------------------------------------------------------------------- 

246 

247@dataclass 

248class MixedPrecisionPolicy: 

249 """ 

250 Configures mixed precision training for HSDP. 

251 

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

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

254 

255 Attributes: 

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

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

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

259 """ 

260 param_dtype: Optional[DType] = None 

261 reduce_dtype: Optional[DType] = None 

262 output_dtype: Optional[DType] = None 

263 cast_forward_inputs: bool = True 

264 apply_grad_on_fp32_main_grad: bool = False 

265 

266 

267@dataclass 

268class OffloadPolicy: 

269 """ 

270 Base class for offload policies. 

271 

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

273 Subclass this to implement custom offload strategies. 

274 """ 

275 

276 

277@dataclass 

278class CPUOffloadPolicy(OffloadPolicy): 

279 """ 

280 Offloads sharded parameters and gradients to CPU memory. 

281 

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

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

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

285 

286 Attributes: 

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

288 and enables overlap with computation. Disable if CPU memory 

289 is constrained. (Default: True) 

290 """ 

291 pin_memory: bool = True 

292 

293 

294@dataclass 

295class CommFusionPolicy(): 

296 enable_comm_fusion: bool = False 

297 comm_fusion_zero_copy: bool = False 

298 

299 

300@dataclass 

301class DataParallelMeshInfo: 

302 mesh: DeviceMesh 

303 shard_mesh_dim: Optional[int] = None 

304 replicate_mesh_dim: Optional[int] = None 

305 

306 def __post_init__(self): 

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

308 raise AssertionError( 

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

310 ) 

311 

312 

313@dataclass 

314class FSDPMeshInfo(DataParallelMeshInfo): 

315 def __post_init__(self): 

316 super().__post_init__() 

317 if self.shard_mesh_dim is None: 

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

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

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

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

322 

323 

324@dataclass 

325class DDPMeshInfo(DataParallelMeshInfo): 

326 def __post_init__(self): 

327 super().__post_init__() 

328 if self.replicate_mesh_dim is None: 

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

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

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

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

333 

334 

335@dataclass 

336class HSDPMeshInfo(FSDPMeshInfo, DDPMeshInfo): 

337 # pylint: disable=W0246 

338 def __post_init__(self): 

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

340 super().__post_init__() 

341 

342 

343@dataclass(frozen=True) 

344class SourceShardMetaInfo: 

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

346 

347 mesh: DeviceMesh 

348 placements: tuple[Placement, ...] 

349 origin_is_dtensor: bool = False 

350 

351 

352__all__ = [ 

353 # Module traversal and tensor plumbing. 

354 "DType", 

355 "get_modules", 

356 "parameters_dict", 

357 "buffers_dict", 

358 "load_into_param", 

359 "cast_fp_tensor", 

360 "apply_to_tensors", 

361 "profiler_record", 

362 # Grad-reduce handle choreography. 

363 "grad_ready_stream", 

364 "set_grad_reduce_handle", 

365 "wait_grad_handle", 

366 # Precision/offload policies and mesh metadata. 

367 "MixedPrecisionPolicy", 

368 "OffloadPolicy", 

369 "CPUOffloadPolicy", 

370 "CommFusionPolicy", 

371 "DataParallelMeshInfo", 

372 "FSDPMeshInfo", 

373 "DDPMeshInfo", 

374 "HSDPMeshInfo", 

375 "SourceShardMetaInfo", 

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

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

378 "get_device_handle", 

379 "get_group_local_rank", 

380]