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

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

16from typing import List 

17import mindspore as ms 

18from mindspore._c_expression import _DisableMsDispatchMode 

19from mindspore.common.api import _pynative_executor 

20from mindspore.utils._pytree import tree_flatten, tree_unflatten 

21from hyper_parallel.tools.logging import get_logger 

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

23from hyper_parallel.platform.mindspore.fully_shard.hook_function import PostBackwardFunction 

24from hyper_parallel.platform.mindspore.fully_shard.state import MindSporeHSDPStateV2 

25from hyper_parallel.platform import get_platform 

26 

27logger = get_logger("FSDP") 

28 

29 

30class MindSporeHSDPSchedulerV2(HSDPSchedulerV2): 

31 """MindSpore HSDP scheduler. 

32 

33 List-unit grouped forward hooks use :class:`HSDPSchedulerV2` defaults for 

34 ``_grouped_forward_pre_hook_skip`` / ``_grouped_forward_post_hook_skip`` (no overrides here). 

35 """ 

36 def zero_grad(self) -> None: 

37 """Zero grad.""" 

38 self.hsdp_state.zero_grad() 

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 from hyper_parallel.platform.mindspore.platform import MindSporePlatform 

47 self.platform = get_platform() 

48 if not isinstance(self.platform, MindSporePlatform): 

49 raise ValueError(f"MindSporeHSDPSchedulerV2 expect MindSporePlatform, but got type: {type(self.platform)}") 

50 

51 def _new_cell_state(self): 

52 """Create a new cell state for mindspore.""" 

53 self.hsdp_state = MindSporeHSDPStateV2( 

54 self.modules, 

55 self.mesh, 

56 self.shard_placement_fn, 

57 self.comm_fusion_policy, 

58 self.mp_policy, 

59 self.offload_policy, 

60 self.ignored_params, 

61 self.replicate_params, 

62 self.platform, 

63 self.scheduler_ctx, 

64 self.device, 

65 ) 

66 

67 def _register_post_backward_hook(self, args, kwargs): 

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

69 if not _pynative_executor.enable_grad(): 

70 return args, kwargs 

71 args_list, args_spec = tree_flatten(args) 

72 kwargs_list, kwargs_spec = tree_flatten(kwargs) 

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

74 inp_tensor_indices: List[int] = [] 

75 inp_tensors: List[ms.Tensor] = [] 

76 for i, obj in enumerate(args_kwargs_list): 

77 if isinstance(obj, ms.Tensor) and obj.requires_grad: 

78 inp_tensor_indices.append(i) 

79 inp_tensors.append(obj) 

80 if len(inp_tensors) == 0: 

81 return args, kwargs # no tensors that require gradients 

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

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

84 args_kwargs_list[inp_tensor_idx] = processed_tensor 

85 args_list = args_kwargs_list[: len(args_list)] 

86 kwargs_list = args_kwargs_list[len(args_list):] 

87 args = tree_unflatten(args_spec, args_list) 

88 kwargs = tree_unflatten(kwargs_spec, kwargs_list) 

89 return args, kwargs 

90 

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

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

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

94 return self._register_post_backward_hook(args, kwargs) 

95 

96 def _register_backward_pre_hook(self, outputs): 

97 """Register output hook to trigger backward pre hook.""" 

98 flat_outputs, _ = tree_flatten(outputs) 

99 for output in flat_outputs: 

100 if isinstance(output, ms.Tensor) and output._requires_grad: 

101 output.register_hook(self._backward_pre_hook) 

102 return outputs 

103 

104 def _forward_hook(self, cell, inputs, outputs): 

105 """Execute forward hook.""" 

106 if self.scheduler_state == FSDPSchedulerState.PRE_BACKWARD: 

107 return 

108 self._register_backward_pre_hook(outputs) 

109 if self.scheduler_ctx.root_bp_state: 

110 self._restore_forward_prefetch_after_recompute() 

111 return 

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

113 

114 # pylint: disable=W0212 

115 def _backward_pre_hook(self, grad): 

116 """Execute backward pre hook.""" 

117 if self.scheduler_state == FSDPSchedulerState.PRE_BACKWARD: 

118 return grad 

119 if self._is_root: 

120 _pynative_executor.queue_backward_final_callback(self._root_backward_hook) 

121 self.scheduler_ctx.root_bp_state = True 

122 self._hsdp_backward_pre_hook(self.cell, None) 

123 return grad 

124 

125 def _root_backward_hook(self): 

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

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

128 for hsdp_scheduler in self.scheduler_ctx.all_hsdp_schedulers: 

129 hsdp_scheduler._backward_hook() 

130 self.scheduler_ctx.root_bp_state = False 

