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

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

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

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15 

16"""Muon optimizer with HSDP shard-group-aware communication.""" 

17 

18import math 

19from collections.abc import Callable 

20from collections import defaultdict 

21import logging 

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

23 

24import torch 

25import torch.distributed as dist 

26 

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

28from hyper_parallel.core.optimizer.dtensor_compat import to_local_if_dtensor 

29from hyper_parallel.core.optimizer.sharding_category import ( 

30 HSDPGroupAssignment, 

31 ParamShardMeta, 

32 build_owner_by_size, 

33) 

34from hyper_parallel.core.optimizer.muon_shard import ( 

35 _debug_param_shard_metadata, 

36 build_pad_ns_inputs, 

37 build_param_shard_metadata_for_group, 

38 chunk_update_by_layout, 

39 fused_allgather_dtensor_params, 

40) 

41 

42logger = logging.getLogger(__name__) 

43 

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

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

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

47 

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

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

50 (4.0848, -6.8946, 2.9270), 

51 (3.9505, -6.3029, 2.6377), 

52 (3.7418, -5.5913, 2.3037), 

53 (2.8769, -3.1427, 1.2046), 

54 (2.8366, -3.0525, 1.2012), 

55) 

56 

57 

58def zeropower_via_newtonschulz5( 

59 ns_inputs: torch.Tensor, 

60 steps: int, 

61 ns_variant: str = "asym5", 

62) -> torch.Tensor: 

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

64 mat_x = ns_inputs 

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

66 if transposed: 

67 mat_x = mat_x.mT 

68 

69 # Normalize input before Newton-Schulz iteration. 

70 mat_x = mat_x / (mat_x.norm(dim=(-2, -1), keepdim=True) + 1e-10) 

71 

72 n_size = mat_x.size(-2) 

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

74 buf_b = torch.empty_like(buf_a) 

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

76 

77 if ns_variant == "legacy": 

78 step_coeffs = [_NS_LEGACY_COEFFS] * steps 

79 elif ns_variant == "asym5": 

80 step_coeffs = _NS_ASYM5_COEFFS[:steps] 

81 else: 

82 raise ValueError( 

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

84 ) 

85 

86 for coeff_a, coeff_b, coeff_c in step_coeffs: 

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

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

89 buf_a.mul_(coeff_b) 

90 buf_b.mul_(coeff_c) 

91 buf_a.add_(buf_b) 

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

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

94 

95 del buf_a, buf_b, buf_c 

96 

97 if transposed: 

98 mat_x = mat_x.mT 

99 

100 return mat_x 

101 

102 

103def compute_muon_slice_scale(slice_tensor: torch.Tensor, matched_adamw_rms: float) -> float: 

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

105 shape = tuple(slice_tensor.shape) 

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

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

108 else: 

109 logical_dims = shape[-2:] 

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

111 

112 

113class Muon(BaseDistributedOptimizer): 

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

115 

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

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

118 sharded parameters. 

119 """ 

120 

121 def __init__( 

122 self, 

123 params, 

124 lr: float = 2e-2, 

125 weight_decay: float = 0.1, 

126 matched_adamw_rms: float = 0.2, 

127 momentum: float = 0.95, 

128 nesterov: bool = True, 

129 ns_steps: int = 5, 

130 ns_variant: str = "asym5", 

131 reshape_fn: Optional[Callable[..., Any]] = None, 

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

133 ): 

134 if ns_variant not in ("legacy", "asym5"): 

135 raise ValueError( 

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

137 ) 

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

139 momentum = [momentum] 

140 if len(momentum) == 1: 

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

142 defaults = { 

143 "lr": lr, 

144 "weight_decay": weight_decay, 

145 "matched_adamw_rms": matched_adamw_rms, 

146 "momentum": momentum, 

147 "nesterov": nesterov, 

148 "ns_steps": ns_steps, 

149 "ns_variant": ns_variant, 

150 } 

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

152 self.reshape_fn = reshape_fn 

153 

154 self._group_dtensor_by_mesh() 

155 self._build_param_shard_metadata() 

156 deduced_count = self._auto_deduce_replica_count() 

157 if deduced_count is None: 

158 self.hsdp_replica_count = None 

159 elif self.hsdp_replica_count is None: 

160 self.hsdp_replica_count = deduced_count 

161 self._split_replicate_groups() 

162 self._build_hsdp_batch() 

163 self._build_param_broadcast_info() 

164 self._classify_parameters_for_step() 

165 

166 def __str__(self): 

167 return super().__repr__() 

168 

169 __repr__ = __str__ 

170 

171 @torch.no_grad() 

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

173 """ 

