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

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 optimizer shared level""" 

16from dataclasses import dataclass, field 

17from enum import auto, Enum 

18from typing import Any, List, Optional, Sequence 

19 

20import numpy as np 

21 

22from hyper_parallel.core.dtensor.device_mesh import DeviceMesh 

23from hyper_parallel.core.dtensor.dtensor import DTensor 

24from hyper_parallel.platform import get_platform 

25from hyper_parallel.platform.platform import PlatformType 

26 

27platform = get_platform() 

28 

29 

30class ShardedState(Enum): 

31 """ 

32 Parameter shard state 

33 """ 

34 SHARDED = auto() 

35 UNSHARDED = auto() 

36 

37 

38class FullyShardParamMode(Enum): 

39 """Internal fully_shard execution modes derived from parameter layout.""" 

40 

41 LOCAL_PARAM = auto() 

42 DTENSOR_COMPAT = auto() 

43 DTENSOR_UNIFIED = auto() 

44 

45 

46@dataclass 

47class GroupInfo: 

48 """Communication group metadata used by fully_shard.""" 

49 

50 group_name: str 

51 group: Any 

52 rank_size: int 

53 

54 

55class FSDPSchedulerState(Enum): 

56 """ 

57 Scheduler state: 

58 - PRE_FORWARD: 

59 already run hook before forward. 

60 - FORWARD: 

61 already run hook after forward. 

62 - PRE_BACKWARD: 

63 already run hook before backward. 

64 - BACKWARD: 

65 already run hook after backward. 

66 """ 

67 PRE_FORWARD = auto() 

68 FORWARD = auto() 

69 PRE_BACKWARD = auto() 

70 BACKWARD = auto() 

71 

72 

73@dataclass 

74class ParamModuleInfo: 

75 """ 

76 Tracks parameter ownership and supports shared weights in HSDP. 

77 

78 This dataclass maintains the mapping between a parameter and its module(s), 

79 enabling parameter swapping during sharding/unsharding transitions. Shared 

80 weights are parameters referenced by multiple modules (e.g., tied embeddings). 

81 

82 This class tracks all references to ensure proper parameter replacement during  

83 sharding/unsharding operations. 

84 

85 Attributes: 

86 module: The module that owns this parameter. 

87 param_name: Attribute name of the parameter in the module (e.g., "weight"). 

88 shared_modules: List of other modules sharing this same parameter object. 

89 shared_param_names: Corresponding parameter names in shared_modules (aligned by index). 

90 """ 

91 module: platform.Module 

92 param_name: str 

93 shared_modules: List[platform.Module] = field(default_factory=list) 

94 shared_param_names: List[str] = field(default_factory=list) 

95 

96 

97def _named_parameters_with_duplicates( 

98 module: platform.Module, **kwargs: Any 

99) -> list[tuple[str, platform.Parameter]]: 

100 """ 

101 This API is required as some modules overwrite `named_parameters()` but do not support 

102 `remove_duplicate`. 

103 """ 

104 if "remove_duplicate" in kwargs: 

105 raise AssertionError( 

106 "_named_parameters_with_duplicates cannot be used with `remove_duplicate` argument." 

107 ) 

108 

109 def get_named_parameters(module, **kwargs): 

110 if platform.platform_type == PlatformType.PYTORCH: 

111 return module.named_parameters(**kwargs) 

112 return module.parameters_and_names(expand=False) 

113 kwargs["remove_duplicate"] = False 

114 try: 

115 ret = list(get_named_parameters(module, **kwargs)) 

116 except AssertionError: 

117 kwargs.pop("remove_duplicate") 

118 ret = list(get_named_parameters(module, **kwargs)) 

119 return ret 

120 

121 

122def _get_param_module_infos( 

123 params: list[platform.Parameter], modules: tuple[platform.Module, ...] 

124) -> list['ParamModuleInfo']: 

125 """ 

126 Shared parameter: lin1.weight = lin2.weight 

127 Shared module: mlp.lin1 = mlp.lin2 

