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

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

7# http://www.apache.org/licenses/LICENSE-2.0

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# ============================================================================

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

17# pylint: disable=forbidden-backend-import

19from collections import defaultdict

20from contextlib import contextmanager, nullcontext

21import logging

22from typing import Any, Dict, Iterator, List, Optional, Tuple

24import torch

25import torch.distributed as dist

26from torch.distributed.checkpoint.state_dict import (

27 StateDictOptions,

28 get_optimizer_state_dict,

29 set_optimizer_state_dict,

30)

32from hyper_parallel.core.optimizer.dtensor_compat import to_local_if_dtensor

33from hyper_parallel.core.optimizer.sharding_category import (

34 HSDPGroupAssignment,

35 build_owner_by_size,

36 get_multi_dim_logical_info,

37 select_owned_records,

38 group_parameters_for_hsdp

39)

40from hyper_parallel.core.optimizer.swap_optimizer import is_swap_optimizer

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

43logger = logging.getLogger(__name__)

46class ChainedOptimizer:

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

49 def __init__(

50 self,

51 model: torch.nn.Module,

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

53 flatten: bool = False,

54 ) -> None:

55 """Compose named optimizers over registered model parameters.

57 Args:

58 model: Module containing the registered model parameters.

59 optimizers: Leaf optimizers keyed by their configured names.

60 flatten: Whether leaf optimizer state uses flattened FQN keys.

61 """

62 self.optimizers_dict = optimizers

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

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

65 self.model = model

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

67 self._is_multi_optimizer = flatten

68 self.model_param_by_optimizer_param: Dict[

69 torch.nn.Parameter, torch.nn.Parameter

70 ] = {}

71 self.optimizer_param_by_model_param: Dict[

72 torch.nn.Parameter, torch.nn.Parameter

73 ] = {}

75 self._rebind_param_attrs()

77 def __iter__(self) -> Iterator[torch.optim.Optimizer]:

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

79 return iter(self.chained_optimizers)

81 def _rebind_param_attrs(self) -> None:

82 """Restore model FQNs and optimizer-family name views."""

83 muon_param_ids: set = set()

84 adamw_param_ids: set = set()

85 self.muon_keys = []

86 self.no_muon_keys = []

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

88 is_muon = hasattr(opt, "is_muon") and opt.is_muon

89 target = muon_param_ids if is_muon else adamw_param_ids

90 for group in opt.param_groups:

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

92 target.add(id(p))

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

95 optimizer_param = self.optimizer_param_by_model_param.get(param, param)

96 setattr(param, "model_name", param_name)

97 setattr(optimizer_param, "model_name", param_name)

98 if id(optimizer_param) in muon_param_ids:

99 self.muon_keys.append(param_name)

100 elif id(optimizer_param) in adamw_param_ids:

101 self.no_muon_keys.append(param_name)

102

103 def reset_optimizer_parameters(

104 self,

105 model_param_by_optimizer_param: Dict[

106 torch.nn.Parameter, torch.nn.Parameter

107 ],

108 ) -> None:

109 """Register replaced optimizer params and rebuild identity-based metadata.

110

111 Args:

112 model_param_by_optimizer_param: Main-to-model parameter identity map.

113 """

114 self.model_param_by_optimizer_param = dict(

115 model_param_by_optimizer_param

116 )

117 self.optimizer_param_by_model_param = {

118 model_param: optimizer_param

119 for optimizer_param, model_param in self.model_param_by_optimizer_param.items()

120 }

121 self._rebind_param_attrs()

122 for optimizer in self.chained_optimizers:

123 reset = getattr(optimizer, "reset_optimizer_parameters", None)

124 if callable(reset):

125 if is_swap_optimizer(optimizer):

126 reset(self.model_param_by_optimizer_param)

127 else:

128 reset()

129

130 @contextmanager

131 def _optimizer_parameter_view(self) -> Iterator[None]:

132 """Temporarily register optimizer params under their model FQNs."""

133 if not self.model_param_by_optimizer_param:

134 yield

135 return

