Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / optimizer / optimizer.py: 0%

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 

16"""Base distributed optimizer and chain optimizer composition.""" 

17 

18from collections import defaultdict 

19import logging 

20from typing import Any, Dict, List, Optional, Tuple, Union 

21 

22import torch 

23import torch.distributed as dist 

24from torch.distributed.checkpoint.state_dict import ( 

25 StateDictOptions, 

26 get_optimizer_state_dict, 

27 set_optimizer_state_dict, 

28) 

29from hyper_parallel.core.optimizer.dtensor_compat import to_local_if_dtensor 

30from hyper_parallel.core.optimizer.sharding_category import ( 

31 HSDPGroupAssignment, 

32 build_owner_by_size, 

33 get_multi_dim_logical_info, 

34 select_owned_records, 

35 group_parameters_for_hsdp 

36) 

37from hyper_parallel.platform import get_platform 

38 

39logging.basicConfig(level=logging.INFO) 

40logger = logging.getLogger(__name__) 

41 

42 

43class ChainedOptimizer: 

44 """Composition wrapper that dispatches step/zero_grad to sub-optimizers.""" 

45 

46 def __init__( 

47 self, 

48 model: torch.nn.Module, 

49 optimizers: Dict[str, torch.optim.Optimizer], 

50 flatten: bool = False 

51 ) -> None: 

52 self.optimizers_dict = optimizers 

53 self.chained_optimizers = list(optimizers.values()) 

54 self.optimizers_keys = list(optimizers.keys()) 

55 self.model = model 

56 self.flatten = flatten # not flatten adamw, flatten for multi-optimizer 

57 self._is_multi_optimizer = flatten 

58 

59 def __iter__(self): 

60 """Allow iteration over the underlying optimizers.""" 

61 return iter(self.chained_optimizers) 

62 

63 def step(self) -> None: 

64 """Call each sub-optimizer's step in order.""" 

65 for opt in self.chained_optimizers: 

66 opt.step() 

67 

68 def zero_grad(self, set_to_none: bool = True) -> None: 

69 """Clear gradients for all sub-optimizers.""" 

70 for opt in self.chained_optimizers: 

71 opt.zero_grad(set_to_none=set_to_none) 

72 

73 @property 

74 def optimizer(self) -> torch.optim.Optimizer: 

75 """Access underlying optimizer when only one optimizer included for backward compatibility.""" 

76 if len(self.chained_optimizers) != 1: 

77 raise ValueError("ChainedOptimizer has more than one optimizer when accessing self.optimizer") 

78 return self.chained_optimizers[0] 

79 

80 @property 

81 def defaults(self) -> Dict[str, Any]: 

82 """Return the defaults of the first sub-optimizer.""" 

83 return self.chained_optimizers[0].defaults 

84 

85 def state_dict(self) -> Dict[str, Any]: 

86 """Return state dicts with DTensor values localized to CPU for serialization. 

87 

88 In HSDP, optimizer state is only populated on the owner rank for replicated 

89 parameters. This method first broadcasts state to all replicate-group peers, 

90 then converts DTensor values to local CPU tensors so ``torch.save`` works. 

91 """ 

92 # Ensure all ranks have consistent optimizer state before snapshotting 

93 for opt in self.chained_optimizers: 

94 if hasattr(opt, "_broadcast_state_fused_for_ckpt"): 

95 opt._broadcast_state_fused_for_ckpt() # pylint: disable=protected-access 

96 

97 merged: Dict[str, Any] = {} 

98 for name, optimizer in self.optimizers_dict.items(): 

99 sd = get_optimizer_state_dict( 

100 self.model, 

101 optimizer, 

102 options=StateDictOptions(flatten_optimizer_state_dict=self.flatten) 

103 ) 

104 overlap = set(merged.keys()) & set(sd.keys()) 

105 if overlap: 

106 raise KeyError( 

107 f"Key clash detected while merging state dict for optimizer '{name}': " 

108 f"{', '.join(sorted(overlap))}" 

109 ) 

110 merged.update(sd) 

111 

112 return merged 

113 

114 def _get_param_groups(self) -> List[Dict[str, Any]]: 

115 """Get param_groups aggregated over underlying optimizers.""" 

116 param_groups: List[Dict[str, Any]] = [] 

117 for optimizer in self.chained_optimizers: 

118 param_groups += optimizer.param_groups 

119 return param_groups 

120 

