Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper

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"""Muon optimizer with HSDP shard-group-aware communication."""

18import math

19from collections.abc import Callable, Iterable

20from collections import defaultdict

21from dataclasses import dataclass

22import logging

23from typing import Any, Dict, List, Optional, Sequence, Tuple, Union

25import torch # pylint: disable=forbidden-backend-import

26import torch.distributed as dist # pylint: disable=forbidden-backend-import

28from hyper_parallel.core.optimizer.optimizer import AsyncReplicateBroadcaster, BaseDistributedOptimizer

29from hyper_parallel.core.optimizer.dtensor_compat import to_local_if_dtensor

30from hyper_parallel.core.optimizer.sharding_category import (

31 HSDPGroupAssignment,

32 ParamShardMeta,

33 build_owner_by_size,

34)

35from hyper_parallel.core.optimizer.muon_shard import (

36 _debug_param_shard_metadata,

37 build_pad_ns_inputs,

38 build_param_shard_metadata_for_group,

39 chunk_update_by_layout,

40 fused_allgather_dtensor_params,

41)

43logger = logging.getLogger(__name__)

45# Legacy: single-coefficient quintic NS (Keller Jordan / Moonlight).

46# Same (a, b, c) applied every step; typically 5 steps.

47_NS_LEGACY_COEFFS: Tuple[float, float, float] = (3.4445, -4.7750, 2.0315)

49# Asym5: 5-step asymmetric NS (dion). Each step uses different coefficients.

50_NS_ASYM5_COEFFS: Tuple[Tuple[float, float, float], ...] = (

51 (4.0848, -6.8946, 2.9270),

52 (3.9505, -6.3029, 2.6377),

53 (3.7418, -5.5913, 2.3037),

54 (2.8769, -3.1427, 1.2046),

55 (2.8366, -3.0525, 1.2012),

56)

59@dataclass

60class NSInputTransform:

61 """Prepared NS tensors and the callback that restores their updates."""

63 tensors: List[torch.Tensor]

64 restore: Callable[[List[torch.Tensor], torch.Tensor], None]

67@dataclass(frozen=True)

68class MuonPostUpdateContext:

69 """Context supplied to an optional parameter post-update callback."""

71 param_fqn: Optional[str]

72 logical_shape: Tuple[int, ...]

73 lr: float

74 weight_decay: float

75 step: int

78_MUON_ADVANCED_DEFAULTS: Dict[str, Any] = {

79 "ns_coefficients": None,

80 "ns_epsilon": 1e-10,

81 "zeropower_fn": None,

82 "momentum_update_fn": None,

83 "reshape_fn": None,

84 "ns_transform_fn": None,

85 "post_update_fn": None,

86 "zero_rms_scale_mode": "zero",

87 "apply_lr_in_update": False,

88}

91@dataclass(frozen=True)

92class _MuonAdvancedOptions:

93 """Validated low-frequency Muon options."""

95 ns_coefficients: Optional[Sequence[Tuple[float, float, float]]]

96 ns_epsilon: float

97 zeropower_fn: Optional[Callable[..., torch.Tensor]]

98 momentum_update_fn: Optional[Callable[..., torch.Tensor]]

99 reshape_fn: Optional[Callable[..., Any]]

100 ns_transform_fn: Optional[Callable[..., Optional[NSInputTransform]]]

101 post_update_fn: Optional[Callable[..., None]]

102 zero_rms_scale_mode: str

103 apply_lr_in_update: bool

104

105

106def _resolve_muon_advanced_options(options: Dict[str, Any]) -> _MuonAdvancedOptions:

107 """Validate and fill low-frequency Muon keyword options."""

108 unknown = options.keys() - _MUON_ADVANCED_DEFAULTS.keys()

109 if unknown:

110 raise TypeError(f"Muon got unexpected keyword arguments: {', '.join(sorted(unknown))}")

111 resolved = dict(_MUON_ADVANCED_DEFAULTS)

112 resolved.update(options)

113 return _MuonAdvancedOptions(**resolved)

114

115

116_TransformNSState = Dict[str, Any]

117

118

119def zeropower_via_newtonschulz5(

120 ns_inputs: torch.Tensor,

121 steps: int,

122 ns_variant: str = "asym5",

123 epsilon: float = 1e-10,

124 ns_coefficients: Optional[Sequence[Tuple[float, float, float]]] = None,

125) -> torch.Tensor:

126 """Newton-Schulz orthogonalization with preallocated matmul buffers."""

127 mat_x = ns_inputs

128 transposed = ns_inputs.size(-2) > ns_inputs.size(-1)

129 if transposed:

130 mat_x = mat_x.mT

131

132 # Normalize input before Newton-Schulz iteration.

133 mat_x = mat_x / (mat_x.norm(dim=(-2, -1), keepdim=True) + epsilon)

134

135 n_size = mat_x.size(-2)

136 buf_a = torch.empty(mat_x.shape[:-2] + (n_size, n_size), dtype=mat_x.dtype, device=mat_x.device)

137 buf_b = torch.empty_like(buf_a)

138 buf_c = torch.empty(mat_x.shape, dtype=mat_x.dtype, device=mat_x.device)

139

140 if ns_variant == "legacy":

141 step_coeffs = [_NS_LEGACY_COEFFS] * steps

142 elif ns_variant == "asym5":

143 step_coeffs = _NS_ASYM5_COEFFS[:steps]

144 elif ns_variant == "custom":

145 if ns_coefficients is None:

146 raise ValueError("ns_coefficients must be provided when ns_variant='custom'")

147 if steps > len(ns_coefficients):

148 raise ValueError(

149 f"ns_steps={steps} exceeds the {len(ns_coefficients)} provided coefficient groups"

150 )

151 step_coeffs = tuple(ns_coefficients[:steps])

152 else:

153 raise ValueError(

154 f"ns_variant must be 'legacy', 'asym5', or 'custom', got {ns_variant!r}"

155 )

156

157 for coeff_a, coeff_b, coeff_c in step_coeffs:

158 torch.matmul(mat_x, mat_x.mT, out=buf_a)

159 torch.matmul(buf_a, buf_a, out=buf_b)

160 buf_a.mul_(coeff_b)

161 buf_b.mul_(coeff_c)

162 buf_a.add_(buf_b)

163 torch.matmul(buf_a, mat_x, out=buf_c)

164 mat_x.mul_(coeff_a).add_(buf_c)

165

166 del buf_a, buf_b, buf_c

