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 

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) 

29 

30from hyper_parallel.core.optimizer.dtensor_compat import to_local_if_dtensor 

31from hyper_parallel.core.optimizer.sharding_category import ( 

32 HSDPGroupAssignment, 

33 build_owner_by_size, 

34 get_multi_dim_logical_info, 

35 select_owned_records, 

36 group_parameters_for_hsdp 

37) 

38from hyper_parallel.core.optimizer.utils import empty_accelerator_cache, get_current_device, get_device_count 

39 

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 self._rebind_param_attrs() 

60 

61 def __iter__(self): 

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

63 return iter(self.chained_optimizers) 

64 

65 def _rebind_param_attrs(self) -> None: 

66 """Restore ``model_name`` and ``is_muon`` on the current parameters.""" 

67 muon_param_ids: set = set() 

68 adamw_param_ids: set = set() 

69 self.muon_keys = [] 

70 self.no_muon_keys = [] 

71 for _, opt in self.optimizers_dict.items(): 

72 target = muon_param_ids if getattr(opt, "is_muon", False) else adamw_param_ids 

73 for group in opt.param_groups: 

74 for p in group.get("params", []): 

75 target.add(id(p)) 

76 

77 for param_name, param in self.model.named_parameters(): 

78 setattr(param, "model_name", param_name) 

79 if id(param) in muon_param_ids: 

80 setattr(param, "is_muon", True) 

81 self.muon_keys.append(param_name) 

82 elif id(param) in adamw_param_ids: 

83 setattr(param, "is_muon", False) 

84 self.no_muon_keys.append(param_name) 

85 

86 def __str__(self): 

87 if not self.optimizers_dict: 

88 return f'{self.__class__.__name__}' 

89 

90 sections = [] 

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

92 sections.append(f'{name}: {repr(optimizer)}') 

93 return "\n\n".join(sections) 

94 

95 __repr__ = __str__ 

96 

97 def step(self, closure=None) -> None: 

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

99 for opt in self.chained_optimizers: 

100 opt.step(closure=closure) 

101 

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

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

104 for opt in self.chained_optimizers: 

105 opt.zero_grad(set_to_none=set_to_none) 

106 

107 @property 

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

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

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

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

112 return self.chained_optimizers[0] 

113 

114 @property 

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

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

117 return self.chained_optimizers[0].defaults 

118 

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

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

121 

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

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

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

125 """ 

126 # Ensure all ranks have consistent optimizer state before snapshotting 

127 for opt in self.chained_optimizers: 

128 if hasattr(opt, "_broadcast_state_fused_for_ckpt"): 

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

130 

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

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

133 sd = get_optimizer_state_dict( 

134 self.model, 

135 optimizer, 

136 options=StateDictOptions(flatten_optimizer_state_dict=self.flatten) 

137 ) 

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

139 if overlap: 

140 raise KeyError( 

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

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

143 ) 

144 merged.update(sd) 

145 

146 return merged 

147 

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

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

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

151 for optimizer in self.chained_optimizers: 

152 param_groups += optimizer.param_groups 

153 return param_groups 

154 

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

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

157 if not isinstance(new_param_groups, list): 

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

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

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

161 

162 start = 0 

163 for optimizer in self.chained_optimizers: 

164 group_len = len(optimizer.param_groups) 

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

166 start += group_len 

167 

168 param_groups = property(_get_param_groups, _set_param_groups) 

169 

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

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

172 for optimizer in self.chained_optimizers: 

173 set_optimizer_state_dict( 

174 self.model, 

175 optimizer, 

176 optim_state_dict=state_dict, 

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

178 ) 

179 

180 self._synchronize_steps() 

181 

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

183 """Synchronize the step of all optimizers. 

184 

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

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

187 """ 

188 steps = [] 

189 for optimizer in self.chained_optimizers: 

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

191 for param_group in actual_opt.param_groups: 

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

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

194 

195 unique_steps = list(set(steps)) 

196 if len(unique_steps) > 1: 

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

198 

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

200 

201 for optimizer in self.chained_optimizers: 

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

203 for param_group in actual_opt.param_groups: 

204 param_group['step'] = step 

205 

206 return step 

207 

208 

209class BaseDistributedOptimizer(torch.optim.Optimizer): 

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

211 

212 Provides fused hierarchical broadcast for parameters and optimizer states. 

213 """ 

