Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / fully_shard / hsdp_state.py: 91%

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"""HSDP cell state""" 

16from typing import List, Tuple, Union 

17 

18from hyper_parallel.platform import get_platform 

19from hyper_parallel.core.fully_shard.hsdp_param import HSDPParamV2 

20from hyper_parallel.core.fully_shard.hsdp_utils import HSDPConfigV2, ShardedState 

21from hyper_parallel.tools.logging import get_logger 

22 

23logger = get_logger("FSDP") 

24 

25platform = get_platform() 

26 

27 

28class HSDPState: 

29 """HSDP state for cell""" 

30 # Record pending per-parameter reduce-scatter/all-reduce work across 

31 # fully_shard states so later backward hooks/root drains can materialize 

32 # gradients launched by earlier states. 

33 pre_reduce_scatter_params = [] 

34 pre_all_reduce_params = [] 

35 

36 def __init__(self, cell: Union[platform.Module, Tuple[platform.Module, ...]], mesh_info, 

37 config: HSDPConfigV2, platform_impl, device=None): 

38 """ 

39 Initialize HSDPState. 

40 

41 Args: 

42 cell (platform.Module or Tuple[platform.Module, ...]): The module(s) whose parameters 

43 are managed by this state. When a tuple is passed, all modules are 

44 treated as one FSDP unit. 

45 mesh_info: Mesh topology for shard/replicate dimensions. 

46 config (HSDPConfigV2): HSDP configuration (mesh, mp_policy, offload_policy, etc.). 

47 platform_impl: Platform abstraction layer (Torch or MindSpore). 

48 device (torch.device, optional): Target device for parameters. 

49 """ 

50 self.modules = (cell,) if isinstance(cell, platform.Module) else tuple(cell) 

51 self.cell = self.modules[0] 

52 self.mesh_info = mesh_info 

53 self.config = config 

54 self.mp_policy = config.mp_policy 

55 self.offload_policy = config.offload_policy 

56 self.platform = platform_impl 

57 self.device = device 

58 self.hsdp_params: List[HSDPParamV2] = [] 

59 self.sharded_hsdp_params: List[HSDPParamV2] = [] 

60 self.replicate_params: List[HSDPParamV2] = [] 

61 self._move_states_to_device() 

62 self._init_hsdp_params() 

63 self.is_shard = True 

64 self.is_replicate_shard = True 

65 self.module_name = None 

66 

67 def __repr__(self) -> str: 

68 """Stable debug name used in log lines. 

69 

70 ``module_name`` is only assigned by the root forward pre-hook, so fall 

71 back to the managed module class and object id before names are set. 

72 Logging's ``%s`` calls this lazily -- only when a record is emitted. 

73 """ 

74 if self.module_name: 

75 return str(self.module_name) 

76 return f"{self.cell.__class__.__name__}@{id(self.cell):x}" 

77 

78 def _init_hsdp_params(self): 

79 """init hsdp parameters for cell""" 

80 raise NotImplementedError("HSDPState subclasses must implement _init_hsdp_params") 

81 

82 def _move_states_to_device(self): 

83 """move states to device""" 

84 raise NotImplementedError("HSDPState subclasses must implement _move_states_to_device") 

85 

86 def _assert_replicate_params_unsharded(self) -> None: 

87 """Validate replicate params are already materialized when state says so.""" 

88 for param in self.replicate_params: 

89 sharded_state = getattr(param, "sharded_state", None) 

90 if sharded_state != ShardedState.UNSHARDED: 

91 param_fqn = getattr(param, "_param_fqn", "<unknown>") 

92 raise AssertionError( 

93 f"Expected replicate parameter {param_fqn} to be " 

94 f"{ShardedState.UNSHARDED}, got {sharded_state}" 

95 ) 

96 

97 def shard(self, shard_replicate: bool = True): 

98 """change parameters to sharded state""" 

99 logger.debug( 

100 "action=reshard module=%s shard_params=%s replicate_params=%s shard_replicate=%s", 

101 self, 

102 self.sharded_hsdp_params, 

103 self.replicate_params, 

104 shard_replicate, 

105 ) 