167

168 if transposed:

169 mat_x = mat_x.mT

170

171 return mat_x

172

173

174def compute_muon_slice_scale(

175 slice_tensor: torch.Tensor,

176 matched_adamw_rms: float,

177 zero_rms_scale_mode: str = "zero",

178) -> float:

179 """Compute Muon scale from the logical matrix dims of a reshaped slice."""

180 if not matched_adamw_rms:

181 return 1.0 if zero_rms_scale_mode == "use_lr" else 0.0

182 shape = tuple(slice_tensor.shape)

183 if len(shape) == 3 and shape[1] == 1:

184 logical_dims = (shape[0], shape[2])

185 else:

186 logical_dims = shape[-2:]

187 return math.sqrt(max(logical_dims)) * matched_adamw_rms

188

189

190class Muon(BaseDistributedOptimizer):

191 """Muon optimizer with HSDP shard-group-aware Newton-Schulz orthogonalization.

192

193 Implements the Muon optimizer which uses Newton-Schulz iteration for matrix

194 orthogonalization of gradient updates, with HSDP-aware communication for

195 sharded parameters.

196 """

197

198 ADDITIONAL_CONFIG_KEYS = set(_MUON_ADVANCED_DEFAULTS)

199

200 def __init__(

201 self,

202 params: Any,

203 lr: float = 2e-2,

204 weight_decay: float = 0.1,

205 matched_adamw_rms: float = 0.2,

206 momentum: Union[float, List[float], Tuple[float, ...]] = 0.95,

207 nesterov: bool = True,

208 ns_steps: int = 5,

209 ns_variant: str = "asym5",

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

211 **advanced_options: Any,

212 ) -> None:

213 """Initialize Muon and build parameter-identity runtime caches.

214

215 Args:

216 params: Parameters or parameter groups optimized by Muon.

217 lr: Learning rate.

218 weight_decay: Decoupled weight-decay coefficient.

219 matched_adamw_rms: RMS scale matched to an AdamW update.

220 momentum: Momentum coefficient or coefficient pair.

221 nesterov: Whether to use Nesterov momentum.

222 ns_steps: Newton-Schulz iteration count.

223 ns_variant: Newton-Schulz coefficient variant.

224 ns_coefficients: Optional custom Newton-Schulz coefficients.

225 ns_epsilon: Numerical stability term used during normalization.

226 zeropower_fn: Optional Newton-Schulz implementation override.

227 momentum_update_fn: Optional momentum update implementation override.

228 reshape_fn: Optional logical-matrix reshape function.

229 ns_transform_fn: Optional Newton-Schulz input transformation.

230 post_update_fn: Optional callback applied after a parameter update.

231 zero_rms_scale_mode: Scaling behavior when matched AdamW RMS is zero.

232 apply_lr_in_update: Whether the update callback applies the learning rate.

233 hsdp_replica_count: Optional optimizer-state replica group size.

234 """

235 advanced = _resolve_muon_advanced_options(advanced_options)

236 if ns_variant not in ("legacy", "asym5", "custom"):

237 raise ValueError(

238 f"ns_variant must be 'legacy', 'asym5', or 'custom', got {ns_variant!r}"

239 )

240 if advanced.ns_epsilon < 0.0 or not math.isfinite(advanced.ns_epsilon):

241 raise ValueError(f"ns_epsilon must be a finite non-negative value, got {advanced.ns_epsilon}")

242 if advanced.zero_rms_scale_mode not in ("zero", "use_lr"):

243 raise ValueError(

244 "zero_rms_scale_mode must be 'zero' or 'use_lr', "

245 f"got {advanced.zero_rms_scale_mode!r}"

246 )

247 normalized_coefficients = self._validate_ns_coefficients(

248 ns_variant, ns_steps, advanced.ns_coefficients

249 )

250 if not isinstance(momentum, (list, tuple)):

251 momentum = [momentum]

252 if len(momentum) == 1:

253 momentum = (momentum[0], momentum[0])

254 defaults = {

255 "lr": lr,

256 "weight_decay": weight_decay,

257 "matched_adamw_rms": matched_adamw_rms,

258 "momentum": momentum,

259 "nesterov": nesterov,

260 "ns_steps": ns_steps,

261 "ns_variant": ns_variant,

262 "ns_coefficients": normalized_coefficients,

263 "ns_epsilon": advanced.ns_epsilon,

264 "zero_rms_scale_mode": advanced.zero_rms_scale_mode,

265 "apply_lr_in_update": advanced.apply_lr_in_update,

266 }

267 super().__init__(params, defaults, is_muon=True, hsdp_replica_count=hsdp_replica_count)

268 self.reshape_fn = advanced.reshape_fn

269 self.zeropower_fn = advanced.zeropower_fn

270 self.momentum_update_fn = advanced.momentum_update_fn

271 self.ns_transform_fn = advanced.ns_transform_fn

272 self.post_update_fn = advanced.post_update_fn

273 self.reset_optimizer_parameters()

274

275 def reset_optimizer_parameters(self) -> None:

276 """Rebuild all parameter-identity caches from current param groups."""

277 self._group_dtensor_by_mesh()

278 self._build_param_shard_metadata()

279 deduced_count = self._auto_deduce_replica_count()

280 if deduced_count is None:

281 self.hsdp_replica_count = None

282 elif self.hsdp_replica_count is None:

283 self.hsdp_replica_count = deduced_count

284 self._split_replicate_groups()

285 self._build_hsdp_batch()

286 self._build_param_broadcast_info()

287 self._classify_parameters_for_step()

288

289 @staticmethod

290 def _validate_ns_coefficients(

291 ns_variant: str,

292 ns_steps: int,

293 ns_coefficients: Optional[Sequence[Tuple[float, float, float]]],

294 ) -> Optional[Tuple[Tuple[float, float, float], ...]]:

295 """Validate and normalize externally supplied Newton-Schulz coefficients."""

296 if ns_variant != "custom":

297 if ns_coefficients is not None:

298 raise ValueError("ns_coefficients can only be provided when ns_variant='custom'")

299 return None

300 if not ns_coefficients:

301 raise ValueError("ns_coefficients must be provided when ns_variant='custom'")

302 if ns_steps > len(ns_coefficients):

303 raise ValueError(

304 f"ns_steps={ns_steps} exceeds the {len(ns_coefficients)} provided coefficient groups"

305 )

306

307 normalized = []

308 for index, coefficients in enumerate(ns_coefficients):