121 def _set_param_groups(self, new_param_groups: List[Dict[str, Any]]) -> None: 

122 """Set param_groups distributed across underlying optimizers.""" 

123 if not isinstance(new_param_groups, list): 

124 raise TypeError("new_param_groups should be a list") 

125 if len(new_param_groups) != len(self.param_groups): 

126 raise ValueError("The size of new_param_groups must be equal to origin param_groups") 

127 

128 start = 0 

129 for optimizer in self.chained_optimizers: 

130 group_len = len(optimizer.param_groups) 

131 optimizer.param_groups = new_param_groups[start: start + group_len] 

132 start += group_len 

133 

134 param_groups = property(_get_param_groups, _set_param_groups) 

135 

136 def load_state_dict(self, state_dict: Dict[str, Any]) -> None: 

137 """Load optimizer state dicts and synchronize steps.""" 

138 for optimizer in self.chained_optimizers: 

139 set_optimizer_state_dict( 

140 self.model, 

141 optimizer, 

142 optim_state_dict=state_dict, 

143 options=StateDictOptions(flatten_optimizer_state_dict=self.flatten), 

144 ) 

145 

146 self._synchronize_steps() 

147 

148 def _synchronize_steps(self) -> Optional[int]: 

149 """Synchronize the step of all optimizers. 

150 

151 TE FusedAdam will not accumulate "step" for empty param groups, 

152 so we need to align the step across param groups before saving and after loading. 

153 """ 

154 steps = [] 

155 for optimizer in self.chained_optimizers: 

156 actual_opt = getattr(optimizer, 'optimizer', optimizer) 

157 for param_group in actual_opt.param_groups: 

158 if len(param_group['params']) > 0 and 'step' in param_group: 

159 steps.append(param_group['step']) 

160 

161 unique_steps = list(set(steps)) 

162 if len(unique_steps) > 1: 

163 raise ValueError(f"steps should <= 1, but got {unique_steps}") 

164 

165 step = unique_steps[0] if len(unique_steps) == 1 else None 

166 

167 for optimizer in self.chained_optimizers: 

168 actual_opt = getattr(optimizer, 'optimizer', optimizer) 

169 for param_group in actual_opt.param_groups: 

170 param_group['step'] = step 

171 

172 return step 

173 

174 

175class BaseDistributedOptimizer(torch.optim.Optimizer): 

176 """Base class for distributed optimizers with HSDP-aware topology communication. 

177 

178 Provides fused hierarchical broadcast for parameters and optimizer states. 

179 """ 

180 

181 def __init__( 

182 self, 

183 params: Any, 

184 defaults: Dict[str, Any], 

185 is_muon: bool, 

186 hsdp_replica_count: Optional[Union[int, Tuple[int, ...]]] = None, 

187 ) -> None: 

188 super().__init__(params, defaults) 

189 self.is_muon = is_muon 

190 self.hsdp_replica_count = hsdp_replica_count 

191 self._param_to_broadcast_info: Dict[ 

192 torch.nn.Parameter, Tuple[Tuple[int, ...], Tuple[dist.ProcessGroup, ...]] 

193 ] = {} 

194 

195 # Cache: (parent_ranks_tuple, sub_size) -> {sub_idx: sub_pg} 

196 self._split_sub_pg_cache: Dict[Tuple[Tuple[int, ...], int], Dict[int, dist.ProcessGroup]] = {} 

197 

198 def _group_dtensor_by_mesh(self): 

199 """Group DTensor parameters by mesh topology and shard layout.""" 

200 self._hsdp_grouping: Dict[int, Tuple[List, List]] = {} 

201 for group_key, group in enumerate(self.param_groups): 

202 no_comm_params, hsdp_groups = group_parameters_for_hsdp(group["params"]) 

203 self._hsdp_grouping[group_key] = (no_comm_params, hsdp_groups) 

204 

205 def _auto_deduce_replica_count(self) -> Optional[Union[int, Tuple[int, ...]]]: 

206 """Deduce hsdp_replica_count based on cluster topology. 

207 

208 - Intra-node PGs: Full dedup (no split), high bandwidth makes broadcast cheap. 

209 - Inter-node PGs: Split at node boundaries to restrict communication domains  

210 within a single node, bypassing cross-node bottlenecks. 

211 """ 

212 devices_per_node = 1 

213 if torch.npu.is_available(): 

214 devices_per_node = torch.npu.device_count() 

215 elif torch.cuda.is_available(): 

