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

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.swap_optimizer_base import _restore_dtensor_metadata

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

44logger = logging.getLogger(__name__)

47class ChainedOptimizer:

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

50 def __init__(

51 self,

52 model: torch.nn.Module,

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

54 flatten: bool = False,

55 ) -> None:

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

58 Args:

59 model: Module containing the registered model parameters.

60 optimizers: Leaf optimizers keyed by their configured names.

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

62 """

63 self.optimizers_dict = optimizers

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

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

66 self.model = model

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

68 self._is_multi_optimizer = flatten

69 self.model_param_by_optimizer_param: Dict[

70 torch.nn.Parameter, torch.nn.Parameter

71 ] = {}

72 self.optimizer_param_by_model_param: Dict[

73 torch.nn.Parameter, torch.nn.Parameter

74 ] = {}

76 self._rebind_param_attrs()

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

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

80 return iter(self.chained_optimizers)

82 def _rebind_param_attrs(self) -> None:

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

84 muon_param_ids: set = set()

85 adamw_param_ids: set = set()

86 self.muon_keys = []

87 self.no_muon_keys = []

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

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

90 target = muon_param_ids if is_muon else adamw_param_ids

91 for group in opt.param_groups:

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

93 target.add(id(p))

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

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

97 setattr(param, "model_name", param_name)

98 setattr(optimizer_param, "model_name", param_name)

99 if id(optimizer_param) in muon_param_ids:

100 self.muon_keys.append(param_name)

101 elif id(optimizer_param) in adamw_param_ids:

102 self.no_muon_keys.append(param_name)

103

104 def reset_optimizer_parameters(

105 self,

106 model_param_by_optimizer_param: Dict[

107 torch.nn.Parameter, torch.nn.Parameter

108 ],

109 ) -> None:

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

111

112 Args:

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

114 """

115 self.model_param_by_optimizer_param = dict(

116 model_param_by_optimizer_param

117 )

118 self.optimizer_param_by_model_param = {

119 model_param: optimizer_param

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

121 }

122 self._rebind_param_attrs()

123 for optimizer in self.chained_optimizers:

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

125 if callable(reset):

126 if is_swap_optimizer(optimizer):

127 reset(self.model_param_by_optimizer_param)

128 else:

129 reset()

130

131 @contextmanager

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

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

134 if not self.model_param_by_optimizer_param:

135 yield

136 return

137

138 registrations = []

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

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

141 if model_param is None:

142 continue

143 optimizer_param = self.optimizer_param_by_model_param.get(model_param)

144 if optimizer_param is None:

145 continue

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

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

148 try:

149 yield

150 finally:

151 for module, parameter_name, model_param in registrations:

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

153

154 def __str__(self) -> str:

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

156 if not self.optimizers_dict:

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

158

159 sections = []

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

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

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

163

164 __repr__ = __str__

165

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

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

168

169 Args:

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

171 """

172 for opt in self.chained_optimizers:

173 opt.step(closure=closure)

174

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

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

177

178 Args:

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

180 """

181 for opt in self.chained_optimizers:

182 opt.zero_grad(set_to_none=set_to_none)

183

184 @property

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

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

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

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

189 return self.chained_optimizers[0]

190

191 @property

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

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

194 return self.chained_optimizers[0].defaults

195

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

197 """Return state dicts with owning checkpoint tensor copies.

198

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

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

201 then clones exported values while preserving every DTensor layout.

202 """

203 prepared = []

204 try:

205 # Swap optimizers must restore released device storage before the

206 # wrapped HSDP implementation reads its state tensors.

207 for opt in self.chained_optimizers:

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

209 if callable(prepare):

210 prepared.append(opt)

211 prepare()

212

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

214 for opt in self.chained_optimizers:

215 if hasattr(opt, "_broadcast_state_fused_for_ckpt"):

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

217

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

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

220 with self._optimizer_parameter_view():

221 sd = get_optimizer_state_dict(

222 self.model,

223 optimizer,

224 options=StateDictOptions(flatten_optimizer_state_dict=self.flatten)

225 )

226 if optimizer in prepared:

227 # The swap lifecycle releases device storage once the

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

229 # their memory instead of aliasing the slot tensors.

230 self._localize_state_dict_tensors(sd)

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

232 if overlap:

233 raise KeyError(

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

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

236 )

237 merged.update(sd)

238

239 return merged

240 finally:

241 for opt in reversed(prepared):

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

