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 import defaultdict

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

22import torch

23import torch.distributed as dist

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

26from hyper_parallel.core.optimizer.dtensor_compat import to_local_if_dtensor

27from hyper_parallel.core.optimizer.sharding_category import (

28 HSDPGroupAssignment,

29 fused_allgather_dtensor_params,

30 build_owner_by_size,

31 chunk_update_by_layout,

32)

33from hyper_parallel.platform import get_platform

36def zeropower_via_newtonschulz5(ns_inputs: torch.Tensor, steps: int) -> torch.Tensor:

37 """Matrix orthogonalization with preallocated matmul buffers."""

38 mat_x = ns_inputs

39 if ns_inputs.size(-2) > ns_inputs.size(-1):

40 mat_x = mat_x.mT

42 # Create a new tensor so later in-place updates are safe.

43 mat_x = mat_x / (torch.norm(mat_x, dim=[-2, -1], keepdim=True) + 1e-7)

45 coeffs = [

46 (4.0848, -6.8946, 2.9270),

47 (3.9505, -6.3029, 2.6377),

48 (3.7418, -5.5913, 2.3037),

49 (2.8769, -3.1427, 1.2046),

50 (2.8366, -3.0525, 1.2012),

51 ]

53 # Preallocate temporary buffers to avoid loop-time allocations.

54 n_size = mat_x.size(-2)

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

56 buf_a2 = torch.empty_like(buf_a)

57 buf_b = torch.empty_like(buf_a)

58 buf_bx = torch.empty_like(mat_x)

60 for coeff_a, coeff_b, coeff_c in coeffs[:steps]:

61 # A = X @ X.T

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

64 # A2 = A @ A

65 torch.matmul(buf_a, buf_a, out=buf_a2)

67 # B = b * A + c * A2

68 buf_b.copy_(buf_a).mul_(coeff_b)

69 buf_a2.mul_(coeff_c)

70 buf_b.add_(buf_a2)

72 # BX = B @ X

73 torch.matmul(buf_b, mat_x, out=buf_bx)

75 # X = a * X + BX

76 mat_x.mul_(coeff_a).add_(buf_bx)

78 if ns_inputs.size(-2) > ns_inputs.size(-1):

79 mat_x = mat_x.mT

81 return mat_x

84def adjust_lr_wd_for_muon(lr: float, matched_adamw_rms: float, param_shape: torch.Size) -> float:

85 """Scale learning rate for 2D Muon parameters based on tensor dimensions."""

86 dim_a, dim_b = param_shape[-2:]

87 adjusted_ratio = math.sqrt(max(dim_a, dim_b)) * matched_adamw_rms

88 return lr * adjusted_ratio

91def adjust_lr_wd_for_muon_conv(lr: float, matched_adamw_rms: float, param_shape: torch.Size) -> float:

92 """Scale learning rate for 3D convolutional Muon parameters."""

93 dim_a, dim_b, dim_c = param_shape[:]

94 adjusted_ratio = math.sqrt(max(dim_a, dim_b, dim_c)) * matched_adamw_rms

95 return lr * adjusted_ratio

98class Muon(BaseDistributedOptimizer):

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

100

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

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

103 sharded parameters.

104 """

105

106 def __init__(

107 self,

108 params,

109 lr: float = 2e-2,

110 weight_decay: float = 0.1,

111 matched_adamw_rms: float = 0.2,

112 momentum: float = 0.95,

113 nesterov: bool = True,

114 ns_steps: int = 5,

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

116 ):

117 defaults = {

118 "lr": lr,

119 "weight_decay": weight_decay,

120 "matched_adamw_rms": matched_adamw_rms,

121 "momentum": momentum,

122 "nesterov": nesterov,

123 "ns_steps": ns_steps,

124 }

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

126

127 self._group_dtensor_by_mesh()

128 deduced_count = self._auto_deduce_replica_count()

129 if deduced_count is None:

130 self.hsdp_replica_count = None

131 elif self.hsdp_replica_count is None:

132 self.hsdp_replica_count = deduced_count

133 self._split_replicate_groups()

134 self._build_hsdp_batch()

135 self._build_param_broadcast_info()

136 self._classify_parameters_for_step()

137

138 @torch.no_grad()

139 def step(self, closure=None) -> Optional[float]:

140 """

141 Perform a single optimization step.

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

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

144 replica peers via ``AsyncReplicateBroadcaster.flush_group``.

