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

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

16import functools 

17from typing import Any, List, Optional, Tuple, Union 

18 

19from hyper_parallel.platform import get_platform 

20from hyper_parallel.core.dtensor.device_mesh import DeviceMesh 

21from hyper_parallel.core.fully_shard.hsdp_utils import ( 

22 FSDPSchedulerState, 

23 HSDPConfigV2, 

24 get_managed_modules_parameters, 

25 get_hsdp_state 

26) 

27from hyper_parallel.tools.logging import get_logger 

28 

29logger = get_logger("FSDP") 

30 

31platform = get_platform() 

32 

33 

34class HSDPSchedulerContext: 

35 """HSDPSchedulerContext""" 

36 

37 def __init__(self) -> None: 

38 # Currently only record is_last_backward flag for scheduler context. 

39 self.is_last_backward: bool = True 

40 # flag to identify "root_module" 

41 self.root_module = None 

42 

43 

44class HSDPSchedulerV2: 

45 """HSDPScheduler is used to scheduler hsdp""" 

46 root_bp_state = False 

47 

48 

49 def __init__(self, cell: Union[platform.Module, Tuple[platform.Module, ...]], mesh, 

50 reshard_after_forward, shard_placement_fn, 

51 mp_policy, offload_policy, ignored_params, replicate_params, device, comm_fusion, 

52 comm_fusion_zero_copy=False): 

53 """init hsdp scheduler. 

54 

55 Args: 

56 cell: A single platform.Module or tuple of platform.Module to manage as one FSDP unit. 

57 """ 

58 self.modules = (cell,) if isinstance(cell, platform.Module) else tuple(cell) 

59 self.cell = self.modules[0] 

60 self.mesh: DeviceMesh = mesh 

61 self.reshard_after_forward = reshard_after_forward 

62 self.shard_placement_fn = shard_placement_fn 

63 self.mp_policy = mp_policy 

64 self.offload_policy = offload_policy 

65 self.ignored_params = ignored_params 

66 self.replicate_params = replicate_params 

67 self.device = device 

68 self.scheduler_state = None 

69 self.forward_prefetch_cells = [] 

70 self.backward_prefetch_cells = [] 

71 self._backup_forward_fetch = None 

72 # Flag to identify root module. 

73 self._is_root = False 

74 # module and its all sub-modules share one same 'HSDPSchedulerContext' 

75 self.scheduler_ctx = HSDPSchedulerContext() 

76 # When ``fully_shard`` is given multiple root modules, forward pre/post hooks coordinate 

77 # so unshard / PostBackward / reshard run once per forward (aligned with PyTorch FSDP2). 

78 self._fsdp_group_post_pending: Optional[set] = set() if len(self.modules) > 1 else None 

79 self.config = HSDPConfigV2( 

80 mesh, 

81 reshard_after_forward, 

82 shard_placement_fn, 

83 mp_policy, 

84 offload_policy, 

85 ignored_params, 

86 replicate_params, 

87 comm_fusion=comm_fusion, 

88 comm_fusion_zero_copy=comm_fusion_zero_copy, 

89 ) 

90 self._init_platform() 

91 self._new_cell_state() 

92 self._register_hooks() 

93 

94 def _init_platform(self): 

95 """Initialize the platform.""" 

96 raise NotImplementedError("HSDPScheduler subclasses must implement _init_platform") 

97 

98 def _new_cell_state(self): 

99 """Create a new cell state.""" 

100 raise NotImplementedError("HSDPScheduler subclasses must implement _new_cell_state") 

101 

102 def _register_hooks(self): 

103 """Register hooks.""" 

104 raise NotImplementedError("HSDPScheduler subclasses must implement _register_hooks.") 

105 

106 def _register_forward_backward_hooks(self): 

107 """Register module forward and backward hook.""" 

108 raise NotImplementedError("HSDPScheduler subclasses must implement _register_forward_backward_hooks.") 

109 

110 def _get_managed_params(self): 

111 """Return deduplicated parameters from all managed modules.""" 

112 return get_managed_modules_parameters(self.modules, self.ignored_params) 

113 

114 def set_reshard_after_forward(self, reshard_after_forward: bool) -> None: 

115 """Set reshard_after_forward flag. 

116 

117 Args: 

118 reshard_after_forward: Whether to reshard parameters after forward. 

119 """ 

