Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / torch / fully_shard / scheduler.py: 63%

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

16import functools 

17import inspect 

18from typing import Callable, List, ParamSpec, TypeVar 

19 

20import torch 

21from torch.autograd import Variable 

22from torch.utils._pytree import tree_flatten, tree_unflatten 

23 

24from hyper_parallel.core.dtensor.dtensor import DTensor 

25from hyper_parallel.tools.logging import get_logger 

26from hyper_parallel.core.fully_shard.hsdp_scheduler import HSDPSchedulerV2, FSDPSchedulerState 

27from hyper_parallel.core.fully_shard.utils import FSDPMeshInfo, DDPMeshInfo, HSDPMeshInfo 

28from hyper_parallel.platform.torch.fully_shard.hook_function import PostBackwardFunction 

29from hyper_parallel.platform.torch.fully_shard.state import TorchHSDPStateV2 

30from hyper_parallel.platform.torch.fully_shard.param_group import get_comm_ctx 

31from hyper_parallel.platform import get_platform 

32 

33logger = get_logger("FSDP") 

34 

35_P = ParamSpec("_P") 

36_R = TypeVar("_R") 

37 

38 

39def _dynamo_disable(func: Callable[_P, _R]) -> Callable[_P, _R]: 

40 """Disable Dynamo tracing while an FSDP runtime hook executes.""" 

41 

42 @functools.wraps(func) 

43 def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: 

44 return torch._dynamo.disable( 

45 func, 

46 recursive=True, 

47 )(*args, **kwargs) 

48 

49 return wrapper 

50 

51 

52class TorchHSDPSchedulerV2(HSDPSchedulerV2): 

53 """TorchHSDPScheduler is used to implement optimizer level.""" 

54 

55 def __init__(self, *args, **kwargs): 

56 """Initialize TorchHSDPSchedulerV2 and register forward/backward hooks.""" 

57 super().__init__(*args, **kwargs) 

58 

59 def _register_hooks(self): 

60 """Register hooks.""" 

61 self._register_forward_backward_hooks() 

62 

63 def _init_platform(self): 

64 """Initialize the platform.""" 

65 # pylint: disable=C0415 

66 from hyper_parallel.platform.torch.platform import TorchPlatform 

67 self.platform = get_platform() 

68 if not isinstance(self.platform, TorchPlatform): 

69 raise ValueError(f"TorchHSDPSchedulerV2 expect TorchPlatform, but got type: {type(self.platform)}") 

70 

71 def _new_cell_state(self): 

72 """Create a new cell state for torch.""" 

73 params = self._get_managed_params() 

74 if self.mesh is None: 

75 compat_meshes = [ 

76 param.device_mesh for param in params if isinstance(param, DTensor) 

77 ] 

78 compat_mesh = compat_meshes[0] if compat_meshes else None 

79 if compat_mesh is None: 

80 raise ValueError( 

81 "Cannot build fully_shard compatibility mesh_info " 

82 "without a DTensor parameter mesh." 

83 ) 

84 compat_mesh_hash = compat_mesh.to_hash() 

85 for param_mesh in compat_meshes[1:]: 

86 if param_mesh.to_hash() != compat_mesh_hash: 

87 raise ValueError( 

88 "fully_shard compatibility mode requires all DTensor parameters to share the same mesh." 

89 ) 

90 self.mesh_info = DDPMeshInfo(mesh=compat_mesh, replicate_mesh_dim=0) 

91 elif self.mesh.ndim == 1: 

92 self.mesh_info = FSDPMeshInfo(mesh=self.mesh, shard_mesh_dim=0) 

93 elif self.mesh.ndim == 2: 

94 self.mesh_info = HSDPMeshInfo(mesh=self.mesh, shard_mesh_dim=1, replicate_mesh_dim=0) 

95 else: 

96 raise ValueError( 

97 "fully_shard only supports explicit 1D DP/FSDP meshes or 2D HSDP meshes. " 

98 f"Got mesh.ndim={self.mesh.ndim}." 

99 ) 

100 self.hsdp_state = TorchHSDPStateV2( 

101 self.modules, self.mesh_info, self.config, self.platform, self.device 

102 ) 

103 

104 def _register_post_backward_hook(self, args, kwargs): 

105 """Wrap forward args/kwargs through PostBackwardFunction to register backward hook.""" 

106 if not torch.is_grad_enabled(): 

107 return args, kwargs 

108 args_list, args_spec = tree_flatten(args) 

109 kwargs_list, kwargs_spec = tree_flatten(kwargs) 

110 args_kwargs_list = list(args_list) + list(kwargs_list) 

111 inp_tensor_indices: List[int] = [] 

112 inp_tensors: List[torch.Tensor] = [] 

113 for i, obj in enumerate(args_kwargs_list): 

114 if torch.is_tensor(obj) and obj.requires_grad: 

115 inp_tensor_indices.append(i) 

