Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / torch / swap_optimizer / adapters.py: 91%
365 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-21 04:29 +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"""Torch Adam/AdamW swap optimizer adapters."""
16# pylint: disable=protected-access
18from __future__ import annotations
20import copy
21from collections import defaultdict
22from typing import Any, Dict, Iterable, List, Optional, Tuple
24import torch
26from hyper_parallel.core.optimizer.adamw import AdamW as HyperAdamW
27from hyper_parallel.core.optimizer.adamw import adamw as hyper_adamw
28from hyper_parallel.core.optimizer.swap_optimizer_base import (
29 OptimizerSwapAdapter,
30 SwapSlot,
31 UpdateUnit,
32)
35class TorchAdamBaseAdapter(OptimizerSwapAdapter):
36 """Common Torch Adam/AdamW adapter logic."""
38 functional_name = "adam"
39 supported_cls = ()
40 decoupled_weight_decay = False
41 is_hyper_adamw = False
42 supports_fused = False
44 @classmethod
45 def matches(cls, optimizer: Any) -> bool:
46 """Return whether this adapter supports ``optimizer``."""
47 return isinstance(optimizer, cls.supported_cls)
49 def __init__(self, optimizer: Any, config: Any, runtime: Any) -> None:
50 super().__init__(optimizer, config, runtime)
51 self._slots: Dict[Tuple[int, str], SwapSlot] = {}
53 def validate(self) -> None:
54 """Validate unsupported optimizer flags."""
55 for group in self.optimizer.param_groups:
56 if group.get("foreach", False) is True:
57 raise ValueError("Swap optimizer does not support foreach=True.")
58 if group.get("fused", False) is True and not self.supports_fused:
59 raise ValueError("Swap optimizer does not support fused=True.")
60 if group.get("differentiable", False):
61 raise ValueError("Swap optimizer does not support differentiable=True.")
62 if group.get("capturable", False):
63 raise ValueError("Swap optimizer does not support capturable=True.")
65 def prepare_step(self, *args: Any, **kwargs: Any) -> Dict[str, Any]:
66 """Initialize lazy state and collect this step's update units."""
67 if args or kwargs:
68 raise ValueError("Torch swap optimizer step does not support closure or extra arguments.")
69 if self.runtime.packed_enabled:
70 return self._prepare_packed_step()
72 units = []
73 for group_index, group in enumerate(self.optimizer.param_groups):
74 if self.is_hyper_adamw:
75 group["step"] = (group.get("step") or 0) + 1
76 for param in group["params"]:
77 grad = getattr(param, "grad", None)
78 if grad is None:
79 continue
80 if getattr(grad, "is_sparse", False):
81 raise ValueError("Swap optimizer only supports dense Adam/AdamW gradients.")
82 state = self.optimizer.state[param]
83 self._init_param_state(param, grad, group)
84 slots = self._build_slots(param, state)
85 units.append(UpdateUnit(
86 adapter_index=group_index,
87 param=param,
88 grad=grad,
89 slots=slots,
90 ))
91 return {"units": units}
93 def _prepare_packed_step(self) -> Dict[str, Any]:
94 """Build a stable packed layout while retaining inactive materialized states."""
95 records = []
96 for group_index, group in enumerate(self.optimizer.param_groups):
97 if self.is_hyper_adamw:
98 group["step"] = (group.get("step") or 0) + 1
99 for param in group["params"]:
100 grad = getattr(param, "grad", None)
101 state = self.optimizer.state.get(param)
102 if grad is not None:
103 if getattr(grad, "is_sparse", False):
104 raise ValueError("Swap optimizer only supports dense Adam/AdamW gradients.")
105 state = self.optimizer.state[param]
106 self._init_param_state(param, grad, group)
107 if state:
108 self._register_present_slots(param, state)
109 has_slots = any((id(param), key) in self._slots for key in self._configured_state_keys())
110 if grad is None and not has_slots:
111 continue
112 records.append((group_index, param, grad))
114 self.runtime.prepare_packed_host(self._ordered_slots())
115 self.publish_packed_state()
116 units = []
117 for group_index, param, grad in records:
118 state = self.optimizer.state[param]
119 slots = self._build_slots(param, state)
120 if grad is None and not any(slot.swappable and slot.packed for slot in slots):
121 continue
122 units.append(UpdateUnit(
123 adapter_index=group_index,
124 param=param,
125 grad=grad,
126 slots=slots,
127 ))
128 return {"units": units}
130 def iter_update_units(self, step_context: Dict[str, Any]) -> List[UpdateUnit]:
131 """Return units collected in ``prepare_step``."""
132 return step_context["units"]
134 def initial_slots(self) -> Iterable[SwapSlot]:
135 """Discover optimizer states materialized before the swap wrapper was created."""
136 slots = []
137 for group in self.optimizer.param_groups:
138 for param in group["params"]:
139 state = self.optimizer.state.get(param)
140 if state:
141 slots.extend(self._build_slots(param, state))
142 return tuple(slots)
144 def step_batch(self, batch: List[UpdateUnit], step_context: Dict[str, Any]) -> None:
145 """Run Torch functional Adam/AdamW for one batch."""
146 del step_context
147 by_group: Dict[int, List[UpdateUnit]] = defaultdict(list)
148 for unit in batch:
149 by_group[unit.adapter_index].append(unit)
150 for group_index, units in by_group.items():
151 group = self.optimizer.param_groups[group_index]
152 state_steps = []
153 params = []
154 grads = []
155 exp_avgs = []
156 exp_avg_sqs = []
157 max_exp_avg_sqs = []
158 for unit in units:
159 if unit.grad is None:
160 continue
161 state = self.optimizer.state[unit.param]
162 params.append(unit.param)
163 grads.append(unit.grad)
164 exp_avgs.append(self._slot_tensor(unit, "exp_avg", state["exp_avg"]))
165 exp_avg_sqs.append(self._slot_tensor(unit, "exp_avg_sq", state["exp_avg_sq"]))
166 if group.get("amsgrad", False):
167 max_exp_avg_sqs.append(
168 self._slot_tensor(unit, "max_exp_avg_sq", state["max_exp_avg_sq"])
169 )
170 if self.is_hyper_adamw:
171 state_steps.append(None)
172 else:
173 state_steps.append(state["step"])
175 if not params:
176 continue
178 if self.is_hyper_adamw:
179 if params and params[0].device.type == "cpu":
180 # torch.optim._functional.adamw increments tensor state_steps
181 # internally. Hyper AdamW already advanced group["step"] in
182 # prepare_step(), so feed step - 1 to preserve outer-step
183 # semantics for CPU-only tests.
184 step_tensor = torch.tensor(float(group["step"] - 1), dtype=torch.float32)
185 torch.optim._functional.adamw(
186 params,
187 grads,
188 exp_avgs,
189 exp_avg_sqs,
190 max_exp_avg_sqs,
191 [step_tensor] * len(params),
192 amsgrad=group["amsgrad"],
193 beta1=group["betas"][0],
194 beta2=group["betas"][1],
195 lr=group["lr"],
196 weight_decay=group["weight_decay"],
197 eps=group["eps"],
198 maximize=group["maximize"],
199 foreach=False,
200 capturable=False,
201 differentiable=False,
202 fused=False,
203 grad_scale=None,
204 found_inf=None,
205 has_complex=False,
206 )
207 else:
208 hyper_adamw(
209 params,
210 grads,
211 exp_avgs,
212 exp_avg_sqs,
213 max_exp_avg_sqs,
214 group["step"],
215 amsgrad=group["amsgrad"],
216 beta1=group["betas"][0],
217 beta2=group["betas"][1],
218 lr=group["lr"],
219 weight_decay=group["weight_decay"],
220 eps=group["eps"],
221 maximize=group["maximize"],
222 )
223 continue
225 func = getattr(torch.optim._functional, self.functional_name)
226 kwargs = {
227 "amsgrad": group["amsgrad"],
228 "beta1": group["betas"][0],
229 "beta2": group["betas"][1],
230 "lr": group["lr"],
231 "weight_decay": group["weight_decay"],
232 "eps": group["eps"],
233 "maximize": group["maximize"],
234 "foreach": False,
235 "capturable": False,
236 "differentiable": False,
237 "fused": bool(group.get("fused", False)),
238 "grad_scale": getattr(self.optimizer, "grad_scale", None),
239 "found_inf": getattr(self.optimizer, "found_inf", None),
240 "has_complex": False,
241 }
242 if self.functional_name == "adam":
243 kwargs["decoupled_weight_decay"] = self.decoupled_weight_decay or group.get(
244 "decoupled_weight_decay", False
245 )
246 func(params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, **kwargs)
248 def all_slots(self):
249 """Iterate known swap slots."""
250 return tuple(self._ordered_slots())
252 def checkpoint_state_dict(self, *args: Any, **kwargs: Any) -> Dict[str, Any]:
253 """Return optimizer checkpoint state using CPU mirrors for swapped slots."""
254 del args, kwargs
255 self.runtime.synchronize_cpu_mirrors(self.all_slots())
256 return self.export_swappable_state(self.optimizer.state_dict())
258 def load_checkpoint_state_dict(
259 self,
260 state_dict: Dict[str, Any],
261 *args: Any,
262 **kwargs: Any,
263 ) -> None:
264 """Load optimizer checkpoint state while restoring swap-managed slots."""
265 del args, kwargs
266 stripped, removed = self.strip_swappable_state(state_dict)
267 self.optimizer.load_state_dict(stripped)
268 self.load_swappable_state(state_dict, removed)
270 def publish_packed_state(self) -> None:
271 """Publish persistent packed CPU mirrors to the wrapped optimizer state."""
272 if not self.runtime.packed_enabled:
273 return
274 for group in self.optimizer.param_groups:
275 for param in group["params"]:
276 state = self.optimizer.state.get(param)
277 for key in self._configured_state_keys():
278 slot = self._slots.get((id(param), key))
279 if slot is None or not slot.packed or slot.cpu_tensor is None:
280 continue
281 if state is None:
282 state = self.optimizer.state[param]
283 state[key] = slot.cpu_tensor
285 def export_swappable_state(self, state_dict: Dict[str, Any]) -> Dict[str, Any]:
286 """Build a checkpoint-safe Torch optimizer state dict.
288 Torch optimizer state dicts are keyed by saved parameter ids, while the
289 adapter tracks live swap slots by the current parameter objects. This
290 method walks both orders together and exports each parameter's optimizer
291 state with the data source that currently owns the valid tensor values.
293 If an Adam state tensor such as ``exp_avg`` or ``exp_avg_sq`` has been
294 offloaded, the live device tensor may only be a placeholder with its
295 storage released. In that case, write a cloned CPU mirror into the
296 exported state dict so checkpoints contain the real optimizer values.
297 Non-swappable state, metadata, and tensors that are still resident on
298 device are deep-copied from the original Torch state dict unchanged.
299 """
300 exported = {
301 key: copy.deepcopy(value)
302 for key, value in state_dict.items()
303 if key not in ("state", "param_groups")
304 }
305 exported["param_groups"] = copy.deepcopy(state_dict.get("param_groups", []))
306 exported["state"] = {}
307 saved_groups = exported.get("param_groups", [])
308 params_in_order = []
309 for group in self.optimizer.param_groups:
310 params_in_order.extend(group["params"])
311 ids_in_order = []
312 for group in saved_groups:
313 ids_in_order.extend(group["params"])
314 for param, param_id in zip(params_in_order, ids_in_order):
315 saved_state = state_dict.get("state", {}).get(param_id)
316 if not saved_state:
317 continue
318 exported_state = {}
319 for key in self._state_keys_for_param(param):
320 slot = self._slots.get((id(param), key))
321 if slot is not None and slot.state == "host" and slot.cpu_tensor is not None and key in saved_state:
322 exported_state[key] = slot.cpu_tensor.detach().clone()
323 for key, value in saved_state.items():
324 if key not in exported_state:
325 exported_state[key] = copy.deepcopy(value)
326 exported["state"][param_id] = exported_state
327 return exported
329 def strip_swappable_state(self, state_dict: Dict[str, Any]) -> Tuple[Dict[str, Any], Dict[int, Dict[str, Any]]]:
330 """Split Adam state tensors out before delegating to Torch loading.
332 PyTorch's ``load_state_dict`` eagerly restores tensors into the
333 optimizer state. For swap-managed Adam buffers, that would bypass the
334 adapter/runtime bookkeeping and can place large tensors directly on the
335 device. This method therefore deep-copies the checkpoint state dict,
336 removes the Adam buffers that may be swap-managed, and returns them in a
337 side table keyed by the checkpoint parameter id.
339 The stripped state dict is safe to pass to the wrapped optimizer's
340 ``load_state_dict`` for ordinary fields such as parameter groups and
341 step counters. The removed tensors must be handed to
342 ``load_swappable_state`` afterwards so they can be restored with the
343 correct CPU mirror/device placeholder layout.
344 """
345 stripped = copy.deepcopy(state_dict)
346 removed: Dict[int, Dict[str, Any]] = {}
347 swappable_keys = self._configured_state_keys()
348 for param_id, saved_state in list(stripped.get("state", {}).items()):
349 if not isinstance(saved_state, dict):
350 continue
351 for key in swappable_keys:
352 if key in saved_state:
353 removed.setdefault(param_id, {})[key] = saved_state.pop(key)
354 return stripped, removed
356 def load_swappable_state(self, original_state_dict: Dict[str, Any], removed: Dict[int, Dict[str, Any]]) -> None:
357 """Restore removed Adam buffers under swap runtime control.
359 This is the second half of checkpoint loading. After Torch has loaded
360 the stripped state dict, this method maps checkpoint parameter ids back
361 to the current parameter objects by walking saved and current parameter
362 groups in order. Each removed Adam buffer is then recreated as an
363 optimizer state entry and registered as a ``SwapSlot``.
365 Packed runtimes place checkpoint values directly in persistent pinned
366 host views. Legacy runtimes retain an empty device placeholder whose
367 storage is restored only during prefetch. Buffers that do not meet the
368 runtime's swappability criteria are materialized directly on the
369 parameter's device and tracked as normal device-resident slots.
370 """
371 saved_groups = original_state_dict.get("param_groups", [])
372 current_groups = self.optimizer.param_groups
373 self._slots = {}
374 saved_ids = []
375 current_params = []
376 for saved_group, current_group in zip(saved_groups, current_groups):
377 saved_ids.extend(saved_group["params"])
378 current_params.extend(current_group["params"])
379 for saved_id, param in zip(saved_ids, current_params):
380 key_to_tensor = removed.get(saved_id, {})
381 if not key_to_tensor:
382 continue
383 state = self.optimizer.state[param]
384 for key, saved_tensor in key_to_tensor.items():
385 cpu_tensor = self._cast_swappable_tensor_to_cpu(param, saved_tensor)
386 if self.runtime.packed_enabled and self.runtime.is_packable_template(param, self.config.min_numel):
387 if self.runtime.is_distributed_tensor(param):
388 logical_tensor = torch.zeros_like(
389 param,
390 memory_format=torch.preserve_format,
391 )
392 slot = self._make_slot(key, logical_tensor)
393 state[key] = logical_tensor
394 self.runtime.release_device_storage(slot)
395 else:
396 slot = self._make_slot(key, None, template=param)
397 state[key] = cpu_tensor
398 slot.tensor = cpu_tensor
399 slot.cpu_tensor = cpu_tensor
400 slot.state = "host"
401 self._slots[(id(param), key)] = slot
402 continue
403 device_tensor = self.runtime.make_empty_device_tensor_like(param, cpu_tensor)
404 slot = self._make_slot(key, device_tensor)
405 if slot.swappable:
406 state[key] = device_tensor
407 slot.cpu_tensor = self.runtime.make_cpu_tensor(cpu_tensor)
408 slot.state = "host"
409 self._slots[(id(param), key)] = slot
410 self.runtime.release_device_storage(slot)
411 else:
412 device_tensor = self._cast_state_tensor_like_torch(
413 param,
414 saved_tensor,
415 saved_id,
416 saved_groups,
417 key,
418 )
419 state[key] = device_tensor
420 self._slots[(id(param), key)] = self._make_slot(key, device_tensor)
421 if self.runtime.packed_enabled:
422 self.runtime.prepare_packed_host(self._ordered_slots())
423 self.publish_packed_state()
425 def _init_param_state(self, param: Any, grad: Any, group: Dict[str, Any]) -> None:
426 """Initialize missing Adam state and swap slots for one parameter."""
427 del grad
428 state = self.optimizer.state[param]
429 if not self.is_hyper_adamw and len(state) == 0:
430 step_device = (
431 param.device
432 if group.get("fused", False)
433 else ("cpu" if self.runtime.packed_enabled else param.device)
434 )
435 state["step"] = torch.zeros((), dtype=torch.float32, device=step_device)
436 state_keys = ["exp_avg", "exp_avg_sq"]
437 if group.get("amsgrad", False):
438 state_keys.append("max_exp_avg_sq")
439 configured_keys = set(self._configured_state_keys())
440 for key in state_keys:
441 if key in state or (id(param), key) in self._slots:
442 continue
443 if (
444 key in configured_keys
445 and not self.runtime.packed_enabled
446 and self.runtime.is_swappable_tensor(param, self.config.min_numel)
447 ):
448 cpu_tensor = self.runtime.make_zero_cpu_tensor_like(param)
449 device_tensor = torch.empty_like(param, memory_format=torch.preserve_format)
450 state[key] = device_tensor
451 slot = self._make_slot(key, device_tensor)
452 slot.cpu_tensor = cpu_tensor
453 slot.state = "host"
454 self._slots[(id(param), key)] = slot
455 self.runtime.release_device_storage(slot)
456 continue
457 if key in configured_keys and self.runtime.is_packable_template(param, self.config.min_numel):
458 self._slots[(id(param), key)] = self._make_slot(key, None, template=param)
459 continue
460 state[key] = torch.zeros_like(param, memory_format=torch.preserve_format)
462 def _register_present_slots(self, param: Any, state: Dict[str, Any]) -> None:
463 """Register configured state tensors that already exist in an optimizer state mapping."""
464 for key in self._configured_state_keys():
465 tensor = state.get(key)
466 if tensor is None or (id(param), key) in self._slots:
467 continue
468 self._slots[(id(param), key)] = self._make_slot(key, tensor)
470 def _build_slots(self, param: Any, state: Dict[str, Any]) -> List[SwapSlot]:
471 """Return swap slots associated with the current parameter state."""
472 slots = []
473 for key in self._state_keys_for_param(param):
474 tensor = state.get(key)
475 if tensor is None:
476 continue
477 slot = self._slots.get((id(param), key))
478 if slot is None:
479 slot = self._make_slot(key, tensor)
480 self._slots[(id(param), key)] = slot
481 elif slot.tensor is not tensor and slot.cpu_tensor is not tensor:
482 slot = self._make_slot(key, tensor)
483 self._slots[(id(param), key)] = slot
484 slots.append(slot)
485 return slots
487 def _make_slot(self, key: str, tensor: Any, template: Optional[Any] = None) -> SwapSlot:
488 """Create a swap slot for a state tensor or a packed-state template."""
489 metadata_tensor = tensor if tensor is not None else template
490 if metadata_tensor is None:
491 raise ValueError(f"Cannot build swap slot {key!r} without a tensor or template.")
492 if tensor is None:
493 swappable = self.runtime.is_packable_template(metadata_tensor, self.config.min_numel)
494 else:
495 swappable = self.runtime.is_swappable_tensor(tensor, self.config.min_numel)
496 packed = bool(self.runtime.packed_enabled and swappable)
497 slot = SwapSlot(
498 name=key,
499 tensor=tensor,
500 cpu_tensor=None,
501 swappable=swappable,
502 state="device" if tensor is not None else "pending",
503 packed=packed,
504 logical_tensor=tensor if packed and self.runtime.is_distributed_tensor(tensor) else None,
505 )
506 self.runtime.populate_slot_metadata(slot, metadata_tensor)
507 return slot
509 def _ordered_slots(self) -> List[SwapSlot]:
510 """Return slots in optimizer parameter and configured state-key order."""
511 slots = []
512 seen_slots = set()
513 for group in self.optimizer.param_groups:
514 for param in group["params"]:
515 for key in self._configured_state_keys():
516 slot = self._slots.get((id(param), key))
517 if slot is not None and id(slot) not in seen_slots:
518 slots.append(slot)
519 seen_slots.add(id(slot))
520 return slots
522 def _state_keys_for_param(self, param: Any) -> Tuple[str, ...]:
523 state = self.optimizer.state[param]
524 keys = self._configured_state_keys()
525 result = []
526 for key in keys:
527 if key in state:
528 result.append(key)
529 elif self.config.state_keys is not None:
530 raise ValueError(f"Requested state key '{key}' is not present for parameter.")
531 return tuple(result)
533 @staticmethod
534 def _slot_tensor(unit: UpdateUnit, key: str, fallback: Any) -> Any:
535 """Return an active swap slot tensor, or the optimizer state fallback."""
536 for slot in unit.slots:
537 if slot.name == key and slot.swappable and slot.state == "device" and slot.tensor is not None:
538 return slot.tensor
539 return fallback
541 def _configured_state_keys(self) -> Tuple[str, ...]:
542 """Return Adam state keys selected for swap by the current config."""
543 keys = self.config.state_keys or self._default_state_keys()
544 result = []
545 for key in keys:
546 if key == "master_param":
547 if self.config.state_keys is not None:
548 raise ValueError(f"Requested state key '{key}' is not available for {type(self.optimizer)!r}.")
549 continue
550 result.append(key)
551 return tuple(result)
553 def _cast_state_tensor_like_torch(
554 self,
555 param: Any,
556 saved_tensor: Any,
557 saved_id: int,
558 saved_groups: List[Dict[str, Any]],
559 key: str,
560 ) -> Any:
561 """Cast a loaded state tensor using PyTorch optimizer load semantics."""
562 if not isinstance(saved_tensor, torch.Tensor):
563 raise ValueError(f"Expected torch.Tensor in optimizer state, got {type(saved_tensor)!r}.")
564 process = getattr(torch.optim.Optimizer, "_process_value_according_to_param_policy", None)
565 if process is not None:
566 return process(param, saved_tensor, saved_id, saved_groups, key).detach().clone()
567 if key == "step":
568 return saved_tensor.detach().clone()
569 if param.is_floating_point():
570 return saved_tensor.detach().to(dtype=param.dtype, device=param.device).clone()
571 return saved_tensor.detach().to(device=param.device).clone()
573 def _cast_swappable_tensor_to_cpu(self, param: Any, saved_tensor: Any) -> Any:
574 """Cast swappable state dtype like PyTorch while keeping values on CPU."""
575 if not isinstance(saved_tensor, torch.Tensor):
576 raise ValueError(f"Expected torch.Tensor in optimizer state, got {type(saved_tensor)!r}.")
577 if param.is_floating_point():
578 return saved_tensor.detach().to(dtype=param.dtype, device="cpu")
579 return saved_tensor.detach().to(device="cpu")
581 @staticmethod
582 def _default_state_keys() -> Tuple[str, ...]:
583 return ("exp_avg", "exp_avg_sq", "max_exp_avg_sq")
586class TorchNativeAdamAdapter(TorchAdamBaseAdapter):
587 """Adapter for ``torch.optim.Adam``."""
589 functional_name = "adam"
591 @classmethod
592 def matches(cls, optimizer: Any) -> bool:
593 # AdamW inherits Adam in PyTorch. Keep Adam subclasses supported, but
594 # let AdamW select its dedicated adapter (which preserves fused=True).
595 return (
596 isinstance(optimizer, torch.optim.Adam)
597 and not isinstance(optimizer, torch.optim.AdamW)
598 )
601class TorchNativeAdamWAdapter(TorchAdamBaseAdapter):
602 """Adapter for ``torch.optim.AdamW``."""
604 functional_name = "adamw"
605 supports_fused = True
607 @classmethod
608 def matches(cls, optimizer: Any) -> bool:
609 return isinstance(optimizer, torch.optim.AdamW)
612class TorchHyperAdamWAdapter(TorchAdamBaseAdapter):
613 """Adapter for hyper-parallel's fused AdamW."""
615 functional_name = "adamw"
616 supported_cls = (HyperAdamW,)
617 is_hyper_adamw = True
619 @classmethod
620 def matches(cls, optimizer: Any) -> bool:
621 return isinstance(optimizer, HyperAdamW)