136

137 registrations = []

138 for module in self.model.modules():

139 for parameter_name, model_param in module._parameters.items(): # pylint: disable=protected-access

140 if model_param is None:

141 continue

142 optimizer_param = self.optimizer_param_by_model_param.get(model_param)

143 if optimizer_param is None:

144 continue

145 registrations.append((module, parameter_name, model_param))

146 module._parameters[parameter_name] = optimizer_param # pylint: disable=protected-access

147 try:

148 yield

149 finally:

150 for module, parameter_name, model_param in registrations:

151 module._parameters[parameter_name] = model_param # pylint: disable=protected-access

152

153 def __str__(self) -> str:

154 """Return each named leaf optimizer's representation."""

155 if not self.optimizers_dict:

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

157

158 sections = []

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

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

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

162

163 __repr__ = __str__

164

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

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

167

168 Args:

169 closure: Optional callable forwarded to every sub-optimizer.

170 """

171 for opt in self.chained_optimizers:

172 opt.step(closure=closure)

173

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

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

176

177 Args:

178 set_to_none: Whether gradients are reset to ``None``.

179 """

180 for opt in self.chained_optimizers:

181 opt.zero_grad(set_to_none=set_to_none)

182

183 @property

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

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

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

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

188 return self.chained_optimizers[0]

189

190 @property

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

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

193 return self.chained_optimizers[0].defaults

194

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

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

197

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

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

200 then converts DTensor values to local CPU tensors for serialization.

201 """

202 prepared = []

203 try:

204 # Swap optimizers must restore released device storage before the

205 # wrapped HSDP implementation reads its state tensors.

206 for opt in self.chained_optimizers:

207 prepare = getattr(opt, "_prepare_swap_checkpoint", None)

208 if callable(prepare):

209 prepared.append(opt)

210 prepare()

211

212 # Ensure all ranks have consistent optimizer state before snapshotting.

213 for opt in self.chained_optimizers:

214 if hasattr(opt, "_broadcast_state_fused_for_ckpt"):

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

216

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

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

219 with self._optimizer_parameter_view():

220 sd = get_optimizer_state_dict(

221 self.model,

222 optimizer,

223 options=StateDictOptions(flatten_optimizer_state_dict=self.flatten)

224 )

225 if optimizer in prepared:

226 # The swap lifecycle releases device storage once the

227 # checkpoint window closes, so the exported values must own

228 # their memory instead of aliasing the slot tensors.

229 self._localize_state_dict_tensors(sd)

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

231 if overlap:

232 raise KeyError(

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

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

235 )

236 merged.update(sd)

237

238 return merged

239 finally:

240 for opt in reversed(prepared):

241 cleanup = getattr(opt, "cleanup_synced_state", None)

242 if callable(cleanup):

243 cleanup()

244 continue

245 finish = getattr(opt, "_finish_swap_checkpoint", None)

246 if callable(finish):

247 finish()

248

249 @staticmethod

250 def _localize_state_dict_tensors(state_dict: Dict[str, Any]) -> None:

251 """Replace exported tensors with CPU copies that own their storage.

252

253 Swap-managed optimizer state lives in CPU mirrors while its device

254 storage is released, so a checkpoint entry must not alias a tensor whose

255 storage the swap runtime is about to free.

256 """

257 def localize(value: Any) -> Any:

258 """Clone tensors in one nested checkpoint value onto CPU.

259

260 Args:

261 value: Checkpoint value or nested container to localize.

262

263 Returns:

264 The value with every nested tensor replaced by an owning CPU copy.