243 if callable(cleanup):

244 cleanup()

245 continue

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

247 if callable(finish):

248 finish()

249

250 @staticmethod

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

252 """Replace exported tensors with owning copies that retain DTensor layouts.

253

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

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

256 storage the swap runtime is about to free.

257 """

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

259 """Clone tensors in one nested checkpoint value.

260

261 Args:

262 value: Checkpoint value or nested container to localize.

263

264 Returns:

265 The value with owning tensor copies and unchanged DTensor metadata.

266 """

267 if isinstance(value, torch.Tensor):

268 local_copy = to_local_if_dtensor(value).detach().cpu().clone()

269 return _restore_dtensor_metadata(value, local_copy)

270 if isinstance(value, dict):

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

272 if isinstance(value, list):

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

274 if isinstance(value, tuple):

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

276 return value

277

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

279 if key == "param_groups":

280 continue

281 state_dict[key] = localize(value)

282

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

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

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

286 for optimizer in self.chained_optimizers:

287 param_groups += optimizer.param_groups

288 return param_groups

289

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

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

292 if not isinstance(new_param_groups, list):

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

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

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

296

297 start = 0

298 for optimizer in self.chained_optimizers:

299 group_len = len(optimizer.param_groups)

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

301 start += group_len

302

303 param_groups = property(_get_param_groups, _set_param_groups)

304

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

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

307

308 Args:

309 state_dict: FQN-keyed optimizer state to restore.

310 """

311 for optimizer in self.chained_optimizers:

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

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

314 with self._optimizer_parameter_view(), context:

315 set_optimizer_state_dict(

316 self.model,

317 optimizer,

318 optim_state_dict=state_dict,

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

320 )

321

322 self._synchronize_steps()

323

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

325 """Synchronize the step of all optimizers.

326

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

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

329 """

330 steps = []

331 for optimizer in self.chained_optimizers:

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

333 for param_group in actual_opt.param_groups:

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

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

336

337 unique_steps = list(set(steps))

338 if len(unique_steps) > 1:

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

340

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

342

343 for optimizer in self.chained_optimizers:

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

345 for param_group in actual_opt.param_groups:

346 param_group['step'] = step

347

348 return step

349

350

351class BaseDistributedOptimizer(torch.optim.Optimizer):

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

353

354 Provides fused hierarchical broadcast for parameters and optimizer states.

355 """

356

357 def __init__(

358 self,

359 params: Any,

360 defaults: Dict[str, Any],

361 is_muon: bool,

362 hsdp_replica_count: Optional[int] = None,

363 ) -> None:

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

365

366 Args:

367 params: Parameter groups optimized by this instance.

368 defaults: Default hyperparameters applied to every group.

369 is_muon: Whether parameters use the Muon update path.

370 hsdp_replica_count: Optional optimizer-state replica group size.

371 """

372 super().__init__(params, defaults)

373 self.is_muon = is_muon

374 self.hsdp_replica_count = hsdp_replica_count

375 self._param_to_broadcast_info: Dict[

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

377 ] = {}

378

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

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

381

382 def _group_dtensor_by_mesh(self):

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

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

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

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

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

388

389 def reset_optimizer_parameters(self) -> None:

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

391 raise NotImplementedError

392

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

394 """Deduce hsdp_replica_count based on cluster topology.

395

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

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

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

