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

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.core.fully_shard.hsdp_utils import get_dtensor_managed_mesh 

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

25from hyper_parallel.platform.mindspore.fully_shard.param_group import get_comm_ctx 

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

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

28from hyper_parallel.platform import get_platform 

29 

30logger = get_logger("FSDP") 

31 

32 

33class MindSporeHSDPSchedulerV2(HSDPSchedulerV2): 

34 """MindSpore HSDP scheduler. 

35 

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

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

38 """ 

39 def zero_grad(self) -> None: 

40 """Zero grad.""" 

41 self.hsdp_state.zero_grad() 

42 

43 def _register_hooks(self): 

44 """Register hooks.""" 

45 self._register_forward_backward_hooks() 

46 

47 def _init_platform(self): 

48 """Initialize the platform.""" 

49 from hyper_parallel.platform.mindspore.platform import MindSporePlatform 

50 self.platform = get_platform() 

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

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

53 

54 def _new_cell_state(self): 

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

56 params = self._get_managed_params() 

57 if self.mesh is None: 

58 compat_meshes = [get_dtensor_managed_mesh(param) for param in params] 

59 compat_meshes = [mesh for mesh in compat_meshes if mesh is not None] 

60 compat_mesh = compat_meshes[0] if compat_meshes else None 

61 if compat_mesh is None: 

62 raise ValueError( 

63 "Cannot build fully_shard compatibility mesh_info " 

64 "without a DTensor parameter mesh." 

65 ) 

66 compat_mesh_hash = compat_mesh.to_hash() 

67 for param_mesh in compat_meshes[1:]: 

68 if param_mesh.to_hash() != compat_mesh_hash: 

69 raise ValueError( 

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

71 ) 

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

73 elif self.mesh.ndim == 1: 

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

75 elif self.mesh.ndim == 2: 

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

77 else: 

78 raise ValueError( 

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

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

81 ) 

82 self.hsdp_state = MindSporeHSDPStateV2( 

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

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 _pynative_executor.enable_grad(): 

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[ms.Tensor] = [] 

95 for i, obj in enumerate(args_kwargs_list): 

96 if isinstance(obj, ms.Tensor) 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_spec, args_list) 

107 kwargs = tree_unflatten(kwargs_spec, kwargs_list) 

108 return args, kwargs 

109 

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

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

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

113 return self._register_post_backward_hook(args, kwargs) 

114 

115 def _register_backward_pre_hook(self, outputs): 

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

117 flat_outputs, _ = tree_flatten(outputs) 

118 for output in flat_outputs: 

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

120 output.register_hook(self._backward_pre_hook) 

121 return outputs 

122 

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

124 """Execute forward hook.""" 

125 if self.scheduler_state == FSDPSchedulerState.PRE_BACKWARD: 

126 return 

127 self._register_backward_pre_hook(outputs) 

128 if HSDPSchedulerV2.root_bp_state: 

129 self._restore_forward_prefetch_after_recompute() 

130 return 

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

132 

133 # pylint: disable=W0212 

134 def _backward_pre_hook(self, grad): 

135 """Execute backward pre hook.""" 

136 _pynative_executor.queue_backward_final_callback(self._root_backward_hook) 

137 if self.scheduler_state == FSDPSchedulerState.PRE_BACKWARD: 

138 return grad 

139 HSDPSchedulerV2.root_bp_state = True 

140 self._hsdp_backward_pre_hook(self.cell, None) 

141 return grad 

142 

143 # pylint: disable=W0613 

144 def _root_backward_hook(self, force_reduce=False): 