214 

215 def __init__( 

216 self, 

217 params: Any, 

218 defaults: Dict[str, Any], 

219 is_muon: bool, 

220 hsdp_replica_count: Optional[int] = None, 

221 ) -> None: 

222 super().__init__(params, defaults) 

223 self.is_muon = is_muon 

224 self.hsdp_replica_count = hsdp_replica_count 

225 self._param_to_broadcast_info: Dict[ 

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

227 ] = {} 

228 

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

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

231 

232 def _group_dtensor_by_mesh(self): 

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

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

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

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

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

238 

239 def _auto_deduce_replica_count(self) -> Optional[int]: 

240 """Deduce hsdp_replica_count based on cluster topology. 

241 

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

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

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

245 """ 

246 devices_per_node = get_device_count() 

247 

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

249 needs_split = False 

250 

251 # group_key, (no_comm_params, hsdp_groups) 

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

253 for hsdp_group in hsdp_groups: 

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

255 if pg is None: 

256 continue 

257 pg_size = dist.get_world_size(pg) 

258 if pg_size <= 1: 

259 continue 

260 

261 if pg_size > devices_per_node: 

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

263 dedup = min(pg_size, devices_per_node) 

264 while pg_size % dedup != 0: 

265 dedup -= 1 

266 needs_split = True 

267 else: 

268 # origin Inter-node 

269 dedup = pg_size 

270 

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

272 if dim_idx not in dedup_per_dim: 

273 dedup_per_dim[dim_idx] = dedup 

274 else: 

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

276 

277 if not needs_split: 

278 return None 

279 

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

281 return min(sorted_dedups) 

282 

283 def _split_replicate_groups(self) -> None: 

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

285 if self.hsdp_replica_count is None: 

286 return 

287 

288 # Argument validation 

289 if not isinstance(self.hsdp_replica_count, int): 

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

291 if self.hsdp_replica_count <= 0: 

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

293 

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

295 for hsdp_group in hsdp_groups: 

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

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

298 if pg is None: 

299 new_replicate_pgs.append(pg) 

300 continue 

301 

302 pg_size = dist.get_world_size(pg) 

303 dedup_size = self.hsdp_replica_count 

304 

305 if pg_size <= dedup_size: 

306 new_replicate_pgs.append(pg) 

307 continue 

308 

309 if pg_size % dedup_size != 0: 

310 raise ValueError( 

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

312 ) 

313 

314 sub_pg = self._get_or_create_sub_pg(pg, dedup_size) 

315 new_replicate_pgs.append(sub_pg) 

316 

317 logger.info_rank0( 

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

319 group_key, dim_idx, pg_size, dedup_size 

320 ) 

321 

322 # replace new sub groups of hsdp groups 

323 hsdp_group.replicate_pgs = tuple(new_replicate_pgs) 

324 

325 def _get_or_create_sub_pg( 

326 self, 

327 parent_pg: dist.ProcessGroup, 

328 sub_size: int, 

329 ) -> dist.ProcessGroup: 

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

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

332 cache_key = (local_parent_ranks, sub_size) 

333 

334 # Cache Hit Check 

335 if cache_key in self._split_sub_pg_cache: 

336 sub_pg_map = self._split_sub_pg_cache[cache_key] 

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

338 if sub_pg is not None: 

339 try: 

340 dist.get_rank(group=sub_pg) 

341 return sub_pg 

342 except RuntimeError: 

343 continue 

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

345 

346 global_rank = dist.get_rank() 

347 world_size = dist.get_world_size() 

348 

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

350 # Far more scalable than all_gather_object for large world_size. 

351 device = get_current_device() 

352 num_parent_ranks = len(local_parent_ranks) 

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

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

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

356 

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

358 # sorted set of tuples for deterministic iteration order. 

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

360 unique_all_parent_ranks = sorted(set( 

361 tuple(row) for row in gathered_cpu_list 

362 )) 

363 

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

365 my_sub_pg: Optional[dist.ProcessGroup] = None 

366 

367 # Synchronized collective creation loop 

