Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / config / optimization.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"""Mixed-precision and optimizer configuration sections. 

16 

17Split from ``auto_models/trainer/config.py`` in stage 7 (05 §15.2.5); 

18class names, fields and defaults are unchanged. 

19""" 

20 

21from dataclasses import dataclass, field 

22from typing import Any, List, Optional 

23 

24from torch.optim import Optimizer # pylint: disable=forbidden-backend-import 

25 

26from hyper_parallel.core.optimizer.swap_optimizer_base import validate_state_keys 

27from hyper_parallel.trainer.config.target import Target, _serialize_config_value 

28 

29 

30@dataclass 

31class MixedPrecisionConfig: 

32 """Mixed-precision parameters exposed by the initial YAML schema.""" 

33 

34 enabled: bool = False 

35 

36 

37@dataclass 

38class OptimizerSwapConfig: 

39 """Optimizer-state swap options for the Trainer-built optimizer. 

40 

41 Swapping moves optimizer state tensors to host memory between updates, so 

42 the peak device footprint drops by the optimizer state size while the 

43 copies overlap the surrounding step. It is independent of 

44 ``fsdp_config.enable_offload``, which offloads weights and gradients. 

45 

46 Args: 

47 enabled: Whether the Trainer wraps the built optimizer with state swap. 

48 swap_times: Number of pipeline partitions the swapped state is split 

49 into; each partition is prefetched one batch ahead of its update. 

50 min_numel: State tensors smaller than this element count stay on device. 

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

52 matching optimizer adapter's defaults. 

53 packed_swap: Whether the backend packs swapped state into two reusable 

54 staging buffers. ``None`` keeps the backend default (enabled on 

55 PyTorch, and for MindFormers AdamW on MindSpore). 

56 """ 

57 

58 enabled: bool = False 

59 swap_times: int = 16 

60 min_numel: int = 1024 

61 state_keys: Optional[List[str]] = None 

62 packed_swap: Optional[bool] = None 

63 

64 def __post_init__(self) -> None: 

65 """Reject a pipeline width or state key set the swap runtime cannot use.""" 

66 if self.swap_times < 1: 

67 raise ValueError( 

68 f"optimizer.swap.swap_times must be at least 1, got {self.swap_times}" 

69 ) 

70 if self.min_numel < 0: 

71 raise ValueError( 

72 f"optimizer.swap.min_numel must be non-negative, got {self.min_numel}" 

73 ) 

74 if self.state_keys is not None: 

75 validate_state_keys(self.state_keys) 

76 

77 

78@dataclass 

79class OptimizerConfig: 

80 """Optimizer target plus Trainer-owned parameter precision policy.""" 

81 

82 target: Target[Optimizer] 

83 fp32_main_params: bool = False 

84 swap: OptimizerSwapConfig = field(default_factory=OptimizerSwapConfig) 

85 

86 def to_dict(self) -> dict[str, Any]: 

87 """Serialize optimizer options in their compact target YAML shape.""" 

88 config = self.target.to_dict() 

89 config["fp32_main_params"] = self.fp32_main_params 

90 config["swap"] = _serialize_config_value(self.swap) 

91 return config