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

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.platform import get_platform 

26from hyper_parallel.tools.logging import get_logger 

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

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

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

30 

31logger = get_logger("FSDP") 

32 

33_P = ParamSpec("_P") 

34_R = TypeVar("_R") 

35 

36 

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

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

39 

40 @functools.wraps(func) 

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

42 return torch._dynamo.disable( 

43 func, 

44 recursive=True, 

45 )(*args, **kwargs) 

46 

47 return wrapper 

48 

49 

50class TorchHSDPSchedulerV2(HSDPSchedulerV2): 

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

52 

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

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

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

56 

57 def _register_hooks(self): 

58 """Register hooks.""" 

59 self._register_forward_backward_hooks() 

60 

61 def _init_platform(self): 

62 """Initialize the platform.""" 

63 # pylint: disable=C0415 

64 from hyper_parallel.platform.torch.platform import TorchPlatform 

65 self.platform = get_platform() 

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

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

68 

69 def _new_cell_state(self): 

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

71 self.hsdp_state = TorchHSDPStateV2( 

72 self.modules, 

73 self.mesh, 

74 self.shard_placement_fn, 

75 self.comm_fusion_policy, 

76 self.mp_policy, 

77 self.offload_policy, 

78 self.ignored_params, 

79 self.replicate_params, 

80 self.platform, 

81 self.scheduler_ctx, 

82 self.device, 

83 source_shard_infos=self.source_shard_infos, 

84 ) 

85 

86 def _register_post_backward_hook(self, args, kwargs): 

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

88 if not torch.is_grad_enabled(): 

89 return args, kwargs 

90 args_list, args_spec = tree_flatten(args) 

91 kwargs_list, kwargs_spec = tree_flatten(kwargs) 

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

93 inp_tensor_indices: List[int] = [] 

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

95 for i, obj in enumerate(args_kwargs_list): 

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

97 inp_tensor_indices.append(i) 

98 inp_tensors.append(obj) 

99 if len(inp_tensors) == 0: 

100 return args, kwargs # no tensors that require gradients 

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

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

103 args_kwargs_list[inp_tensor_idx] = processed_tensor 

104 args_list = args_kwargs_list[: len(args_list)] 

105 kwargs_list = args_kwargs_list[len(args_list) :] 

106 args = tree_unflatten(args_list, args_spec) 

107 kwargs = tree_unflatten(kwargs_list, kwargs_spec) 

108 return args, kwargs 

109 

110 @_dynamo_disable 

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

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

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

114 return self._register_post_backward_hook(args, kwargs) 

115 

116 def _register_backward_pre_hook(self, outputs): 

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

118 flat_outputs, _ = tree_flatten(outputs) 

119 for output in flat_outputs: 

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

121 handle_ref = [None] 

122 # pylint: disable=C0103, W0102 

123 

124 def wrapper_for_backward_pre_hook(grad, _handle_ref=handle_ref): 

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

126 handle = _handle_ref[0] 

127 if handle is not None: 

128 handle.remove() 

129 return self._backward_pre_hook(grad) 

130 # pylint: enable=C0103, W0102 

131 handle = output.register_hook(wrapper_for_backward_pre_hook) 

132 handle_ref[0] = handle 

133 return outputs 

134 

135 @_dynamo_disable 

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

137 """Execute forward hook.""" 

138 if self.scheduler_state == FSDPSchedulerState.PRE_BACKWARD: 

139 return 

140 self._register_backward_pre_hook(outputs) 

141 if self.scheduler_ctx.root_bp_state: 

142 self._restore_forward_prefetch_after_recompute() 

143 return 

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

145 

146 # pylint: disable=W0212 

147 @_dynamo_disable 

148 def _backward_pre_hook(self, grad): 

149 """Execute backward pre hook.""" 

150 if self.scheduler_state == FSDPSchedulerState.PRE_BACKWARD: 

151 return grad 

152 if self._is_root: 

153 Variable._execution_engine.queue_callback(self._root_backward_hook) 

154 self.scheduler_ctx.root_bp_state = True 

155 self._hsdp_backward_pre_hook(self.cell, None) 

156 return grad 

157 

158 @_dynamo_disable 

159 def _root_backward_hook(self): 

160 """Drain all DP pipelines, then run final TP reduction and apply gradients.""" 

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

162 for hsdp_scheduler in self.scheduler_ctx.all_hsdp_schedulers: 

163 # let modules which are not triggered backward_hook launch backward communication. 

164 hsdp_scheduler._backward_hook() 

165 self.scheduler_ctx.root_bp_state = False 

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

167 logger.debug( 

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

169 self.hsdp_state, 

170 ) 