399 """

400 devices_per_node = get_device_count()

401

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

403 needs_split = False

404

405 # group_key, (no_comm_params, hsdp_groups)

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

407 for hsdp_group in hsdp_groups:

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

409 if pg is None:

410 continue

411 pg_size = dist.get_world_size(pg)

412 if pg_size <= 1:

413 continue

414

415 if pg_size > devices_per_node:

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

417 dedup = min(pg_size, devices_per_node)

418 while pg_size % dedup != 0:

419 dedup -= 1

420 needs_split = True

421 else:

422 # origin Inter-node

423 dedup = pg_size

424

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

426 if dim_idx not in dedup_per_dim:

427 dedup_per_dim[dim_idx] = dedup

428 else:

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

430

431 if not needs_split:

432 return None

433

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

435 return min(sorted_dedups)

436

437 def _split_replicate_groups(self) -> None:

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

439 if self.hsdp_replica_count is None:

440 return

441

442 # Argument validation

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

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

445 if self.hsdp_replica_count <= 0:

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

447

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

449 for hsdp_group in hsdp_groups:

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

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

452 if pg is None:

453 new_replicate_pgs.append(pg)

454 continue

455

456 pg_size = dist.get_world_size(pg)

457 dedup_size = self.hsdp_replica_count

458

459 if pg_size <= dedup_size:

460 new_replicate_pgs.append(pg)

461 continue

462

463 if pg_size % dedup_size != 0:

464 raise ValueError(

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

466 )

467

468 sub_pg = self._get_or_create_sub_pg(pg, dedup_size)

469 new_replicate_pgs.append(sub_pg)

470

471 logger.info_rank0(

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

473 group_key, dim_idx, pg_size, dedup_size

474 )

475

476 # replace new sub groups of hsdp groups

477 hsdp_group.replicate_pgs = tuple(new_replicate_pgs)

478

479 def _get_or_create_sub_pg(

480 self,

481 parent_pg: dist.ProcessGroup,

482 sub_size: int,

483 ) -> dist.ProcessGroup:

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

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

486 cache_key = (local_parent_ranks, sub_size)

487

488 # Cache Hit Check

489 if cache_key in self._split_sub_pg_cache:

490 sub_pg_map = self._split_sub_pg_cache[cache_key]

491 for sub_pg in sub_pg_map.values():

492 if sub_pg is not None:

493 try:

494 dist.get_rank(group=sub_pg)

495 return sub_pg

496 except RuntimeError:

497 continue

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

499

500 global_rank = dist.get_rank()

501 world_size = dist.get_world_size()

502

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

504 # Far more scalable than all_gather_object for large world_size.

505 device = get_current_device()

506 num_parent_ranks = len(local_parent_ranks)

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

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

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

510

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

512 # sorted set of tuples for deterministic iteration order.

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

514 unique_all_parent_ranks = sorted(set(

515 tuple(row) for row in gathered_cpu_list

516 ))

517

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

519 my_sub_pg: Optional[dist.ProcessGroup] = None

520

521 # Synchronized collective creation loop

522 for parent_ranks in unique_all_parent_ranks:

523 num_sub_groups = len(parent_ranks) // sub_size

524

525 for sub_idx in range(num_sub_groups):

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

527 sub_pg = dist.new_group(sub_ranks)

528

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

530 if parent_ranks == local_parent_ranks:

531 if global_rank in sub_ranks:

532 my_sub_pg = sub_pg

533 sub_pg_map[sub_idx] = sub_pg

534 else:

535 sub_pg_map[sub_idx] = None

536

537 self._split_sub_pg_cache[cache_key] = sub_pg_map

538

539 if my_sub_pg is None:

540 raise RuntimeError(

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

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

543 )

544

545 return my_sub_pg

546

547 def _build_hsdp_batch(

548 self,

549 max_batch_numel: Optional[int] = None,

550 ) -> None:

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

552 if max_batch_numel is None:

553 broadcast_max_bytes = getattr(

554 self, "replicate_broadcast_max_bytes", 512 * 1024 * 1024,

555 )

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

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

558

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

560

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

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

563 sorted_groups = sorted(

564 hsdp_groups,

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

566 )

567

568 batch_groups: List[Dict] = []

569

570 for hsdp_group in sorted_groups:

571 # Sort records within group by numel descending.

572 sorted_records = sorted(

573 hsdp_group.records,

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

575 )

576

577 sub_batches: List[List[Dict]] = []

578 current_batch: List[Dict] = []

579 current_numel = 0

580

581 for record in sorted_records:

582 p_numel = record.param.numel()

583

584 current_batch.append({

585 "record": record,

586 "hsdp_group": hsdp_group,

587 })

588 current_numel += p_numel

589

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

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

592 if current_numel >= max_batch_numel:

593 sub_batches.append(current_batch)

594 current_batch = []

595 current_numel = 0

596

597 if current_batch:

598 sub_batches.append(current_batch)

599

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

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

602

603 batch_groups.append({

604 "hsdp_group": hsdp_group,

605 "sub_batches": sub_batches,

606 })

607

608 # log batch split info.

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

610 logger.info_rank0(

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

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

613 group_key,

614 len(hsdp_groups),

615 len(batch_groups),

616 total_sub_batches,

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

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

619 max_batch_numel

620 )

621 self._hsdp_batches[group_key] = batch_groups

622

623 def _build_sub_batch_assignment(

624 self,

625 sub_batch_entries: List[Dict],

626 hsdp_group: Any,

627 ) -> Optional[HSDPGroupAssignment]:

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

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

630 if not records:

631 return None

632

633 device_mesh = records[0].param.device_mesh

634

635 replicate_group_ranks, replicate_sizes = get_multi_dim_logical_info(

636 device_mesh,

637 hsdp_group.comm_key.replicate_mesh_dims,

638 )

639

640 # When hsdp_replica_count is set, remap coordinates into

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

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

643 if self.hsdp_replica_count is not None:

644 replicate_group_ranks = tuple(

645 r % self.hsdp_replica_count for r in replicate_group_ranks

646 )

647 replicate_sizes = tuple(

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

649 )

650

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

652 owner_by_index = build_owner_by_size(

653 records=records,

654 replicate_sizes=replicate_sizes,

655 )

656

657 owned_records = select_owned_records(

658 records=records,

659 owner_by_index=owner_by_index,

660 replicate_group_ranks=replicate_group_ranks,

661 )

662

663 is_shard_for_ns = (

664 hsdp_group.comm_key.has_shard_group

665 and hsdp_group.layout_spec.is_last2d_sharded

666 )

667

668 hsdp_assign = HSDPGroupAssignment(

669 owned_records=owned_records,

670 all_records=records,

671 owner_by_index=owner_by_index,

672 replicate_group_ranks=replicate_group_ranks,

673 replicate_sizes=replicate_sizes,

674 replicate_pgs=hsdp_group.replicate_pgs,

675 shard_pgs=hsdp_group.shard_pgs,

676 is_shard=is_shard_for_ns,

677 layout_spec=hsdp_group.layout_spec,

678 )

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

680

681 # Build broadcast reverse mapping for replicated groups

682 if hsdp_assign.is_replicated and hsdp_assign.replicate_pgs:

683 for record in hsdp_assign.all_records:

684 src_coord = hsdp_assign.owner_rank_coord(record)

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

686 continue

687 self._param_to_broadcast_info[record.param] = (

688 src_coord, hsdp_assign.replicate_pgs,

689 )

690

691 return hsdp_assign

692

693 def _build_param_broadcast_info(self) -> None:

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

695

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

697 are remapped so that owner assignment and broadcast happen within

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

699 replicate group.

700

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

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

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

704 """

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