131 logger.debug( 

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

133 self.hsdp_state, 

134 ) 

135 self._finalize_comm_fusion_reductions() 

136 self._finalize_per_param_reductions() 

137 self.launch_tp_replicate_reduce_and_apply() 

138 

139 def _finalize_comm_fusion_reductions(self) -> None: 

140 """Drain the comm_fusion=True reduce-scatter/all-reduce pipeline.""" 

141 comm_ctx = self.scheduler_ctx.param_group_comm_ctx 

142 if comm_ctx.all_reduce_param_group is not None: 

143 logger.debug( 

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

145 self.hsdp_state, 

146 ) 

147 comm_ctx.all_reduce_param_group.wait_all_reduce_and_save_grad() 

148 comm_ctx.all_reduce_param_group = None 

149 if comm_ctx.pre_param_group is not None: 

150 logger.debug( 

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

152 self.hsdp_state, 

153 ) 

154 comm_ctx.pre_param_group.wait_reduce_scatter_and_issue_all_reduce() 

155 comm_ctx.pre_param_group = None 

156 if comm_ctx.all_reduce_param_group is not None: 

157 comm_ctx.all_reduce_param_group.wait_all_reduce_and_save_grad() 

158 comm_ctx.all_reduce_param_group = None 

159 

160 def _finalize_per_param_reductions(self) -> None: 

161 """Drain the module-tree-local comm_fusion=False communication queues.""" 

162 previous_groups = self.hsdp_state._wait_prev_reduce_scatter() 

163 self.hsdp_state._wait_prev_reduce_scatter_without_all_reduce() 

164 self.hsdp_state._issue_prev_fused_all_reduce(previous_groups) 

165 self.hsdp_state.wait_and_split_all_reduce_work_groups() 

166 

167 def launch_tp_replicate_reduce_and_apply(self) -> None: 

168 """Run final source-layout reductions and apply gradients for all states.""" 

169 for hsdp_scheduler in self.scheduler_ctx.all_hsdp_schedulers: 

170 hsdp_state = hsdp_scheduler.hsdp_state 

171 if hsdp_state is None: 

172 continue 

173 need_synchronize = False 

174 for hsdp_param in hsdp_state.hsdp_params: 

175 reduced_grad = hsdp_param.all_reduce_comm_ctx.all_reduce_output 

176 if reduced_grad is None: 

177 reduced_grad = hsdp_param.reduce_scatter_comm_ctx.reduce_scatter_output 

178 if reduced_grad is None: 

179 continue 

180 hsdp_param.all_reduce_tp_replicate_grad_inplace( 

181 reduced_grad, 

182 hsdp_state.reduce_op_type, 

183 ) 

184 need_synchronize = hsdp_param.apply_reduced_grad(reduced_grad) or need_synchronize 

185 hsdp_param.clear_all_reduce_output() 

186 hsdp_param.clear_reduce_scatter_output() 

187 hsdp_state._sync_current_stream_if_needed(need_synchronize) 

188 

189 def reset_iter_state(self) -> None: 

190 """Reset MindSpore fully_shard iteration state after communication completes.""" 

191 super().reset_iter_state() 

192 self.hsdp_state.reset_iter_state() 

193 comm_ctx = self.scheduler_ctx.param_group_comm_ctx 

194 comm_ctx.pre_param_group = None 

195 comm_ctx.all_reduce_param_group = None 

196 

197 def _backward_hook(self): 

198 """Execute backward hook.""" 

199 if self.scheduler_state == FSDPSchedulerState.BACKWARD: 

200 return 

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

202 

203 @staticmethod 

204 def _without_ms_dispatch_mode(hook): 

205 """Run HSDP hook internals outside any outer MsDispatchMode.""" 

206 def wrapped_hook(*args, **kwargs): 

207 with _DisableMsDispatchMode(): 

208 return hook(*args, **kwargs) 

209 return wrapped_hook 

210 

211 def _register_forward_backward_hooks(self): 

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

213 if self._fsdp_group_post_pending is None: 

214 for mod in self.modules: 

215 mod.register_forward_pre_hook( 

216 self._without_ms_dispatch_mode(self._forward_pre_hook), 

217 with_kwargs=True, 

218 ) 

219 mod.register_forward_hook(self._without_ms_dispatch_mode(self._forward_hook)) 

220 return 

221 for mod in self.modules: 

222 mod.register_forward_pre_hook( 

223 self._without_ms_dispatch_mode(self._grouped_forward_pre_hook), 

224 with_kwargs=True, 

225 ) 

226 mod.register_forward_hook( 

227 self._without_ms_dispatch_mode(self._make_grouped_forward_post_hook(mod)) 

228 )