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

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 inspect 

17import torch 

18from typing import List 

19from torch.autograd import Variable 

20from torch.utils._pytree import tree_flatten, tree_unflatten 

21from hyper_parallel.core.dtensor.dtensor import DTensor 

22from hyper_parallel.tools.logging import get_logger 

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

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

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

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

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

28from hyper_parallel.platform import get_platform 

29 

30logger = get_logger("FSDP") 

31 

32 

33class TorchHSDPSchedulerV2(HSDPSchedulerV2): 

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

35 

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

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

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

39 

40 def _register_hooks(self): 

41 """Register hooks.""" 

42 self._register_forward_backward_hooks() 

43 

44 def _init_platform(self): 

45 """Initialize the platform.""" 

46 # pylint: disable=C0415 

47 from hyper_parallel.platform.torch.platform import TorchPlatform 

48 self.platform = get_platform() 

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

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

51 

52 def _new_cell_state(self): 

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

54 params = self._get_managed_params() 

55 if self.mesh is None: 

56 compat_meshes = [ 

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

58 ] 

59 compat_mesh = compat_meshes[0] if compat_meshes else None 

60 if compat_mesh is None: 

61 raise ValueError( 

62 "Cannot build fully_shard compatibility mesh_info " 

63 "without a DTensor parameter mesh." 

64 ) 

65 compat_mesh_hash = compat_mesh.to_hash() 

66 for param_mesh in compat_meshes[1:]: 

67 if param_mesh.to_hash() != compat_mesh_hash: 

68 raise ValueError( 

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

70 ) 

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

72 elif self.mesh.ndim == 1: 

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

74 elif self.mesh.ndim == 2: 

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

76 else: 

77 raise ValueError( 

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

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

80 ) 

81 self.hsdp_state = TorchHSDPStateV2( 

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

83 ) 

84 

85 def _register_post_backward_hook(self, args, kwargs): 

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

87 if not torch.is_grad_enabled(): 

88 return args, kwargs 

89 args_list, args_spec = tree_flatten(args) 

90 kwargs_list, kwargs_spec = tree_flatten(kwargs) 

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

92 inp_tensor_indices: List[int] = [] 

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

94 for i, obj in enumerate(args_kwargs_list): 

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

96 inp_tensor_indices.append(i) 

97 inp_tensors.append(obj) 

98 if len(inp_tensors) == 0: 

99 return args, kwargs # no tensors that require gradients 

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

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

102 args_kwargs_list[inp_tensor_idx] = processed_tensor 

103 args_list = args_kwargs_list[: len(args_list)] 

104 kwargs_list = args_kwargs_list[len(args_list) :] 

105 args = tree_unflatten(args_list, args_spec) 

106 kwargs = tree_unflatten(kwargs_list, kwargs_spec) 

107 return args, kwargs 

108 

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

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

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

112 return self._register_post_backward_hook(args, kwargs) 

113 

114 def _register_backward_pre_hook(self, outputs): 

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

116 flat_outputs, _ = tree_flatten(outputs) 

117 for output in flat_outputs: 

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

119 handle_ref = [None] 

120 # pylint: disable=C0103, W0102 

121 

122 def wrapper_for_backward_pre_hook(grad, _handle_ref=handle_ref): 

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

124 handle = _handle_ref[0] 

125 if handle is not None: 

126 handle.remove() 

127 return self._backward_pre_hook(grad) 

128 # pylint: enable=C0103, W0102 

129 handle = output.register_hook(wrapper_for_backward_pre_hook) 

130 handle_ref[0] = handle 

131 return outputs 

132 

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

134 """Execute forward hook.""" 

135 if self.scheduler_state == FSDPSchedulerState.PRE_BACKWARD: 

136 return 

137 self._register_backward_pre_hook(outputs) 

138 if HSDPSchedulerV2.root_bp_state: 

139 self._restore_forward_prefetch_after_recompute() 

140 return 

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

142 

143 # pylint: disable=W0212 