265 """

266 if isinstance(value, torch.Tensor):

267 return to_local_if_dtensor(value).detach().cpu().clone()

268 if isinstance(value, dict):

269 return {key: localize(item) for key, item in value.items()}

270 if isinstance(value, list):

271 return [localize(item) for item in value]

272 if isinstance(value, tuple):

273 return tuple(localize(item) for item in value)

274 return value

275

276 for key, value in list(state_dict.items()):

277 if key == "param_groups":

278 continue

279 state_dict[key] = localize(value)

280

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

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

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

284 for optimizer in self.chained_optimizers:

285 param_groups += optimizer.param_groups

286 return param_groups

287

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

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

290 if not isinstance(new_param_groups, list):

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

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

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

294

295 start = 0

296 for optimizer in self.chained_optimizers:

297 group_len = len(optimizer.param_groups)

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

299 start += group_len

300

301 param_groups = property(_get_param_groups, _set_param_groups)

302

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

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

305

306 Args:

307 state_dict: FQN-keyed optimizer state to restore.

308 """

309 for optimizer in self.chained_optimizers:

310 load_context = getattr(optimizer, "_swap_checkpoint_load_context", None)

311 context = load_context(state_dict, self.flatten) if callable(load_context) else nullcontext()

312 with self._optimizer_parameter_view(), context:

313 set_optimizer_state_dict(

314 self.model,

315 optimizer,

316 optim_state_dict=state_dict,

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

318 )

319

320 self._synchronize_steps()

321

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

323 """Synchronize the step of all optimizers.

324

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

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

327 """

328 steps = []

329 for optimizer in self.chained_optimizers:

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

331 for param_group in actual_opt.param_groups:

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

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

334

335 unique_steps = list(set(steps))

336 if len(unique_steps) > 1:

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

338

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

340

341 for optimizer in self.chained_optimizers:

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

343 for param_group in actual_opt.param_groups:

344 param_group['step'] = step

345

346 return step

347

348

349class BaseDistributedOptimizer(torch.optim.Optimizer):

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

351

352 Provides fused hierarchical broadcast for parameters and optimizer states.

353 """

354

355 def __init__(

356 self,

357 params: Any,

358 defaults: Dict[str, Any],

359 is_muon: bool,

360 hsdp_replica_count: Optional[int] = None,

361 ) -> None:

362 """Initialize topology-aware optimizer state and parameter groups.

363

364 Args:

365 params: Parameter groups optimized by this instance.

366 defaults: Default hyperparameters applied to every group.

367 is_muon: Whether parameters use the Muon update path.

368 hsdp_replica_count: Optional optimizer-state replica group size.

369 """

370 super().__init__(params, defaults)

371 self.is_muon = is_muon

372 self.hsdp_replica_count = hsdp_replica_count

373 self._param_to_broadcast_info: Dict[

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

375 ] = {}

376

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

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

379

380 def _group_dtensor_by_mesh(self):

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

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

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

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

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

386

387 def reset_optimizer_parameters(self) -> None:

388 """Rebuild subclass metadata after optimizer parameters are replaced."""

389 raise NotImplementedError

390

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

392 """Deduce hsdp_replica_count based on cluster topology.

393

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

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

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