106 if not self.is_shard: 

107 for param in self.sharded_hsdp_params: 

108 param.to_sharded() 

109 self.is_shard = True 

110 if shard_replicate and not self.is_replicate_shard: 

111 for param in self.replicate_params: 

112 param.to_sharded() 

113 self.is_replicate_shard = True 

114 

115 def unshard(self, async_op=False, unshard_replicate: bool = True): 

116 """change parameters to unsharded state""" 

117 logger.debug( 

118 "action=unshard module=%s async_op=%s shard_params=%s replicate_params=%s unshard_replicate=%s", 

119 self, 

120 async_op, 

121 self.sharded_hsdp_params, 

122 self.replicate_params, 

123 unshard_replicate, 

124 ) 

125 if not self.is_shard and (not unshard_replicate or not self.is_replicate_shard): 

126 if unshard_replicate: 

127 self._assert_replicate_params_unsharded() 

128 return 

129 

130 if unshard_replicate: 

131 if self.is_replicate_shard: 

132 for param in self.replicate_params: 

133 param.unshard(async_op) 

134 else: 

135 self._assert_replicate_params_unsharded() 

136 if self.is_shard: 

137 if self.config.comm_fusion and self.param_group is not None: 

138 self.param_group.unshard(async_op) 

139 else: 

140 for param in self.sharded_hsdp_params: 

141 param.unshard(async_op) 

142 if not async_op: 

143 self.wait_for_unshard(unshard_replicate) 

144 

145 def prefetch(self, unshard_replicate: bool = True): 

146 """prefetch unsharded parameters""" 

147 logger.debug( 

148 "action=prefetch module=%s shard_params=%s replicate_params=%s unshard_replicate=%s", 

149 self, 

150 self.sharded_hsdp_params, 

151 self.replicate_params, 

152 unshard_replicate, 

153 ) 

154 self.unshard(async_op=True, unshard_replicate=unshard_replicate) 

155 

156 def wait_for_unshard(self, wait_for_replicate: bool = True): 

157 """wait for all unshard parameters""" 

158 logger.debug( 

159 "action=wait_unshard module=%s shard_params=%s replicate_params=%s wait_for_replicate=%s", 

160 self, 

161 self.sharded_hsdp_params, 

162 self.replicate_params, 

163 wait_for_replicate, 

164 ) 

165 if not self.is_shard and (not wait_for_replicate or not self.is_replicate_shard): 

166 if wait_for_replicate: 

167 self._assert_replicate_params_unsharded() 

168 return 

169 if wait_for_replicate: 

170 if self.is_replicate_shard: 

171 for param in self.replicate_params: 

172 param.wait_for_unshard() 

173 self.is_replicate_shard = False 

174 else: 

175 self._assert_replicate_params_unsharded() 

176 if self.is_shard: 

177 if self.config.comm_fusion and self.param_group is not None: 

178 self.param_group.wait_for_unshard() 

179 else: 

180 for param in self.sharded_hsdp_params: 

181 param.wait_for_unshard() 

182 self.is_shard = False 

183 

184 def set_gradient_scaling_factor(self, factor): 

185 """Propagate the gradient scaling factor to the layer that applies it. 

186 

187 The factor is consumed on the reduce input: ``param_group.foreach_reduce`` 

188 for the fused (comm_fusion) path, or per-parameter ``reduce_scatter_grad`` 

189 / ``all_reduce_grad`` otherwise. The state does not hold a copy. 

190 """ 

191 param_group = getattr(self, "param_group", None) 

192 if param_group is not None: 

193 param_group.gradient_scaling_factor = factor 

194 else: 

195 for hsdp_param in self._iter_managed_params(): 

196 hsdp_param.gradient_scaling_factor = factor 

197 

198 def _iter_managed_params(self): 

199 """Return all fully_shard-managed parameters, including replicate_params.""" 

200 return [*self.hsdp_params, *self.replicate_params]