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

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, field, replace 

20from typing import Any, Optional, Sequence 

21 

22from hyper_parallel.core.optimizer.swap_optimizer_base import validate_state_keys 

23from hyper_parallel.platform import get_platform 

24from hyper_parallel.platform.platform import PlatformType 

25 

26 

27def _default_packed_swap() -> bool: 

28 """Return the packed swap default for the active backend.""" 

29 return get_platform().platform_type == PlatformType.PYTORCH 

30 

31 

32def _is_mindformers_adamw(optimizer: Any) -> bool: 

33 """Return whether an optimizer is the supported MindFormers AdamW type.""" 

34 optimizer_type = type(optimizer) 

35 return ( 

36 optimizer_type.__name__ == "AdamW" 

37 and optimizer_type.__module__ == "mindformers.pynative.optimizer.adamw" 

38 ) 

39 

40 

41def _resolve_packed_swap_default( 

42 optimizer: Any, 

43 config: "SwapOptimizerConfig", 

44 platform_type: PlatformType, 

45) -> "SwapOptimizerConfig": 

46 """Enable packed swap by default for MindFormers AdamW on MindSpore.""" 

47 if ( 

48 platform_type == PlatformType.MINDSPORE 

49 and not config.packed_swap_was_explicit 

50 and _is_mindformers_adamw(optimizer) 

51 ): 

52 return replace(config, packed_swap=True) 

53 return config 

54 

55 

56@dataclass(frozen=True) 

57class SwapOptimizerConfig: 

58 """Configuration for Adam/AdamW optimizer state swap. 

59 

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

61 

62 Args: 

63 swap_times: Number of pipeline partitions. 

64 state_keys: Logical state keys to swap. ``None`` uses adapter defaults. 

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

66 include_master_params: Whether optimizer-owned fp32 master params are swapped. 

67 packed_swap: Whether supported backends use two packed A/B staging buffers. 

68 Defaults to ``True`` on PyTorch and for MindFormers AdamW on 

69 MindSpore; other MindSpore optimizers default to ``False``. When 

70 ``False``, optimizer states are swapped tensor by tensor. An 

71 explicit value always takes precedence over these defaults. 

72 """ 

73 

74 swap_times: int = 16 

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

76 min_numel: int = 1024 

77 include_master_params: bool = False 

78 packed_swap: Optional[bool] = None 

79 _packed_swap_explicit: bool = field(init=False, repr=False, compare=False) 

80 

81 def __post_init__(self) -> None: 

82 packed_swap_explicit = self.packed_swap is not None 

83 object.__setattr__(self, "_packed_swap_explicit", packed_swap_explicit) 

84 if not packed_swap_explicit: 

85 object.__setattr__(self, "packed_swap", _default_packed_swap()) 

86 if self.swap_times <= 0: 

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

88 if self.min_numel < 0: 

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

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

91 

92 @property 

93 def packed_swap_was_explicit(self) -> bool: 

94 """Return whether ``packed_swap`` was explicitly supplied by the caller.""" 

95 return self._packed_swap_explicit 

96 

97 

98class SwapOptimizer: 

99 """Core facade that dispatches to the active backend implementation.""" 

100 

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

102 return swap_optimizer(optimizer, config) 

103 

104 

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

106 """Wrap a supported Adam/AdamW optimizer with optimizer-state swap. 

107 

108 Args: 

109 optimizer: Base optimizer instance. 

110 config: Swap optimizer configuration. 

111 

112 Returns: 

113 Backend-specific swap optimizer wrapper. 

114 

115 Raises: 

116 ValueError: If the active backend or optimizer type is unsupported. 

117 """ 

118 platform = get_platform() 

119 cfg = config or SwapOptimizerConfig() 

120 cfg = _resolve_packed_swap_default(optimizer, cfg, platform.platform_type) 

121 return platform.get_swap_optimizer()(optimizer, cfg) 

122 

123 

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

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

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