145 """Finalize the outermost backward: drain pending reductions and apply grads. 

146 

147 The drain is unconditional. Every step below is self-limiting -- the fused 

148 groups are ``None``-guarded and ``reduce_params`` is an empty-queue no-op 

149 when there is no pending work -- so running them on every invocation never 

150 double-applies and preserves the invariant that a parameter's ``.grad`` is 

151 either ``None`` or a fully reduced value. This mirrors torch FSDP2, whose 

152 wait in ``_root_post_backward_final_callback`` is likewise not gated on the 

153 per-group post-backward training state. 

154 

155 Gating the drain on ``scheduler_state != BACKWARD`` was unsafe for any unit that 

156 acts as a root while being fed a differentiable activation: its input's 

157 ``PostBackwardFunction`` drives ``scheduler_state == BACKWARD``, so the gate would 

158 skip the drain and leak the last module's reduce-scatter into the next optimizer 

159 step. This happens when ``fully_shard`` wraps only inner layers and not the root 

160 module (each layer becomes its own root yet is fed a grad-requiring activation). 

161 PP hit the same boundary and worked around it with ``force_reduce=True`` from 

162 ``PipelineStage.execute_reduce_grad``; that call site keeps working -- the drain is 

163 simply always performed now. PP per-micro-batch accumulation is unaffected because 

164 each chunk's backward sets ``requires_gradient_sync=False``, leaving the reduce 

165 queue empty here so this drain is a no-op until the explicit reduce step. 

166 

167 ``root_bp_state`` (top-level root backward in flight; gates forward prefetch during 

168 activation recompute) is independent of the drain and is cleared only by the root 

169 module's own hook, keyed on ``_is_root``. 

170 """ 

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

172 self._backward_hook() 

173 if self._is_root: 

174 HSDPSchedulerV2.root_bp_state = False 

175 comm_ctx = get_comm_ctx() 

176 if comm_ctx.all_reduce_param_group is not None: 

177 logger.debug( 

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

179 self.hsdp_state, 

180 ) 

181 comm_ctx.all_reduce_param_group.wait_all_reduce_and_apply_grad() 

182 comm_ctx.all_reduce_param_group = None 

183 if comm_ctx.pre_param_group is not None: 

184 logger.debug( 

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

186 self.hsdp_state, 

187 ) 

188 comm_ctx.pre_param_group.apply_fusion_reduced_grad() 

189 comm_ctx.pre_param_group = None 

190 # Step 1: Wait for previous reduce-scatter groups and get them for all-reduce 

191 prev_groups = self.hsdp_state._wait_prev_reduce_scatter() 

192 # Step 2: Accumulate and issue async all-reduce for previous groups 

193 for group in prev_groups: 

194 group.accumulate_existing_grads_to_buffer() 

195 group.issue_async_allreduce() 

196 MindSporeHSDPStateV2.pending_all_reduce_groups.append(group) 

197 # Step 3: Wait/apply any remaining reduce-scatter for pure FSDP params 

198 self.hsdp_state.reduce_scattered_params() 

199 # Step 4: Wait for pending all-reduce groups and apply grads 

200 MindSporeHSDPStateV2.delay_apply_reduce_grads() 

201 # Step 5: Process any remaining all-reduce params (without fusion) 

202 logger.debug( 

203 "hook=root_backward_hook action=reduce_params module=%s", 

204 self.hsdp_state, 

205 ) 

206 self.hsdp_state.reduce_params() 

207 

208 def _backward_hook(self): 

209 """Execute backward hook.""" 

210 if self.scheduler_state == FSDPSchedulerState.BACKWARD: 

211 return 

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

213 

214 @staticmethod 

215 def _without_ms_dispatch_mode(hook): 

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

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

218 with _DisableMsDispatchMode(): 

219 return hook(*args, **kwargs) 

220 return wrapped_hook 

221 

222 def _register_forward_backward_hooks(self): 

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

224 if self._fsdp_group_post_pending is None: 

225 for mod in self.modules: 

226 mod.register_forward_pre_hook( 

227 self._without_ms_dispatch_mode(self._forward_pre_hook), 

228 with_kwargs=True, 

229 ) 

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

231 return 

232 for mod in self.modules: 

233 mod.register_forward_pre_hook( 

234 self._without_ms_dispatch_mode(self._grouped_forward_pre_hook), 

235 with_kwargs=True, 

236 ) 

237 mod.register_forward_hook( 

238 self._without_ms_dispatch_mode(self._make_grouped_forward_post_hook(mod)) 

239 )