145 """

146 loss = None

147 if closure is not None:

148 with torch.enable_grad():

149 loss = closure()

150 broadcaster = AsyncReplicateBroadcaster(self)

151

152 num_groups = len(self.param_groups)

153

154 for group in self.param_groups:

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

156

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

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

159 for group_idx in range(num_groups):

160 info = self._hsdp_assignment_batches.get(group_idx)

161 if not info:

162 # No HSDP info — all params are no_comm.

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

164 no_comm_ns[group_idx] = self._update_muon_momentum(

165 self.param_groups[group_idx], unshard_params

166 )

167 else:

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

169 if no_comm_params:

170 no_comm_ns[group_idx] = self._update_muon_momentum(

171 self.param_groups[group_idx], no_comm_params

172 )

173 else:

174 no_comm_ns[group_idx] = {}

175

176 # Process no-comm params.

177 for group_idx in range(num_groups):

178 if no_comm_ns[group_idx]:

179 self._process_unshard_params(self.param_groups[group_idx], no_comm_ns[group_idx])

180

181 # Flatten nested batch into a linear schedule.

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

183 max_num_batches = 0

184 for group_idx in range(num_groups):

185 info = self._hsdp_assignment_batches.get(group_idx)

186 linear_batches = []

187 if info:

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

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

190 group_linear_batches[group_idx] = linear_batches

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

192

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

194 for batch_idx in range(max_num_batches):

195 for group_idx in range(num_groups):

196 linear_batches = group_linear_batches[group_idx]

197 if batch_idx >= len(linear_batches):

198 continue

199

200 hsdp_assign = linear_batches[batch_idx]

201 group = self.param_groups[group_idx]

202

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

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

205 if not hsdp_assign.is_shard:

206 if ns_inputs:

207 self._process_unshard_params(group, ns_inputs)

208 else:

209 self._process_shard_params(

210 group, ns_inputs, [hsdp_assign], group_idx,

211 buffer_cache={},

212 )

213

214 # Flush broadcasts after each assignment.

215 broadcaster.flush_group(hsdp_assign)

216

217 broadcaster.wait_all()

218 return loss

219

220 def _classify_parameters_for_step(self) -> None:

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

222

223 unshard: run Newton-Schulz locally.

224 shard: all-gather before Newton-Schulz.

225

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

227 """

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

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

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

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

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

233

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

235 assignment_info = self._hsdp_assignment_batches.get(group_idx)

236

237 unshard_params = []

238 shard_hsdp_assignments: List[HSDPGroupAssignment] = []

239

240 if assignment_info:

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

242 for bg in assignment_info["batch_groups"]:

243 for hsdp_assign in bg["sub_batches"]:

244 if hsdp_assign.is_shard:

245 shard_hsdp_assignments.append(hsdp_assign)

246 else:

247 unshard_params.extend(hsdp_assign.owned_params)

248 else:

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

250

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

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

253

254 # Greedily assign NS compute across shard ranks.

255 self._shard_compute_coord[group_idx] = {}

256 for hsdp_assign in shard_hsdp_assignments:

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

258 compute_by_index = build_owner_by_size(

259 records=hsdp_assign.owned_records,

260 replicate_sizes=shard_sizes,

261 )

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

263

264 self.unshard_params_by_group[group_idx] = unshard_params

265 self.shard_assignments_by_group[group_idx] = shard_hsdp_assignments

266 self.shard_params_by_group[group_idx] = shard_params

267

