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

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"""Public API for optimizer state swap.""" 

16 

17from __future__ import annotations 

18 

19from dataclasses import dataclass 

20from typing import Any, Optional, Sequence 

21 

22from hyper_parallel.core.optimizer.swap_optimizer_base import validate_state_keys 

23 

24 

25@dataclass(frozen=True) 

26class SwapOptimizerConfig: 

27 """Configuration for optimizer state swap. 

28 

29 The runtime uses a fixed one-batch-ahead prefetch pipeline. 

30 

31 Args: 

32 swap_times: Number of pipeline partitions. 

33 state_keys: Logical optimizer state keys to swap. ``None`` uses the 

34 defaults of the adapter matching the wrapped optimizer; an explicit 

35 value must only name keys that optimizer family actually owns. 

36 min_numel: Tensor states smaller than this element count are not swapped. 

37 packed_swap: Whether to use two packed A/B staging buffers. Defaults to 

38 ``True``; when ``False``, optimizer states are swapped tensor by 

39 tensor. 

40 """ 

41 

42 swap_times: int = 16 

43 state_keys: Optional[Sequence[str]] = None 

44 min_numel: int = 1024 

45 packed_swap: bool = True 

46 

47 def __post_init__(self) -> None: 

48 if self.swap_times <= 0: 

49 raise ValueError("SwapOptimizerConfig.swap_times must be positive.") 

50 if self.min_numel < 0: 

51 raise ValueError("SwapOptimizerConfig.min_numel must be non-negative.") 

52 object.__setattr__(self, "state_keys", validate_state_keys(self.state_keys)) 

53 

54 

55class SwapOptimizer: 

56 """Core facade that creates the Torch optimizer-state swap wrapper.""" 

57 

58 def __new__(cls, optimizer: Any, config: Optional[SwapOptimizerConfig] = None): 

59 return swap_optimizer(optimizer, config) 

60 

61 

62def swap_optimizer(optimizer: Any, config: Optional[SwapOptimizerConfig] = None) -> Any: 

63 """Wrap a supported optimizer with optimizer-state swap. 

64 

65 Public dispatcher: it resolves the optimizer family and defers to that 

66 family's factory, which owns adapter selection and the resulting error 

67 message. Optimizer modules are imported lazily so loading this facade does 

68 not require every algorithm implementation. 

69 

70 Args: 

71 optimizer: Base optimizer instance. 

72 config: Swap optimizer configuration. When omitted, every supported 

73 optimizer family uses packed transfers. 

74 

75 Returns: 

76 Swap optimizer wrapper. 

77 

78 Raises: 

79 ValueError: If the optimizer type is unsupported. 

80 """ 

81 resolved_config = config or SwapOptimizerConfig() 

82 if is_muon_optimizer(optimizer): 

83 from hyper_parallel.core.optimizer.swap_muon import ( # pylint: disable=import-outside-toplevel 

84 swap_muon, 

85 ) 

86 return swap_muon(optimizer, resolved_config) 

87 

88 from hyper_parallel.core.optimizer.swap_adam import ( # pylint: disable=import-outside-toplevel 

89 swap_adam, 

90 ) 

91 return swap_adam(optimizer, resolved_config) 

92 

93 

94def is_swap_optimizer(optimizer: Any) -> bool: 

95 """Return whether ``optimizer`` is a swap optimizer wrapper.""" 

96 return bool(getattr(optimizer, "_is_swap_optimizer", False)) 

97 

98 

99def is_muon_optimizer(optimizer: Any) -> bool: 

100 """Return whether ``optimizer`` is a Muon leaf optimizer. 

101 

102 The check is by concrete leaf type, never by class name or by inspecting the 

103 shape of a ``ChainedOptimizer``: a chain is not itself a Muon optimizer, and 

104 its leaves must each be wrapped on their own. 

105 """ 

106 from hyper_parallel.core.optimizer.muon import ( # pylint: disable=import-outside-toplevel 

107 Muon, 

108 ) 

109 return isinstance(optimizer, Muon)