706 self._param_to_broadcast_info: Dict[

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

708 ] = {}

709

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

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

712 assignment_batch_groups = []

713

714 for bg in batch_groups:

715 hsdp_group = bg["hsdp_group"]

716 sub_batch_assigns: List[HSDPGroupAssignment] = []

717

718 for sub_batch_entries in bg["sub_batches"]:

719 hsdp_assign = self._build_sub_batch_assignment(sub_batch_entries, hsdp_group)

720 if hsdp_assign is not None:

721 sub_batch_assigns.append(hsdp_assign)

722

723 assignment_batch_groups.append({

724 "hsdp_group": hsdp_group,

725 "sub_batches": sub_batch_assigns,

726 })

727

728 self._hsdp_assignment_batches[group_key] = {

729 "no_comm": no_comm_params,

730 "batch_groups": assignment_batch_groups,

731 }

732

733 def _broadcast_replicate_params_after_step(self) -> None:

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

735 self._broadcast_op_fused(target="param")

736

737 def _broadcast_state_fused_for_ckpt(self) -> None:

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

739 self._broadcast_op_fused(

740 target="state",

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

742 )

743

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

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

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

747

748 def _collect_broadcast_tensors(

749 self,

750 target: str,

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

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

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

754

755 Args:

756 target: "param" or "state".

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

758

759 Returns:

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

761 """

762 rank_dtype_tensors = defaultdict(list)

763

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

765 if target == "param":

766 local_tensor = to_local_if_dtensor(p.data)

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

768

769 elif target == "state" and state_keys:

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

771 for key in state_keys:

772 if key in param_state:

773 state_tensor = param_state[key]

774 local_tensor = to_local_if_dtensor(state_tensor)

775 else:

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

777 param_state[key] = local_tensor

778 local_tensor = to_local_if_dtensor(local_tensor)

779

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

781

782 return rank_dtype_tensors

783

784 @staticmethod

785 def _compute_broadcast_batches(

786 tensors: List[torch.Tensor],

787 alignment_elements: int,

788 max_broadcast_elements: int

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

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

791

792 Args:

793 tensors: List of tensors to batch.

794 alignment_elements: Alignment granularity in elements.

795 max_broadcast_elements: Maximum elements per batch.

796

797 Returns:

798 List of (batch_offsets, batch_total_size) tuples.

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

800 """

801 batches = []

802 current_batch = []

803 current_total_size = 0

804

805 for t in tensors:

806 actual_numel = t.numel()

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

808

809 if current_batch and current_total_size + padded_numel > max_broadcast_elements:

810 batches.append((current_batch, current_total_size))

811 current_batch = []

812 current_total_size = 0

813

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

815 current_total_size += padded_numel

816

817 if current_batch:

818 batches.append((current_batch, current_total_size))

819

820 return batches

821

822 @staticmethod

823 def _hierarchical_broadcast_buffer(

824 batch_buffer: torch.Tensor,

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

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

827 local_coord: Tuple[int, ...]

828 ) -> None:

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

830

831 Args:

832 batch_buffer: Contiguous buffer to broadcast.

833 src_coord: Source coordinate tuple.

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

835 local_coord: Local rank coordinate tuple.

836 """

837

838 for dim_idx, pg in enumerate(replicate_pgs):

839 if pg is None:

840 continue

841

842 participate = True

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

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

845 participate = False

846 break

847

848 if not participate:

849 continue

850

851 src_rank_in_pg = src_coord[dim_idx]

852 global_src_rank = dist.get_global_rank(pg, src_rank_in_pg)

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

854

855 @staticmethod

856 def _hierarchical_broadcast_buffer_async(

857 batch_buffer: torch.Tensor,

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

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

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

861 ) -> List[dist.Work]:

862 """Async version of _hierarchical_broadcast_buffer.

863

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

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

866

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

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

869 """

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

871

872 for dim_idx, pg in enumerate(replicate_pgs):

873 if pg is None:

874 continue

875

876 participate = True

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

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

879 participate = False

880 break

881

882 if not participate:

883 continue

884

885 src_rank_in_pg = src_coord[dim_idx]

886 global_src_rank = dist.get_global_rank(pg, src_rank_in_pg)

887 broadcast_ops.append((pg, global_src_rank))

888

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

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

891 handle = dist.broadcast(

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

893 )

894 if op_idx + 1 < len(broadcast_ops):

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

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

897 handle.wait()

898 else:

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

900 handles.append(handle)

901

902 return handles

903

904 def _broadcast_op_fused(

905 self,

906 target: str,

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

908 ) -> None:

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

910

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

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

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

914 the accelerator while packing and unpacking each batch.

915

916 Args:

917 target: "param" or "state".

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

919 """

920 device = get_current_device()

921

922 alignment = 512 # bytes

923 rank_dtype_tensors = self._collect_broadcast_tensors(target, state_keys)

924

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

926 if not tensors:

927 continue

928

929 local_coord = tuple(

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

931 for pg in replicate_pgs

932 )

933

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

935 alignment_elements = alignment // element_size

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

937 max_broadcast_elements = max_broadcast_bytes // element_size

938 max_broadcast_elements = max(

939 alignment_elements,

940 (max_broadcast_elements // alignment_elements) * alignment_elements,

941 )

942

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

944 if not batches:

945 continue

946

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

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

949

950 for batch_tensor_offsets, batch_total_size in batches:

951 batch_buffer = buffer[:batch_total_size]

952

953 # Pack: owner rank

954 if local_coord == src_coord:

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

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

957 if padded_numel > actual_numel:

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

959

960 # Hierarchical Broadcast

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

962

963 # Unpack: copy buffer data back to individual tensors

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

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

966

967 buffer.untyped_storage().resize_(0)

968 del buffer

969

970 def cleanup_synced_state(self) -> None:

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

972

973 In HSDP mode, _broadcast_state_fused_for_ckpt broadcasts state to all

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

975 the non-owned state to restore per-rank memory savings. Call once all

976 checkpoint consumers have finished reading the live state, either after

977 persistence or after creating an independent snapshot for async I/O.

978 """

979 params_to_remove = []

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

981 for bg in assignment_info["batch_groups"]:

982 for hsdp_info in bg["sub_batches"]:

983 if not hsdp_info.is_replicated:

984 continue

985 for record in hsdp_info.all_records:

986 # Records without an owner are never broadcast, so their state is not a replica copy.

987 if record.index not in hsdp_info.owner_by_index:

988 continue

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

990 params_to_remove.append(record.param)

991

992 for p in params_to_remove:

993 self.state[p].clear()

994 del self.state[p]

995

996 empty_accelerator_cache()

997

998

999class AsyncReplicateBroadcaster:

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

1001

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

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

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

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

1006

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

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

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

1010

1011 Usage::

1012

1013 broadcaster = AsyncReplicateBroadcaster(optimizer)

1014 for hsdp_assign in hsdp_assignments:

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

1016 broadcaster.flush_group(hsdp_assign)

1017 broadcaster.wait_all()

1018 """

1019

1020 def __init__(

1021 self,

1022 optimizer: BaseDistributedOptimizer,

1023 ) -> None:

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

1025

1026 Args:

1027 optimizer: Distributed optimizer owning the replica assignments.

1028 """

1029 self._optimizer = optimizer

1030

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

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

1033 self._inflight: List[

1034 Tuple[

1035 torch.Tensor,

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

1037 List[dist.Work],

1038 ]

1039 ] = []

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

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

1042 self._max_inflight_batches = 1

1043

1044 def flush_group(

1045 self,

1046 hsdp_assign: Any,

1047 records: Optional[List] = None,

1048 ) -> None:

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

1050

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

1052 sub_batch) so that the async replicate broadcast overlaps with

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

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

1055

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

1057 to ensure collective communication consistency.

1058 """

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

1060 return

1061

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

1063

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

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

1066 # _collect_broadcast_tensors but scoped to the provided records.

1067 rank_dtype_tensors: Dict[

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

1069 List[torch.Tensor],

1070 ] = defaultdict(list)

1071

1072 for record in flush_records:

1073 src_coord = hsdp_assign.owner_rank_coord(record)

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

1075 continue

1076 local_tensor = to_local_if_dtensor(record.param.data)

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

1078 rank_dtype_tensors[key].append(local_tensor)

1079

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

1081 if tensors:

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

1083

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

1085 self._wait_and_release_oldest()

1086

1087 def _flush_key(

1088 self,

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

1090 tensors: List[torch.Tensor],

1091 async_op: bool = True,

1092 ) -> None:

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

1094 device = get_current_device()

1095 src_coord, dtype, replicate_pgs = key

1096 alignment = 512 # bytes

1097

1098 local_coord = tuple(

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

1100 for pg in replicate_pgs

1101 )

1102

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

1104 alignment_elements = alignment // element_size

1105 max_broadcast_bytes = getattr(

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

1107 )

1108 max_broadcast_elements = max_broadcast_bytes // element_size

1109 max_broadcast_elements = max(

1110 alignment_elements,

1111 (max_broadcast_elements // alignment_elements) * alignment_elements,

1112 )

1113

1114 # pylint: disable=protected-access

1115 batches = BaseDistributedOptimizer._compute_broadcast_batches(

1116 tensors, alignment_elements, max_broadcast_elements,

1117 )

1118 if not batches:

1119 return

1120

1121 buffer: torch.Tensor

1122 if not async_op:

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

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

1125

1126 for batch_tensor_offsets, batch_total_size in batches:

1127 if async_op:

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

1129 else:

1130 batch_buffer = buffer[:batch_total_size]

1131

1132 # Pack: owner rank

1133 if local_coord == src_coord:

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

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

1136 if padded_numel > actual_numel:

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

1138

1139 if async_op:

1140 handles = BaseDistributedOptimizer._hierarchical_broadcast_buffer_async(

1141 batch_buffer, src_coord, replicate_pgs, local_coord,

1142 )

1143 if handles:

1144 # Pin buffer + offsets until wait_all unpacks them.

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

1146 else:

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

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

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

1150 batch_buffer.untyped_storage().resize_(0)

1151 del batch_buffer

1152 else:

1153 BaseDistributedOptimizer._hierarchical_broadcast_buffer(

1154 batch_buffer, src_coord, replicate_pgs, local_coord,

1155 )

1156 # Unpack immediately for sync path

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

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

1159

1160 if not async_op:

1161 # Sync path: buffer can be freed immediately

1162 buffer.untyped_storage().resize_(0)

1163 del buffer

1164

1165 def _wait_and_release_oldest(self) -> None:

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

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

1168 for handle in handles:

1169 handle.wait()

1170

1171 # Unpack after the async broadcast completes.

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

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

1174

1175 batch_buffer.untyped_storage().resize_(0)

1176 del batch_buffer

1177

1178 def wait_all(self) -> None:

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

1180 while self._inflight:

1181 self._wait_and_release_oldest()

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

545 statements