268 def _update_muon_momentum(

269 self,

270 group: Dict[str, Any],

271 params: List[torch.Tensor],

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

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

274 momentum = group['momentum']

275 nesterov = group['nesterov']

276

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

278 valid_params = []

279 grads = []

280 bufs = []

281 for p in params:

282 g = p.grad

283 if g is None:

284 continue

285 state = self.state[p]

286 if "momentum_buffer" not in state:

287 state["momentum_buffer"] = torch.zeros_like(g)

288 valid_params.append(p)

289 grads.append(g)

290 bufs.append(state["momentum_buffer"])

291

292 if not valid_params:

293 return {}

294

295 # Strip DTensor wrappers for elementwise foreach ops

296 local_grads = [to_local_if_dtensor(g) for g in grads]

297 local_bufs = [to_local_if_dtensor(b) for b in bufs]

298

299 # Fused momentum update: buf = momentum * buf + grad

300 # pylint: disable=protected-access

301 torch._foreach_mul_(local_bufs, momentum)

302 torch._foreach_add_(local_bufs, local_grads)

303

304 # Nesterov: u = momentum * buf + grad (out-of-place, keeps buf intact)

305 if nesterov:

306 local_us = torch._foreach_mul(local_bufs, momentum)

307 torch._foreach_add_(local_us, local_grads)

308 else:

309 local_us = list(local_bufs)

310

311 # Cast to bfloat16 for NS iteration

312 if local_us[0].dtype == torch.bfloat16:

313 local_us_bf = local_us

314 else:

315 local_us_bf = [u.to(torch.bfloat16) for u in local_us]

316

317 return dict(zip(valid_params, local_us_bf))

318

319 def _process_unshard_params(

320 self,

321 group: Dict[str, Any],

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

323 ) -> None:

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

325 lr = group["lr"]

326 weight_decay = group["weight_decay"]

327 rms = group["matched_adamw_rms"]

328 ns_steps = group["ns_steps"]

329

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

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

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

333

334 for sub_batch in safe_batches:

335 updates_dict, adjusted_lr = self._compute_batched_ns_updates(

336 sub_batch, param_to_ns_input, lr, rms, ns_steps, no_shard=True

337 )

338

339 # Fused batched apply — all params in the same sub_batch share

340 # the same adjusted_lr, so we can use foreach ops.

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

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

343

344 if weight_decay != 0.0:

345 # pylint: disable=protected-access

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

347 # pylint: disable=protected-access

348 torch._foreach_add_(local_params, local_updates, alpha=-adjusted_lr)

349

350 def _gather_and_compute_shard_updates(

351 self,

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

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

354 hsdp_assign: HSDPGroupAssignment,

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

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

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

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

359 total_shard_size: int,

360 lr: float,

361 rms: float,

362 ns_steps: int,

363 buffer_cache: Optional[Dict],

364 ) -> Tuple[

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

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

367 ]:

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

369

370 Returns:

371 (my_updates, param_compute_coord) for this HSDP assignment.

372 """

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

374

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

376 my_param_ids: set = set()

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

378 my_indices: set = set()

379

380 for idx, p in enumerate(valid_params):

381 p_index = param_to_index[p]

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

383 param_compute_coord[p] = compute_coord

384 if compute_coord == local_coords:

385 my_params.append(p)

386 my_param_ids.add(id(p))

387 my_indices.add(idx)

388

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

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

391 local_inputs = [param_to_ns_input[p] for p in valid_params]

392 gathered_list = fused_allgather_dtensor_params(

393 local_inputs, shard_pgs, hsdp_assign.layout_spec, buffer_cache=buffer_cache,

394 keep_indices=my_indices,

395 )

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

397 if id(p) in my_param_ids:

398 gathered_inputs[p] = full_inp

399 local_inputs.clear()

400 gathered_list.clear()

401

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

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

404 if my_params:

405 shape_groups = self._group_by_shape(my_params)

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

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

408 for sub_batch in safe_batches:

409 updates_dict, _ = self._compute_batched_ns_updates(

410 sub_batch, gathered_inputs, lr, rms, ns_steps

411 )

412 for p in sub_batch:

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

414 del updates_dict

415 gathered_inputs.clear()

416

417 return my_updates, param_compute_coord

418

419 def _process_shard_params(

420 self,

421 group: Dict[str, Any],

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

423 hsdp_assignments: List[HSDPGroupAssignment],

424 group_idx: int,

425 buffer_cache: Optional[Dict] = None,

426 ) -> None:

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

428 platform = get_platform()

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

430

431 lr = group["lr"]

432 weight_decay = group["weight_decay"]

433 rms = group["matched_adamw_rms"]

434 ns_steps = group["ns_steps"]

435

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

437

438 for hsdp_assign in hsdp_assignments:

439 owned_params = hsdp_assign.owned_params

440 valid_params = [p for p in owned_params if p in param_to_ns_input]

441

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

443

444 # Gather NS inputs and compute updates assigned to this shard coordinate.

445 if valid_params:

446 my_updates, param_compute_coord = self._gather_and_compute_shard_updates(

447 valid_params, param_to_ns_input, hsdp_assign,

448 shard_compute_coord, shard_sizes, local_coords,

449 shard_pgs, total_shard_size,

450 lr, rms, ns_steps, buffer_cache,

451 )

452

453 # Fused broadcast full updates within the shard group, then batched apply.

454 self._fused_broadcast_and_apply(

455 valid_params, my_updates, param_compute_coord,

456 lr, weight_decay, rms,

457 shard_pgs, shard_sizes, local_coords, total_shard_size,

458 hsdp_assign, platform, device,

459 )

460

461 def _group_by_shape(

462 self,

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

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

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

466

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

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

469 """

470 groups = defaultdict(list)

471

472 for p in params:

473 shape = tuple(p.shape)

474

475 if len(shape) == 2:

476 core_shape = (shape[0], shape[1])

477 elif len(shape) == 3 and shape[1] == 1:

478 core_shape = (shape[0], shape[2])

479 elif len(shape) >= 3:

480 core_shape = (shape[-2], shape[-1])

481 else:

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

483

484 groups[core_shape].append(p)

485

486 return groups

487

488 def _split_into_memory_safe_batches(

489 self,

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

491 shard_size: int = 1,

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

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

494

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

496 for the memory amplification from allgather.

497 """

498 max_numel_per_batch = 512 * 1024 * 1024 // shard_size

499

500 batches = []

501 current_batch = []

502 current_count = 0

503

504 for p in p_list:

505 p_count = p.numel()

506 if current_batch and current_count + p_count > max_numel_per_batch:

507 batches.append(current_batch)

508 current_batch = [p]

509 current_count = p_count

510 else:

511 current_batch.append(p)

512 current_count += p_count

513

514 if current_batch:

515 batches.append(current_batch)

516

517 return batches

518

519 def _compute_batched_ns_updates(

520 self,

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

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

523 lr: float,

524 rms: float,

525 ns_steps: int,

526 no_shard: bool = False

527 ) -> Tuple[Dict[torch.nn.Parameter, torch.Tensor], float]:

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

529

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

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

532

533 Returns:

534 updates_dict: per-parameter NS-orthogonalized updates.

535 adjusted_lr: a single scalar — all params in the same batch share

536 the same shape group and optimizer hyper-params, so their

537 adjusted_lr is identical.

538 """

539 updates_dict = {}

540

541 if not p_list:

542 return updates_dict, 0.0

543

544 inputs_3d = []

545 slice_sizes = []

546 shapes_info = []

547

548 for p in p_list:

549 origin_shape = tuple(getattr(p, 'local_shape', None) or p.to_local().shape) if no_shard else tuple(p.shape)

550 ns_input = ns_inputs[p].view(origin_shape)

551

552 is_conv = False

553 if len(origin_shape) == 2:

554 inp_3d = ns_input.unsqueeze(0)

555 n_dim = 1

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

557 inp_3d = ns_input.squeeze(1).unsqueeze(0)

558 is_conv = True

559 n_dim = 1

560 else:

561 inp_3d = ns_input

562 n_dim = origin_shape[0]

563

564 inputs_3d.append(inp_3d)

565 slice_sizes.append(n_dim)

566 shapes_info.append((origin_shape, is_conv))

567

568 merged_input = torch.cat(inputs_3d, dim=0)

569 merged_update = zeropower_via_newtonschulz5(merged_input, steps=ns_steps)

570 del merged_input

571

572 current_idx = 0

573 for i, p in enumerate(p_list):

574 n_dim = slice_sizes[i]

575 origin_shape, is_conv = shapes_info[i]

576

577 update = merged_update[current_idx: current_idx + n_dim]

578 current_idx += n_dim

579

580 if is_conv:

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

582 elif len(origin_shape) == 2:

583 update = update.squeeze(0)

584

585 updates_dict[p] = update

586 del merged_update

587

588 # Compute adjusted_lr once — all params share the same shape group

589 ref_shape, is_conv = shapes_info[0]

590 if is_conv:

591 adjusted_lr = adjust_lr_wd_for_muon_conv(lr, rms, ref_shape)

592 else:

593 adjusted_lr = adjust_lr_wd_for_muon(lr, rms, ref_shape)

594

595 return updates_dict, adjusted_lr

596

597 def _fused_broadcast_and_apply(

598 self,

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

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

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

602 lr: float,

603 weight_decay: float,

604 rms: float,

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

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

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

608 total_shard_size: int,

609 hsdp_assign: HSDPGroupAssignment,

610 platform: Any,

611 device: torch.device,

612 ) -> None:

613 """Fused broadcast + batched apply for shard-group updates (Optimized for low Free time)."""

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

615 for p in valid_params:

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

617

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

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

620 all_adjusted_lrs: List[float] = []

621

622 # Phase 1: Batched Pack

623 alignment_bytes = 512

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

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

626

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

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

629

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

631 is_compute_rank = coord == local_coords

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

633 total_padded_numel = 0

634

635 for p in coord_params:

636 actual_numel = p.numel()

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

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

639 total_padded_numel += padded_numel

640

641 coord_param_offsets[coord] = param_offsets

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

643

644 if is_compute_rank:

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

646 update = my_updates[p]

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

648 if padded_numel > actual_numel:

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

650

651 pack_buffers[coord] = pack_buffer

652

653 # Phase 2: Async Batched Relay Broadcast

654 if total_shard_size > 1:

655 self._batched_relay_broadcast(

656 pack_buffers, shard_pgs, shard_sizes, local_coords, platform

657 )

658

659 # Phase 3: Batched Unpack & Apply

660 layout_spec = hsdp_assign.layout_spec

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

662 pack_buffer = pack_buffers[coord]

663 param_offsets = coord_param_offsets[coord]

664

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

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

667 update_to_apply = chunk_update_by_layout(full_update, p, layout_spec)

668

669 origin_shape = tuple(p.shape)

670 if len(origin_shape) == 3 and origin_shape[1] == 1:

671 adjusted_lr = adjust_lr_wd_for_muon_conv(lr, rms, origin_shape)

672 else:

673 adjusted_lr = adjust_lr_wd_for_muon(lr, rms, origin_shape)

674

675 all_local_params.append(to_local_if_dtensor(p.data))

676 all_update_shards.append(update_to_apply.view(to_local_if_dtensor(p.data).shape))

677 all_adjusted_lrs.append(adjusted_lr)

678

679 # Batched Apply

680 if not all_local_params:

681 return

682

683 if weight_decay != 0.0:

684 coeff = 1.0 - lr * weight_decay

685 # pylint: disable=protected-access

686 torch._foreach_mul_(all_local_params, coeff)

687

688 lr_group_map: Dict[float, Tuple[List[torch.Tensor], List[torch.Tensor]]] = defaultdict(lambda: ([], []))

689 for local_p, update_shard, adj_lr in zip(all_local_params, all_update_shards, all_adjusted_lrs):

690 params_list, updates_list = lr_group_map[adj_lr]

691 params_list.append(local_p)

692 updates_list.append(update_shard)

693

694 for adj_lr, (params_list, updates_list) in lr_group_map.items():

695 if params_list:

696 # pylint: disable=protected-access

697 torch._foreach_add_(params_list, updates_list, alpha=-adj_lr)

698

699 @staticmethod

700 def _get_shard_info(

701 hsdp_assign: HSDPGroupAssignment,

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

703 """Extract shard topology from HSDPGroupAssignment.

704

705 Returns:

706 shard_sizes: Size of each shard mesh dimension.

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

708 shard_pgs: ProcessGroup for each shard mesh dim.

709 total_shard_size: Product of all shard_sizes.

710 """

711 shard_pgs = hsdp_assign.shard_pgs

712 shard_sizes = tuple(

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

714 for pg in shard_pgs

715 )

716 local_coords = tuple(

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

718 for pg in shard_pgs

719 )

720 total_shard_size = 1

721 for s in shard_sizes:

722 total_shard_size *= s

723

724 return shard_sizes, local_coords, shard_pgs, total_shard_size

725

726 @staticmethod

727 def _batched_relay_broadcast(

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

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

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

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

732 platform: Any,

733 ) -> None:

734 """

735 Batched asynchronous multi-dimensional relay broadcast.

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

737 while strictly preserving the multidimensional relay dependency.

738 """

739 for dim_idx, pg in enumerate(shard_pgs):

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

741 continue

742

743 work_handles = []

744

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

746 aligned = all(

747 local_coords[sub_dim] == coord[sub_dim]

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

749 )

750 if not aligned:

751 continue

752

753 src_rank_in_pg = coord[dim_idx]

754 global_src_rank = platform.get_global_rank(pg, src_rank_in_pg)

755

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

757 if work is not None:

758 work_handles.append(work)

759

760 for work in work_handles:

761 work.wait()

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

378 statements