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

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, Set, 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.utils import MixedPrecisionPolicy, OffloadPolicy, CommFusionPolicy 

21from hyper_parallel.tools.logging import get_logger 

22 

23logger = get_logger("FSDP") 

24 

25platform = get_platform() 

26 

27ModuleClass = platform.Module 

28ParameterClass = platform.Parameter 

29 

30 

31class HSDPState: 

32 """HSDP state for cell""" 

33 

34 def __init__( 

35 self, 

36 cell: Union[ModuleClass, Tuple[ModuleClass, ...]], 

37 mesh, 

38 shard_placement_fn, 

39 comm_fusion_policy: CommFusionPolicy, 

40 mp_policy: MixedPrecisionPolicy, 

41 offload_policy: OffloadPolicy, 

42 raw_ignored_params: Set[ParameterClass], 

43 raw_replicate_params: set[ParameterClass], 

44 platform_impl, 

45 scheduler_ctx, 

46 device=None, 

47 ): 

48 """ 

49 Initialize HSDPState. 

50 

51 Args: 

52 cell: Module or modules managed as one fully_shard unit. 

53 mesh: Explicit data-parallel device mesh. 

54 shard_placement_fn: Optional function selecting the parameter shard dimension. 

55 comm_fusion_policy: Communication fusion configuration. 

56 mp_policy: Mixed-precision policy. 

57 offload_policy: Parameter offload policy. 

58 raw_ignored_params: Parameters excluded from fully_shard management. 

59 raw_replicate_params: Managed parameters that remain replicated. 

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

61 scheduler_ctx: Scheduler context shared by this module tree. 

62 device: Optional target device for parameters. 

63 """ 

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

65 self.cell = self.modules[0] 

66 self.mesh = mesh 

67 self.shard_placement_fn = shard_placement_fn 

68 self.mp_policy = mp_policy 

69 self.offload_policy = offload_policy 

70 self.comm_fusion_policy = comm_fusion_policy 

71 self.raw_ignored_params = set(raw_ignored_params or ()) 

72 self.raw_replicate_params = set(raw_replicate_params or ()) 

73 self.platform = platform_impl 

74 self.scheduler_ctx = scheduler_ctx 

75 self.device = device 

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

77 self.param_group = None 

78 self._move_states_to_device() 

79 self._init_hsdp_params() 

80 self.is_shard = True 

81 self.module_name = None 

82 # requires_gradient_sync 

83 self.reduce_grads = True 

84 # Reshard parameter after backward 

85 self.reshard_after_backward = True 

86 # Requires AllReduce for grad When HSDP 

87 self.requires_all_reduce = True 

88 self.set_reduce_op_type("avg") 

89 self._reset_sharded_params = False 

90 

91 def __repr__(self) -> str: 

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

93 

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

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

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

97 """ 

98 if self.module_name: 

99 return str(self.module_name) 

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

101 

102 def _init_hsdp_params(self): 

103 """init hsdp parameters for cell""" 

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

105 

106 def _move_states_to_device(self): 

107 """move states to device""" 

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

109 

110 def set_reduce_op_type(self, reduce_op_type: str) -> None: 

111 """Set the gradient reduction operation for the current backend.""" 

112 raise NotImplementedError("HSDPState subclasses must implement set_reduce_op_type") 

113 

114 def shard(self) -> None: 

115 """change parameters to sharded state""" 

116 logger.debug( 

117 "action=reshard module=%s params=%s", 

118 self, 

119 self.hsdp_params, 

120 ) 

121 if self.is_shard: 

122 return 

123 for param in self.hsdp_params: 

124 param.to_sharded() 

125 self.is_shard = True 

126 

127 def unshard(self, async_op: bool = False) -> None: 

128 """change parameters to unsharded state""" 

129 logger.debug( 

130 "action=unshard module=%s async_op=%s params=%s", 

131 self, 

132 async_op, 

133 self.hsdp_params, 

134 ) 

135 if not self.is_shard: 

136 return 

137 

138 if self.comm_fusion_policy.enable_comm_fusion and self.param_group is not None: 

139 self.param_group.unshard(async_op) 

140 else: 

141 for param in self.hsdp_params: 

142 param.unshard(async_op) 

143 if not async_op: 

144 self.wait_for_unshard() 

145 

146 def prefetch(self) -> None: 

147 """prefetch unsharded parameters""" 

148 logger.debug( 

149 "action=prefetch module=%s params=%s", 

150 self, 

151 self.hsdp_params, 

152 ) 

153 self.unshard(async_op=True) 

154 

155 def wait_for_unshard(self) -> None: 

156 """wait for all unshard parameters""" 

157 logger.debug( 

158 "action=wait_unshard module=%s params=%s", 

159 self, 

160 self.hsdp_params, 

161 ) 

162 if not self.is_shard: 

163 return 

164 

165 if self.comm_fusion_policy.enable_comm_fusion and self.param_group is not None: 

166 self.param_group.wait_for_unshard() 

167 else: 

168 for param in self.hsdp_params: 

169 param.wait_for_unshard() 

170 self.is_shard = False 

171 

172 def set_gradient_scaling_factor(self, factor): 

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

174 

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

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

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

178 """ 

179 if self.param_group is not None: 

180 self.param_group.gradient_scaling_factor = factor 

181 else: 

182 for hsdp_param in self.hsdp_params: 

183 hsdp_param.gradient_scaling_factor = factor 

184 

185 def set_requires_all_reduce(self, requires_all_reduce: bool) -> None: 

186 """Propagate the HSDP all-reduce switch to the active communication path.""" 

187 self.requires_all_reduce = requires_all_reduce 

188 if self.param_group is not None: 

189 self.param_group.requires_all_reduce = requires_all_reduce