216 devices_per_node = torch.cuda.device_count() 

217 

218 dedup_per_dim: Dict[int, int] = {} 

219 needs_split = False 

220 

221 # group_key, (no_comm_params, hsdp_groups) 

222 for _, (_, hsdp_groups) in self._hsdp_grouping.items(): 

223 for hsdp_group in hsdp_groups: 

224 for dim_idx, pg in enumerate(hsdp_group.replicate_pgs): 

225 if pg is None: 

226 continue 

227 pg_size = dist.get_world_size(pg) 

228 if pg_size <= 1: 

229 continue 

230 

231 if pg_size > devices_per_node: 

232 # Inter-node: Find largest divisor safe for node boundary 

233 dedup = min(pg_size, devices_per_node) 

234 while pg_size % dedup != 0: 

235 dedup -= 1 

236 needs_split = True 

237 else: 

238 # origin Inter-node 

239 dedup = pg_size 

240 

241 # Enforce conservative (smallest) dedup across shared mesh axes 

242 if dim_idx not in dedup_per_dim: 

243 dedup_per_dim[dim_idx] = dedup 

244 else: 

245 dedup_per_dim[dim_idx] = min(dedup_per_dim[dim_idx], dedup) 

246 

247 if not needs_split: 

248 return None 

249 

250 sorted_dedups = [dedup_per_dim[k] for k in sorted(dedup_per_dim.keys())] 

251 if len(sorted_dedups) == 1: 

252 return sorted_dedups[0] 

253 

254 return tuple(sorted_dedups) 

255 

256 def _split_replicate_groups(self) -> None: 

257 """Split replicate ProcessGroups into smaller sub-groups based on hsdp_replica_count.""" 

258 if self.hsdp_replica_count is None: 

259 return 

260 

261 # Argument validation 

262 if isinstance(self.hsdp_replica_count, int): 

263 if self.hsdp_replica_count <= 0: 

264 raise ValueError(f"hsdp_replica_count must be positive, got {self.hsdp_replica_count}") 

265 elif isinstance(self.hsdp_replica_count, tuple): 

266 if not self.hsdp_replica_count: 

267 raise ValueError("hsdp_replica_count tuple must not be empty") 

268 for i, v in enumerate(self.hsdp_replica_count): 

269 if not isinstance(v, int) or v <= 0: 

270 raise ValueError(f"hsdp_replica_count[{i}] must be positive, got {v}") 

271 else: 

272 raise TypeError(f"Unsupported hsdp_replica_count type: {type(self.hsdp_replica_count).__name__}") 

273 

274 for group_key, (_, hsdp_groups) in self._hsdp_grouping.items(): 

275 for hsdp_group in hsdp_groups: 

276 new_replicate_pgs: List[dist.ProcessGroup] = [] 

277 for dim_idx, pg in enumerate(hsdp_group.replicate_pgs): 

278 if pg is None: 

279 new_replicate_pgs.append(pg) 

280 continue 

281 

282 pg_size = dist.get_world_size(pg) 

283 dedup_size = self._get_dedup_size_for_dim(dim_idx, pg_size) 

284 

285 if pg_size <= dedup_size: 

286 new_replicate_pgs.append(pg) 

287 continue 

288 

289 if pg_size % dedup_size != 0: 

290 raise ValueError( 

291 f"hsdp_replica_count {dedup_size} must evenly divide replicate group size {pg_size}" 

292 ) 

293 

294 sub_pg = self._get_or_create_sub_pg(pg, dedup_size) 

295 new_replicate_pgs.append(sub_pg) 

296 

297 logger.info_rank0( 

298 "[HSDP Split] group_key=%s, dim_idx=%s, original_size=%s, sub_size=%s", 

299 group_key, dim_idx, pg_size, dedup_size 

300 ) 

301 

302 # replace new sub groups of hsdp groups 

303 hsdp_group.replicate_pgs = tuple(new_replicate_pgs) 

304 

305 def _get_dedup_size_for_dim(self, dim_idx: int, pg_size: int) -> int: 

306 """Get the effective dedup size for a specific replicate dimension.""" 

307 if isinstance(self.hsdp_replica_count, int): 

308 return self.hsdp_replica_count 

309 

310 if dim_idx < len(self.hsdp_replica_count): 

311 return self.hsdp_replica_count[dim_idx] 

312 

313 return pg_size 

314 