120 if not isinstance(reshard_after_forward, bool): 

121 raise ValueError(f"reshard_after_forward should be a bool, got {type(reshard_after_forward)}") 

122 self.reshard_after_forward = reshard_after_forward 

123 self.config.reshard_after_forward = reshard_after_forward 

124 

125 def set_reshard_after_backward(self, reshard_after_backward: bool) -> None: 

126 """Set reshard_after_backward flag. 

127 

128 Args: 

129 reshard_after_backward: Whether to reshard after backward completes. 

130 """ 

131 if not isinstance(reshard_after_backward, bool): 

132 raise ValueError(f"reshard_after_backward should be a bool, got {type(reshard_after_backward)}") 

133 if self.hsdp_state is not None: 

134 self.hsdp_state.reshard_after_backward = reshard_after_backward 

135 

136 def set_requires_all_reduce(self, requires_all_reduce: bool) -> None: 

137 """Set requires_all_reduce flag. 

138 

139 Args: 

140 requires_all_reduce: Whether this unit participates in all-reduce. 

141 """ 

142 if not isinstance(requires_all_reduce, bool): 

143 raise ValueError(f"requires_all_reduce should be a bool, got {type(requires_all_reduce)}") 

144 if self.hsdp_state is not None: 

145 self.hsdp_state.requires_all_reduce = requires_all_reduce 

146 

147 def set_requires_grad_sync(self, requires_grad_sync: bool) -> None: 

148 """Set flag controlling whether gradients are synchronized. 

149 

150 Args: 

151 requires_grad_sync: When True, enable grad sync for this scheduler. 

152 """ 

153 if not isinstance(requires_grad_sync, bool): 

154 raise ValueError(f"requires_grad_sync should be a bool, got {type(requires_grad_sync)}") 

155 self.hsdp_state.set_requires_grad_sync(requires_grad_sync) 

156 

157 # pylint: disable=W0613 

158 def _hsdp_forward_pre_hook(self, cell, args, kwargs): 

159 """Forward pre hook to unsharded parameter for forward process.""" 

160 logger.debug("hook=forward_pre enter module=%s", self.hsdp_state) 

161 if self.scheduler_state == FSDPSchedulerState.PRE_BACKWARD: 

162 logger.debug("hook=forward_pre skip module=%s reason=pre_backward", self.hsdp_state) 

163 return args, kwargs 

164 if HSDPSchedulerV2.root_bp_state: 

165 self._disable_forward_prefetch_for_recompute() 

166 if self.scheduler_ctx.root_module is None: 

167 self.scheduler_ctx.root_module = self.cell 

168 self._is_root = True 

169 for _, module in platform.get_cells_and_names(self.scheduler_ctx.root_module): 

170 from hyper_parallel.core.fully_shard.api import HSDPModule # pylint: disable=C0415 

171 if isinstance(module, HSDPModule): 

172 submod_scheduler = getattr(module, "hsdp_scheduler", None) 

173 if submod_scheduler and submod_scheduler.scheduler_ctx is not self.scheduler_ctx: 

174 submod_scheduler.scheduler_ctx = self.scheduler_ctx 

175 

176 if not self._is_root and not self.hsdp_state.module_name: 

177 for module_name, module in platform.get_cells_and_names(self.scheduler_ctx.root_module): 

178 if module == self.cell: 

179 self.hsdp_state.module_name = module_name 

180 break 

181 self.scheduler_state = FSDPSchedulerState.PRE_FORWARD 

182 self._init_params_fqn() 

183 self._lazy_init_all_states() 

184 if self.mp_policy.cast_forward_inputs and self.mp_policy.param_dtype: 

185 cast_fn = functools.partial(self.platform.cast_fp_tensor, self.mp_policy.param_dtype) 

186 args = self.platform.apply_to_tensors(cast_fn, args) 

187 kwargs = self.platform.apply_to_tensors(cast_fn, kwargs) 

188 with self.platform.profiler_record(f"pre_forward unshard:{self.hsdp_state.module_name}"): 

189 logger.debug("hook=forward_pre action=unshard module=%s", self.hsdp_state) 

190 self.hsdp_state.unshard() 

191 for prefetch_cell in self.forward_prefetch_cells: 