368 for parent_ranks in unique_all_parent_ranks: 

369 num_sub_groups = len(parent_ranks) // sub_size 

370 

371 for sub_idx in range(num_sub_groups): 

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

373 sub_pg = dist.new_group(sub_ranks) 

374 

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

376 if parent_ranks == local_parent_ranks: 

377 if global_rank in sub_ranks: 

378 my_sub_pg = sub_pg 

379 sub_pg_map[sub_idx] = sub_pg 

380 else: 

381 sub_pg_map[sub_idx] = None 

382 

383 self._split_sub_pg_cache[cache_key] = sub_pg_map 

384 

385 if my_sub_pg is None: 

386 raise RuntimeError( 

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

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

389 ) 

390 

391 return my_sub_pg 

392 

393 def _build_hsdp_batch( 

394 self, 

395 max_batch_numel: Optional[int] = None, 

396 ) -> None: 

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

398 if max_batch_numel is None: 

399 broadcast_max_bytes = getattr( 

400 self, "replicate_broadcast_max_bytes", 512 * 1024 * 1024, 

401 ) 

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

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

404 

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

406 

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

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

409 sorted_groups = sorted( 

410 hsdp_groups, 

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

412 ) 

413 

414 batch_groups: List[Dict] = [] 

415 

416 for hsdp_group in sorted_groups: 

417 # Sort records within group by numel descending. 

418 sorted_records = sorted( 

419 hsdp_group.records, 

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

421 ) 

422 

423 sub_batches: List[List[Dict]] = [] 

424 current_batch: List[Dict] = [] 

425 current_numel = 0 

426 

427 for record in sorted_records: 

428 p_numel = record.param.numel() 

429 

430 current_batch.append({ 

431 "record": record, 

432 "hsdp_group": hsdp_group, 

433 }) 

434 current_numel += p_numel 

435 

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

437 # so that symmetric structures stay together and fragmentation is reduced. 

438 if current_numel >= max_batch_numel: 

439 sub_batches.append(current_batch) 

440 current_batch = [] 

441 current_numel = 0 

442 

443 if current_batch: 

444 sub_batches.append(current_batch) 

445 

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

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

448 

449 batch_groups.append({ 

450 "hsdp_group": hsdp_group, 

451 "sub_batches": sub_batches, 

452 }) 

453 

454 # log batch split info. 

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

456 logger.info_rank0( 

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

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

459 group_key, 

460 len(hsdp_groups), 

461 len(batch_groups), 

462 total_sub_batches, 

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

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

465 max_batch_numel 

466 ) 

467 self._hsdp_batches[group_key] = batch_groups 

468 

469 def _build_sub_batch_assignment( 

470 self, 

471 sub_batch_entries: List[Dict], 

472 hsdp_group: Any, 

473 ) -> Optional[HSDPGroupAssignment]: 

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

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

476 if not records: 

477 return None 

478 

479 device_mesh = records[0].param.device_mesh 

480 

481 replicate_group_ranks, replicate_sizes = get_multi_dim_logical_info( 

482 device_mesh, 

483 hsdp_group.comm_key.replicate_mesh_dims, 

484 ) 

485 

486 # When hsdp_replica_count is set, remap coordinates into 

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

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

489 if self.hsdp_replica_count is not None: 

490 replicate_group_ranks = tuple( 

491 r % self.hsdp_replica_count for r in replicate_group_ranks 

492 ) 

493 replicate_sizes = tuple( 

494 min(s, self.hsdp_replica_count) for s in replicate_sizes 

495 ) 

496 

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

498 owner_by_index = build_owner_by_size( 

499 records=records, 

500 replicate_sizes=replicate_sizes, 

501 ) 

502 

503 owned_records = select_owned_records( 

504 records=records, 

505 owner_by_index=owner_by_index, 

506 replicate_group_ranks=replicate_group_ranks, 

507 ) 

508 

509 is_shard_for_ns = ( 

510 hsdp_group.comm_key.has_shard_group 

511 and hsdp_group.layout_spec.is_last2d_sharded 

512 ) 

513 