315 def _get_or_create_sub_pg( 

316 self, 

317 parent_pg: dist.ProcessGroup, 

318 sub_size: int, 

319 ) -> dist.ProcessGroup: 

320 """Retrieve or collectively create a synchronized sub-ProcessGroup.""" 

321 local_parent_ranks = tuple(sorted(list(dist.get_process_group_ranks(parent_pg)))) 

322 cache_key = (local_parent_ranks, sub_size) 

323 

324 # Cache Hit Check 

325 if cache_key in self._split_sub_pg_cache: 

326 sub_pg_map = self._split_sub_pg_cache[cache_key] 

327 for sub_idx, sub_pg in sub_pg_map.items(): 

328 if sub_pg is not None: 

329 try: 

330 dist.get_rank(group=sub_pg) 

331 return sub_pg 

332 except RuntimeError: 

333 continue 

334 raise RuntimeError(f"Current rank not found in cached sub-groups for parent_pg={local_parent_ranks}") 

335 

336 global_rank = dist.get_rank() 

337 world_size = dist.get_world_size() 

338 

339 # Global Rendezvous via GPU all_gather_into_tensor (NCCL/HCCL fast-path). 

340 # Far more scalable than all_gather_object for large world_size. 

341 device = torch.npu.current_device() if torch.npu.is_available() else torch.cuda.current_device() 

342 num_parent_ranks = len(local_parent_ranks) 

343 local_tensor = torch.tensor(local_parent_ranks, dtype=torch.long, device=device) 

344 gathered_tensor = torch.empty((world_size, num_parent_ranks), dtype=torch.long, device=device) 

345 dist.all_gather_into_tensor(gathered_tensor.view(-1), local_tensor) 

346 

347 # Deduplicate: each row is one rank's parent_ranks; convert to 

348 # sorted set of tuples for deterministic iteration order. 

349 gathered_cpu_list = gathered_tensor.cpu().tolist() 

350 unique_all_parent_ranks = sorted(set( 

351 tuple(row) for row in gathered_cpu_list 

352 )) 

353 

354 sub_pg_map: Dict[int, dist.ProcessGroup] = {} 

355 my_sub_pg: Optional[dist.ProcessGroup] = None 

356 

357 # Synchronized collective creation loop 

358 for parent_ranks in unique_all_parent_ranks: 

359 num_sub_groups = len(parent_ranks) // sub_size 

360 

361 for sub_idx in range(num_sub_groups): 

362 sub_ranks = parent_ranks[sub_idx * sub_size: (sub_idx + 1) * sub_size] 

363 sub_pg = dist.new_group(sub_ranks) 

364 

365 # Map results only if this parent ranks list matches the current rank's context 

366 if parent_ranks == local_parent_ranks: 

367 if global_rank in sub_ranks: 

368 my_sub_pg = sub_pg 

369 sub_pg_map[sub_idx] = sub_pg 

370 else: 

371 sub_pg_map[sub_idx] = None 

372 

373 self._split_sub_pg_cache[cache_key] = sub_pg_map 

374 

375 if my_sub_pg is None: 

376 raise RuntimeError( 

377 f"Rank {global_rank} not found in any sub-group of parent_pg " 

378 f"with ranks {local_parent_ranks} and sub_size={sub_size}" 

379 ) 

380 

381 return my_sub_pg 

382 

383 def _build_hsdp_batch( 

384 self, 

385 max_batch_numel: Optional[int] = None, 

386 ) -> None: 

387 """Split HSDP groups into memory-capped batches for compute-broadcast overlap.""" 

388 if max_batch_numel is None: 

389 broadcast_max_bytes = getattr( 

390 self, "replicate_broadcast_max_bytes", 512 * 1024 * 1024, 

391 ) 

392 hsdp_size = self.hsdp_replica_count if self.hsdp_replica_count is not None else 1 

393 max_batch_numel = broadcast_max_bytes * hsdp_size if hsdp_size > 1 else float('inf') 

394 

395 self._hsdp_batches: Dict[int, List[Dict]] = {} 

396 

397 for group_key, (_, hsdp_groups) in self._hsdp_grouping.items(): 

398 # Sort groups by total numel descending — large groups first. 

399 sorted_groups = sorted( 

400 hsdp_groups, 

401 key=lambda g: -sum(r.param.numel() for r in g.records), 

402 ) 

403 

404 batch_groups: List[Dict] = [] 

405 

406 for hsdp_group in sorted_groups: 