192 prefetch_state = prefetch_cell.hsdp_scheduler.hsdp_state 

193 with self.platform.profiler_record(f"pre_forward prefetch:" 

194 f"{prefetch_state.module_name}"): 

195 logger.debug( 

196 "hook=forward_pre action=prefetch module=%s target=%s", 

197 self.hsdp_state, 

198 prefetch_state, 

199 ) 

200 prefetch_state.prefetch() 

201 return args, kwargs 

202 

203 def _lazy_init_all_states(self): 

204 if self._is_root and self.scheduler_ctx.root_module is not None: 

205 for _, module in platform.get_cells_and_names(self.scheduler_ctx.root_module): 

206 hsdp_state = get_hsdp_state(module) 

207 if hsdp_state: 

208 hsdp_state.lazy_init() 

209 

210 def _init_params_fqn(self): # pylint: disable=W0212 

211 if not self._is_root or self.scheduler_ctx.root_module is None: 

212 return 

213 # Build a map from original (sharded) parameter tensor → hsdp_param wrapper, 

214 # covering both sharded hsdp_params and replicate_params. 

215 param_to_hsdp_param = {} 

216 for _, module in platform.get_cells_and_names(self.scheduler_ctx.root_module): 

217 hsdp_state = get_hsdp_state(module) 

218 if hsdp_state is None: 

219 continue 

220 for hsdp_param in hsdp_state._iter_managed_params(): # pylint: disable=W0212 

221 orig_param = hsdp_param.sharded_param 

222 # Shared parameters: keep only the first mapping to preserve the 

223 # first-seen FQN (consistent with the deduplication in _init_hsdp_params). 

224 if orig_param not in param_to_hsdp_param: 

225 param_to_hsdp_param[orig_param] = hsdp_param 

226 

227 # Walk the full parameter tree and assign FQNs; skip params already seen 

228 # (shared-parameter deduplication: first name wins). 

229 visited_params = set() 

230 for param_name, parameter in platform.parameters_dict(self.scheduler_ctx.root_module): 

231 if parameter in visited_params: 

232 continue 

233 visited_params.add(parameter) 

234 hsdp_param = param_to_hsdp_param.get(parameter) 

235 if hsdp_param is not None: 

236 hsdp_param._param_fqn = param_name # pylint: disable=W0212 

237 

238 # pylint: disable=W0613, R1710 

239 def _hsdp_forward_hook(self, cell, inputs, outputs): 

240 """Forward hook to shard parameter for saving memory.""" 

241 logger.debug("hook=forward enter module=%s", self.hsdp_state) 

242 if self.scheduler_state == FSDPSchedulerState.PRE_BACKWARD: 

243 logger.debug("hook=forward skip module=%s reason=pre_backward", self.hsdp_state) 

244 return 

245 self.scheduler_state = FSDPSchedulerState.FORWARD 

246 if self.reshard_after_forward: 

247 with self.platform.profiler_record(f"forward reshard:{self.hsdp_state.module_name}"): 

248 logger.debug("hook=forward action=reshard module=%s", self.hsdp_state) 

249 self.hsdp_state.shard(shard_replicate=False) 

250 if self.mp_policy.output_dtype is not None: 

251 outputs = self.platform.apply_to_tensors( 

252 functools.partial(self.platform.cast_fp_tensor, self.mp_policy.output_dtype), 

253 outputs, 

254 ) 

255 return outputs 

256 

257 # pylint: disable=W0613 

258 def _hsdp_backward_pre_hook(self, cell, grad_outputs): 

259 """Backward pre hook to unsharded parameter for backward process.""" 

260 logger.debug("hook=backward_pre enter module=%s", self.hsdp_state) 

261 self.scheduler_state = FSDPSchedulerState.PRE_BACKWARD 

262 if self.reshard_after_forward: 

263 with self.platform.profiler_record(f"pre_backward unshard:{self.hsdp_state.module_name}"): 

264 logger.debug("hook=backward_pre action=unshard module=%s", self.hsdp_state) 

265 self.hsdp_state.unshard(unshard_replicate=False) 

266 for prefetch_cell in self.backward_prefetch_cells: 

267 prefetch_state = prefetch_cell.hsdp_scheduler.hsdp_state 