397 """

398 devices_per_node = get_device_count()

399

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

401 needs_split = False

402

403 # group_key, (no_comm_params, hsdp_groups)

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

405 for hsdp_group in hsdp_groups:

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

407 if pg is None:

408 continue

409 pg_size = dist.get_world_size(pg)

410 if pg_size <= 1:

411 continue

412

413 if pg_size > devices_per_node:

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

415 dedup = min(pg_size, devices_per_node)

416 while pg_size % dedup != 0:

417 dedup -= 1

418 needs_split = True

419 else:

420 # origin Inter-node

421 dedup = pg_size

422

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

424 if dim_idx not in dedup_per_dim:

425 dedup_per_dim[dim_idx] = dedup

426 else:

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

428

429 if not needs_split:

430 return None

431

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

433 return min(sorted_dedups)

434

435 def _split_replicate_groups(self) -> None:

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

437 if self.hsdp_replica_count is None:

438 return

439

440 # Argument validation

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

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

443 if self.hsdp_replica_count <= 0:

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

445

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

447 for hsdp_group in hsdp_groups:

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

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

450 if pg is None:

451 new_replicate_pgs.append(pg)

452 continue

453

454 pg_size = dist.get_world_size(pg)

455 dedup_size = self.hsdp_replica_count

456

457 if pg_size <= dedup_size:

458 new_replicate_pgs.append(pg)

459 continue

460

461 if pg_size % dedup_size != 0:

462 raise ValueError(

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

464 )

465

466 sub_pg = self._get_or_create_sub_pg(pg, dedup_size)

467 new_replicate_pgs.append(sub_pg)

468

469 logger.info_rank0(

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

471 group_key, dim_idx, pg_size, dedup_size

472 )

473

474 # replace new sub groups of hsdp groups

475 hsdp_group.replicate_pgs = tuple(new_replicate_pgs)

476

477 def _get_or_create_sub_pg(

478 self,

479 parent_pg: dist.ProcessGroup,

480 sub_size: int,

481 ) -> dist.ProcessGroup:

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

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

484 cache_key = (local_parent_ranks, sub_size)

485

486 # Cache Hit Check

487 if cache_key in self._split_sub_pg_cache:

488 sub_pg_map = self._split_sub_pg_cache[cache_key]

489 for sub_pg in sub_pg_map.values():

490 if sub_pg is not None:

491 try:

492 dist.get_rank(group=sub_pg)

493 return sub_pg

494 except RuntimeError:

495 continue

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

497

498 global_rank = dist.get_rank()

499 world_size = dist.get_world_size()

500

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

502 # Far more scalable than all_gather_object for large world_size.

503 device = get_current_device()

504 num_parent_ranks = len(local_parent_ranks)

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

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

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

508

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

510 # sorted set of tuples for deterministic iteration order.

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

512 unique_all_parent_ranks = sorted(set(

513 tuple(row) for row in gathered_cpu_list

514 ))

515

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

517 my_sub_pg: Optional[dist.ProcessGroup] = None

518

519 # Synchronized collective creation loop

520 for parent_ranks in unique_all_parent_ranks:

521 num_sub_groups = len(parent_ranks) // sub_size

522

523 for sub_idx in range(num_sub_groups):

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

525 sub_pg = dist.new_group(sub_ranks)

526

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

528 if parent_ranks == local_parent_ranks:

529 if global_rank in sub_ranks:

530 my_sub_pg = sub_pg

531 sub_pg_map[sub_idx] = sub_pg

532 else:

533 sub_pg_map[sub_idx] = None

534

535 self._split_sub_pg_cache[cache_key] = sub_pg_map

536

537 if my_sub_pg is None:

538 raise RuntimeError(

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

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

541 )

542

543 return my_sub_pg

544

545 def _build_hsdp_batch(

546 self,

547 max_batch_numel: Optional[int] = None,

548 ) -> None:

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

550 if max_batch_numel is None:

551 broadcast_max_bytes = getattr(

552 self, "replicate_broadcast_max_bytes", 512 * 1024 * 1024,

553 )

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

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

556

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

558

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

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

561 sorted_groups = sorted(

562 hsdp_groups,

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

564 )

565

566 batch_groups: List[Dict] = []

567

568 for hsdp_group in sorted_groups:

569 # Sort records within group by numel descending.

570 sorted_records = sorted(

571 hsdp_group.records,

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

573 )

574

575 sub_batches: List[List[Dict]] = []

576 current_batch: List[Dict] = []

577 current_numel = 0

578

579 for record in sorted_records:

580 p_numel = record.param.numel()

581

582 current_batch.append({

583 "record": record,

584 "hsdp_group": hsdp_group,

585 })

586 current_numel += p_numel

587

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

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

590 if current_numel >= max_batch_numel:

591 sub_batches.append(current_batch)

592 current_batch = []

593 current_numel = 0

594

595 if current_batch:

596 sub_batches.append(current_batch)

597

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

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

600

601 batch_groups.append({

602 "hsdp_group": hsdp_group,

603 "sub_batches": sub_batches,

604 })

605

606 # log batch split info.

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

608 logger.info_rank0(

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

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

611 group_key,

612 len(hsdp_groups),

613 len(batch_groups),

614 total_sub_batches,

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

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

617 max_batch_numel

618 )

619 self._hsdp_batches[group_key] = batch_groups

620

621 def _build_sub_batch_assignment(

622 self,

623 sub_batch_entries: List[Dict],

624 hsdp_group: Any,

625 ) -> Optional[HSDPGroupAssignment]:

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

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

628 if not records:

629 return None

630

631 device_mesh = records[0].param.device_mesh

632

633 replicate_group_ranks, replicate_sizes = get_multi_dim_logical_info(

634 device_mesh,

635 hsdp_group.comm_key.replicate_mesh_dims,

636 )

637

638 # When hsdp_replica_count is set, remap coordinates into

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

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

641 if self.hsdp_replica_count is not None:

642 replicate_group_ranks = tuple(

643 r % self.hsdp_replica_count for r in replicate_group_ranks

644 )

645 replicate_sizes = tuple(

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

647 )

648

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

650 owner_by_index = build_owner_by_size(

651 records=records,

652 replicate_sizes=replicate_sizes,

653 )

654

655 owned_records = select_owned_records(

656 records=records,

657 owner_by_index=owner_by_index,

658 replicate_group_ranks=replicate_group_ranks,

659 )

660

661 is_shard_for_ns = (

662 hsdp_group.comm_key.has_shard_group

663 and hsdp_group.layout_spec.is_last2d_sharded

664 )

665

666 hsdp_assign = HSDPGroupAssignment(

667 owned_records=owned_records,

668 all_records=records,

669 owner_by_index=owner_by_index,

670 replicate_group_ranks=replicate_group_ranks,

671 replicate_sizes=replicate_sizes,

672 replicate_pgs=hsdp_group.replicate_pgs,

673 shard_pgs=hsdp_group.shard_pgs,

674 is_shard=is_shard_for_ns,

675 layout_spec=hsdp_group.layout_spec,

676 )

677 logger.debug_rank0(f'[Hyper-optimizer] hsdp_assign: {hsdp_assign}')

678

679 # Build broadcast reverse mapping for replicated groups

680 if hsdp_assign.is_replicated and hsdp_assign.replicate_pgs:

681 for record in hsdp_assign.all_records:

682 src_coord = hsdp_assign.owner_rank_coord(record)

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

684 continue

685 self._param_to_broadcast_info[record.param] = (

686 src_coord, hsdp_assign.replicate_pgs,

687 )

688

689 return hsdp_assign

690

691 def _build_param_broadcast_info(self) -> None:

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

693

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

695 are remapped so that owner assignment and broadcast happen within

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

697 replicate group.

698

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

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

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

702 """

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