514 hsdp_assign = HSDPGroupAssignment( 

515 owned_records=owned_records, 

516 all_records=records, 

517 owner_by_index=owner_by_index, 

518 replicate_group_ranks=replicate_group_ranks, 

519 replicate_sizes=replicate_sizes, 

520 replicate_pgs=hsdp_group.replicate_pgs, 

521 shard_pgs=hsdp_group.shard_pgs, 

522 is_shard=is_shard_for_ns, 

523 layout_spec=hsdp_group.layout_spec, 

524 ) 

525 logger.info_rank0(f'[Hyper-optimizer] hsdp_assign: {hsdp_assign}') 

526 

527 # Build broadcast reverse mapping for replicated groups 

528 if hsdp_assign.is_replicated and hsdp_assign.replicate_pgs: 

529 for record in hsdp_assign.all_records: 

530 src_coord = hsdp_assign.owner_rank_coord(record) 

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

532 continue 

533 self._param_to_broadcast_info[record.param] = ( 

534 src_coord, hsdp_assign.replicate_pgs, 

535 ) 

536 

537 return hsdp_assign 

538 

539 def _build_param_broadcast_info(self) -> None: 

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

541 

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

543 are remapped so that owner assignment and broadcast happen within 

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

545 replicate group. 

546 

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

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

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

550 """ 

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

552 self._param_to_broadcast_info: Dict[ 

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

554 ] = {} 

555 

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

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

558 assignment_batch_groups = [] 

559 

560 for bg in batch_groups: 

561 hsdp_group = bg["hsdp_group"] 

562 sub_batch_assigns: List[HSDPGroupAssignment] = [] 

563 

564 for sub_batch_entries in bg["sub_batches"]: 

565 hsdp_assign = self._build_sub_batch_assignment(sub_batch_entries, hsdp_group) 

566 if hsdp_assign is not None: 

567 sub_batch_assigns.append(hsdp_assign) 

568 

569 assignment_batch_groups.append({ 

570 "hsdp_group": hsdp_group, 

571 "sub_batches": sub_batch_assigns, 

572 }) 

573 

574 self._hsdp_assignment_batches[group_key] = { 

575 "no_comm": no_comm_params, 

576 "batch_groups": assignment_batch_groups, 

577 } 

578 

579 def _broadcast_replicate_params_after_step(self) -> None: 

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

581 self._broadcast_op_fused(target="param") 

582 

583 def _broadcast_state_fused_for_ckpt(self) -> None: 

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

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

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

587 

588 def _collect_broadcast_tensors( 

589 self, 

590 target: str, 

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

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

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

594 

595 Args: 

596 target: "param" or "state". 

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

598 

599 Returns: 

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

601 """ 

602 rank_dtype_tensors = defaultdict(list) 

603 

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

605 if target == "param": 

606 local_tensor = to_local_if_dtensor(p.data) 

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

608 

609 elif target == "state" and state_keys: 

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

611 for key in state_keys: 

612 if key in param_state: 

613 state_tensor = param_state[key] 

614 local_tensor = to_local_if_dtensor(state_tensor) 

615 else: 

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

617 param_state[key] = local_tensor 

618 local_tensor = to_local_if_dtensor(local_tensor) 

619 

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

621 

622 return rank_dtype_tensors 

623 

624 @staticmethod 

