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

1# Copyright 2025 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 parameter""" 

16 

17from hyper_parallel.core.dtensor.device_mesh import DeviceMesh 

18from hyper_parallel.core.dtensor.dtensor import DTensor 

19from hyper_parallel.core.dtensor.placement_types import Replicate 

20from hyper_parallel.core.fully_shard.hsdp_utils import ( 

21 FullyShardParamMode, 

22 GroupInfo, 

23 get_rank_list_for_axes, 

24 get_split_rank_lists_for_axes, 

25) 

26from hyper_parallel.core.fully_shard.utils import DDPMeshInfo, FSDPMeshInfo 

27from hyper_parallel.platform import get_platform 

28 

29platform = get_platform() 

30_GROUP_INFO_CACHE = {} 

31 

32 

33def _build_group_info_from_rank_list(group_name: str, rank_list) -> GroupInfo: 

34 """Create group metadata from an explicit rank list.""" 

35 normalized_rank_list = tuple(sorted(int(rank) for rank in rank_list)) 

36 if len(normalized_rank_list) <= 1: 

37 return GroupInfo(f"{group_name}_invalid", None, 1) 

38 if normalized_rank_list in _GROUP_INFO_CACHE: 

39 cached_group = _GROUP_INFO_CACHE[normalized_rank_list] 

40 return GroupInfo(str(normalized_rank_list), cached_group, len(normalized_rank_list)) 

41 try: 

42 group = platform.create_group(list(normalized_rank_list)) 

43 except (RuntimeError, ValueError): # pragma: no cover - UT may run without dist init 

44 group = None 

45 _GROUP_INFO_CACHE[normalized_rank_list] = group 

46 return GroupInfo(str(normalized_rank_list), group, len(normalized_rank_list)) 

47 

48 

49def _build_group_info_from_process_group( 

50 group_name: str, 

51 process_group, 

52 rank_size: int, 

53 *, 

54 resolved_group_name: str | None = None, 

55) -> GroupInfo: 

56 """Create group metadata from an existing process group.""" 

57 if process_group is None or rank_size <= 1: 

58 return GroupInfo(f"{group_name}_invalid", None, 1) 

59 return GroupInfo(resolved_group_name or group_name, process_group, rank_size) 

60 

61 

62class HSDPParamV2: 

63 """ 

64 HSDP parameter. 

65 """ 

66 

67 def __repr__(self) -> str: 

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

69 

70 Prefers the parameter FQN (assigned by the root forward pre-hook); before 

71 that, falls back to the owning module class plus param name and object id. 

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

73 """ 

74 fqn = getattr(self, "_param_fqn", None) 

75 if fqn: 

76 return str(fqn) 

77 module_info = getattr(self, "_module_info", None) 

78 if module_info is not None: 

79 return f"{module_info.module.__class__.__name__}.{module_info.param_name}@{id(self):x}" 

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

81 

82 def __init__( 

83 self, 

84 param, 

85 module_info, 

86 mesh_info, 

87 post_forward_mesh_info, 

88 shard_placement_fn, 

89 mp_policy, 

90 offload_policy, 

91 threshold, 

92 ): 

93 """ 

94 Initialize HSDPParamV2. 

95 

96 Args: 

97 param (nn.Parameter): The original parameter to shard. 

98 module_info (ParamModuleInfo): Ownership and shared-weight metadata for the parameter. 

99 mesh_info (FSDPMeshInfo): Mesh topology describing shard/replicate dimensions. 

100 post_forward_mesh_info: Mesh info used after forward (reserved for subclass use). 

101 shard_placement_fn (Callable, optional): Returns a Shard placement for the parameter, 

102 or None to use default (Shard(0)). 

103 mp_policy (MixedPrecisionPolicy, optional): Mixed precision dtype policy. 

104 offload_policy (OffloadPolicy, optional): CPU offload policy. 

105 threshold: Minimum parameter size to enable sharding (reserved for subclass use). 

