Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / optimizer / swap_adam.py: 88%
190 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"""Adam/AdamW optimizer-state swap adapters."""
17# pylint: disable=protected-access
18# pylint: disable=forbidden-backend-import
20from __future__ import annotations
22import inspect
23from collections import defaultdict
24from dataclasses import dataclass
25from typing import Any, Callable, Dict, List, NamedTuple, Sequence, Tuple
27import torch
29from hyper_parallel.core.optimizer.swap_optimizer_base import (
30 StateSwapAdapter,
31 SwapOptimizer as CoreSwapOptimizer,
32 SwapSlot,
33 _slot_tensor,
34)
36#: Logical state keys an Adam/AdamW optimizer owns.
37ADAM_STATE_KEYS = ("exp_avg", "exp_avg_sq", "max_exp_avg_sq")
40class GroupArgs(NamedTuple):
41 """Per-parameter argument lists one functional Adam/AdamW step consumes."""
43 params: List[Any]
44 grads: List[Any]
45 exp_avgs: List[Any]
46 exp_avg_sqs: List[Any]
47 max_exp_avg_sqs: List[Any]
48 state_steps: List[Any]
51@dataclass
52class AdamUpdateUnit:
53 """Per-parameter Adam/AdamW update unit consumed by the pipeline runtime.
55 ``adapter_index`` identifies the Torch parameter group owning ``param``.
56 """
58 adapter_index: int
59 param: Any
60 grad: Any
61 slots: List[SwapSlot]
64class AdamSwapAdapter(StateSwapAdapter):
65 """Adam/AdamW optimizer-state swap adapter."""
67 functional_name = "adam"
68 decoupled_weight_decay = False
69 is_new_adamw = False
70 supports_fused = False
72 def validate(self) -> None:
73 """Validate unsupported optimizer flags."""
74 for group in self.optimizer.param_groups:
75 if group.get("foreach", False) is True:
76 raise ValueError("Swap optimizer does not support foreach=True.")
77 if group.get("fused", False) is True and not self.supports_fused:
78 raise ValueError("Swap optimizer does not support fused=True.")
79 if group.get("differentiable", False):
80 raise ValueError("Swap optimizer does not support differentiable=True.")
81 if group.get("capturable", False):
82 raise ValueError("Swap optimizer does not support capturable=True.")
84 def prepare_step(self, *args: Any, **kwargs: Any) -> Dict[str, Any]:
85 """Initialize lazy state and collect this step's update units."""
86 if args or kwargs:
87 raise ValueError("Torch swap optimizer step does not support closure or extra arguments.")
88 if self.runtime.packed_enabled:
89 return self._prepare_packed_step()
91 units = []
92 for group_index, group in enumerate(self.optimizer.param_groups):
93 if self.is_new_adamw:
94 group["step"] = (group.get("step") or 0) + 1
95 for param in group["params"]:
96 grad = getattr(param, "grad", None)
97 if grad is None:
98 continue
99 if getattr(grad, "is_sparse", False):
100 raise ValueError("Swap optimizer only supports dense Adam/AdamW gradients.")
101 state = self.optimizer.state[param]
102 self._init_param_state(param, group)
103 slots = self._build_slots(param, state)
104 units.append(AdamUpdateUnit(
105 adapter_index=group_index,
106 param=param,
107 grad=grad,
108 slots=slots,
109 ))
110 return {"units": units}
112 def _prepare_packed_step(self) -> Dict[str, Any]:
113 """Build a stable packed layout while retaining inactive materialized states."""
114 records = []
115 for group_index, group in enumerate(self.optimizer.param_groups):
116 if self.is_new_adamw:
117 group["step"] = (group.get("step") or 0) + 1
118 for param in group["params"]:
119 grad = getattr(param, "grad", None)
120 state = self.optimizer.state.get(param)
121 if grad is not None:
122 if getattr(grad, "is_sparse", False):
123 raise ValueError("Swap optimizer only supports dense Adam/AdamW gradients.")
124 state = self.optimizer.state[param]
125 self._init_param_state(param, group)
126 if state:
127 self._register_present_slots(param, state)
128 has_slots = any((id(param), key) in self._slots for key in self._configured_state_keys())
129 if grad is None and not has_slots:
130 continue
131 records.append((group_index, param, grad))
133 self.runtime.prepare_packed_host(self._ordered_slots())
134 self.publish_packed_state()
135 units = []
136 for group_index, param, grad in records:
137 state = self.optimizer.state[param]
138 slots = self._build_slots(param, state)
139 if grad is None and not any(slot.swappable and slot.packed for slot in slots):
140 continue
141 units.append(AdamUpdateUnit(
142 adapter_index=group_index,
143 param=param,
144 grad=grad,
145 slots=slots,
146 ))
147 return {"units": units}
149 def _collect_group_args(
150 self,
151 units: Sequence[AdamUpdateUnit],
152 group: Dict[str, Any],
153 ) -> GroupArgs:
154 """Gather the per-parameter argument lists one optimizer step needs.
156 Units without a gradient are skipped. ``max_exp_avg_sqs`` stays empty
157 unless the group is amsgrad, and ``state_steps`` holds ``None`` for new
158 AdamW, which advances the step counter itself.
159 """
160 args = GroupArgs([], [], [], [], [], [])
161 for unit in units:
162 if unit.grad is None:
163 continue
164 state = self.optimizer.state[unit.param]
165 args.params.append(unit.param)
166 args.grads.append(unit.grad)
167 args.exp_avgs.append(_slot_tensor(unit, "exp_avg", state["exp_avg"]))
168 args.exp_avg_sqs.append(_slot_tensor(unit, "exp_avg_sq", state["exp_avg_sq"]))
169 if group.get("amsgrad", False):
170 args.max_exp_avg_sqs.append(
171 _slot_tensor(unit, "max_exp_avg_sq", state["max_exp_avg_sq"])
172 )
173 if self.is_new_adamw:
174 args.state_steps.append(None)
175 else:
176 args.state_steps.append(state["step"])
177 return args
179 def step_batch(self, batch: List[AdamUpdateUnit], step_context: Dict[str, Any]) -> None:
180 """Run Torch functional Adam/AdamW for one batch."""
181 del step_context
182 by_group: Dict[int, List[AdamUpdateUnit]] = defaultdict(list)
183 for unit in batch:
184 by_group[unit.adapter_index].append(unit)
185 for group_index, units in by_group.items():
186 self._step_group(self.optimizer.param_groups[group_index], units)
188 def _step_group(self, group: Dict[str, Any], units: List[AdamUpdateUnit]) -> None:
189 """Run one group's parameters through the matching functional Adam/AdamW."""
190 args = self._collect_group_args(units, group)
191 params = args.params
193 if not params:
194 return
196 if self.is_new_adamw:
197 self._step_new_adamw(
198 group,
199 args.params,
200 args.grads,
201 args.exp_avgs,
202 args.exp_avg_sqs,
203 args.max_exp_avg_sqs,
204 )
205 return
207 func = getattr(torch.optim._functional, self.functional_name)
208 kwargs = {
209 "amsgrad": group["amsgrad"],
210 "beta1": group["betas"][0],
211 "beta2": group["betas"][1],
212 "lr": group["lr"],
213 "weight_decay": group["weight_decay"],
214 "eps": group["eps"],
215 "maximize": group["maximize"],
216 "foreach": False,
217 "capturable": False,
218 "differentiable": False,
219 "fused": bool(group.get("fused", False)),
220 "grad_scale": getattr(self.optimizer, "grad_scale", None),
221 "found_inf": getattr(self.optimizer, "found_inf", None),
222 "has_complex": False,
223 }
224 if self.functional_name == "adam":
225 if "decoupled_weight_decay" in inspect.signature(func).parameters:
226 kwargs["decoupled_weight_decay"] = self._decoupled_weight_decay(group)
227 func(
228 args.params,
229 args.grads,
230 args.exp_avgs,
231 args.exp_avg_sqs,
232 args.max_exp_avg_sqs,
233 args.state_steps,
234 **kwargs,
235 )
237 def _decoupled_weight_decay(self, group: Dict[str, Any]) -> bool:
238 """Resolve the decoupled weight decay flag for one parameter group."""
239 return self.decoupled_weight_decay or group.get("decoupled_weight_decay", False)
241 def _step_new_adamw(
242 self,
243 group: Dict[str, Any],
244 params: Sequence[Any],
245 grads: Sequence[Any],
246 exp_avgs: Sequence[Any],
247 exp_avg_sqs: Sequence[Any],
248 max_exp_avg_sqs: Sequence[Any],
249 ) -> None:
250 """Run the new AdamW functional for one group."""
251 if params and params[0].device.type == "cpu":
252 # torch.optim._functional.adamw increments tensor state_steps
253 # internally. New AdamW already advanced group["step"] in
254 # prepare_step(), so feed step - 1 to preserve outer-step
255 # semantics for CPU-only tests.
256 step_tensor = torch.tensor(float(group["step"] - 1), dtype=torch.float32)
257 torch.optim._functional.adamw(
258 params,
259 grads,
260 exp_avgs,
261 exp_avg_sqs,
262 max_exp_avg_sqs,
263 [step_tensor] * len(params),
264 amsgrad=group["amsgrad"],
265 beta1=group["betas"][0],
266 beta2=group["betas"][1],
267 lr=group["lr"],
268 weight_decay=group["weight_decay"],
269 eps=group["eps"],
270 maximize=group["maximize"],
271 foreach=False,
272 capturable=False,
273 differentiable=False,
274 fused=False,
275 grad_scale=None,
276 found_inf=None,
277 has_complex=False,
278 )
279 return
280 _new_adamw_func()(
281 params,
282 grads,
283 exp_avgs,
284 exp_avg_sqs,
285 max_exp_avg_sqs,
286 group["step"],
287 amsgrad=group["amsgrad"],
288 beta1=group["betas"][0],
289 beta2=group["betas"][1],
290 lr=group["lr"],
291 weight_decay=group["weight_decay"],
292 eps=group["eps"],
293 maximize=group["maximize"],
294 )
296 def _init_param_state(self, param: Any, group: Dict[str, Any]) -> None:
297 """Initialize missing Adam state and swap slots for one parameter."""
298 state = self.optimizer.state[param]
299 if not self.is_new_adamw and len(state) == 0:
300 step_device = (
301 param.device
302 if group.get("fused", False)
303 else ("cpu" if self.runtime.packed_enabled else param.device)
304 )
305 state["step"] = torch.zeros((), dtype=torch.float32, device=step_device)
306 state_keys = ["exp_avg", "exp_avg_sq"]
307 if group.get("amsgrad", False):
308 state_keys.append("max_exp_avg_sq")
309 configured_keys = set(self._configured_state_keys())
310 for key in state_keys:
311 if key in state or (id(param), key) in self._slots:
312 continue
313 if (
314 key in configured_keys
315 and not self.runtime.packed_enabled
316 and self.runtime.is_swappable_tensor(param, self.config.min_numel)
317 ):
318 cpu_tensor = self.runtime.make_zero_cpu_tensor_like(param)
319 device_tensor = torch.empty_like(param, memory_format=torch.preserve_format)
320 state[key] = device_tensor
321 slot = self._make_slot(key, device_tensor)
322 slot.cpu_tensor = cpu_tensor
323 slot.state = "host"
324 self._slots[(id(param), key)] = slot
325 self.runtime.release_device_storage(slot)
326 continue
327 if key in configured_keys and self.runtime.is_packable_template(param, self.config.min_numel):
328 self._slots[(id(param), key)] = self._make_slot(key, None, template=param)
329 continue
330 state[key] = torch.zeros_like(param, memory_format=torch.preserve_format)
332 @staticmethod
333 def default_state_keys() -> Tuple[str, ...]:
334 """Return the swap state keys an Adam/AdamW optimizer owns."""
335 return ADAM_STATE_KEYS
338class TorchNativeAdamAdapter(AdamSwapAdapter):
339 """Adapter for ``torch.optim.Adam``."""
341 functional_name = "adam"
343 @classmethod
344 def matches(cls, optimizer: Any) -> bool:
345 """Match ``torch.optim.Adam``, excluding its AdamW subclass."""
346 # AdamW inherits Adam in PyTorch. Keep Adam subclasses supported, but
347 # let AdamW select its dedicated adapter (which preserves fused=True).
348 return (
349 isinstance(optimizer, torch.optim.Adam)
350 and not isinstance(optimizer, torch.optim.AdamW)
351 )
354class TorchNativeAdamWAdapter(AdamSwapAdapter):
355 """Adapter for ``torch.optim.AdamW``."""
357 functional_name = "adamw"
358 supports_fused = True
360 @classmethod
361 def matches(cls, optimizer: Any) -> bool:
362 """Match ``torch.optim.AdamW``."""
363 return isinstance(optimizer, torch.optim.AdamW)
366class TorchNewAdamWAdapter(AdamSwapAdapter):
367 """Adapter for hyper-parallel's fused AdamW."""
369 functional_name = "adamw"
370 is_new_adamw = True
372 @classmethod
373 def matches(cls, optimizer: Any) -> bool:
374 """Match hyper-parallel's lazily imported AdamW."""
375 return isinstance(optimizer, _new_adamw_cls())
377def _new_adamw_cls() -> type:
378 """Return hyper-parallel's AdamW class without importing it at module load."""
379 from hyper_parallel.core.optimizer.adamw import ( # pylint: disable=import-outside-toplevel
380 AdamW as new_adamw_cls,
381 )
382 return new_adamw_cls
385def _new_adamw_func() -> Callable[..., Any]:
386 """Return hyper-parallel's functional AdamW without importing it at module load."""
387 from hyper_parallel.core.optimizer.adamw import ( # pylint: disable=import-outside-toplevel
388 adamw as new_adamw_func,
389 )
390 return new_adamw_func
393ADAM_ADAPTERS = (
394 TorchNewAdamWAdapter,
395 TorchNativeAdamAdapter,
396 TorchNativeAdamWAdapter,
397)
400def build_adam_swap_adapter(optimizer: Any, config: Any, runtime: Any) -> AdamSwapAdapter:
401 """Build the Adam/AdamW swap adapter matching ``optimizer``.
403 Args:
404 optimizer: Wrapped Torch or HyperParallel Adam/AdamW optimizer.
405 config: Resolved ``SwapOptimizerConfig``.
406 runtime: Shared ``PipelineSwapRuntime``.
408 Returns:
409 The adapter instance for ``optimizer``.
411 Raises:
412 ValueError: If ``optimizer`` is not a supported Adam/AdamW optimizer.
413 """
414 for adapter_cls in ADAM_ADAPTERS:
415 if adapter_cls.matches(optimizer):
416 return adapter_cls(optimizer, config, runtime)
417 raise ValueError(
418 "Swap optimizer only supports torch.optim.Adam, torch.optim.AdamW, "
419 "and hyper_parallel.core.optimizer.adamw.AdamW on the Torch backend. "
420 f"Got {type(optimizer)!r}."
421 )
424def swap_adam(optimizer: Any, config: Any) -> Any:
425 """Wrap an Adam/AdamW optimizer with optimizer-state swap.
427 Args:
428 optimizer: Supported Adam/AdamW optimizer instance.
429 config: Resolved ``SwapOptimizerConfig``.
431 Returns:
432 The shared core swap-optimizer wrapper.
433 """
434 return CoreSwapOptimizer(
435 optimizer,
436 config,
437 adapter_factory=build_adam_swap_adapter,
438 )