309 if len(coefficients) != 3:

310 raise ValueError(

311 f"ns_coefficients[{index}] must contain exactly three values, got {coefficients!r}"

312 )

313 normalized_coefficients = tuple(float(value) for value in coefficients)

314 if not all(math.isfinite(value) for value in normalized_coefficients):

315 raise ValueError(f"ns_coefficients[{index}] must contain only finite values")

316 normalized.append(normalized_coefficients)

317 return tuple(normalized)

318

319 def __str__(self) -> str:

320 """Return the optimizer representation."""

321 return super().__repr__()

322

323 __repr__ = __str__

324

325 @torch.no_grad()

326 def step(

327 self,

328 closure: Optional[Callable[[], Any]] = None,

329 ) -> Optional[float]:

330 """Perform a single optimization step.

331

332 De-duplication is controlled by the caller: ``param_to_ns_input`` should already contain only the owned

333 params (via ``hsdp_assign.owned_params``). The caller is responsible for broadcasting the updated params to

334 replica peers via ``AsyncReplicateBroadcaster.flush_group``.

335

336 Args:

337 closure: Optional callable that reevaluates the model and returns loss.

338

339 Returns:

340 The closure loss when provided, otherwise ``None``.

341 """

342 loss = None

343 if closure is not None:

344 with torch.enable_grad():

345 loss = closure()

346 broadcaster = AsyncReplicateBroadcaster(self)

347

348 num_groups = len(self.param_groups)

349

350 for group in self.param_groups:

351 group['step'] = (group.get('step') or 0) + 1

352

353 # Compute momentum only for no-comm params upfront.

354 no_comm_ns: Dict[int, Dict] = {}

355 for group_idx in range(num_groups):

356 info = self._hsdp_assignment_batches.get(group_idx)

357 if not info:

358 # No HSDP info — all params are no_comm.

359 unshard_params = self.unshard_params_by_group.get(group_idx, [])

360 no_comm_ns[group_idx] = self._update_muon_momentum(

361 self.param_groups[group_idx], unshard_params

362 )

363 else:

364 no_comm_params = info.get("no_comm", [])

365 if no_comm_params:

366 no_comm_ns[group_idx] = self._update_muon_momentum(

367 self.param_groups[group_idx], no_comm_params

368 )

369 else:

370 no_comm_ns[group_idx] = {}

371

372 # Process no-comm params.

373 for group_idx in range(num_groups):

374 if no_comm_ns[group_idx]:

375 group = self.param_groups[group_idx]

376 self._process_unshard_params(group, no_comm_ns[group_idx])

377 if self.post_update_fn is not None:

378 self._run_post_update_fn(group, no_comm_ns[group_idx].keys())

379

380 # Flatten nested batch into a linear schedule.

381 group_linear_batches: Dict[int, List[HSDPGroupAssignment]] = {}

382 max_num_batches = 0

383 for group_idx in range(num_groups):

384 info = self._hsdp_assignment_batches.get(group_idx)

385 linear_batches = []

386 if info:

387 for bg in info.get("batch_groups", []):

388 linear_batches.extend(bg.get("sub_batches", []))

389 group_linear_batches[group_idx] = linear_batches

390 max_num_batches = max(max_num_batches, len(linear_batches))

391

392 # Process batches: compute momentum per-batch for owned params only (HSDP de-duplication).

393 for batch_idx in range(max_num_batches):

394 for group_idx in range(num_groups):

395 linear_batches = group_linear_batches[group_idx]

396 if batch_idx >= len(linear_batches):

397 continue

398

399 hsdp_assign = linear_batches[batch_idx]

400 group = self.param_groups[group_idx]

401

402 # Compute momentum for this assignment's owned params only.

403 ns_inputs = self._update_muon_momentum(group, hsdp_assign.owned_params)

404 if not hsdp_assign.is_shard:

405 if ns_inputs:

406 self._process_unshard_params(group, ns_inputs)

407 else:

408 self._process_shard_params(

409 group, ns_inputs, [hsdp_assign], group_idx,

410 buffer_cache={},

411 )

412 if self.post_update_fn is not None:

413 self._run_post_update_fn(group, ns_inputs.keys())

414

415 # Flush broadcasts after each assignment.

416 broadcaster.flush_group(hsdp_assign)

417

418 broadcaster.wait_all()

419 return loss

420

421 def _run_post_update_fn(

422 self,

423 group: Dict[str, Any],

424 params: Iterable[torch.nn.Parameter],

425 ) -> None:

426 """Run the optional model callback on updated owner parameters before broadcast."""

427 if self.post_update_fn is None:

428 return

429 context_values = {

430 "lr": group["lr"],

431 "weight_decay": group["weight_decay"],

432 "step": group["step"],

433 }

434 for param in params:

435 local_tensor = to_local_if_dtensor(param.data)

436 context = MuonPostUpdateContext(

437 param_fqn=getattr(param, "model_name", None),

438 logical_shape=tuple(param.shape),

439 **context_values,

440 )

441 self.post_update_fn(param, local_tensor, context)

442

443 def _classify_parameters_for_step(self) -> None:

444 """Classify params by whether the last two dims are sharded.

445

446 unshard: run Newton-Schulz locally.

447 shard: all-gather before Newton-Schulz.

448

449 Reads from self._hsdp_assignment_batches which is organized by batch.

450 """

451 self.unshard_params_by_group: Dict[int, List] = {}

452 self.shard_params_by_group: Dict[int, List] = {}

453 self.shard_assignments_by_group: Dict[int, List[HSDPGroupAssignment]] = {}

454 # Per group: record.index -> shard coord that computes NS.

455 self._shard_compute_coord: Dict[int, Dict[int, Tuple[int, ...]]] = {}

456

457 for group_idx, group in enumerate(self.param_groups):

458 assignment_info = self._hsdp_assignment_batches.get(group_idx)

459

460 unshard_params = []

461 shard_hsdp_assignments: List[HSDPGroupAssignment] = []

462

463 if assignment_info:

464 unshard_params.extend(assignment_info["no_comm"])

465 for bg in assignment_info["batch_groups"]:

466 for hsdp_assign in bg["sub_batches"]:

467 if hsdp_assign.is_shard:

468 shard_hsdp_assignments.append(hsdp_assign)

469 else:

470 unshard_params.extend(hsdp_assign.owned_params)

471 else:

472 unshard_params.extend(group["params"])

473