407 # Sort records within group by numel descending. 

408 sorted_records = sorted( 

409 hsdp_group.records, 

410 key=lambda r: -r.param.numel(), 

411 ) 

412 

413 sub_batches: List[List[Dict]] = [] 

414 current_batch: List[Dict] = [] 

415 current_numel = 0 

416 

417 for record in sorted_records: 

418 p_numel = record.param.numel() 

419 

420 current_batch.append({ 

421 "record": record, 

422 "hsdp_group": hsdp_group, 

423 }) 

424 current_numel += p_numel 

425 

426 # Soft limit: allow the bucket to slightly exceed the cap 

427 # so that symmetric structures stay together and fragmentation 

428 # is reduced (same approach as PyTorch FSDP/DDP bucketing). 

429 if current_numel >= max_batch_numel: 

430 sub_batches.append(current_batch) 

431 current_batch = [] 

432 current_numel = 0 

433 

434 if current_batch: 

435 sub_batches.append(current_batch) 

436 

437 # Sort sub-batches by numel descending within this group. 

438 sub_batches.sort(key=lambda b: -sum(e["record"].param.numel() for e in b)) 

439 

440 batch_groups.append({ 

441 "hsdp_group": hsdp_group, 

442 "sub_batches": sub_batches, 

443 }) 

444 

445 # log batch split info. 

446 total_sub_batches = sum(len(bg["sub_batches"]) for bg in batch_groups) 

447 logger.info_rank0( 

448 "[HSDP Batch] group_key=%s, num_hsdp_groups=%s, num_batch_groups=%s, " 

449 "total_sub_batches=%s, group_numels=%s, max_batch_numel=%s", 

450 group_key, 

451 len(hsdp_groups), 

452 len(batch_groups), 

453 total_sub_batches, 

454 [sum(bg['hsdp_group'].records_numel) if hasattr(bg['hsdp_group'], 'records_numel') else sum( 

455 r.param.numel() for r in bg['hsdp_group'].records) for bg in batch_groups], 

456 max_batch_numel 

457 ) 

458 self._hsdp_batches[group_key] = batch_groups 

459 

460 def _build_sub_batch_assignment( 

461 self, 

462 sub_batch_entries: List[Dict], 

463 hsdp_group: Any, 

464 ) -> Optional[HSDPGroupAssignment]: 

465 """Build an HSDPGroupAssignment for one sub-batch and update broadcast info.""" 

466 records = [e["record"] for e in sub_batch_entries] 

467 if not records: 

468 return None 

469 

470 device_mesh = records[0].param.device_mesh 

471 

472 replicate_group_ranks, replicate_sizes = get_multi_dim_logical_info( 

473 device_mesh, 

474 hsdp_group.comm_key.replicate_mesh_dims, 

475 ) 

476 

477 # When hsdp_replica_count is set, remap coordinates into 

478 # sub-groups: each original coord maps to (coord % sub_size) 

479 # within its sub-group, and the effective group size shrinks. 

480 # Supports per-dimension control via Tuple[int, ...]. 

481 if self.hsdp_replica_count is not None: 

482 if isinstance(self.hsdp_replica_count, int): 

483 dedup_per_dim = (self.hsdp_replica_count,) * len(replicate_group_ranks) 

484 else: 

485 dedup_per_dim = self.hsdp_replica_count 

486 replicate_group_ranks = tuple( 

487 r % dedup_per_dim[i] for i, r in enumerate(replicate_group_ranks) 

488 ) 

489 replicate_sizes = tuple( 

490 min(s, dedup_per_dim[i]) for i, s in enumerate(replicate_sizes) 

491 ) 

492 

493 # Greedy owner assignment on this sub-batch's records. 

494 owner_by_index = build_owner_by_size( 

495 records=records, 

496 replicate_sizes=replicate_sizes, 

497 ) 

498 

499 owned_records = select_owned_records( 

500 records=records, 

501 owner_by_index=owner_by_index, 

502 replicate_group_ranks=replicate_group_ranks, 

503 ) 

504 

505 is_shard_for_ns = ( 

506 hsdp_group.comm_key.has_shard_group 

507 and hsdp_group.layout_spec.is_last2d_sharded 

508 ) 

509 