171 self._finalize_comm_fusion_reductions() 

172 self._finalize_per_param_reductions() 

173 self.launch_tp_replicate_reduce_and_apply() 

174 

175 def _finalize_comm_fusion_reductions(self) -> None: 

176 """Drain the comm_fusion=True RS/AR pipeline.""" 

177 comm_ctx = self.scheduler_ctx.param_group_comm_ctx 

178 if comm_ctx.all_reduce_param_group is not None: 

179 logger.debug( 

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

181 self.hsdp_state, 

182 ) 

183 comm_ctx.all_reduce_param_group.wait_all_reduce_and_save_grad() 

184 comm_ctx.all_reduce_param_group = None 

185 if comm_ctx.pre_param_group is not None: 

186 logger.debug( 

187 "hook=root_backward_hook wait=comm_fusion_reduce_scatter module=%s", 

188 self.hsdp_state, 

189 ) 

190 comm_ctx.pre_param_group.wait_reduce_scatter_and_issue_all_reduce() 

191 comm_ctx.pre_param_group = None 

192 if comm_ctx.all_reduce_param_group is not None: 

193 comm_ctx.all_reduce_param_group.wait_all_reduce_and_save_grad() 

194 comm_ctx.all_reduce_param_group = None 

195 

196 def _finalize_per_param_reductions(self) -> None: 

197 """Drain the module-tree-local comm_fusion=False RS/AR queues.""" 

198 # A fused root may own non-fused children, so always drain the tree queues. 

199 last_all_reduce_groups = self.hsdp_state._wait_prev_reduce_scatter() 

200 self.hsdp_state._wait_prev_reduce_scatter_without_all_reduce() 

201 self.hsdp_state._issue_prev_fused_all_reduce(last_all_reduce_groups) 

202 self.hsdp_state.wait_and_split_all_reduce_work_groups() 

203 

204 def launch_tp_replicate_reduce_and_apply(self) -> None: 

205 """Run final TP replicate reductions and apply gradients for all states.""" 

206 for hsdp_scheduler in self.scheduler_ctx.all_hsdp_schedulers: 

207 hsdp_state = hsdp_scheduler.hsdp_state 

208 if hsdp_state is None: 

209 continue 

210 need_synchronize = False 

211 for hsdp_param in hsdp_state.hsdp_params: 

212 reduced_grad = hsdp_param.all_reduce_comm_ctx.all_reduce_output 

213 if reduced_grad is None: 

214 reduced_grad = hsdp_param.reduce_scatter_comm_ctx.reduce_scatter_output 

215 if reduced_grad is None: 

216 continue 

217 hsdp_param.all_reduce_source_replicate_grad_inplace( 

218 reduced_grad, 

219 hsdp_state.reduce_op_type, 

220 ) 

221 need_synchronize = hsdp_param.apply_reduced_grad(reduced_grad) or need_synchronize 

222 hsdp_param.clear_all_reduce_output() 

223 hsdp_param.clear_reduce_scatter_output() 

224 hsdp_state._sync_current_stream_if_needed(need_synchronize) 

225 

226 @_dynamo_disable 

227 def reset_iter_state(self) -> None: 

228 """Reset Torch fully_shard iteration state after communication is complete.""" 

229 super().reset_iter_state() 

230 self.hsdp_state.reset_iter_state() 

231 comm_ctx = self.scheduler_ctx.param_group_comm_ctx 

232 comm_ctx.pre_param_group = None 

233 comm_ctx.all_reduce_param_group = None 

234 

235 @_dynamo_disable 

236 def _backward_hook(self): 

237 """Execute backward hook.""" 

238 if self.scheduler_state == FSDPSchedulerState.BACKWARD: 

239 return 

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

241 

242 # pylint: disable=W0613 

243 @staticmethod 

244 def _grouped_forward_pre_hook_skip(cell, args, kwargs) -> None: 

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

246 return None 

247 

248 @staticmethod 

249 def _grouped_forward_post_hook_skip(outputs) -> None: 

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

251 return None 

252 

253 @_dynamo_disable 

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

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

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

257 

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

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

260 sig = inspect.signature(mod.register_forward_hook) 

261 if "always_call" in sig.parameters: 

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

263 else: 

264 mod.register_forward_hook(hook, prepend=False) 

265 

266 def _register_forward_backward_hooks(self): 

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

268 if self._fsdp_group_post_pending is None: 

269 for mod in self.modules: 

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

271 mod.register_forward_hook(self._forward_hook) 

272 return 

273 for mod in self.modules: 

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

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

276 self._register_forward_module_hook(mod, grouped_forward_hook)