268 with self.platform.profiler_record(f"pre_backward prefetch:" 

269 f"{prefetch_state.module_name}"): 

270 logger.debug( 

271 "hook=backward_pre action=prefetch module=%s target=%s", 

272 self.hsdp_state, 

273 prefetch_state, 

274 ) 

275 prefetch_state.prefetch(unshard_replicate=False) 

276 

277 # pylint: disable=W0613 

278 def _hsdp_backward_hook(self, cell, grad_inputs, grad_outputs): 

279 """Backward hook to shard parameter for optimizer process or saving memory.""" 

280 logger.debug("hook=backward_hook enter module=%s", self.hsdp_state) 

281 self.scheduler_state = FSDPSchedulerState.BACKWARD 

282 with self.platform.profiler_record(f"post_backward:{self.hsdp_state.module_name}"): 

283 logger.debug("hook=backward_hook action=post_backward module=%s", self.hsdp_state) 

284 self.hsdp_state.post_backward() 

285 if self._fsdp_group_post_pending is not None: 

286 self._fsdp_group_post_pending.clear() 

287 

288 # pylint: disable=W0613 

289 @staticmethod 

290 def _grouped_forward_pre_hook_skip(cell, args, kwargs): 

291 """Return value when grouped pre-forward should not run (first module already did). 

292 

293 Default matches MindSpore Cell forward pre-hooks (explicit ``(args, kwargs)``). 

294 ``TorchHSDPSchedulerV2`` overrides this to return ``None`` (``nn.Module`` idiom). 

295 """ 

296 return args, kwargs 

297 

298 @staticmethod 

299 def _grouped_forward_post_hook_skip(outputs): 

300 """Return value when grouped post-forward is deferred to a later module in the group. 

301 

302 Default returns ``outputs`` (MindSpore). ``TorchHSDPSchedulerV2`` overrides to ``None``. 

303 """ 

304 return outputs 

305 

306 def _grouped_forward_pre_hook(self, cell, args, kwargs): 

307 """Run FSDP pre-forward only for the first module in the group (PyTorch FSDP2-aligned).""" 

308 pending = self._fsdp_group_post_pending 

309 if pending is None: 

310 return self._forward_pre_hook(cell, args, kwargs) 

311 if len(pending) == 0: 

312 pending.update(self.modules) 

313 return self._forward_pre_hook(cell, args, kwargs) 

314 return self._grouped_forward_pre_hook_skip(cell, args, kwargs) 

315 

316 def _make_grouped_forward_post_hook(self, mod): 

317 """Build post-forward hook: last module in the group runs reshard + output backward hooks.""" 

318 

319 def grouped_post_hook(cell, inputs, outputs): 

320 pending = self._fsdp_group_post_pending 

321 if pending is None: 

322 return self._forward_hook(cell, inputs, outputs) 

323 pending.discard(mod) 

324 if len(pending) == 0: 

325 return self._forward_hook(cell, inputs, outputs) 

326 return self._grouped_forward_post_hook_skip(outputs) 

327 

328 return grouped_post_hook 

329 

330 def set_forward_prefetch_cells(self, hsdp_cell_list: List[Any]) -> None: 

331 """Set cells prefetched during forward. 

332 

333 Args: 

334 hsdp_cell_list: HSDP cells to prefetch ahead of forward. 

335 """ 

336 self.forward_prefetch_cells = hsdp_cell_list 

337 

338 def set_backward_prefetch_cells(self, hsdp_cell_list: List[Any]) -> None: 

339 """Set cells prefetched during backward. 

340 

341 Args: 

342 hsdp_cell_list: HSDP cells to prefetch ahead of backward. 

343 """ 

344 self.backward_prefetch_cells = hsdp_cell_list 

345 

346 def _disable_forward_prefetch_for_recompute(self) -> None: 

347 """Temporarily disable forward prefetch during activation recompute.""" 

348 self._backup_forward_fetch = self.forward_prefetch_cells 

349 self.forward_prefetch_cells = [] 

350 

351 def _restore_forward_prefetch_after_recompute(self) -> bool: 

352 """Restore forward prefetch list after a recompute forward hook finishes.""" 

353 if self._backup_forward_fetch is None: 

354 return False 

355 self.forward_prefetch_cells = self._backup_forward_fetch 

356 self._backup_forward_fetch = None 

357 return True