474 # Shard params only include locally owned params; replica ownership is enforced at apply time.

475 shard_params = [p for a in shard_hsdp_assignments for p in a.owned_params]

476

477 # Greedily assign NS compute across shard ranks.

478 self._shard_compute_coord[group_idx] = {}

479 for hsdp_assign in shard_hsdp_assignments:

480 shard_sizes, _, _, _ = self._get_shard_info(hsdp_assign)

481 compute_by_index = build_owner_by_size(

482 records=hsdp_assign.owned_records,

483 replicate_sizes=shard_sizes,

484 )

485 self._shard_compute_coord[group_idx].update(compute_by_index)

486

487 self.unshard_params_by_group[group_idx] = unshard_params

488 self.shard_assignments_by_group[group_idx] = shard_hsdp_assignments

489 self.shard_params_by_group[group_idx] = shard_params

490

491 def _update_muon_momentum(

492 self,

493 group: Dict[str, Any],

494 params: List[torch.Tensor],

495 ) -> Dict[torch.nn.Parameter, torch.Tensor]:

496 """Compute first-order momentum and return bfloat16 NS inputs."""

497 momentum = group["momentum"]

498 if not isinstance(momentum, (list, tuple)):

499 momentum = [momentum]

500 if len(momentum) == 1:

501 momentum = (momentum[0], momentum[0])

502 momentum1, momentum2 = momentum

503 # Pre-filter params with valid grads and ensure momentum buffers exist

504 valid_params = []

505 grads = []

506 bufs = []

507 for p in params:

508 g = p.grad

509 if g is None:

510 continue

511 if "momentum_buffer" not in self.state[p]:

512 self.state[p]["momentum_buffer"] = torch.zeros_like(g)

513 valid_params.append(p)

514 grads.append(g)

515 bufs.append(self.state[p]["momentum_buffer"])

516

517 if not valid_params:

518 return {}

519

520 grads = [to_local_if_dtensor(grad) for grad in grads]

521 bufs = [to_local_if_dtensor(buffer) for buffer in bufs]

522 if self.momentum_update_fn is not None:

523 custom_updates = list(self.momentum_update_fn(grads, bufs, momentum1, momentum2, group["nesterov"]))

524 if len(custom_updates) != len(valid_params):

525 raise ValueError(

526 "momentum_update_fn must return one update tensor per input gradient, "

527 f"got {len(custom_updates)} updates for {len(valid_params)} gradients"

528 )

529 custom_updates = [update.to(torch.bfloat16) for update in custom_updates]

530 param_updates = dict(zip(valid_params, custom_updates))

531 return param_updates

532

533 # Match muon_update_core():

534 # m_for_update = grad + momentum1 * m_old

535 # m_new = grad + momentum2 * m_old

536 # pylint: disable=protected-access

537 local_us = torch._foreach_add(grads, bufs, alpha=momentum1)

538

539 torch._foreach_mul_(bufs, momentum2)

540 torch._foreach_add_(bufs, grads)

541

542 if group["nesterov"]:

543 torch._foreach_mul_(local_us, momentum1)

544 torch._foreach_add_(local_us, grads)

545

546 if local_us[0].dtype != torch.bfloat16:

547 local_us = [update.to(torch.bfloat16) for update in local_us]

548

549 param_updates = dict(zip(valid_params, local_us))

550 return param_updates

551

552 def _process_unshard_params(

553 self,

554 group: Dict[str, Any],

555 param_to_ns_input: Dict[torch.nn.Parameter, torch.Tensor],

556 ) -> None:

557 """Process un-sharded params: shape-group -> memory-batch -> NS -> local update."""

558 lr = group["lr"]

559 weight_decay = group["weight_decay"]

560

561 # empty params are not processed in muon, for uneven shard

562 param_to_ns_input = {

563 p: ns_input for p, ns_input in param_to_ns_input.items() if ns_input.numel() > 0

564 }

565 if not param_to_ns_input:

566 return

567

568 shape_groups = self._group_by_shape(list(param_to_ns_input.keys()))

569 for _, p_list in shape_groups.items():

570 safe_batches = self._split_into_memory_safe_batches(p_list, shard_size=1)

571

572 for sub_batch in safe_batches:

573 updates_dict = self._compute_batched_ns_updates(

574 sub_batch, param_to_ns_input, group, no_shard=True

575 )

576

577 local_params = [to_local_if_dtensor(p.data) for p in sub_batch]

578 local_updates = [updates_dict[p].view(lp.shape) for p, lp in zip(sub_batch, local_params)]

579

580 if weight_decay != 0.0:

581 # pylint: disable=protected-access

582 torch._foreach_mul_(local_params, 1 - lr * weight_decay)

583 # pylint: disable=protected-access

584 apply_alpha = 1.0 if group["apply_lr_in_update"] else -lr

585 torch._foreach_add_(local_params, local_updates, alpha=apply_alpha)

586

587 def _gather_and_compute_shard_updates(

588 self,

589 valid_params: List[torch.nn.Parameter],

590 param_to_ns_input: Dict[torch.nn.Parameter, torch.Tensor],

591 hsdp_assign: HSDPGroupAssignment,

592 shard_compute_coord: Dict[int, Tuple[int, ...]],

593 group: Dict[str, Any],

594 buffer_cache: Optional[Dict],

595 ) -> Tuple[

596 Dict[torch.nn.Parameter, torch.Tensor],

597 Dict[torch.nn.Parameter, Tuple[int, ...]],

598 ]:

599 """Gather NS inputs and compute updates for locally-assigned shard params.

600

601 Returns:

602 (my_updates, param_compute_coord) for this HSDP assignment.

603 """

604 shard_sizes, local_coords, shard_pgs, total_shard_size = self._get_shard_info(hsdp_assign)

605 param_to_index = {record.param: record.index for record in hsdp_assign.owned_records}

606

607 my_params: List[torch.nn.Parameter] = []

608 my_param_ids: set = set()

609 param_compute_coord: Dict[torch.nn.Parameter, Tuple[int, ...]] = {}

610 my_indices: set = set()

611

612 for idx, p in enumerate(valid_params):

613 p_index = param_to_index[p]

614 compute_coord = shard_compute_coord.get(p_index, (0,) * len(shard_sizes))

615 param_compute_coord[p] = compute_coord

616 if compute_coord == local_coords:

617 my_params.append(p)

618 my_param_ids.add(id(p))

619 my_indices.add(idx)

620

621 # Fused all-gather full NS inputs; keep only tensors computed locally.