510 hsdp_assign = HSDPGroupAssignment( 

511 owned_records=owned_records, 

512 all_records=records, 

513 owner_by_index=owner_by_index, 

514 replicate_group_ranks=replicate_group_ranks, 

515 replicate_sizes=replicate_sizes, 

516 replicate_pgs=hsdp_group.replicate_pgs, 

517 shard_pgs=hsdp_group.shard_pgs, 

518 is_shard=is_shard_for_ns, 

519 layout_spec=hsdp_group.layout_spec, 

520 ) 

521 

522 # Build broadcast reverse mapping for replicated groups 

523 if hsdp_assign.is_replicated and hsdp_assign.replicate_pgs: 

524 for record in hsdp_assign.all_records: 

525 src_coord = hsdp_assign.owner_rank_coord(record) 

526 if not src_coord or any(c < 0 for c in src_coord): 

527 continue 

528 self._param_to_broadcast_info[record.param] = ( 

529 src_coord, hsdp_assign.replicate_pgs, 

530 ) 

531 

532 return hsdp_assign 

533 

534 def _build_param_broadcast_info(self) -> None: 

535 """Build per-batch HSDP assignments and param broadcast reverse mapping. 

536 

537 When ``hsdp_replica_count`` is set, replicate coordinates and sizes 

538 are remapped so that owner assignment and broadcast happen within 

539 sub-groups of size ``hsdp_replica_count`` instead of the full 

540 replicate group. 

541 

542 Sub-group broadcast already covers all ranks (each rank belongs to 

543 exactly one sub-group), so param and state broadcast use the same 

544 sub-group mapping — no separate full-group path is needed. 

545 """ 

546 self._hsdp_assignment_batches: Dict[int, Dict] = {} 

547 self._param_to_broadcast_info: Dict[ 

548 torch.nn.Parameter, Tuple[Tuple[int, ...], Tuple[dist.ProcessGroup, ...]] 

549 ] = {} 

550 

551 for group_key, batch_groups in self._hsdp_batches.items(): 

552 no_comm_params = self._hsdp_grouping[group_key][0] 

553 assignment_batch_groups = [] 

554 

555 for bg in batch_groups: 

556 hsdp_group = bg["hsdp_group"] 

557 sub_batch_assigns: List[HSDPGroupAssignment] = [] 

558 

559 for sub_batch_entries in bg["sub_batches"]: 

560 hsdp_assign = self._build_sub_batch_assignment(sub_batch_entries, hsdp_group) 

561 if hsdp_assign is not None: 

562 sub_batch_assigns.append(hsdp_assign) 

563 

564 assignment_batch_groups.append({ 

565 "hsdp_group": hsdp_group, 

566 "sub_batches": sub_batch_assigns, 

567 }) 

568 

569 self._hsdp_assignment_batches[group_key] = { 

570 "no_comm": no_comm_params, 

571 "batch_groups": assignment_batch_groups, 

572 } 

573 

574 def _broadcast_replicate_params_after_step(self) -> None: 

575 """Broadcast updated params from assigned rank to replicate-group peers.""" 

576 self._broadcast_op_fused(target="param") 

577 

578 def _broadcast_state_fused_for_ckpt(self) -> None: 

579 """Broadcast optimizer state before checkpoint save.""" 

580 state_keys = ["momentum_buffer"] if self.is_muon else ["exp_avg", "exp_avg_sq"] 

581 self._broadcast_op_fused(target="state", state_keys=state_keys) 

582 

583 def _collect_broadcast_tensors( 

584 self, 

585 target: str, 

586 state_keys: Optional[List[str]] = None, 

587 ) -> Dict[Tuple, List[torch.Tensor]]: 

588 """Collect tensors to broadcast, grouped by (src_coord, dtype, replicate_pgs). 

589 

590 Args: 

591 target: "param" or "state". 

592 state_keys: State dict keys to broadcast when target="state". 

593 

594 Returns: 

595 Dict mapping (src_coord, dtype, replicate_pgs) to list of local tensors. 

596 """ 

597 rank_dtype_tensors = defaultdict(list) 

598 

599 for p, (src_coord, replicate_pgs) in self._param_to_broadcast_info.items(): 

600 if target == "param": 

601 local_tensor = to_local_if_dtensor(p.data) 

602 rank_dtype_tensors[(src_coord, local_tensor.dtype, replicate_pgs)].append(local_tensor) 

603 

604 elif target == "state" and state_keys: 

605 param_state = self.state.setdefault(p, {}) 

606 for key in state_keys: 

607 if key in param_state: 

608 state_tensor = param_state[key] 

609 local_tensor = to_local_if_dtensor(state_tensor) 