704 self._param_to_broadcast_info: Dict[

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

706 ] = {}

707

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

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

710 assignment_batch_groups = []

711

712 for bg in batch_groups:

713 hsdp_group = bg["hsdp_group"]

714 sub_batch_assigns: List[HSDPGroupAssignment] = []

715

716 for sub_batch_entries in bg["sub_batches"]:

717 hsdp_assign = self._build_sub_batch_assignment(sub_batch_entries, hsdp_group)

718 if hsdp_assign is not None:

719 sub_batch_assigns.append(hsdp_assign)

720

721 assignment_batch_groups.append({

722 "hsdp_group": hsdp_group,

723 "sub_batches": sub_batch_assigns,

724 })

725

726 self._hsdp_assignment_batches[group_key] = {

727 "no_comm": no_comm_params,

728 "batch_groups": assignment_batch_groups,

729 }

730

731 def _broadcast_replicate_params_after_step(self) -> None:

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

733 self._broadcast_op_fused(target="param")

734

735 def _broadcast_state_fused_for_ckpt(self) -> None:

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

737 self._broadcast_op_fused(

738 target="state",

739 state_keys=list(self._state_keys_for_checkpoint_broadcast()),

740 )

741

742 def _state_keys_for_checkpoint_broadcast(self) -> Tuple[str, ...]:

743 """Return optimizer state fields synchronized for an HSDP checkpoint."""

744 return ("momentum_buffer",) if self.is_muon else ("exp_avg", "exp_avg_sq")