622 gathered_inputs: Dict[torch.nn.Parameter, torch.Tensor] = {}

623 local_inputs = build_pad_ns_inputs(

624 valid_params,

625 param_to_ns_input,

626 self._param_shard_metadata,

627 )

628 param_shard_metadata = [self._param_shard_metadata.get(p) for p in valid_params]

629 gathered_list = fused_allgather_dtensor_params(

630 local_inputs, shard_pgs, hsdp_assign.layout_spec,

631 param_shard_metadata=param_shard_metadata,

632 buffer_cache=buffer_cache,

633 keep_indices=my_indices,

634 )

635 for p, full_inp in zip(valid_params, gathered_list):

636 if id(p) in my_param_ids:

637 gathered_inputs[p] = full_inp

638 local_inputs.clear()

639 gathered_list.clear()

640

641 # Compute NS updates in shape groups and memory-safe batches.

642 my_updates: Dict[torch.nn.Parameter, torch.Tensor] = {}

643 if my_params:

644 shape_groups = self._group_by_shape(my_params)

645 for _, p_list in shape_groups.items():

646 safe_batches = self._split_into_memory_safe_batches(p_list, shard_size=total_shard_size)

647 for sub_batch in safe_batches:

648 updates_dict = self._compute_batched_ns_updates(

649 sub_batch, gathered_inputs, group

650 )

651 for p in sub_batch:

652 my_updates[p] = updates_dict[p].contiguous()

653 del updates_dict

654 gathered_inputs.clear()

655

656 return my_updates, param_compute_coord

657

658 def _process_shard_params(

659 self,

660 group: Dict[str, Any],

661 param_to_ns_input: Dict[torch.nn.Parameter, torch.Tensor],

662 hsdp_assignments: List[HSDPGroupAssignment],

663 group_idx: int,

664 buffer_cache: Optional[Dict] = None,

665 ) -> None:

666 """Process sharded params with greedy shard-group compute assignment."""

667 shard_compute_coord = self._shard_compute_coord.get(group_idx, {})

668

669 for hsdp_assign in hsdp_assignments:

670 owned_params = hsdp_assign.owned_params

671 # Communication membership is static per assignment. Ranks with empty

672 # local shards still participate via build_pad_ns_inputs().

673 valid_params = owned_params

674

675 shard_sizes, _, shard_pgs, total_shard_size = self._get_shard_info(hsdp_assign)

676

677 # uneven shard check

678 shard_pgs_rank = []

679 for pg in shard_pgs:

680 pg_rank = torch.distributed.get_process_group_ranks(pg)

681 shard_pgs_rank.append(pg_rank)

682

683 for p in valid_params:

684 if tuple(p.shape)[0] % shard_sizes[0] != 0:

685 logger.debug_rank0(

686 "[hyper-optimizer uneven-shard] p %s placement %s device_mesh %s "

687 "local_shape %s global_shape %s pg_rank %s shard_sizes %s",

688 p.model_name,

689 p.placements,

690 p.device_mesh,

691 p.to_local().shape,

692 p.shape,

693 shard_pgs_rank,

694 shard_sizes,

695 )

696

697 if valid_params:

698 safe_batches = self._split_into_memory_safe_batches(

699 valid_params, shard_size=total_shard_size, min_batch_size=total_shard_size

700 )

701 for sub_batch in safe_batches:

702 my_updates, param_compute_coord = self._gather_and_compute_shard_updates(

703 sub_batch, param_to_ns_input, hsdp_assign,

704 shard_compute_coord, group, buffer_cache,

705 )

706

707 self._fused_broadcast_and_apply(

708 sub_batch, my_updates, param_compute_coord,

709 group, hsdp_assign,

710 )

711

712 def _group_by_shape(

713 self,

714 params: List[torch.nn.Parameter],

715 ) -> Dict[tuple, List[torch.nn.Parameter]]:

716 """Group parameters by their last-2-dim shape (A, B) for batched NS.

717

718 [1024, 1024], [1024, 1, 1024], and [3, 1024, 1024] all map to

719 key (1024, 1024) for maximum batch merging.

720 """

721 groups = defaultdict(list)

722

723 for p in params:

724 core_shape = self._shape_to_core_shape(tuple(p.shape))

725 groups[core_shape].append(p)

726

727 return groups

728

729 @staticmethod

730 def _shape_to_core_shape(shape: Tuple[int, ...]) -> Tuple[int, int]:

731 if len(shape) == 2:

732 return (shape[0], shape[1])

733 if len(shape) == 3 and shape[1] == 1:

734 return (shape[0], shape[2])

735 if len(shape) >= 3:

736 return (shape[-2], shape[-1])

737 raise ValueError('1D parameters are not supported in Muon')

738