116 inp_tensors.append(obj) 

117 if len(inp_tensors) == 0: 

118 return args, kwargs # no tensors that require gradients 

119 processed_tensors = PostBackwardFunction.apply(self, *inp_tensors) 

120 for inp_tensor_idx, processed_tensor in zip(inp_tensor_indices, processed_tensors): 

121 args_kwargs_list[inp_tensor_idx] = processed_tensor 

122 args_list = args_kwargs_list[: len(args_list)] 

123 kwargs_list = args_kwargs_list[len(args_list) :] 

124 args = tree_unflatten(args_list, args_spec) 

125 kwargs = tree_unflatten(kwargs_list, kwargs_spec) 

126 return args, kwargs 

127 

128 @_dynamo_disable 

129 def _forward_pre_hook(self, cell, args, kwargs): 

130 """Execute forward pre hook and set up backward hook.""" 

131 args, kwargs = self._hsdp_forward_pre_hook(cell, args, kwargs) 

132 return self._register_post_backward_hook(args, kwargs) 

133 

134 def _register_backward_pre_hook(self, outputs): 

135 """Register gradient hooks on all requires-grad outputs to trigger backward pre hook.""" 

136 flat_outputs, _ = tree_flatten(outputs) 

137 for output in flat_outputs: 

138 if isinstance(output, torch.Tensor) and output.requires_grad: 

139 handle_ref = [None] 

140 # pylint: disable=C0103, W0102 

141 

142 def wrapper_for_backward_pre_hook(grad, _handle_ref=handle_ref): 

143 """Remove this hook after it fires to prevent accmulation""" 

144 handle = _handle_ref[0] 

145 if handle is not None: 

146 handle.remove() 

147 return self._backward_pre_hook(grad) 

148 # pylint: enable=C0103, W0102 

149 handle = output.register_hook(wrapper_for_backward_pre_hook) 

150 handle_ref[0] = handle 

151 return outputs 

152 

153 @_dynamo_disable 

154 def _forward_hook(self, cell, inputs, outputs): # pylint: disable=R1710 

155 """Execute forward hook.""" 

156 if self.scheduler_state == FSDPSchedulerState.PRE_BACKWARD: 

157 return 

158 self._register_backward_pre_hook(outputs) 

159 if HSDPSchedulerV2.root_bp_state: 

160 self._restore_forward_prefetch_after_recompute() 

161 return 

162 return self._hsdp_forward_hook(cell, inputs, outputs) 

163 

164 # pylint: disable=W0212 

165 @_dynamo_disable 

166 def _backward_pre_hook(self, grad): 

167 """Execute backward pre hook.""" 

168 Variable._execution_engine.queue_callback(self._root_backward_hook) 

169 if self.scheduler_state == FSDPSchedulerState.PRE_BACKWARD: 

170 return grad 

171 HSDPSchedulerV2.root_bp_state = True 

172 self._hsdp_backward_pre_hook(self.cell, None) 

173 return grad 

174 

175 @_dynamo_disable 

176 def _root_backward_hook(self, force_reduce=False): 

177 """Finalize gradient reduction for the outermost HSDP module after backward. 

178 

179 ``apply_final_reduce`` selects between two cases, distinguished by whether 

180 this unit's forward input was differentiable: 

181 

182 * input ``requires_grad=False`` (the common case): no ``PostBackwardFunction`` 

183 is inserted on the input, so ``scheduler_state != BACKWARD`` and this hook 

184 owns the finalization -- it drains the pending reductions (comm_fusion=True 

185 via ``CommContext``; comm_fusion=False via the last module's reduce_scatter 

186 + allreduce) and applies the per-parameter gradients. 

187 * input ``requires_grad=True`` (a boundary case where the unit is fed a 

188 differentiable activation from an enclosing graph): the input's 

189 ``PostBackwardFunction`` drives ``scheduler_state == BACKWARD``, so the 

190 natural path does not finalize here. ``force_reduce`` lets a caller that 

191 owns that backward boundary demand the drain happen now rather than have it 

192 deferred a step. 

193 

194 ``root_bp_state`` -- whether the top-level root module's backward is still in 

195 flight, used to gate forward prefetch during activation recompute -- is 

196 independent of the reduce branch: it is cleared only by the root module's own 

197 hook, keyed on ``_is_root``. 

198 """ 

199 logger.debug("hook=root_backward_hook enter module=%s", self.hsdp_state) 

200 apply_final_reduce = self.scheduler_state != FSDPSchedulerState.BACKWARD 

201 self._backward_hook() 

202 if self._is_root: 

203 HSDPSchedulerV2.root_bp_state = False 

204 if apply_final_reduce or force_reduce: 

205 with torch.profiler.record_function(f"root_backward reduce:{self.hsdp_state.module_name}"): 

206 logger.debug( 

207 "hook=root_backward_hook action=final_reduce module=%s", 

208 self.hsdp_state, 

209 ) 