106 """ 

107 raise NotImplementedError("HSDP param subclasses must implement __init__") 

108 

109 def _init_sharded_param(self, param, shard_placement_fn): 

110 """add and init sharded param""" 

111 raise NotImplementedError("HSDP param subclasses must implement _init_sharded_param") 

112 

113 def init_dtype_attrs(self, mp_policy): 

114 """Initialize dtype attributes from mixed precision policy.""" 

115 raise NotImplementedError("HSDP param subclasses must implement init_dtype_attrs") 

116 

117 def init_all_gather_outputs( 

118 self, all_gather_input_numels, all_gather_input_dtypes, world_size, device, force_recreate=False 

119 ): 

120 """Allocate or reuse output buffers for all-gather communication.""" 

121 raise NotImplementedError("HSDP param subclasses must implement init_all_gather_outputs") 

122 

123 def init_unsharded_param(self): 

124 """Reconstruct the full unsharded parameter from all-gather outputs.""" 

125 raise NotImplementedError("HSDP param subclasses must implement init_unsharded_param") 

126 

127 def to_sharded(self): 

128 """Transition parameter from unsharded back to sharded state and free unsharded storage.""" 

129 raise NotImplementedError("HSDP param subclasses must implement to_sharded") 

130 

131 def to_unsharded(self): 

132 """Transition parameter to unsharded state after all-gather completes.""" 

133 raise NotImplementedError("HSDP param subclasses must implement to_unsharded") 

134 

135 def to_sharded_dtensor(self, tensor): 

136 """Wrap a local sharded tensor as a DTensor with the correct mesh and placements.""" 

137 raise NotImplementedError("HSDP param subclasses must implement to_sharded_dtensor") 

138 

139 def to_accumulated_grad_if_needed(self): 

140 """Move unsharded grad to accumulated grad buffer if dtype conversion is required.""" 

141 raise NotImplementedError("HSDP param subclasses must implement to_accumulated_grad_if_needed") 

142 

143 def accumulate_unsharded_grad_if_needed(self): 

144 """Accumulate unsharded param grad into accumulated grad buffer if both exist.""" 

145 raise NotImplementedError("HSDP param subclasses must implement accumulate_unsharded_grad_if_needed") 

146 

147 def alloc_all_gather_outputs(self): 

148 """Resize all-gather output buffers to their full capacity for communication.""" 

149 raise NotImplementedError("HSDP param subclasses must implement alloc_all_gather_outputs") 

150 

151 def free_unsharded_param(self): 

152 """Release storage of all-gather outputs and inner tensors to free device memory.""" 

153 raise NotImplementedError("HSDP param subclasses must implement free_unsharded_param") 

154 

155 @property 

156 def all_gather_inputs(self): 

157 """Return the local sharded tensor(s) to use as input for all-gather communication.""" 

158 raise NotImplementedError("HSDP param subclasses must implement all_gather_inputs") 

159 

160 @property 

161 def unsharded_param(self): 

162 """Return the full unsharded parameter after all-gather.""" 

163 raise NotImplementedError("HSDP param subclasses must implement unsharded_param") 

164 

165 @property 

166 def unsharded_grad_data(self): 

167 """Return the unsharded_param.grad.""" 

168 raise NotImplementedError("HSDP param subclasses must implement unsharded_grad_data") 

169 

170 @property 

171 def unsharded_accumulated_grad_data(self): 

172 """Return the unsharded accumulated gradient buffer.""" 

173 raise NotImplementedError("HSDP param subclasses must implement unsharded_accumulated_grad_data") 

174 

175 @property 

176 def _sharded_local_tensor(self): 

177 """Return the underlying local tensor of the sharded DTensor parameter.""" 

178 raise NotImplementedError("HSDP param subclasses must implement _sharded_local_tensor") 

179 

180 def _get_unsharded_param_data(self, async_op=False): 

181 """Perform all-gather to obtain unsharded parameter data, returning (tensor, handle).""" 

182 raise NotImplementedError("HSDP param subclasses must implement _get_unsharded_param_data") 

183 

184 def unshard(self, async_op=False): 

185 """Trigger all-gather to unshard the parameter, optionally asynchronously.""" 

186 raise NotImplementedError("HSDP param subclasses must implement unshard") 

187 

188 def wait_for_unshard(self): 

189 """Wait for all-gather to complete and transition parameter to unsharded state.""" 

190 raise NotImplementedError("HSDP param subclasses must implement wait_for_unshard") 

191 

192 def shard(self): 

193 """Transition parameter from unsharded back to sharded state.""" 

194 raise NotImplementedError("HSDP param subclasses must implement shard") 

195 

196 def reduce_scatter_grad(self): 

197 """Perform reduce-scatter on the unsharded gradient to produce a sharded gradient.""" 

198 raise NotImplementedError("HSDP param subclasses must implement reduce_scatter_grad") 

199 

200 def all_reduce_grad(self): 

201 """Perform all-reduce on gradient across the replicate dimension (HSDP mode only).""" 

202 raise NotImplementedError("HSDP param subclasses must implement all_reduce_grad") 

203 

204 def _resolve_process_group_name(self, group_name: str, process_group) -> str: 

205 """Resolve the name recorded in GroupInfo for an existing process group.""" 

206 del process_group 

207 return group_name 

208 

209 def _get_base_spmd_placements(self) -> tuple: 

210 """Return placements before explicit data-parallel semantics are applied.""" 

211 if ( 

212 getattr(self, "param_mode", None) == FullyShardParamMode.DTENSOR_UNIFIED 

213 and getattr(self, "_orig_param_is_dtensor", False) 

214 ): 

215 self._spmd_mesh = DeviceMesh.concatenate([self.mesh_info.mesh, self._orig_dtensor_mesh]) 

216 dp_prefix_placements = tuple(Replicate() for _ in range(self.mesh_info.mesh.ndim)) 

217 return dp_prefix_placements + tuple(self._orig_dtensor_placements) 

218 

219 if ( 

220 getattr(self, "param_mode", None) == FullyShardParamMode.DTENSOR_COMPAT 

221 and getattr(self, "_orig_param_is_dtensor", False) 

222 ): 

223 self._spmd_mesh = self._orig_dtensor_mesh 

224 return tuple(self._orig_dtensor_placements) 

225 

226 self._spmd_mesh = self.mesh_info.mesh 

227 return tuple(Replicate() for _ in range(self._spmd_mesh.ndim)) 

228 

229 def _get_data_parallel_shard_placement(self, placements: list, shard_placement): 

230 """Return the placement to apply on the explicit fully_shard dimension.""" 

231 del placements 

232 return shard_placement 

233 

234 def _apply_data_parallel_placements(self, placements: list, shard_placement) -> tuple: 

235 """Apply explicit DDP/FSDP placements on top of the base SPMD layout.""" 

236 if len(placements) != self._spmd_mesh.ndim: 

237 raise AssertionError( 

238 f"Expected {self._spmd_mesh.ndim} unified placements, got {len(placements)}: {placements}" 

239 ) 

240 

241 spmd_replicate_mesh_dim = getattr(self, "_spmd_replicate_mesh_dim", None) 

242 if ( 

243 isinstance(self.mesh_info, DDPMeshInfo) 

244 and spmd_replicate_mesh_dim is not None 

245 and not getattr(self, "_orig_param_is_dtensor", False) 

246 ): 

247 placements[spmd_replicate_mesh_dim] = Replicate() 

248 

249 spmd_shard_mesh_dim = getattr(self, "_spmd_shard_mesh_dim", None) 

250 if ( 

251 getattr(self, "uses_param_shard", False) 

252 and isinstance(self.mesh_info, FSDPMeshInfo) 

253 and spmd_shard_mesh_dim is not None 

254 ): 

255 placements[spmd_shard_mesh_dim] = self._get_data_parallel_shard_placement( 

256 placements, shard_placement 

257 ) 

258 return tuple(placements) 

259 

260 def _init_group_infos(self) -> None: 

261 """Initialize sharded/unsharded communication groups from the current layout.""" 

262 if ( 

263 getattr(self, "uses_param_shard", False) 

264 and getattr(self, "is_sharded", False) 

265 and isinstance(self.mesh_info, FSDPMeshInfo) 

266 ): 

267 resolved_group_name = self._resolve_process_group_name( 

268 "fully_shard_sharded_group", 

269 self.mesh_info.shard_process_group, 

270 ) 

271 self.sharded_group_info = _build_group_info_from_process_group( 

272 "fully_shard_sharded_group", 

273 self.mesh_info.shard_process_group, 

274 self.mesh_info.shard_mesh_size, 

275 resolved_group_name=resolved_group_name, 

276 ) 

277 else: 

278 self.sharded_group_info = GroupInfo("fully_shard_sharded_group_invalid", None, 1) 

279 

280 self.unsharded_group_info = self._build_layout_driven_group_info() 

281 self.shard_size = self.sharded_group_info.rank_size 

282 self.dp_size = self.unsharded_group_info.rank_size 

283 self.rank_size = max(1, self.shard_size * self.dp_size) 

284 

285 def _build_layout_driven_group_info(self) -> GroupInfo: 

286 """Build the group that should all-reduce an unsharded gradient from the final layout.""" 

287 group_axes = [ 

288 axis 

289 for axis, placement in enumerate(self._spmd_placements) 

290 if placement.is_replicate() 

291 ] 

292 spmd_shard_mesh_dim = getattr(self, "_spmd_shard_mesh_dim", None) 

293 if getattr(self, "uses_param_shard", False) and spmd_shard_mesh_dim is not None: 

294 group_axes = [axis for axis in group_axes if axis != spmd_shard_mesh_dim] 

295 if not group_axes: 

296 return GroupInfo("fully_shard_unsharded_group_invalid", None, 1) 

297 

298 group_dim_names = getattr(self._spmd_mesh, "mesh_dim_names", None) 

299 if group_dim_names: 

300 try: 

301 mesh_axis_names = tuple(group_dim_names[axis] for axis in group_axes) 

302 if len(mesh_axis_names) == 1: 

303 axis_name = mesh_axis_names[0] 

304 process_group = self._spmd_mesh.get_group(axis_name) 

305 if process_group is not None: 

306 rank_size = self._spmd_mesh.mesh_shape[group_dim_names.index(axis_name)] 

307 resolved_group_name = self._resolve_process_group_name( 

308 "fully_shard_unsharded_group", 

309 process_group, 

310 ) 

311 return _build_group_info_from_process_group( 

312 "fully_shard_unsharded_group", 

313 process_group, 

314 rank_size, 

315 resolved_group_name=resolved_group_name, 

316 ) 

317 

318 split_rank_lists = get_split_rank_lists_for_axes(self._spmd_mesh, group_axes) 

319 process_group = platform.split_group(split_ranks=split_rank_lists) 

320 if process_group is not None: 

321 rank_size = 1 

322 for axis in group_axes: 

323 rank_size *= self._spmd_mesh.mesh_shape[axis] 

324 resolved_group_name = self._resolve_process_group_name( 

325 "fully_shard_unsharded_group", 

326 process_group, 

327 ) 

328 return _build_group_info_from_process_group( 

329 "fully_shard_unsharded_group", 

330 process_group, 

331 rank_size, 

332 resolved_group_name=resolved_group_name, 

333 ) 

334 except ( 

335 AssertionError, 

336 AttributeError, 

337 KeyError, 

338 RuntimeError, 

339 TypeError, 

340 ValueError, 

341 ): 

342 pass 

343 

344 rank_list = get_rank_list_for_axes(self._spmd_mesh, group_axes) 

345 return _build_group_info_from_rank_list("fully_shard_unsharded_group", rank_list) 

346 

347 def _normalize_unsharded_grad_to_local(self, grad, *, reduce_partial_dtensor: bool = True): 

348 """Normalize a pending gradient to the local tensor expected by fully_shard collectives.""" 

349 if not isinstance(grad, DTensor): 

350 return grad 

351 

352 if reduce_partial_dtensor and any(placement.is_partial() for placement in grad.placements): 

353 grad = grad.reduce_partial() 

354 

355 orig_dtensor_mesh = getattr(self, "_orig_dtensor_mesh", None) 

356 orig_dtensor_placements = getattr(self, "_orig_dtensor_placements", None) 

357 mesh_mismatch = ( 

358 orig_dtensor_mesh is not None 

359 and grad.device_mesh.to_hash() != orig_dtensor_mesh.to_hash() 

360 ) 

361 placement_mismatch = ( 

362 orig_dtensor_placements is not None 

363 and tuple(grad.placements) != tuple(orig_dtensor_placements) 

364 ) 

365 if mesh_mismatch or placement_mismatch: 

366 grad = grad.redistribute(orig_dtensor_mesh, orig_dtensor_placements) 

367 return grad.to_local()