625 def _compute_broadcast_batches( 

626 tensors: List[torch.Tensor], 

627 alignment_elements: int, 

628 max_broadcast_elements: int 

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

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

631 

632 Args: 

633 tensors: List of tensors to batch. 

634 alignment_elements: Alignment granularity in elements. 

635 max_broadcast_elements: Maximum elements per batch. 

636 

637 Returns: 

638 List of (batch_offsets, batch_total_size) tuples. 

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

640 """ 

641 batches = [] 

642 current_batch = [] 

643 current_total_size = 0 

644 

645 for t in tensors: 

646 actual_numel = t.numel() 

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

648 

649 if current_batch and current_total_size + padded_numel > max_broadcast_elements: 

650 batches.append((current_batch, current_total_size)) 

651 current_batch = [] 

652 current_total_size = 0 

653 

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

655 current_total_size += padded_numel 

656 

657 if current_batch: 

658 batches.append((current_batch, current_total_size)) 

659 

660 return batches 

661 

662 @staticmethod 

663 def _hierarchical_broadcast_buffer( 

664 batch_buffer: torch.Tensor, 

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

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

667 local_coord: Tuple[int, ...] 

668 ) -> None: 

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

670 

671 Args: 

672 batch_buffer: Contiguous buffer to broadcast. 

673 src_coord: Source coordinate tuple. 

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

675 local_coord: Local rank coordinate tuple. 

676 """ 

677 

678 for dim_idx, pg in enumerate(replicate_pgs): 

679 if pg is None: 

680 continue 

681 

682 participate = True 

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

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

685 participate = False 

686 break 

687 

688 if not participate: 

689 continue 

690 

691 src_rank_in_pg = src_coord[dim_idx] 

692 global_src_rank = dist.get_global_rank(pg, src_rank_in_pg) 

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

694 

695 @staticmethod 

696 def _hierarchical_broadcast_buffer_async( 

697 batch_buffer: torch.Tensor, 

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

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

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

701 ) -> List[dist.Work]: 

702 """Async version of _hierarchical_broadcast_buffer. 

703 

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

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

706 

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

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

709 """ 

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

711 

712 for dim_idx, pg in enumerate(replicate_pgs): 

713 if pg is None: 

714 continue 

715 

716 participate = True 

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

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

719 participate = False 

720 break 

721 

722 if not participate: 

723 continue 

724 

725 src_rank_in_pg = src_coord[dim_idx] 

726 global_src_rank = dist.get_global_rank(pg, src_rank_in_pg) 

727 handle = dist.broadcast( 

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

729 ) 

730 handles.append(handle) 

731 

732 return handles 

733 

734 def _broadcast_op_fused( 

735 self, 

736 target: str, 

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

738 ) -> None: 

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

740 

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

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

743 

744 Args: 

745 target: "param" or "state". 

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

747 """ 

748 device = get_current_device() 

749 

750 alignment = 512 # bytes 

751 rank_dtype_tensors = self._collect_broadcast_tensors(target, state_keys) 

752 

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

754 if not tensors: 

755 continue 

756 

757 local_coord = tuple( 

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

759 for pg in replicate_pgs 

760 ) 

761 

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

763 alignment_elements = alignment // element_size 

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

765 max_broadcast_elements = max_broadcast_bytes // element_size 

766 max_broadcast_elements = max( 

767 alignment_elements, 

768 (max_broadcast_elements // alignment_elements) * alignment_elements, 

769 ) 

770 

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

772 if not batches: 

773 continue 

774 

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

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

777 

778 for batch_tensor_offsets, batch_total_size in batches: 

779 batch_buffer = buffer[:batch_total_size] 

780 

781 # Pack: owner rank 

782 if local_coord == src_coord: 

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

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

785 if padded_numel > actual_numel: 

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

787 

788 # Hierarchical Broadcast 

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

790 

791 # Unpack: copy buffer data back to individual tensors 

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

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

794 

795 buffer.untyped_storage().resize_(0) 

796 del buffer 

797 

798 def cleanup_synced_state(self) -> None: 

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

800 

801 In HSDP mode, _broadcast_state_fused_for_ckpt broadcasts state to all 

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

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

804 after the checkpoint has been fully written to disk. 

805 """ 

806 params_to_remove = [] 

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

808 for bg in assignment_info["batch_groups"]: 

809 for hsdp_info in bg["sub_batches"]: 

810 if not hsdp_info.is_replicated: 

811 continue 

812 for record in hsdp_info.all_records: 

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

814 params_to_remove.append(record.param) 

815 

816 for p in params_to_remove: 

817 self.state[p].clear() 

818 del self.state[p] 

819 

820 empty_accelerator_cache() 

821 

822 

823class AsyncReplicateBroadcaster: 

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

825 

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

827 computation. After all sub_batches for an HSDP group are done, the 

828 caller invokes flush_group(hsdp_assign) to issue an async hierarchical 

829 broadcast — overlapping with the next HSDP group's NS computation. 

830 

831 Flush is driven by the caller at HSDP-group boundaries (not by a 

832 threshold) so that all ranks issue the same collective operations in 

833 the same order, which is required by NCCL/HCCL. 

834