174 Perform a single optimization step. 

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

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

177 replica peers via ``AsyncReplicateBroadcaster.flush_group``. 

178 """ 

179 loss = None 

180 if closure is not None: 

181 with torch.enable_grad(): 

182 loss = closure() 

183 broadcaster = AsyncReplicateBroadcaster(self) 

184 

185 num_groups = len(self.param_groups) 

186 

187 for group in self.param_groups: 

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

189 

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

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

192 for group_idx in range(num_groups): 

193 info = self._hsdp_assignment_batches.get(group_idx) 

194 if not info: 

195 # No HSDP info — all params are no_comm. 

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

197 no_comm_ns[group_idx] = self._update_muon_momentum( 

198 self.param_groups[group_idx], unshard_params 

199 ) 

200 else: 

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

202 if no_comm_params: 

203 no_comm_ns[group_idx] = self._update_muon_momentum( 

204 self.param_groups[group_idx], no_comm_params 

205 ) 

206 else: 

207 no_comm_ns[group_idx] = {} 

208 

209 # Process no-comm params. 

210 for group_idx in range(num_groups): 

211 if no_comm_ns[group_idx]: 

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

213 

214 # Flatten nested batch into a linear schedule. 

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

216 max_num_batches = 0 

217 for group_idx in range(num_groups): 

218 info = self._hsdp_assignment_batches.get(group_idx) 

219 linear_batches = [] 

220 if info: 

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

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

223 group_linear_batches[group_idx] = linear_batches 

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

225 

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

227 for batch_idx in range(max_num_batches): 

228 for group_idx in range(num_groups): 

229 linear_batches = group_linear_batches[group_idx] 

230 if batch_idx >= len(linear_batches): 

231 continue 

232 

233 hsdp_assign = linear_batches[batch_idx] 

234 group = self.param_groups[group_idx] 

235 

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

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

238 if not hsdp_assign.is_shard: 

239 if ns_inputs: 

240 self._process_unshard_params(group, ns_inputs) 

241 else: 

242 self._process_shard_params( 

243 group, ns_inputs, [hsdp_assign], group_idx, 

244 buffer_cache={}, 

245 ) 

246 

247 # Flush broadcasts after each assignment. 

248 broadcaster.flush_group(hsdp_assign) 

249 

250 broadcaster.wait_all() 

251 return loss 

252 

253 def _classify_parameters_for_step(self) -> None: 

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

255 

256 unshard: run Newton-Schulz locally. 

257 shard: all-gather before Newton-Schulz. 

258 

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

260 """ 

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

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

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

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

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

266 

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

268 assignment_info = self._hsdp_assignment_batches.get(group_idx) 

269 

270 unshard_params = [] 

271 shard_hsdp_assignments: List[HSDPGroupAssignment] = [] 

272 

273 if assignment_info: 

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

275 for bg in assignment_info["batch_groups"]: 

276 for hsdp_assign in bg["sub_batches"]: 

277 if hsdp_assign.is_shard: 

278 shard_hsdp_assignments.append(hsdp_assign) 

279 else: 

280 unshard_params.extend(hsdp_assign.owned_params) 

281 else: 

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

283 

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

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

286 

287 # Greedily assign NS compute across shard ranks. 

288 self._shard_compute_coord[group_idx] = {} 

289 for hsdp_assign in shard_hsdp_assignments: 

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

291 compute_by_index = build_owner_by_size( 

292 records=hsdp_assign.owned_records, 

293 replicate_sizes=shard_sizes, 

294 ) 

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

296 

297 self.unshard_params_by_group[group_idx] = unshard_params 

298 self.shard_assignments_by_group[group_idx] = shard_hsdp_assignments 

299 self.shard_params_by_group[group_idx] = shard_params 

300 