610 else: 

611 local_tensor = torch.empty_like(p, dtype=torch.float32) 

612 param_state[key] = local_tensor 

613 local_tensor = to_local_if_dtensor(local_tensor) 

614 

615 rank_dtype_tensors[(src_coord, local_tensor.dtype, replicate_pgs)].append(local_tensor) 

616 

617 return rank_dtype_tensors 

618 

619 @staticmethod 

620 def _compute_broadcast_batches( 

621 tensors: List[torch.Tensor], 

622 alignment_elements: int, 

623 max_broadcast_elements: int 

624 ) -> List[Tuple[List[Tuple[torch.Tensor, int, int, int]], int]]: 

625 """Split tensors into memory-capped batches with alignment padding. 

626 

627 Args: 

628 tensors: List of tensors to batch. 

629 alignment_elements: Alignment granularity in elements. 

630 max_broadcast_elements: Maximum elements per batch. 

631 

632 Returns: 

633 List of (batch_offsets, batch_total_size) tuples. 

634 Each batch_offsets entry is (tensor, offset, actual_numel, padded_numel). 

635 """ 

636 batches = [] 

637 current_batch = [] 

638 current_total_size = 0 

639 

640 for t in tensors: 

641 actual_numel = t.numel() 

642 padded_numel = ((actual_numel + alignment_elements - 1) // alignment_elements) * alignment_elements 

643 

644 if current_batch and current_total_size + padded_numel > max_broadcast_elements: 

645 batches.append((current_batch, current_total_size)) 

646 current_batch = [] 

647 current_total_size = 0 

648 

649 current_batch.append((t, current_total_size, actual_numel, padded_numel)) 

650 current_total_size += padded_numel 

651 

652 if current_batch: 

653 batches.append((current_batch, current_total_size)) 

654 

655 return batches 

656 

657 @staticmethod 

658 def _hierarchical_broadcast_buffer( 

659 batch_buffer: torch.Tensor, 

660 src_coord: Tuple[int, ...], 

661 replicate_pgs: Tuple[dist.ProcessGroup, ...], 

662 local_coord: Tuple[int, ...] 

663 ) -> None: 

664 """Broadcast a buffer dimension-by-dimension across replicate groups. 

665 

666 Args: 

667 batch_buffer: Contiguous buffer to broadcast. 

668 src_coord: Source coordinate tuple. 

669 replicate_pgs: Tuple of ProcessGroups (one per dimension). 

670 local_coord: Local rank coordinate tuple. 

671 """ 

672 platform = get_platform() 

673 

674 for dim_idx, pg in enumerate(replicate_pgs): 

675 if pg is None: 

676 continue 

677 

678 participate = True 

679 for subsequent_dim in range(dim_idx + 1, len(replicate_pgs)): 

680 if local_coord[subsequent_dim] != src_coord[subsequent_dim]: 

681 participate = False 

682 break 

683 

684 if not participate: 

685 continue 

686 

687 src_rank_in_pg = src_coord[dim_idx] 

688 global_src_rank = platform.get_global_rank(pg, src_rank_in_pg) 

689 dist.broadcast(batch_buffer, src=global_src_rank, group=pg) 

690 

691 @staticmethod 

692 def _hierarchical_broadcast_buffer_async( 

693 batch_buffer: torch.Tensor, 

694 src_coord: Tuple[int, ...], 

695 replicate_pgs: Tuple[dist.ProcessGroup, ...], 

696 local_coord: Tuple[int, ...], 

697 ) -> List[dist.Work]: 

698 """Async version of _hierarchical_broadcast_buffer. 

699 

700 Same dimension-by-dimension relay logic, but each dist.broadcast uses 

701 async_op=True. Returns a list of Work handles to wait on later. 

702 

703 Note: dimensions within a single buffer are still sequential (dim N+1 

704 depends on dim N completing), but different buffers can overlap. 

705 """ 

706 platform = get_platform() 

707 handles: List[dist.Work] = [] 

708 

709 for dim_idx, pg in enumerate(replicate_pgs): 

710 if pg is None: 

711 continue 

712 

713 participate = True 

714 for subsequent_dim in range(dim_idx + 1, len(replicate_pgs)): 

715 if local_coord[subsequent_dim] != src_coord[subsequent_dim]: 

716 participate = False 

717 break 

718 

719 if not participate: 

720 continue 

721 

722 src_rank_in_pg = src_coord[dim_idx] 

723 global_src_rank = platform.get_global_rank(pg, src_rank_in_pg) 

724 handle = dist.broadcast( 

725 batch_buffer, src=global_src_rank, group=pg, async_op=True, 

726 ) 

727 handles.append(handle) 

728 

729 return handles 

730 

731 def _broadcast_op_fused( 

732 self, 

733 target: str, 

734 state_keys: Optional[List[str]] = None, 

735 ) -> None: 

736 """Fused hierarchical broadcast for param or state across replicate groups. 

737 

738 Groups tensors by (src_coord, dtype, replicate_pgs), packs into contiguous buffers 

739 with 512-byte alignment, and broadcasts dimension-by-dimension in memory-capped batches. 

740 

741 Args: 

742 target: "param" or "state". 

743 state_keys: State dict keys to broadcast when target="state". 

744 """ 

745 device = torch.npu.current_device() if torch.npu.is_available() else torch.cuda.current_device() 

746 

747 alignment = 512 # bytes 

748 rank_dtype_tensors = self._collect_broadcast_tensors(target, state_keys) 

749 

750 for (src_coord, dtype, replicate_pgs), tensors in rank_dtype_tensors.items(): 

751 if not tensors: 

752 continue 

753 

754 local_coord = tuple( 

755 dist.get_rank(group=pg) if pg is not None else -1 

756 for pg in replicate_pgs 

757 ) 

758 

759 element_size = torch.empty(0, dtype=dtype, device=device).element_size() 

760 alignment_elements = alignment // element_size 

761 max_broadcast_bytes = getattr(self, "replicate_broadcast_max_bytes", 512 * 1024 * 1024) 

762 max_broadcast_elements = max_broadcast_bytes // element_size 

763 max_broadcast_elements = max( 

764 alignment_elements, 

765 (max_broadcast_elements // alignment_elements) * alignment_elements, 

766 ) 

767 

768 batches = self._compute_broadcast_batches(tensors, alignment_elements, max_broadcast_elements) 

769 if not batches: 

770 continue 

771 

772 max_batch_size = max(batch_size for _, batch_size in batches) 

773 buffer = torch.empty(max_batch_size, dtype=dtype, device=device) 

774 

775 for batch_tensor_offsets, batch_total_size in batches: 

776 batch_buffer = buffer[:batch_total_size] 

777 

778 # Pack: owner rank 

779 if local_coord == src_coord: 

780 for t, offset, actual_numel, padded_numel in batch_tensor_offsets: 

781 batch_buffer[offset:offset + actual_numel].copy_(t.view(-1)) 

782 if padded_numel > actual_numel: 

783 batch_buffer[offset + actual_numel:offset + padded_numel].zero_() 

784 

785 # Hierarchical Broadcast 

786 self._hierarchical_broadcast_buffer(batch_buffer, src_coord, replicate_pgs, local_coord) 

787 

788 # Unpack: copy buffer data back to individual tensors 

789 for t, offset, actual_numel, _ in batch_tensor_offsets: 

790 t.view(-1).copy_(batch_buffer[offset:offset + actual_numel]) 

791 

792 buffer.untyped_storage().resize_(0) 

793 del buffer 

794 

795 def cleanup_synced_state(self) -> None: 

796 """Release optimizer state for non-owned params after checkpoint saving. 

797 

798 In HSDP mode, _broadcast_state_fused_for_ckpt broadcasts state to all 

799 ranks so every rank can save a complete checkpoint. This method removes 

800 the non-owned state to restore per-rank memory savings. Must be called 

801 after the checkpoint has been fully written to disk. 

802 """ 

803 params_to_remove = [] 

804 for assignment_info in self._hsdp_assignment_batches.values(): 

805 for bg in assignment_info["batch_groups"]: 

806 for hsdp_info in bg["sub_batches"]: 

807 if not hsdp_info.is_replicated: 

808 continue 

809 for record in hsdp_info.all_records: 

810 if not hsdp_info.is_owned(record) and record.param in self.state: 

811 params_to_remove.append(record.param) 

812 

813 for p in params_to_remove: 

814 self.state[p].clear() 

815 del self.state[p] 

816 

817 torch.npu.empty_cache() 

818 

819 

820class AsyncReplicateBroadcaster: 

821 """Incremental async replicate broadcast with HSDP-group-level flushing. 

822 

823 Designed to overlap HSDP replicate-group broadcasts with Muon NS