739 def _reshape_ns_input(

740 self,

741 param: torch.nn.Parameter,

742 ns_input: torch.Tensor,

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

744 """Return the contiguous NS input and any reshape views used for NS."""

745 working_input = ns_input if ns_input.is_contiguous() else ns_input.contiguous()

746 if self.reshape_fn is None:

747 return working_input, [working_input]

748

749 param_fqn = param.model_name

750 if param_fqn is None:

751 return working_input, [working_input]

752

753 reshaped_inputs = list(self.reshape_fn(param_fqn, working_input))

754 if not reshaped_inputs:

755 return working_input, [working_input]

756

757 if reshaped_inputs[0].shape != working_input.shape:

758 logger.info_rank0(

759 "Reshape %s from %s to %s",

760 param_fqn,

761 working_input.shape,

762 reshaped_inputs[0].shape,

763 )

764

765 for reshaped_input in reshaped_inputs:

766 if reshaped_input.untyped_storage().data_ptr() != working_input.untyped_storage().data_ptr():

767 raise ValueError("reshape_fn must return views that share storage with the working NS input tensor")

768

769 return working_input, reshaped_inputs

770

771 def _prepare_ns_transform(

772 self,

773 param: torch.nn.Parameter,

774 ns_input: torch.Tensor,

775 ) -> Tuple[torch.Tensor, NSInputTransform]:

776 """Route one parameter through custom, zero-copy, or identity NS preparation."""

777 working_input = ns_input if ns_input.is_contiguous() else ns_input.contiguous()

778 param_fqn = getattr(param, "model_name", None)

779 if self.ns_transform_fn is not None and param_fqn is not None:

780 transform = self.ns_transform_fn(param_fqn, working_input)

781 if transform is not None:

782 self._validate_ns_transform(transform, working_input)

783 return working_input, transform

784

785 _, reshaped_inputs = self._reshape_ns_input(param, working_input)

786

787 def restore_view_updates(updates: List[torch.Tensor], output: torch.Tensor) -> None:

788 """Write reshaped NS updates back into the views and the working input."""

789 for reshaped_input, update in zip(reshaped_inputs, updates):

790 reshaped_input.copy_(update.contiguous().view_as(reshaped_input))

791 if output.untyped_storage().data_ptr() != working_input.untyped_storage().data_ptr():

792 output.copy_(working_input)

793

794 return working_input, NSInputTransform(tensors=reshaped_inputs, restore=restore_view_updates)

795

796 @staticmethod

797 def _validate_ns_transform(transform: NSInputTransform, working_input: torch.Tensor) -> None:

798 """Validate a model-provided reversible NS transform."""

799 if not isinstance(transform, NSInputTransform):

800 raise ValueError(

801 "ns_transform_fn must return NSInputTransform or None, "

802 f"got {type(transform).__name__}"

803 )

804 if not transform.tensors:

805 raise ValueError("NSInputTransform.tensors must not be empty")

806 for index, tensor in enumerate(transform.tensors):

807 if not isinstance(tensor, torch.Tensor):

808 raise ValueError(f"NSInputTransform.tensors[{index}] must be a torch.Tensor")

809 if tensor.device != working_input.device:

810 raise ValueError("NSInputTransform tensors must be on the same device as the NS input")

811 if tensor.dtype != working_input.dtype:

812 raise ValueError("NSInputTransform tensors must have the same dtype as the NS input")

813 if tensor.dim() < 2:

814 raise ValueError("NSInputTransform tensors must have at least two dimensions")

815 if not callable(transform.restore):

816 raise ValueError("NSInputTransform.restore must be callable")

817

818 def _compute_batched_ns_outputs_for_tensors(

819 self,

820 tensor_list: List[torch.Tensor],

821 ns_steps: int,

822 ns_variant: str = "asym5",

823 ns_coefficients: Optional[Sequence[Tuple[float, float, float]]] = None,

824 ns_epsilon: float = 1e-10,

825 ) -> List[torch.Tensor]:

826 """Run batched NS on mixed-shape tensors and restore their original shapes."""

827

828 if not tensor_list:

829 return []

830

831 state = {"inputs_3d": [], "slice_sizes": [], "shapes_info": [], "restored_updates": [], "current_idx": 0}

832

833 for tensor in tensor_list:

834 origin_shape = tuple(tensor.shape)

835 is_conv = False

836

837 if len(origin_shape) == 2:

838 inp_3d = tensor.unsqueeze(0)

839 n_dim = 1

840 elif len(origin_shape) == 3 and origin_shape[1] == 1:

841 inp_3d = tensor.squeeze(1).unsqueeze(0)

842 is_conv = True

843 n_dim = 1

844 elif len(origin_shape) == 3:

845 inp_3d = tensor

846 n_dim = origin_shape[0]

847 else:

848 inp_3d = tensor.reshape(-1, origin_shape[-2], origin_shape[-1])

849 n_dim = inp_3d.shape[0]

850

851 state["inputs_3d"].append(inp_3d)

852 state["slice_sizes"].append(n_dim)

853 state["shapes_info"].append((origin_shape, is_conv))

854

855 merged_input = torch.cat(state["inputs_3d"], dim=0)

856

857 if merged_input.shape[0] == 1:

858 merged_input = merged_input.squeeze(0)

859

860 if self.zeropower_fn is None:

861 merged_update = zeropower_via_newtonschulz5(

862 merged_input,

863 steps=ns_steps,

864 ns_variant=ns_variant,

865 epsilon=ns_epsilon,

866 ns_coefficients=ns_coefficients,

867 )

868 else:

869 merged_update = self.zeropower_fn(merged_input, steps=ns_steps)

870

871 if merged_input.dim() == 2:

872 merged_update = merged_update.unsqueeze(0)

873 del merged_input

874

875 for n_dim, (origin_shape, is_conv) in zip(state["slice_sizes"], state["shapes_info"]):

876 update = merged_update[state["current_idx"]:state["current_idx"] + n_dim]

877 state["current_idx"] += n_dim

878

879 if is_conv:

880 update = update.squeeze(0).unsqueeze(1)

881 elif len(origin_shape) == 2:

882 update = update.squeeze(0)

883 elif len(origin_shape) >= 4:

884 update = update.reshape(origin_shape)

885

886 state["restored_updates"].append(update)

887

888 del merged_update

889 return state["restored_updates"]

890

891 def _split_into_memory_safe_batches(

892 self,

893 p_list: List[torch.nn.Parameter],

894 shard_size: int = 1,

895 min_batch_size: int = 1,

896 ) -> List[List[torch.nn.Parameter]]:

897 """Split parameters into memory-safe batches to prevent OOM during NS.

898

899 The per-batch element limit is scaled down by shard_size to account

900 for the memory amplification from allgather. Sharded paths prioritize

901 filling the compute-owner ranks before applying this soft limit.

902 """

903 max_numel_per_batch = 512 * 1024 * 1024 // shard_size

904 min_batch_size = max(1, min_batch_size)

905

906 batches = []

907 current_batch = []

908 current_count = 0

909

910 for p in p_list:

911 p_count = p.numel()

912 if (

913 len(current_batch) >= min_batch_size

914 and current_count + p_count > max_numel_per_batch

915 ):

916 batches.append(current_batch)

917 current_batch = [p]

918 current_count = p_count

919 else:

920 current_batch.append(p)

921 current_count += p_count

922

923 if current_batch:

924 batches.append(current_batch)

925

926 return batches

927

928 def _compute_batched_ns_updates(

929 self,

930 p_list: List[torch.nn.Parameter],

931 ns_inputs: Dict[torch.nn.Parameter, torch.Tensor],

932 group: Dict[str, Any],

933 no_shard: bool = False

934 ) -> Dict[torch.nn.Parameter, torch.Tensor]:

935 """Route batched NS through the native or reversible-transform path."""

936 if self.ns_transform_fn is None:

937 return self._compute_batched_ns_updates_fast(p_list, ns_inputs, group, no_shard)

938 return self._compute_batched_ns_updates_with_transform(p_list, ns_inputs, group, no_shard)

939

940 def _compute_batched_ns_updates_fast(

941 self,

942 p_list: List[torch.nn.Parameter],

943 ns_inputs: Dict[torch.nn.Parameter, torch.Tensor],

944 group: Dict[str, Any],

945 no_shard: bool = False,

946 ) -> Dict[torch.nn.Parameter, torch.Tensor]:

947 """Compute native reshape/view NS updates without transform bookkeeping."""

948 if not p_list:

949 return {}

950 state = {"reshape_groups": defaultdict(list), "origin_shapes": {}, "working_inputs": {}, "restored": {}}

951 for param in p_list:

952 local_shape = getattr(param, "local_shape", None)

953 if local_shape is None:

954 local_shape = to_local_if_dtensor(param.data).shape

955 state["origin_shapes"][param] = tuple(local_shape) if no_shard else tuple(param.shape)

956 working_input, reshaped_inputs = self._reshape_ns_input(

957 param, ns_inputs[param].view(state["origin_shapes"][param])

958 )

959 state["working_inputs"][param] = working_input

960 for reshaped_input in reshaped_inputs:

961 state["reshape_groups"][self._shape_to_core_shape(tuple(reshaped_input.shape))].append(

962 reshaped_input

963 )

964 for tensor_list in state["reshape_groups"].values():

965 updates = self._compute_batched_ns_outputs_for_tensors(

966 tensor_list, group["ns_steps"], group["ns_variant"], group["ns_coefficients"], group["ns_epsilon"]

967 )

968 for reshaped_input, update in zip(tensor_list, updates):

969 scale = compute_muon_slice_scale(

970 update, group["matched_adamw_rms"], zero_rms_scale_mode=group["zero_rms_scale_mode"]

971 )

972 if group["apply_lr_in_update"]:

973 scale *= -group["lr"]

974 update.mul_(scale)

975 reshaped_input.copy_(update.contiguous().view_as(reshaped_input))

976

977 for param in p_list:

978 update = ns_inputs[param].view(state["origin_shapes"][param])

979 working_input = state["working_inputs"][param]

980 if working_input.untyped_storage().data_ptr() != update.untyped_storage().data_ptr():

981 update.copy_(working_input)

982 state["restored"][param] = update

983 return state["restored"]

984

985 def _compute_batched_ns_updates_with_transform(

986 self,

987 p_list: List[torch.nn.Parameter],

988 ns_inputs: Dict[torch.nn.Parameter, torch.Tensor],

989 group: Dict[str, Any],

990 no_shard: bool = False,

991 ) -> Dict[torch.nn.Parameter, torch.Tensor]:

992 """Batched Newton-Schulz update for mixed 2D / Conv3D / 3D parameters.

993

994 Normalizes all inputs to 3D, concatenates along dim 0, runs a single

995 NS iteration, then slices results back to original shapes.

996

997 """

998

999 if not p_list:

1000 return {}

1001

1002 state: _TransformNSState = {

1003 "reshape_groups": defaultdict(list),

1004 "origin_shapes": {},

1005 "working_inputs": {},

1006 "transforms": {},

1007 "transform_updates": [],

1008 "restored": {},

1009 }

1010

1011 for p in p_list:

1012 state["origin_shapes"][p] = (

1013 tuple(to_local_if_dtensor(p.data).shape) if no_shard else tuple(p.shape)

1014 )

1015 prepared = self._prepare_ns_transform(p, ns_inputs[p].view(state["origin_shapes"][p]))

1016 state["working_inputs"][p] = prepared[0]

1017 state["transforms"][p] = prepared[1]

1018 state["transform_updates"].append([None] * len(prepared[1].tensors))

1019

1020 for tensor_index, transformed_input in enumerate(prepared[1].tensors):

1021 state["reshape_groups"][self._shape_to_core_shape(tuple(transformed_input.shape))].append(

1022 (transformed_input, len(state["transform_updates"]) - 1, tensor_index)

1023 )

1024

1025 for tensor_records in state["reshape_groups"].values():

1026 reshaped_updates = self._compute_batched_ns_outputs_for_tensors(

1027 [record[0] for record in tensor_records],

1028 group["ns_steps"],

1029 ns_variant=group["ns_variant"],

1030 ns_coefficients=group["ns_coefficients"],

1031 ns_epsilon=group["ns_epsilon"],

1032 )

1033

1034 # Scale updates.

1035 for record, reshaped_update in zip(tensor_records, reshaped_updates):

1036 slice_scale = compute_muon_slice_scale(

1037 reshaped_update,

1038 group["matched_adamw_rms"],

1039 zero_rms_scale_mode=group["zero_rms_scale_mode"],

1040 )

1041

1042 if group["apply_lr_in_update"]:

1043 slice_scale *= -group["lr"]

1044 reshaped_update.mul_(slice_scale)

1045 state["transform_updates"][record[1]][record[2]] = reshaped_update

1046

1047 for p, reshaped_updates in zip(p_list, state["transform_updates"]):

1048

1049 if any(update is None for update in reshaped_updates):

1050 raise RuntimeError("Missing Newton-Schulz update for a transformed input")

1051 state["transforms"][p].restore(reshaped_updates, state["working_inputs"][p])

1052

1053 for p in p_list:

1054 reshaped_update = ns_inputs[p].view(state["origin_shapes"][p])

1055

1056 if (state["working_inputs"][p].untyped_storage().data_ptr()

1057 != reshaped_update.untyped_storage().data_ptr()):

1058 reshaped_update.copy_(state["working_inputs"][p])

1059 state["restored"][p] = reshaped_update

1060

1061 return state["restored"]

1062

1063 def _fused_broadcast_and_apply(

1064 self,

1065 valid_params: List[torch.nn.Parameter],

1066 my_updates: Dict[torch.nn.Parameter, torch.Tensor],

1067 param_compute_coord: Dict[torch.nn.Parameter, Tuple[int, ...]],

1068 group: Dict[str, Any],

1069 hsdp_assign: HSDPGroupAssignment,

1070 ) -> None:

1071 """Fused broadcast and apply for shard-group updates."""

1072 lr = group["lr"]

1073 weight_decay = group["weight_decay"]

1074 shard_sizes, local_coords, shard_pgs, total_shard_size = self._get_shard_info(hsdp_assign)

1075 device = to_local_if_dtensor(valid_params[0].data).device

1076

1077 coord_groups: Dict[Tuple[int, ...], List[torch.nn.Parameter]] = defaultdict(list)

1078 for p in valid_params:

1079 coord_groups[param_compute_coord[p]].append(p)

1080

1081 all_local_params: List[torch.Tensor] = []

1082 all_update_shards: List[torch.Tensor] = []

1083

1084 alignment_bytes = 512

1085 element_size = torch.empty(0, dtype=torch.bfloat16, device=device).element_size()

1086 alignment_elements = max(1, alignment_bytes // element_size)

1087

1088 pack_buffers: Dict[Tuple[int, ...], torch.Tensor] = {}

1089 coord_param_offsets: Dict[Tuple[int, ...], List[Tuple[int, int, int]]] = {}

1090

1091 for coord, coord_params in coord_groups.items():

1092 is_compute_rank = coord == local_coords

1093 param_offsets: List[Tuple[int, int, int]] = []

1094 total_padded_numel = 0

1095

1096 for p in coord_params:

1097 actual_numel = p.numel()

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

1099 param_offsets.append((total_padded_numel, actual_numel, padded_numel))

1100 total_padded_numel += padded_numel

1101

1102 coord_param_offsets[coord] = param_offsets

1103 pack_buffer = torch.empty(total_padded_numel, dtype=torch.bfloat16, device=device)

1104

1105 if is_compute_rank:

1106 for p, (offset, actual_numel, padded_numel) in zip(coord_params, param_offsets):

1107 update = my_updates[p]

1108 pack_buffer[offset:offset + actual_numel].copy_(update.reshape(-1))

1109 if padded_numel > actual_numel:

1110 pack_buffer[offset + actual_numel:offset + padded_numel].zero_()

1111

1112 pack_buffers[coord] = pack_buffer

1113

1114 if total_shard_size > 1:

1115 self._batched_relay_broadcast(

1116 pack_buffers, shard_pgs, shard_sizes, local_coords

1117 )

1118

1119 layout_spec = hsdp_assign.layout_spec

1120 for coord, coord_params in coord_groups.items():

1121 pack_buffer = pack_buffers[coord]

1122 param_offsets = coord_param_offsets[coord]

1123

1124 for p, (offset, actual_numel, _) in zip(coord_params, param_offsets):

1125 full_update = pack_buffer[offset:offset + actual_numel].view(tuple(p.shape))

1126 update_to_apply = chunk_update_by_layout(

1127 full_update,

1128 p,

1129 layout_spec,

1130 self._param_shard_metadata.get(p),

1131 )

1132

1133 local_param = to_local_if_dtensor(p.data)

1134 all_local_params.append(local_param)

1135 all_update_shards.append(update_to_apply.view(local_param.shape))

1136

1137 if not all_local_params:

1138 return

1139

1140 if weight_decay != 0.0:

1141 coeff = 1.0 - lr * weight_decay

1142 # pylint: disable=protected-access

1143 torch._foreach_mul_(all_local_params, coeff)

1144

1145 # Slice-wise Muon scaling has already been applied during NS postprocess.

1146 # pylint: disable=protected-access

1147 apply_alpha = 1.0 if group["apply_lr_in_update"] else -lr

1148 torch._foreach_add_(all_local_params, all_update_shards, alpha=apply_alpha)

1149

1150 def _build_param_shard_metadata(self) -> None:

1151 """Build shard metadata once during optimizer init."""

1152 self._param_shard_metadata: Dict[torch.nn.Parameter, ParamShardMeta] = {}

1153

1154 for _, hsdp_groups in self._hsdp_grouping.values():

1155 for hsdp_group in hsdp_groups:

1156 if hsdp_group.layout_spec is None or not hsdp_group.layout_spec.shard_axes:

1157 continue

1158 group_param_to_meta = build_param_shard_metadata_for_group(hsdp_group)

1159 for param, shard_meta in group_param_to_meta.items():

1160 self._param_shard_metadata[param] = shard_meta

1161 _debug_param_shard_metadata(hsdp_group, group_param_to_meta)

1162

1163 @staticmethod

1164 def _get_shard_info(

1165 hsdp_assign: HSDPGroupAssignment,

1166 ) -> Tuple[Tuple[int, ...], Tuple[int, ...], Tuple[dist.ProcessGroup, ...], int]:

1167 """Extract shard topology from HSDPGroupAssignment.

1168

1169 Returns:

1170 shard_sizes: Size of each shard mesh dimension.

1171 local_coords: Current rank's coordinate in each shard mesh dim.

1172 shard_pgs: ProcessGroup for each shard mesh dim.

1173 total_shard_size: Product of all shard_sizes.

1174 """

1175 shard_pgs = hsdp_assign.shard_pgs

1176 shard_sizes = tuple(

1177 dist.get_world_size(pg) if pg is not None else 1

1178 for pg in shard_pgs

1179 )

1180 local_coords = tuple(

1181 dist.get_rank(pg) if pg is not None else 0

1182 for pg in shard_pgs

1183 )

1184 total_shard_size = 1

1185 for s in shard_sizes:

1186 total_shard_size *= s

1187

1188 return shard_sizes, local_coords, shard_pgs, total_shard_size

1189

1190 @staticmethod

1191 def _batched_relay_broadcast(

1192 tensor_dict: Dict[Tuple[int, ...], torch.Tensor],

1193 shard_pgs: Tuple[dist.ProcessGroup, ...],

1194 shard_sizes: Tuple[int, ...],

1195 local_coords: Tuple[int, ...],

1196 ) -> None:

1197 """

1198 Batched asynchronous multi-dimensional relay broadcast.

1199 By operating asynchronously within each dimension, we eliminate CPU overhead bubbles

1200 while strictly preserving the multidimensional relay dependency.

1201 """

1202 for dim_idx, pg in enumerate(shard_pgs):

1203 if pg is None or shard_sizes[dim_idx] <= 1:

1204 continue

1205

1206 work_handles = []

1207

1208 for coord, tensor in tensor_dict.items():

1209 aligned = all(

1210 local_coords[sub_dim] == coord[sub_dim]

1211 for sub_dim in range(dim_idx + 1, len(shard_pgs))

1212 )

1213 if not aligned:

1214 continue

1215

1216 src_rank_in_pg = coord[dim_idx]

1217 global_src_rank = dist.get_global_rank(pg, src_rank_in_pg)

1218

1219 work = dist.broadcast(tensor, src=global_src_rank, group=pg, async_op=True)

1220 if work is not None:

1221 work_handles.append(work)

1222

1223 for work in work_handles:

1224 work.wait()

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

615 statements