128 We do not remove duplicates when traversing both modules and parameters to 

129 find shared modules' parameters and shared parameters within a module. 

130 """ 

131 params_set = set(params) 

132 param_to_module_info: dict[platform.Parameter, ParamModuleInfo] = {} 

133 

134 def get_named_modules(module): 

135 if platform.platform_type == PlatformType.PYTORCH: 

136 return module.named_modules(remove_duplicate=False) 

137 return module.cells_and_names() 

138 

139 for module in modules: 

140 for _, submodule in get_named_modules(module): 

141 for param_name, param in _named_parameters_with_duplicates( 

142 submodule, recurse=False 

143 ): 

144 if param in params_set: 

145 if param not in param_to_module_info: 

146 param_to_module_info[param] = ParamModuleInfo( 

147 submodule, param_name 

148 ) 

149 else: 

150 param_to_module_info[param].shared_modules.append(submodule) 

151 param_to_module_info[param].shared_param_names.append( 

152 param_name 

153 ) 

154 if len(param_to_module_info) != len(params): 

155 raise AssertionError(f"Some parameters are not in the module tree of {modules}") 

156 return [param_to_module_info[param] for param in params] 

157 

158 

159def get_managed_modules_parameters( 

160 modules: Sequence[platform.Module], 

161 ignored_params: Optional[Sequence[platform.Parameter]] = None, 

162) -> list[platform.Parameter]: 

163 """Collect deduplicated parameters from ``modules`` while skipping ignored params. 

164 

165 Parameters that were already initialized by an inner ``fully_shard`` instance 

166 are intentionally excluded so nested ``fully_shard(mesh=None)`` resolves mesh 

167 mode from the parameters that the current wrapper will actually manage. 