210 # Drain any pending async fused reduction from the last module's backward 

211 comm_ctx = get_comm_ctx() 

212 # Drain any pending pipelined HSDP reductions (comm_fusion=True) 

213 if comm_ctx.all_reduce_param_group is not None: 

214 logger.debug( 

215 "hook=root_backward_hook wait=comm_fusion_all_reduce module=%s", 

216 self.hsdp_state, 

217 ) 

218 comm_ctx.all_reduce_param_group.wait_all_reduce_and_apply_grad() 

219 comm_ctx.all_reduce_param_group = None 

220 if comm_ctx.pre_param_group is not None: 

221 logger.debug( 

222 "hook=root_backward_hook apply=comm_fusion_reduce_scatter module=%s", 

223 self.hsdp_state, 

224 ) 

225 comm_ctx.pre_param_group.apply_fusion_reduced_grad() 

226 comm_ctx.pre_param_group = None 

227 

228 # Process the last module's reduce_scatter and allreduce (comm_fusion=False) 

229 if TorchHSDPStateV2.pre_all_reduce_groups: 

230 for group in TorchHSDPStateV2.pre_all_reduce_groups: 

231 logger.debug( 

232 "hook=root_backward_hook wait=pre_reduce_scatter group_size=%s module=%s", 

233 len(group.hsdp_params), 

234 self.hsdp_state, 

235 ) 

236 # Wait reduce_scatter 

237 for hsdp_param in group.hsdp_params: 

238 hsdp_param.reduce_scatter_output() 

239 hsdp_param.clear_reduce_scatter_output() 

240 # Accumulate existing gradients (from previous mini steps) to fused_buffer 

241 # This is for gradient accumulation scenario 

242 # where previous mini steps used pre_reduce_scatter_params. 

243 # The gradients in sharded_param.grad are reduce_scatter results (not allreduced) 

244 group.accumulate_existing_grads_to_buffer() 

245 # Issue allreduce 

246 logger.debug( 

247 "hook=root_backward_hook launch=fused_all_reduce group_size=%s module=%s", 

248 len(group.hsdp_params), 

249 self.hsdp_state, 

250 ) 

251 group.issue_async_allreduce() 

252 TorchHSDPStateV2.pending_all_reduce_groups.append(group) 

253 TorchHSDPStateV2.pre_all_reduce_groups.clear() 

254 

255 # Apply gradients for params without all_reduce needs 

256 self.hsdp_state.reduce_scattered_params() 

257 # Finally, wait all allreduce and apply gradients 

258 TorchHSDPStateV2.delay_apply_reduce_grads(self.hsdp_state.device) 

259 

260 # Handle user config replicated_param 

261 self.hsdp_state.reduce_params() 

262 

263 

264 @_dynamo_disable 

265 def _backward_hook(self): 

266 """Execute backward hook.""" 

267 if self.scheduler_state == FSDPSchedulerState.BACKWARD: 

268 return 

269 self._hsdp_backward_hook(self.cell, None, None) 

270 

271 # pylint: disable=W0613 

272 def _grouped_forward_pre_hook_skip(self, cell, args, kwargs) -> None: # pylint: disable=arguments-differ 

273 """Override base ``(args, kwargs)`` return; ``nn.Module`` pre-hook uses ``None`` for no-op.""" 

274 return None 

275 

276 def _grouped_forward_post_hook_skip(self, outputs) -> None: # pylint: disable=arguments-differ 

277 """Override base output pass-through; forward hook uses ``None`` for no-op.""" 

278 return None 

279 

280 @_dynamo_disable 

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

282 """Run the grouped FSDP pre-forward hook outside Dynamo tracing.""" 

283 return super()._grouped_forward_pre_hook(cell, args, kwargs) 

284 

285 def _register_forward_module_hook(self, mod, hook) -> None: 

286 """Register forward hook; use ``always_call=True`` when supported (matches PyTorch FSDP).""" 

287 sig = inspect.signature(mod.register_forward_hook) 

288 if "always_call" in sig.parameters: 

289 mod.register_forward_hook(hook, prepend=False, always_call=True) 

290 else: 

291 mod.register_forward_hook(hook, prepend=False) 

292 

293 def _register_forward_backward_hooks(self): 

294 """Register module forward and backward hook on all managed modules.""" 

295 if self._fsdp_group_post_pending is None: 

296 for mod in self.modules: 

297 mod.register_forward_pre_hook(self._forward_pre_hook, with_kwargs=True) 

298 mod.register_forward_hook(self._forward_hook) 

299 return 

300 for mod in self.modules: 

301 mod.register_forward_pre_hook(self._grouped_forward_pre_hook, with_kwargs=True) 

302 grouped_forward_hook = _dynamo_disable(self._make_grouped_forward_post_hook(mod)) 

303 self._register_forward_module_hook(mod, grouped_forward_hook)