Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / optimizer / swap_muon.py: 95%
234 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
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"""Muon momentum-only optimizer-state swap adapter."""
17# pylint: disable=protected-access
18# pylint: disable=forbidden-backend-import
20from __future__ import annotations
22from dataclasses import dataclass
23import math
24from typing import Any, Dict, Iterable, Iterator, List, Optional, Sequence, Tuple
26import torch
28from hyper_parallel.core.optimizer.muon import Muon
29from hyper_parallel.core.optimizer.optimizer import AsyncReplicateBroadcaster
30from hyper_parallel.core.optimizer.sharding_category import HSDPGroupAssignment
31from hyper_parallel.core.optimizer.swap_optimizer_base import (
32 StateSwapAdapter,
33 SwapOptimizer as CoreSwapOptimizer,
34 SwapSlot,
35)
37#: Muon's only persistent optimizer state that participates in swap.
38MUON_STATE_KEYS = ("momentum_buffer",)
40_NO_COMM = "no_comm"
41_HSDP = "hsdp"
44@dataclass
45class MuonSwapUnit:
46 """One atomic Muon update unit and the momentum slots it needs.
48 A unit is never split by the pipeline: its parameters run through momentum,
49 Newton-Schulz and the parameter update as one indivisible piece of work, so
50 a single matrix is never partially swapped. ``kind`` only selects which
51 Muon execution path consumes the unit, and ``hsdp_assign`` keeps the
52 original HSDP communication boundary intact.
54 Attributes:
55 group_index: Index of the owning parameter group, derived from
56 ``schedule_key``.
57 kind: ``"no_comm"`` for locally orthonormalized params, ``"hsdp"`` for
58 an HSDP shard/replicate assignment. Derived from ``schedule_key``.
59 params: Parameters this unit updates, in schedule order.
60 slots: Swappable momentum-buffer slots for ``params``.
61 schedule_key: Rank-independent group, kind, assignment and sub-batch indices.
62 logical_nbytes: Full logical momentum bytes, independent of local ownership.
63 hsdp_assign: The original assignment for ``"hsdp"`` units, else ``None``.
64 post_update_params: Parameters whose post-update callbacks run after this
65 unit. Only the final unit of the original Muon update boundary owns
66 this tuple, in the original parameter order.
67 flushes_broadcast: Whether this unit commits its assignment's parameter
68 broadcast. Only the last unit of an assignment sets it, so an
69 assignment split into several NS batches still issues the single
70 flush the bare schedule issues.
71 """
73 params: List[torch.nn.Parameter]
74 slots: List[SwapSlot]
75 schedule_key: Tuple[int, str, int, int]
76 logical_nbytes: int
77 hsdp_assign: Optional[HSDPGroupAssignment] = None
78 post_update_params: Tuple[torch.nn.Parameter, ...] = ()
79 flushes_broadcast: bool = False
81 @property
82 def group_index(self) -> int:
83 """Return the owning parameter-group index from the schedule identity."""
84 return self.schedule_key[0]
86 @property
87 def kind(self) -> str:
88 """Return the execution-path kind from the schedule identity."""
89 return self.schedule_key[1]
92class MuonSwapAdapter(StateSwapAdapter):
93 """Swap adapter that offloads and prefetches Muon's momentum buffers.
95 The adapter is the swap *coordinator* for Muon, never a second Muon
96 implementation. It replays the schedule of the unmodified ``Muon`` and
97 calls that optimizer's existing execution methods -- momentum,
98 Newton-Schulz, all-gather, broadcast and the parameter write-back all stay
99 in ``muon.py`` as the single source of truth. What the adapter owns is
100 (a) building the atomic update-unit schedule in the original order, and
101 (b) making sure the momentum buffer each unit reads is materialized on the
102 device only for the duration of that unit.
103 """
105 @staticmethod
106 def default_state_keys() -> Tuple[str, ...]:
107 """Return the swap state keys a Muon optimizer owns."""
108 return MUON_STATE_KEYS
110 def validate(self) -> None:
111 """Reject configurations the momentum-only swap path cannot preserve.
113 Raises:
114 ValueError: If a configured state key is not one Muon owns, or a
115 parameter holds momentum state the runtime cannot rebind, or a
116 custom momentum update would lose its original call boundary.
117 """
118 self._validate_configured_state_keys()
119 if self.optimizer.momentum_update_fn is not None:
120 raise ValueError(
121 "Muon swap optimizer does not support momentum_update_fn because "
122 "state partitioning changes its call boundary."
123 )
124 for group in self.optimizer.param_groups:
125 for param in group["params"]:
126 self._validate_param(param)
128 def _validate_configured_state_keys(self) -> None:
129 """Reject a configured swap key that is not Muon state.
131 ``_configured_state_keys`` shares a vocabulary with the other families
132 and silently drops keys none of them implement, so Muon has to check its
133 own selection explicitly: asking to swap ``exp_avg`` on a Muon optimizer
134 must fail rather than quietly swap nothing.
136 Raises:
137 ValueError: If the configured selection names a foreign state key.
138 """
139 if self.config.state_keys is None:
140 return
141 unavailable = sorted(set(self.config.state_keys) - set(self.default_state_keys()))
142 if unavailable:
143 raise ValueError(
144 f"Requested state keys {unavailable} are not available for {type(self.optimizer)!r}; "
145 f"Muon only swaps {self.default_state_keys()}."
146 )
148 def _validate_param(self, param: torch.nn.Parameter) -> None:
149 """Validate one parameter's momentum state is swap-compatible."""
150 state = self.optimizer.state.get(param)
151 if not state:
152 return
153 buffer = state.get("momentum_buffer")
154 if buffer is None:
155 return
156 if getattr(param, "grad", None) is not None and getattr(param.grad, "is_sparse", False):
157 raise ValueError("Muon swap optimizer only supports dense gradients.")
158 local = self._runtime_local_tensor(buffer)
159 if not isinstance(local, torch.Tensor):
160 raise ValueError(
161 f"Muon swap optimizer requires tensor momentum state, got {type(buffer)!r}."
162 )
163 if local.layout != torch.strided:
164 raise ValueError("Muon swap optimizer requires dense momentum state.")
165 if not local.is_contiguous():
166 raise ValueError("Muon swap optimizer requires contiguous momentum state.")
168 @staticmethod
169 def _runtime_local_tensor(tensor: Any) -> Any:
170 """Return the local shard of a possibly distributed tensor."""
171 to_local = getattr(tensor, "to_local", None)
172 return to_local() if callable(to_local) else tensor
174 def _scheduled_param_ids(self) -> Optional[set[int]]:
175 """Return parameters this rank owns in the current Muon schedule.
177 ``None`` means the optimizer has no HSDP schedule and therefore all
178 parameters are local. With HSDP metadata present, only no-comm
179 parameters and assignment owners may carry Muon state on this rank.
180 """
181 schedule = getattr(self.optimizer, "_hsdp_assignment_batches", None)
182 # ``Muon`` initialises this as an empty dict in non-HSDP runs, so an
183 # empty schedule carries no ownership metadata at all: every parameter
184 # is local and pre-existing momentum state must be kept.
185 if not schedule:
186 return None
187 scheduled: set[int] = set()
188 for group_index, group in enumerate(self.optimizer.param_groups):
189 info = schedule.get(group_index)
190 if info is None:
191 scheduled.update(id(param) for param in group["params"])
192 continue
193 scheduled.update(id(param) for param in info.get("no_comm", []))
194 for batch_group in info.get("batch_groups", []):
195 for assignment in batch_group.get("sub_batches", []):
196 scheduled.update(id(param) for param in assignment.owned_params)
197 return scheduled
199 def _params_with_local_state(self, params: Iterable[Any]) -> Iterable[Any]:
200 """Keep checkpoint-restored momentum only for local schedule owners."""
201 scheduled_ids = self._scheduled_param_ids()
202 if scheduled_ids is None:
203 return params
204 return tuple(param for param in params if id(param) in scheduled_ids)
206 def initial_slots(self) -> Iterable[SwapSlot]:
207 """Discover only local-owner momentum state before the first step."""
208 scheduled_ids = self._scheduled_param_ids()
209 slots = []
210 for group in self.optimizer.param_groups:
211 for param in group["params"]:
212 if scheduled_ids is not None and id(param) not in scheduled_ids:
213 # State may have been primed on every rank before this
214 # wrapper was created. Drop non-owner entries before the
215 # swap runtime publishes any slots.
216 self.optimizer.state.pop(param, None)
217 continue
218 state = self.optimizer.state.get(param)
219 if not state:
220 continue
221 for key in self._configured_state_keys():
222 if key in state and (id(param), key) not in self._slots:
223 self._slots[(id(param), key)] = self._make_slot(key, state[key])
224 slots.extend(self._slots_for_params([param]))
225 return tuple(slots)
227 # ------------------------------------------------------------------ schedule
228 def prepare_step(self, *args: Any, **kwargs: Any) -> Dict[str, Any]:
229 """Advance steps, register momentum slots and build the unit schedule.
231 The unit schedule mirrors ``Muon.step()`` exactly: the no-comm params of
232 every group first, then the flattened HSDP assignment schedule, keeping
233 the per-group-then-per-batch nesting so collective order across ranks is
234 unchanged.
236 Returns:
237 Step context holding the ordered units and the step's broadcaster.
239 Raises:
240 ValueError: If extra step arguments are passed.
241 """
242 if args or kwargs:
243 raise ValueError("Muon swap optimizer step does not support closure or extra arguments.")
244 for group in self.optimizer.param_groups:
245 group['step'] = (group.get('step') or 0) + 1
247 units: List[MuonSwapUnit] = []
248 num_groups = len(self.optimizer.param_groups)
249 for group_index in range(num_groups):
250 no_comm_params = self._no_comm_params(group_index)
251 no_comm_batches = list(self._iter_unshard_batches(no_comm_params))
252 for batch_index, params in enumerate(no_comm_batches):
253 post_update_params = no_comm_params if batch_index == len(no_comm_batches) - 1 else ()
254 units.append(self._make_unit(
255 (group_index, _NO_COMM, -1, batch_index), params, params,
256 post_update_params=post_update_params,
257 ))
259 linear_batches = [
260 self._hsdp_linear_batches(group_index) for group_index in range(num_groups)
261 ]
262 max_num_batches = max((len(batches) for batches in linear_batches), default=0)
263 for batch_index in range(max_num_batches):
264 for group_index in range(num_groups):
265 group_batches = linear_batches[group_index]
266 if batch_index >= len(group_batches):
267 continue
268 units.extend(self._make_hsdp_units(group_index, group_batches[batch_index], batch_index))
270 return {
271 "units": units,
272 "broadcaster": AsyncReplicateBroadcaster(self.optimizer),
273 }
275 def _no_comm_params(self, group_index: int) -> List[torch.nn.Parameter]:
276 """Return the params of ``group_index`` that need no HSDP collective.
278 This reads the same private schedule metadata ``Muon.step()`` reads, and
279 falls back to every parameter in the group when no HSDP assignment
280 metadata was built for it -- the same "no HSDP info means all params are
281 no_comm" rule the bare step applies.
282 """
283 info = getattr(self.optimizer, "_hsdp_assignment_batches", {}).get(group_index)
284 if not info:
285 return list(self.optimizer.unshard_params_by_group.get(group_index, []))
286 return list(info.get("no_comm", []))
288 def _hsdp_linear_batches(self, group_index: int) -> List[HSDPGroupAssignment]:
289 """Flatten one group's nested HSDP schedule into a linear batch list."""
290 info = getattr(self.optimizer, "_hsdp_assignment_batches", {}).get(group_index)
291 linear_batches: List[HSDPGroupAssignment] = []
292 if info:
293 for batch_group in info.get("batch_groups", []):
294 linear_batches.extend(batch_group.get("sub_batches", []))
295 return linear_batches
297 def _iter_unshard_batches(
298 self,
299 params: List[torch.nn.Parameter],
300 ) -> Iterator[List[torch.nn.Parameter]]:
301 """Plan logical batches using Muon's existing shape and memory rules."""
302 for shape_group in self.optimizer._group_by_shape(params).values():
303 yield from self.optimizer._split_into_memory_safe_batches(shape_group, shard_size=1)
305 def _make_unit(
306 self,
307 schedule_key: Tuple[int, str, int, int],
308 logical_params: Iterable[torch.nn.Parameter],
309 params: Iterable[torch.nn.Parameter],
310 hsdp_assign: Optional[HSDPGroupAssignment] = None,
311 post_update_params: Iterable[torch.nn.Parameter] = (),
312 flushes_broadcast: bool = False,
313 ) -> MuonSwapUnit:
314 """Project a logical unit onto local gradient-bearing owners and slots."""
315 logical_nbytes = sum(
316 math.prod(param.shape) * self._runtime_local_tensor(param).element_size()
317 for param in logical_params
318 )
319 params = [param for param in params if param.grad is not None]
320 post_update_params = tuple(
321 param for param in post_update_params if param.grad is not None
322 )
323 # HSDP assignments contain every logical parameter, but only the local
324 # owner participates in this rank's momentum update. Register state
325 # only for parameters that are actually present in this atomic unit;
326 # otherwise a non-owner would retain a device/host slot indefinitely.
327 for param in params:
328 self._ensure_momentum_state(param)
329 return MuonSwapUnit(
330 params=params,
331 slots=self._slots_for_params(params),
332 schedule_key=schedule_key,
333 logical_nbytes=logical_nbytes,
334 hsdp_assign=hsdp_assign,
335 post_update_params=post_update_params,
336 flushes_broadcast=flushes_broadcast,
337 )
339 def _make_hsdp_units(
340 self,
341 group_index: int,
342 hsdp_assign: HSDPGroupAssignment,
343 assignment_index: int,
344 ) -> List[MuonSwapUnit]:
345 """Build the units for one HSDP assignment, in assignment order.
347 Sharded assignments stay atomic because the shard executor consumes the
348 entire assignment. For unsharded assignments, every rank derives all
349 owners' NS batches before selecting its own local batch.
350 """
351 if hsdp_assign.is_shard:
352 return [
353 self._make_unit(
354 (group_index, _HSDP, assignment_index, 0),
355 hsdp_assign.all_params, hsdp_assign.owned_params,
356 hsdp_assign, post_update_params=hsdp_assign.owned_params,
357 flushes_broadcast=True,
358 )
359 ]
360 params_by_owner: Dict[Tuple[int, ...], List[torch.nn.Parameter]] = {}
361 for record in hsdp_assign.all_records:
362 owner = hsdp_assign.owner_by_index.get(record.index)
363 # Frozen records remain in all_records but have no assigned owner.
364 # () is a valid owner when there are no replicate dimensions.
365 if owner is None:
366 continue
367 params_by_owner.setdefault(owner, []).append(record.param)
368 batches_by_owner = {
369 owner: list(self._iter_unshard_batches(params_by_owner[owner]))
370 for owner in sorted(params_by_owner)
371 }
372 batch_count = max((len(batches) for batches in batches_by_owner.values()), default=1)
373 local_batches = batches_by_owner.get(hsdp_assign.replicate_group_ranks, [])
374 units = []
375 for index in range(batch_count):
376 logical_params = [
377 param for batches in batches_by_owner.values() if index < len(batches)
378 for param in batches[index]
379 ]
380 local_params = local_batches[index] if index < len(local_batches) else []
381 post_update_params = hsdp_assign.owned_params if index == batch_count - 1 else ()
382 units.append(self._make_unit(
383 (group_index, _HSDP, assignment_index, index),
384 logical_params, local_params, hsdp_assign,
385 post_update_params=post_update_params,
386 flushes_broadcast=index == batch_count - 1,
387 ))
388 return units
390 def _ensure_momentum_state(self, param: torch.nn.Parameter) -> None:
391 """Create the host-resident momentum slot for one parameter.
393 First-step state is created host-side on purpose: the design requires
394 that a freshly created momentum buffer never becomes device-resident
395 outside its own update unit. Parameters without a gradient keep Muon's
396 original skip semantics and get no slot at all.
398 The device shell published into ``state[param]`` is what
399 ``Muon._update_muon_momentum`` reads during this step. That method only
400 allocates momentum when the key is missing, so a registered slot is
401 written in place and never replaced.
402 """
403 if self._runtime_local_tensor(param.grad).layout != torch.strided:
404 raise ValueError("Muon swap optimizer only supports dense gradients.")
405 state = self.optimizer.state[param]
406 for key in self._configured_state_keys():
407 if (id(param), key) in self._slots:
408 continue
409 device_tensor = self.runtime.make_empty_state_like(param)
410 if not self._runtime_local_tensor(device_tensor).is_contiguous():
411 raise ValueError("Muon swap optimizer requires contiguous momentum state.")
412 state[key] = device_tensor
413 slot = self._make_slot(key, device_tensor)
414 if not slot.swappable:
415 # Too small or otherwise ineligible: keep ordinary device state.
416 device_tensor.zero_()
417 self._slots[(id(param), key)] = slot
418 continue
419 slot.cpu_tensor = self.runtime.make_zero_cpu_tensor_like(param)
420 slot.state = "host"
421 self._slots[(id(param), key)] = slot
422 self.runtime.release_device_storage(slot)
424 def _slots_for_params(self, params: Iterable[torch.nn.Parameter]) -> List[SwapSlot]:
425 """Return the registered momentum slots for ``params``, in order.
427 Only slots the adapter already owns are returned: they carry the swap
428 state (host mirror, released device storage, pending event) that the
429 runtime manipulates, so rebuilding a slot here would drop that state.
430 Momentum state materialized before the wrapper was created is registered
431 separately by :meth:`initial_slots`.
432 """
433 keys = self._configured_state_keys()
434 slots = []
435 seen = set()
436 for param in params:
437 for key in keys:
438 slot = self._slots.get((id(param), key))
439 if slot is None or id(slot) in seen:
440 continue
441 slots.append(slot)
442 seen.add(id(slot))
443 return slots
445 # ------------------------------------------------------------------ execution
446 def iter_update_units(self, step_context: Dict[str, Any]) -> List[MuonSwapUnit]:
447 """Return the units built by :meth:`prepare_step`."""
448 return step_context["units"]
450 def partition_costs(self, units: Sequence[MuonSwapUnit], step_context: Any) -> Sequence[int]:
451 """Return logical momentum bytes without consulting local slots or gradients."""
452 del step_context
453 return [unit.logical_nbytes for unit in units]
455 def step_batch(self, batch: List[MuonSwapUnit], step_context: Dict[str, Any]) -> None:
456 """Run one prefetched batch of Muon update units.
458 The runtime has already waited for this batch's H2D before calling in,
459 so every slot the batch reads is device-resident here.
460 """
461 broadcaster = step_context["broadcaster"]
462 for unit in batch:
463 self._run_unit(unit, broadcaster)
465 def _run_unit(self, unit: MuonSwapUnit, broadcaster: AsyncReplicateBroadcaster) -> None:
466 """Run one atomic update unit with the matching Muon method."""
467 kind = unit.kind
468 if kind not in (_NO_COMM, _HSDP):
469 raise ValueError(f"Unknown Muon swap unit kind {kind!r}.")
470 group = self.optimizer.param_groups[unit.group_index]
471 ns_inputs = self.optimizer._update_muon_momentum(group, unit.params)
472 if kind == _NO_COMM:
473 self._apply_unshard_update(group, ns_inputs)
474 else:
475 hsdp_assign = unit.hsdp_assign
476 if hsdp_assign is None:
477 raise RuntimeError("Muon swap HSDP unit is missing its group assignment.")
478 if hsdp_assign.is_shard:
479 self.optimizer._process_shard_params(
480 group, ns_inputs, [hsdp_assign], unit.group_index,
481 buffer_cache={},
482 )
483 else:
484 self._apply_unshard_update(group, ns_inputs)
485 if self.optimizer.post_update_fn is not None and unit.post_update_params:
486 self.optimizer._run_post_update_fn(group, unit.post_update_params)
487 # Flush at the global assignment boundary, after owner callbacks, even
488 # when this rank has no local updates in the last logical unit.
489 if unit.flushes_broadcast:
490 broadcaster.flush_group(unit.hsdp_assign)
492 def _apply_unshard_update(
493 self,
494 group: Dict[str, Any],
495 ns_inputs: Dict[torch.nn.Parameter, torch.Tensor],
496 ) -> None:
497 """Run the unsharded NS and parameter update for one unit."""
498 if not ns_inputs:
499 return
500 # Muon's own helper re-groups by shape and re-splits into memory-safe
501 # batches; a unit's params are already one such batch, so the split is a
502 # single group and the boundary cannot drift.
503 self.optimizer._process_unshard_params(group, ns_inputs)
505 def finish_step(self, step_context: Any) -> None:
506 """Wait for every parameter broadcast issued during this step."""
507 step_context["broadcaster"].wait_all()
510def build_muon_swap_adapter(optimizer: Any, config: Any, runtime: Any) -> MuonSwapAdapter:
511 """Build the Muon swap adapter for ``optimizer``.
513 Args:
514 optimizer: Muon optimizer instance.
515 config: Resolved ``SwapOptimizerConfig``.
516 runtime: Shared ``PipelineSwapRuntime``.
518 Returns:
519 The Muon swap adapter.
521 Raises:
522 ValueError: If ``optimizer`` is not a Muon optimizer.
523 """
524 if not isinstance(optimizer, Muon):
525 raise ValueError(
526 "Swap muon only supports hyper_parallel.core.optimizer.muon.Muon. "
527 f"Got {type(optimizer)!r}."
528 )
529 return MuonSwapAdapter(optimizer, config, runtime)
532def swap_muon(optimizer: Any, config: Any) -> Any:
533 """Wrap a Muon optimizer with momentum-only state swap.
535 Args:
536 optimizer: Muon optimizer instance.
537 config: Resolved ``SwapOptimizerConfig``. Packed swap preserves Muon's
538 logical units and only coalesces their momentum transfers.
540 Returns:
541 The shared core swap-optimizer wrapper.
542 """
543 return CoreSwapOptimizer(
544 optimizer,
545 config,
546 adapter_factory=build_muon_swap_adapter,
547 )