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

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"""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 

21import torch 

22from torch import nn 

23 

24from hyper_parallel.core.dtensor.device_mesh import DeviceMesh 

25from hyper_parallel.core.dtensor.dtensor import DTensor 

26 

27 

28class ShardedState(Enum): 

29 """ 

30 Parameter shard state 

31 """ 

32 SHARDED = auto() 

33 UNSHARDED = auto() 

34 

35 

36class FullyShardParamMode(Enum): 

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

38 

39 LOCAL_PARAM = auto() 

40 DTENSOR_COMPAT = auto() 

41 DTENSOR_UNIFIED = auto() 

42 

43 

44@dataclass 

45class GroupInfo: 

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

47 

48 group_name: str 

49 group: Any 

50 rank_size: int 

51 

52 

53class FSDPSchedulerState(Enum): 

54 """ 

55 Scheduler state: 

56 - PRE_FORWARD: 

57 already run hook before forward. 

58 - FORWARD: 

59 already run hook after forward. 

60 - PRE_BACKWARD: 

61 already run hook before backward. 

62 - BACKWARD: 

63 already run hook after backward. 

64 """ 

65 PRE_FORWARD = auto() 

66 FORWARD = auto() 

67 PRE_BACKWARD = auto() 

68 BACKWARD = auto() 

69 

70 

71@dataclass 

72class ParamModuleInfo: 

73 """ 

74 Tracks parameter ownership and supports shared weights in HSDP. 

75 

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

77 enabling parameter swapping during sharding/unsharding transitions. Shared 

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

79 

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

81 sharding/unsharding operations. 

82 

83 Attributes: 

84 module: The module that owns this parameter. 

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

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

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

88 """ 

89 module: nn.Module 

90 param_name: str 

91 shared_modules: List[nn.Module] = field(default_factory=list) 

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

93 

94 

95def _named_parameters_with_duplicates( 

96 module: nn.Module, **kwargs: Any 

97) -> list[tuple[str, nn.Parameter]]: 

98 """ 

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

100 `remove_duplicate`. 

101 """ 

102 if "remove_duplicate" in kwargs: 

103 raise AssertionError( 

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

105 ) 

106 kwargs["remove_duplicate"] = False 

107 try: 

108 ret = list(module.named_parameters(**kwargs)) 

109 except AssertionError: 

110 kwargs.pop("remove_duplicate") 

111 ret = list(module.named_parameters(**kwargs)) 

112 return ret 

113 

114 

115def _get_param_module_infos( 

116 params: list[nn.Parameter], modules: tuple[nn.Module, ...] 

117) -> list['ParamModuleInfo']: 

118 """ 

119 Shared parameter: lin1.weight = lin2.weight 

120 Shared module: mlp.lin1 = mlp.lin2 

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

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

123 """ 

124 params_set = set(params) 

125 param_to_module_info: dict[nn.Parameter, ParamModuleInfo] = {} 

126 

127 for module in modules: 

128 for _, submodule in module.named_modules(remove_duplicate=False): 

129 for param_name, param in _named_parameters_with_duplicates( 

130 submodule, recurse=False 

131 ): 

132 if param in params_set: 

133 if param not in param_to_module_info: 

134 param_to_module_info[param] = ParamModuleInfo( 

135 submodule, param_name 

136 ) 

137 else: 

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

139 param_to_module_info[param].shared_param_names.append( 

140 param_name 

141 ) 

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

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

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

145 

146 

147def get_managed_modules_parameters( 

148 modules: Sequence[nn.Module], 

149 ignored_params: Optional[Sequence[nn.Parameter]] = None, 

150) -> list[nn.Parameter]: 

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

152 

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

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

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

156 """ 

157 params: list[nn.Parameter] = [] 

158 ignored_params_set = set(ignored_params or ()) 

159 visited_params: set[nn.Parameter] = set() 

160 for mod in modules: 

161 for _, param in mod.named_parameters(): 

162 if param in ignored_params_set or param in visited_params: 

163 continue 

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

165 continue 

166 visited_params.add(param) 

167 params.append(param) 

168 return params 

169 

170 

171def infer_fully_shard_param_mode( 

172 mesh: Optional[DeviceMesh], 

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

174) -> FullyShardParamMode: 

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

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

177 if not has_dtensor_param: 

178 return FullyShardParamMode.LOCAL_PARAM 

179 if mesh is None: 

180 return FullyShardParamMode.DTENSOR_COMPAT 

181 return FullyShardParamMode.DTENSOR_UNIFIED 

182 

183 

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

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

186 if isinstance(param, DTensor): 

187 return param 

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

189 if isinstance(param_data, DTensor): 

190 return param_data 

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

192 return param 

193 return None 

194 

195 

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

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

198 return unwrap_dtensor_param(param) is not None 

199 

200 

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

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

203 payload = unwrap_dtensor_param(param) 

204 if payload is None: 

205 return None 

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

207 

208 

209def get_rank_list_for_axes( 

210 mesh: DeviceMesh, 

211 axes: Sequence[int], 

212 rank: Optional[int] = None, 

213) -> list[int]: 

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

215 if rank is None: 

216 rank = mesh.rank 

217 if rank not in mesh.rank_list: 

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

219 

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

221 if len(normalized_axes) == 0: 

222 return [rank] 

223 

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

225 rank_index = mesh.rank_list.index(rank) 

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

227 temp = rank_index 

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

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

230 temp //= mesh.mesh_shape[i] 

231 mesh_slice = [] 

232 for axis, axis_coord in enumerate(coord): 

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

234 selected = mesh_tensor[tuple(mesh_slice)] 

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

236 

237 

238def get_split_rank_lists_for_axes( 

239 mesh: DeviceMesh, 

240 axes: Sequence[int], 

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

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

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

244 if len(normalized_axes) == 0: 

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

246 

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

248 complementary_axes = tuple( 

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

250 ) 

251 if len(complementary_axes) == 0: 

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

253 

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

255 split_rank_lists: list[list[int]] = [] 

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

257 mesh_slice = [] 

258 coord_idx = 0 

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

260 if axis in normalized_axes: 

261 mesh_slice.append(slice(None)) 

262 else: 

263 mesh_slice.append(complementary_coord[coord_idx]) 

264 coord_idx += 1 

265 selected = mesh_tensor[tuple(mesh_slice)] 

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

267 return split_rank_lists 

268 

269 

270def get_hsdp_state(module): 

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

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

273 if isinstance(module, HSDPModule): 

274 if module.hsdp_scheduler is None: 

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

276 return module.hsdp_scheduler.hsdp_state 

277 

278 

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

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

281 

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

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

284 """ 

285 if factor is None: 

286 return 

287 if isinstance(factor, torch.Tensor) and factor.dtype != reduced_grad.dtype: 

288 factor = factor.to(reduced_grad.dtype) 

289 reduced_grad.mul_(factor) 

290 return None