745

746 def _collect_broadcast_tensors(

747 self,

748 target: str,

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

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

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

752

753 Args:

754 target: "param" or "state".

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

756

757 Returns:

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

759 """

760 rank_dtype_tensors = defaultdict(list)

761

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

763 if target == "param":

764 local_tensor = to_local_if_dtensor(p.data)

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

766

767 elif target == "state" and state_keys:

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

769 for key in state_keys:

770 if key in param_state:

771 state_tensor = param_state[key]

772 local_tensor = to_local_if_dtensor(state_tensor)

773 else:

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

775 param_state[key] = local_tensor

776 local_tensor = to_local_if_dtensor(local_tensor)

777

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

779

780 return rank_dtype_tensors

781

782 @staticmethod

783 def _compute_broadcast_batches(

784 tensors: List[torch.Tensor],

785 alignment_elements: int,

786 max_broadcast_elements: int

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

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

789

790 Args:

791 tensors: List of tensors to batch.

792 alignment_elements: Alignment granularity in elements.

793 max_broadcast_elements: Maximum elements per batch.

794

795 Returns:

796 List of (batch_offsets, batch_total_size) tuples.

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

798 """

799 batches = []

800 current_batch = []

801 current_total_size = 0

802

803 for t in tensors:

804 actual_numel = t.numel()

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

806

807 if current_batch and current_total_size + padded_numel > max_broadcast_elements:

808 batches.append((current_batch, current_total_size))

809 current_batch = []

810 current_total_size = 0

811

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

813 current_total_size += padded_numel

814

815 if current_batch:

816 batches.append((current_batch, current_total_size))

817

818 return batches

819

820 @staticmethod

821 def _hierarchical_broadcast_buffer(

822 batch_buffer: torch.Tensor,

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

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

825 local_coord: Tuple[int, ...]

826 ) -> None:

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

828

829 Args:

830 batch_buffer: Contiguous buffer to broadcast.

831 src_coord: Source coordinate tuple.

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

833 local_coord: Local rank coordinate tuple.

834 """

835

836 for dim_idx, pg in enumerate(replicate_pgs):

837 if pg is None:

838 continue

839

840 participate = True

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

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

843 participate = False

844 break

845

846 if not participate:

847 continue

848

849 src_rank_in_pg = src_coord[dim_idx]

850 global_src_rank = dist.get_global_rank(pg, src_rank_in_pg)

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

852

853 @staticmethod

854 def _hierarchical_broadcast_buffer_async(

855 batch_buffer: torch.Tensor,

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

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

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

859 ) -> List[dist.Work]:

860 """Async version of _hierarchical_broadcast_buffer.

861

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

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

864

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

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

867 """

868 broadcast_ops: List[Tuple[dist.ProcessGroup, int]] = []

869

870 for dim_idx, pg in enumerate(replicate_pgs):

871 if pg is None:

872 continue

873

874 participate = True

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

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

877 participate = False

878 break

879

880 if not participate:

881 continue

882

883 src_rank_in_pg = src_coord[dim_idx]

884 global_src_rank = dist.get_global_rank(pg, src_rank_in_pg)

885 broadcast_ops.append((pg, global_src_rank))

886

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

888 for op_idx, (pg, global_src_rank) in enumerate(broadcast_ops):

889 handle = dist.broadcast(

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

891 )

892 if op_idx + 1 < len(broadcast_ops):

893 # The next mesh dimension relays the data received here.

894 # It must not read the shared buffer until this hop finishes.

895 handle.wait()

896 else:

897 # Only the final hop may overlap with subsequent computation.

898 handles.append(handle)

899

900 return handles

901

902 def _broadcast_op_fused(

903 self,

904 target: str,

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

906 ) -> None:

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

908

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

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

911 State tensors may be CPU swap mirrors; only the current fused buffer is then resident on

912 the accelerator while packing and unpacking each batch.

913

914 Args:

915 target: "param" or "state".

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

917 """

918 device = get_current_device()

919

920 alignment = 512 # bytes

921 rank_dtype_tensors = self._collect_broadcast_tensors(target, state_keys)

922

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

924 if not tensors:

925 continue

926

927 local_coord = tuple(

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

929 for pg in replicate_pgs

930 )

931

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

933 alignment_elements = alignment // element_size

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

935 max_broadcast_elements = max_broadcast_bytes // element_size

936 max_broadcast_elements = max(

937 alignment_elements,

938 (max_broadcast_elements // alignment_elements) * alignment_elements,

939 )

940

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

942 if not batches:

943 continue

944

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

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

947

948 for batch_tensor_offsets, batch_total_size in batches:

949 batch_buffer = buffer[:batch_total_size]

950

951 # Pack: owner rank

952 if local_coord == src_coord:

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

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

955 if padded_numel > actual_numel:

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

957

958 # Hierarchical Broadcast

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

960

961 # Unpack: copy buffer data back to individual tensors

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

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

964

965 buffer.untyped_storage().resize_(0)

966 del buffer

967

968 def cleanup_synced_state(self) -> None:

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

970

971 In HSDP mode, _broadcast_state_fused_for_ckpt broadcasts state to all

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

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

974 after the checkpoint has been fully written to disk.

975 """

976 params_to_remove = []

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

978 for bg in assignment_info["batch_groups"]:

979 for hsdp_info in bg["sub_batches"]:

980 if not hsdp_info.is_replicated:

981 continue

982 for record in hsdp_info.all_records:

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

984 params_to_remove.append(record.param)

985

986 for p in params_to_remove:

987 self.state[p].clear()

988 del self.state[p]

989

990 empty_accelerator_cache()

991

992

993class AsyncReplicateBroadcaster:

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

995

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

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

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

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

1000

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

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

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

1004

1005 Usage::

1006

1007 broadcaster = AsyncReplicateBroadcaster(optimizer)

1008 for hsdp_assign in hsdp_assignments:

1009 # ... NS compute + apply for this group ...

1010 broadcaster.flush_group(hsdp_assign)

1011 broadcaster.wait_all()

1012 """

1013

1014 def __init__(

1015 self,

1016 optimizer: BaseDistributedOptimizer,

1017 ) -> None:

1018 """Bind the distributed optimizer whose broadcasts are batched.

1019

1020 Args:

1021 optimizer: Distributed optimizer owning the replica assignments.

1022 """

1023 self._optimizer = optimizer

1024

1025 # Inflight async broadcasts: list of (buffer, batch_offsets, handles)

1026 # Buffer must stay alive until handles are waited on.

1027 self._inflight: List[

1028 Tuple[

1029 torch.Tensor,

1030 List[Tuple[torch.Tensor, int, int, int]],

1031 List[dist.Work],

1032 ]

1033 ] = []

1034 # Allow one batch of replicate communication to overlap with the

1035 # next batch's compute, but avoid unbounded inflight buffer growth.

1036 self._max_inflight_batches = 1

1037

1038 def flush_group(

1039 self,

1040 hsdp_assign: Any,

1041 records: Optional[List] = None,

1042 ) -> None:

1043 """Flush params belonging to the given HSDP group.

1044

1045 Can be called per sub_batch (passing only the records in that

1046 sub_batch) so that the async replicate broadcast overlaps with

1047 the next sub_batch's NS computation. When *records* is None,

1048 all records in the group are flushed (backward compatible).

1049

1050 All ranks must call this at the same point in the execution flow

1051 to ensure collective communication consistency.

1052 """

1053 if not hsdp_assign.is_replicated or not hsdp_assign.replicate_pgs:

1054 return

1055

1056 flush_records = records if records is not None else hsdp_assign.all_records

1057

1058 # Collect local tensors for the given records, grouped by

1059 # (src_coord, dtype, replicate_pgs) — same logic as

1060 # _collect_broadcast_tensors but scoped to the provided records.

1061 rank_dtype_tensors: Dict[

1062 Tuple[Tuple[int, ...], torch.dtype, Tuple[dist.ProcessGroup, ...]],

1063 List[torch.Tensor],

1064 ] = defaultdict(list)

1065

1066 for record in flush_records:

1067 src_coord = hsdp_assign.owner_rank_coord(record)

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

1069 continue

1070 local_tensor = to_local_if_dtensor(record.param.data)

1071 key = (src_coord, local_tensor.dtype, hsdp_assign.replicate_pgs)

1072 rank_dtype_tensors[key].append(local_tensor)

1073

1074 for key, tensors in rank_dtype_tensors.items():

1075 if tensors:

1076 self._flush_key(key, tensors, async_op=True)

1077

1078 while len(self._inflight) > self._max_inflight_batches:

1079 self._wait_and_release_oldest()

1080

1081 def _flush_key(

1082 self,

1083 key: Tuple[Tuple[int, ...], torch.dtype, Tuple[dist.ProcessGroup, ...]],

1084 tensors: List[torch.Tensor],

1085 async_op: bool = True,

1086 ) -> None:

1087 """Pack and broadcast tensors for one broadcast key."""

1088 device = get_current_device()

1089 src_coord, dtype, replicate_pgs = key

1090 alignment = 512 # bytes

1091

1092 local_coord = tuple(

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

1094 for pg in replicate_pgs

1095 )

1096

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

1098 alignment_elements = alignment // element_size

1099 max_broadcast_bytes = getattr(

1100 self._optimizer, "replicate_broadcast_max_bytes", 512 * 1024 * 1024,

1101 )

1102 max_broadcast_elements = max_broadcast_bytes // element_size

1103 max_broadcast_elements = max(

1104 alignment_elements,

1105 (max_broadcast_elements // alignment_elements) * alignment_elements,

1106 )

1107

1108 # pylint: disable=protected-access

1109 batches = BaseDistributedOptimizer._compute_broadcast_batches(

1110 tensors, alignment_elements, max_broadcast_elements,

1111 )

1112 if not batches:

1113 return

1114

1115 buffer: torch.Tensor

1116 if not async_op:

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

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

1119

1120 for batch_tensor_offsets, batch_total_size in batches:

1121 if async_op:

1122 batch_buffer = torch.empty(batch_total_size, dtype=dtype, device=device)

1123 else:

1124 batch_buffer = buffer[:batch_total_size]

1125

1126 # Pack: owner rank

1127 if local_coord == src_coord:

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

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

1130 if padded_numel > actual_numel:

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

1132

1133 if async_op:

1134 handles = BaseDistributedOptimizer._hierarchical_broadcast_buffer_async(

1135 batch_buffer, src_coord, replicate_pgs, local_coord,

1136 )

1137 if handles:

1138 # Pin buffer + offsets until wait_all unpacks them.

1139 self._inflight.append((batch_buffer, batch_tensor_offsets, handles))

1140 else:

1141 # No async work was enqueued on this rank, so unpack and free now.

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

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

1144 batch_buffer.untyped_storage().resize_(0)

1145 del batch_buffer

1146 else:

1147 BaseDistributedOptimizer._hierarchical_broadcast_buffer(

1148 batch_buffer, src_coord, replicate_pgs, local_coord,

1149 )

1150 # Unpack immediately for sync path

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

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

1153

1154 if not async_op:

1155 # Sync path: buffer can be freed immediately

1156 buffer.untyped_storage().resize_(0)

1157 del buffer

1158

1159 def _wait_and_release_oldest(self) -> None:

1160 """Wait, unpack, and release the oldest inflight async batch."""

1161 batch_buffer, batch_tensor_offsets, handles = self._inflight.pop(0)

1162 for handle in handles:

1163 handle.wait()

1164

1165 # Unpack after the async broadcast completes.

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

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

1168

1169 batch_buffer.untyped_storage().resize_(0)

1170 del batch_buffer

1171

1172 def wait_all(self) -> None:

1173 """Wait for all inflight async broadcasts and unpack results."""

1174 while self._inflight:

1175 self._wait_and_release_oldest()

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

541 statements