144 def _backward_pre_hook(self, grad): 

145 """Execute backward pre hook.""" 

146 Variable._execution_engine.queue_callback(self._root_backward_hook) 

147 if self.scheduler_state == FSDPSchedulerState.PRE_BACKWARD: 

148 return grad 

149 HSDPSchedulerV2.root_bp_state = True 

150 self._hsdp_backward_pre_hook(self.cell, None) 

151 return grad 

152 

153 def _root_backward_hook(self, force_reduce=False): 

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

155 

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

157 this unit's forward input was differentiable: 

158 

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

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

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

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

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

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

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

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

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

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

169 deferred a step. 

170 

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

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

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

174 hook, keyed on ``_is_root``. 

175 """ 

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

177 apply_final_reduce = self.scheduler_state != FSDPSchedulerState.BACKWARD 

178 self._backward_hook() 

179 if self._is_root: 

180 HSDPSchedulerV2.root_bp_state = False 

181 if apply_final_reduce or force_reduce: 

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

183 logger.debug( 

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

185 self.hsdp_state, 

186 ) 

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

188 comm_ctx = get_comm_ctx() 

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

190 if comm_ctx.all_reduce_param_group is not None: 

191 logger.debug( 

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

193 self.hsdp_state, 

194 ) 

195 comm_ctx.all_reduce_param_group.wait_all_reduce_and_apply_grad() 

196 comm_ctx.all_reduce_param_group = None 

197 if comm_ctx.pre_param_group is not None: 

198 logger.debug( 

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

200 self.hsdp_state, 

201 ) 

202 comm_ctx.pre_param_group.apply_fusion_reduced_grad() 

203 comm_ctx.pre_param_group = None 

204 

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

206 if TorchHSDPStateV2.pre_all_reduce_groups: 

207 for group in TorchHSDPStateV2.pre_all_reduce_groups: 

208 logger.debug( 

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

210 len(group.hsdp_params), 

211 self.hsdp_state, 

212 ) 

213 # Wait reduce_scatter 

214 for hsdp_param in group.hsdp_params: 

215 hsdp_param.reduce_scatter_output() 

216 hsdp_param.clear_reduce_scatter_output() 

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

218 # This is for gradient accumulation scenario 

219 # where previous mini steps used pre_reduce_scatter_params. 

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

221 group.accumulate_existing_grads_to_buffer() 

222 # Issue allreduce 

223 logger.debug( 

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

225 len(group.hsdp_params), 

226 self.hsdp_state, 

227 ) 

228 group.issue_async_allreduce() 

229 TorchHSDPStateV2.pending_all_reduce_groups.append(group) 

230 TorchHSDPStateV2.pre_all_reduce_groups.clear() 

231 

232 # Apply gradients for params without all_reduce needs 

233 self.hsdp_state.reduce_scattered_params() 

234 # Finally, wait all allreduce and apply gradients 

235 TorchHSDPStateV2.delay_apply_reduce_grads(self.hsdp_state.device) 

236 

237 # Handle user config replicated_param 

238 self.hsdp_state.reduce_params() 

239 

240 

241 def _backward_hook(self): 

242 """Execute backward hook.""" 

243 if self.scheduler_state == FSDPSchedulerState.BACKWARD: 

244 return 

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

246 

247 # pylint: disable=W0613 

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

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

250 return None 

251 

252 def _grouped_forward_post_hook_skip(self, outputs) -> None: 

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

254 return None 

255 

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

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

258 sig = inspect.signature(mod.register_forward_hook) 

259 if "always_call" in sig.parameters: 

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

261 else: 

262 mod.register_forward_hook(hook, prepend=False) 

263 

264 def _register_forward_backward_hooks(self): 

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

266 if self._fsdp_group_post_pending is None: 

267 for mod in self.modules: 

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

269 mod.register_forward_hook(self._forward_hook) 

270 return 

271 for mod in self.modules: 

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

273 self._register_forward_module_hook(mod, self._make_grouped_forward_post_hook(mod))