301 def _update_muon_momentum( 

302 self, 

303 group: Dict[str, Any], 

304 params: List[torch.Tensor], 

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

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

307 momentum1, momentum2 = group["momentum"] 

308 nesterov = group['nesterov'] 

309 

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

311 valid_params = [] 

312 grads = [] 

313 bufs = [] 

314 for p in params: 

315 g = p.grad 

316 if g is None: 

317 continue 

318 state = self.state[p] 

319 if "momentum_buffer" not in state: 

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

321 valid_params.append(p) 

322 grads.append(g) 

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

324 

325 if not valid_params: 

326 return {} 

327 

328 # Match muon_update_core(): 

329 # m_for_update = grad + momentum1 * m_old 

330 # m_new = grad + momentum2 * m_old 

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

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

333 # pylint: disable=protected-access 

334 local_us = torch._foreach_add(local_grads, local_bufs, alpha=momentum1) 

335 

336 torch._foreach_mul_(local_bufs, momentum2) 

337 torch._foreach_add_(local_bufs, local_grads) 

338 

339 if nesterov: 

340 torch._foreach_mul_(local_us, momentum1) 

341 torch._foreach_add_(local_us, local_grads) 

342 

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

344 local_us_bf = local_us 

345 else: 

346 local_us = list(local_us) 

347 for i, u in enumerate(local_us): 

348 local_us[i] = u.to(torch.bfloat16) 

349 local_us_bf = local_us 

350 

351 return dict(zip(valid_params, local_us_bf)) 

352 

353 def _process_unshard_params( 

354 self, 

355 group: Dict[str, Any], 

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

357 ) -> None: 

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

359 lr = group["lr"] 

360 weight_decay = group["weight_decay"] 

361 

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

363 param_to_ns_input = { 

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

365 } 

366 if not param_to_ns_input: 

367 return 

368 

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

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

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

372 

373 for sub_batch in safe_batches: 

374 updates_dict = self._compute_batched_ns_updates( 

375 sub_batch, param_to_ns_input, group, no_shard=True 

376 ) 

377 

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

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

380 

381 if weight_decay != 0.0: 

382 # pylint: disable=protected-access 

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

384 # pylint: disable=protected-access 

385 torch._foreach_add_(local_params, local_updates, alpha=-lr) 

386 

387 def _gather_and_compute_shard_updates( 

388 self, 

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

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

391 hsdp_assign: HSDPGroupAssignment, 

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

393 group: Dict[str, Any], 

394 buffer_cache: Optional[Dict], 

395 ) -> Tuple[ 

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

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

398 ]: 

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

400 

401 Returns: 

402 (my_updates, param_compute_coord) for this HSDP assignment. 

403 """ 

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

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

406 

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

408 my_param_ids: set = set() 

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

410 my_indices: set = set() 

411 

412 for idx, p in enumerate(valid_params): 

413 p_index = param_to_index[p] 

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

415 param_compute_coord[p] = compute_coord 

416 if compute_coord == local_coords: 

417 my_params.append(p) 

418 my_param_ids.add(id(p)) 

419 my_indices.add(idx) 

420 

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

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

423 local_inputs = build_pad_ns_inputs( 

424 valid_params, 

425 param_to_ns_input, 

426 self._param_shard_metadata, 

427 ) 

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

429 gathered_list = fused_allgather_dtensor_params( 

430 local_inputs, shard_pgs, hsdp_assign.layout_spec, 

431 param_shard_metadata=param_shard_metadata, 

432 buffer_cache=buffer_cache, 

433 keep_indices=my_indices, 

434 ) 

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

436 if id(p) in my_param_ids: 

437 gathered_inputs[p] = full_inp 

438 local_inputs.clear() 

439 gathered_list.clear() 

440 

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

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

443 if my_params: 

444 shape_groups = self._group_by_shape(my_params) 

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

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

447 for sub_batch in safe_batches: 

448 updates_dict = self._compute_batched_ns_updates( 

449 sub_batch, gathered_inputs, group 

450 ) 

451 for p in sub_batch: 

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

453 del updates_dict 

454 gathered_inputs.clear() 

455 

456 return my_updates, param_compute_coord 

457 

458 def _process_shard_params( 

459 self, 

460 group: Dict[str, Any], 

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

462 hsdp_assignments: List[HSDPGroupAssignment], 

463 group_idx: int, 

464 buffer_cache: Optional[Dict] = None, 

465 ) -> None: 

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

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

468 

469 for hsdp_assign in hsdp_assignments: 

470 owned_params = hsdp_assign.owned_params 

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

472 # local shards still participate via build_pad_ns_inputs(). 

473 valid_params = owned_params 

474 

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

476 

477 # uneven shard check 

478 shard_pgs_rank = [] 

479 for pg in shard_pgs: 

480 pg_rank = torch.distributed.get_process_group_ranks(pg) 

481 shard_pgs_rank.append(pg_rank) 

482 

483 for p in valid_params: 

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

485 logger.debug_rank0( 

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

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

488 p.model_name, 

489 p.placements, 

490 p.device_mesh, 

491 p.to_local().shape, 

492 p.shape, 

493 shard_pgs_rank, 

494 shard_sizes, 

495 ) 

496 

497 if valid_params: 

498 safe_batches = self._split_into_memory_safe_batches( 

499 valid_params, shard_size=total_shard_size, 

500 ) 

501 for sub_batch in safe_batches: 

502 my_updates, param_compute_coord = self._gather_and_compute_shard_updates( 

503 sub_batch, param_to_ns_input, hsdp_assign, 

504 shard_compute_coord, group, buffer_cache, 

505 ) 

506 

507 self._fused_broadcast_and_apply( 

508 sub_batch, my_updates, param_compute_coord, 

509 group, hsdp_assign, 

510 ) 

511 

512 def _group_by_shape( 

513 self, 

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

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

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

517 

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

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

520 """ 

521 groups = defaultdict(list) 

522 

523 for p in params: 

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

525 groups[core_shape].append(p) 

526 

527 return groups 

528 

529 @staticmethod 

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

531 if len(shape) == 2: 

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

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

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

535 if len(shape) >= 3: 

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

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

538 

539 def _reshape_ns_input( 

540 self, 

541 param: torch.nn.Parameter, 

542 ns_input: torch.Tensor, 

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

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

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

546 if self.reshape_fn is None: 

547 return working_input, [working_input] 

548 

549 param_fqn = param.model_name 

550 if param_fqn is None: 

551 return working_input, [working_input] 

552 

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

554 if not reshaped_inputs: 

555 return working_input, [working_input] 

556 

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

558 logger.info_rank0( 

559 "Reshape %s from %s to %s", 

560 param_fqn, 

561 working_input.shape, 

562 reshaped_inputs[0].shape, 

563 ) 

564 

565 for reshaped_input in reshaped_inputs: 

566 assert ( 

567 reshaped_input.untyped_storage().data_ptr() == working_input.untyped_storage().data_ptr() 

568 ), "reshape_fn must return views that share storage with the working NS input tensor." 

569 

570 return working_input, reshaped_inputs 

571 

572 def _compute_batched_ns_outputs_for_tensors( 

573 self, 

574 tensor_list: List[torch.Tensor], 

575 ns_steps: int, 

576 ns_variant: str = "asym5", 

577 ) -> List[torch.Tensor]: 

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

579 if not tensor_list: 

580 return [] 

581 

582 inputs_3d = [] 

583 slice_sizes = [] 

584 shapes_info = [] 

585 

586 for tensor in tensor_list: 

587 origin_shape = tuple(tensor.shape) 

588 is_conv = False 

589 if len(origin_shape) == 2: 

590 inp_3d = tensor.unsqueeze(0) 

591 n_dim = 1 

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

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

594 is_conv = True 

595 n_dim = 1 

596 elif len(origin_shape) == 3: 

597 inp_3d = tensor 

598 n_dim = origin_shape[0] 

599 else: 

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

601 n_dim = inp_3d.shape[0] 

602 

603 inputs_3d.append(inp_3d) 

604 slice_sizes.append(n_dim) 

605 shapes_info.append((origin_shape, is_conv)) 

606 

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

608 squeeze_batch = merged_input.shape[0] == 1 

609 if squeeze_batch: 

610 merged_input = merged_input.squeeze(0) 

611 

612 merged_update = zeropower_via_newtonschulz5(merged_input, steps=ns_steps, ns_variant=ns_variant) 

613 del merged_input 

614 

615 if squeeze_batch: 

616 merged_update = merged_update.unsqueeze(0) 

617 

618 outputs = [] 

619 current_idx = 0 

620 for n_dim, (origin_shape, is_conv) in zip(slice_sizes, shapes_info): 

621 update = merged_update[current_idx: current_idx + n_dim] 

622 current_idx += n_dim 

623 

624 if is_conv: 

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

626 elif len(origin_shape) == 2: 

627 update = update.squeeze(0) 

628 elif len(origin_shape) >= 4: 

629 update = update.reshape(origin_shape) 

630 

631 outputs.append(update) 

632 

633 del merged_update 

634 return outputs 

635 

636 def _split_into_memory_safe_batches( 

637 self, 

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

639 shard_size: int = 1, 

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

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

642 

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

644 for the memory amplification from allgather. 

645 """ 

646 max_numel_per_batch = 512 * 1024 * 1024 // shard_size 

647 

648 batches = [] 

649 current_batch = [] 

650 current_count = 0 

651 

652 for p in p_list: 

653 p_count = p.numel() 

654 if current_batch and current_count + p_count > max_numel_per_batch: 

655 batches.append(current_batch) 

656 current_batch = [p] 

657 current_count = p_count 

658 else: 

659 current_batch.append(p) 

660 current_count += p_count 

661 

662 if current_batch: 

663 batches.append(current_batch) 

664 

665 return batches 

666 

667 def _compute_batched_ns_updates( 

668 self, 

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

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

671 group: Dict[str, Any], 

672 no_shard: bool = False 

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

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

675 

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

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

678 

679 """ 

680 updates_dict = {} 

681 

682 if not p_list: 

683 return updates_dict 

684 

685 rms = group["matched_adamw_rms"] 

686 ns_steps = group["ns_steps"] 

687 ns_variant = group["ns_variant"] 

688 

689 reshape_groups: Dict[Tuple[int, int], List[torch.Tensor]] = defaultdict(list) 

690 origin_shapes: Dict[torch.nn.Parameter, Tuple[int, ...]] = {} 

691 working_inputs: Dict[torch.nn.Parameter, torch.Tensor] = {} 

692 

693 for p in p_list: 

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

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

696 origin_shapes[p] = origin_shape 

697 

698 working_input, reshaped_inputs = self._reshape_ns_input(p, ns_input) 

699 working_inputs[p] = working_input 

700 for reshaped_input in reshaped_inputs: 

701 core_shape = self._shape_to_core_shape(tuple(reshaped_input.shape)) 

702 reshape_groups[core_shape].append(reshaped_input) 

703 

704 for _, tensor_list in reshape_groups.items(): 

705 reshaped_updates = self._compute_batched_ns_outputs_for_tensors( 

706 tensor_list, 

707 ns_steps, 

708 ns_variant=ns_variant, 

709 ) 

710 

711 # scale updates 

712 for reshaped_input, reshaped_update in zip(tensor_list, reshaped_updates): 

713 slice_scale = compute_muon_slice_scale(reshaped_update, rms) 

714 reshaped_update.mul_(slice_scale) 

715 reshaped_input.copy_(reshaped_update.contiguous().view_as(reshaped_input)) 

716 

717 for p in p_list: 

718 ns_input = ns_inputs[p].view(origin_shapes[p]) 

719 working_input = working_inputs[p] 

720 if working_input.untyped_storage().data_ptr() != ns_input.untyped_storage().data_ptr(): 

721 ns_input.copy_(working_input) 

722 updates_dict[p] = ns_input 

723 

724 return updates_dict 

725 

726 def _fused_broadcast_and_apply( 

727 self, 

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

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

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

731 group: Dict[str, Any], 

732 hsdp_assign: HSDPGroupAssignment, 

733 ) -> None: 

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

735 lr = group["lr"] 

736 weight_decay = group["weight_decay"] 

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

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

739 

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

741 for p in valid_params: 

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

743 

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

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

746 

747 alignment_bytes = 512 

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

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

750 

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

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

753 

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

755 is_compute_rank = coord == local_coords 

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

757 total_padded_numel = 0 

758 

759 for p in coord_params: 

760 actual_numel = p.numel() 

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

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