168 """ 

169 params: list[platform.Parameter] = [] 

170 ignored_params_set = set(ignored_params or ()) 

171 visited_params: set[platform.Parameter] = set() 

172 for mod in modules: 

173 for _, param in platform.parameters_dict(mod): 

174 if param in ignored_params_set or param in visited_params: 

175 continue 

176 if getattr(param, "_hsdp_param_initialized", False): 

177 continue 

178 visited_params.add(param) 

179 params.append(param) 

180 return params 

181 

182 

183def infer_fully_shard_param_mode( 

184 mesh: Optional[DeviceMesh], 

185 params: Optional[Sequence[Any]] = None, 

186) -> FullyShardParamMode: 

187 """Infer the compatibility parameter mode from parameter layout and mesh.""" 

188 has_dtensor_param = any(is_dtensor_managed_param(param) for param in params or ()) 

189 if not has_dtensor_param: 

190 return FullyShardParamMode.LOCAL_PARAM 

191 if mesh is None: 

192 return FullyShardParamMode.DTENSOR_COMPAT 

193 return FullyShardParamMode.DTENSOR_UNIFIED 

194 

195 

196def unwrap_dtensor_param(param: Any) -> Optional[DTensor]: 

197 """Return the DTensor payload carried by ``param`` if one exists.""" 

198 if isinstance(param, DTensor): 

199 return param 

200 param_data = getattr(param, "data", None) 

201 if isinstance(param_data, DTensor): 

202 return param_data 

203 if all(hasattr(param, attr) for attr in ("_device_mesh", "_placements", "_local_tensor")): 

204 return param 

205 return None 

206 

207 

208def is_dtensor_managed_param(param: Any) -> bool: 

209 """Return whether a parameter already carries DTensor layout metadata.""" 

210 return unwrap_dtensor_param(param) is not None 

211 

212 

213def get_dtensor_managed_mesh(param: Any) -> Optional[DeviceMesh]: 

214 """Return the DTensor mesh carried by ``param`` if one exists.""" 

215 payload = unwrap_dtensor_param(param) 

216 if payload is None: 

217 return None 

218 return getattr(payload, "device_mesh", getattr(payload, "_device_mesh", None)) 

219 

220 

221def get_rank_list_for_axes( 

222 mesh: DeviceMesh, 

223 axes: Sequence[int], 

224 rank: Optional[int] = None, 

225) -> list[int]: 

226 """Return ranks that vary along ``axes`` and keep all other coordinates fixed.""" 

227 if rank is None: 

228 rank = mesh.rank 

229 if rank not in mesh.rank_list: 

230 raise ValueError(f"Rank {rank} not found in mesh rank list {mesh.rank_list}.") 

231 

232 normalized_axes = tuple(sorted(set(axes))) 

233 if len(normalized_axes) == 0: 

234 return [rank] 

235 

236 mesh_tensor = np.array(mesh.rank_list).reshape(mesh.mesh_shape) 

237 rank_index = mesh.rank_list.index(rank) 

238 coord = [0] * len(mesh.mesh_shape) 

239 temp = rank_index 

240 for i in range(len(mesh.mesh_shape) - 1, -1, -1): 

241 coord[i] = temp % mesh.mesh_shape[i] 

242 temp //= mesh.mesh_shape[i] 

243 mesh_slice = [] 

244 for axis, axis_coord in enumerate(coord): 

245 mesh_slice.append(slice(None) if axis in normalized_axes else axis_coord) 

246 selected = mesh_tensor[tuple(mesh_slice)] 

247 return [int(item) for item in np.array(selected).reshape(-1).tolist()] 

248 

249 

250def get_split_rank_lists_for_axes( 

251 mesh: DeviceMesh, 

252 axes: Sequence[int], 

253) -> list[list[int]]: 

254 """Return all rank lists induced by varying ``axes`` and fixing the complementary axes.""" 

255 normalized_axes = tuple(sorted(set(axes))) 

256 if len(normalized_axes) == 0: 

257 return [[int(rank) for rank in mesh.rank_list]] 

258 

259 mesh_tensor = np.array(mesh.rank_list).reshape(mesh.mesh_shape) 

260 complementary_axes = tuple( 

261 axis for axis in range(len(mesh.mesh_shape)) if axis not in normalized_axes 

262 ) 

263 if len(complementary_axes) == 0: 

264 return [[int(item) for item in np.array(mesh_tensor).reshape(-1).tolist()]] 

265 

266 complementary_shape = tuple(mesh.mesh_shape[axis] for axis in complementary_axes) 

267 split_rank_lists: list[list[int]] = [] 

268 for complementary_coord in np.ndindex(*complementary_shape): 

269 mesh_slice = [] 

270 coord_idx = 0 

271 for axis in range(len(mesh.mesh_shape)): 

272 if axis in normalized_axes: 

273 mesh_slice.append(slice(None)) 

274 else: 

275 mesh_slice.append(complementary_coord[coord_idx]) 

276 coord_idx += 1 

277 selected = mesh_tensor[tuple(mesh_slice)] 

278 split_rank_lists.append([int(item) for item in np.array(selected).reshape(-1).tolist()]) 

279 return split_rank_lists 

280 

281 

282def get_hsdp_state(module): 

283 """Return the HSDPState for a fully_shard-managed module, or None.""" 

284 from hyper_parallel.core.fully_shard.api import HSDPModule # pylint: disable=C0415 

285 if isinstance(module, HSDPModule): 

286 if module.hsdp_scheduler is None: 

287 raise AssertionError("Expect HSDPModule contains 'hsdp_scheduler'.") 

288 return module.hsdp_scheduler.hsdp_state 

289 

290 

291def apply_gradient_scaling_factor(reduced_grad: Any, factor: Any) -> None: 

292 """In-place scale ``reduced_grad`` by ``factor`` (no-op when ``factor`` is ``None``). 

293 

294 Tensor factors are cast to ``reduced_grad``'s dtype so the multiply stays 

295 in the reduction precision (e.g. fp32 accumulation, bf16 reduce). 

296 """ 

297 if factor is None: 

298 return 

299 if isinstance(factor, platform.Tensor) and factor.dtype != reduced_grad.dtype: 

300 factor = factor.to(reduced_grad.dtype) 

301 reduced_grad.mul_(factor) 

302 return None