Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / config / trainer.py: 94%

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"""Aggregate trainer configuration tree. 

16 

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

18class name, fields and defaults are unchanged. 

19""" 

20 

21from dataclasses import dataclass, field, fields 

22from typing import Any, List, Literal, Optional 

23 

24from hyper_parallel.models.build_options import CompileConfig, FSDP2Config 

25from hyper_parallel.components.checkpoint.config import CheckpointingConfig 

26 

27from hyper_parallel.trainer.config.data import DataLoaderConfig, DatasetConfig 

28from hyper_parallel.trainer.config.optimization import ( 

29 MixedPrecisionConfig, 

30 OptimizerConfig, 

31) 

32from hyper_parallel.trainer.config.parallelism import ( 

33 AcceleratorConfig, 

34 ActivationCheckpointConfig, 

35 PlanOverride, 

36) 

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

38from hyper_parallel.trainer.config.stress_detect import StressDetectConfig 

39from hyper_parallel.trainer.config.training import ( 

40 DebugConfig, 

41 MemoryConfig, 

42 ModelIntegrationConfig, 

43 ProfilingConfig, 

44 TrainingConfig, 

45 WandbConfig, 

46) 

47 

48 

49@dataclass 

50class TrainerConfig: 

51 """Resolved component tree; runtime objects are built by the task trainer.""" 

52 

53 model: Target[Any] 

54 optimizer: OptimizerConfig 

55 

56 lr_scheduler: Optional[Target[Any]] = None 

57 loss_fn: Optional[Target[Any]] = None 

58 # Final floating-point dtype after model weights are loaded or initialized 

59 # from scratch. None preserves the dtype produced by the initialization path. 

60 model_init_dtype: Optional[Literal["float16", "bfloat16", "float32"]] = None 

61 training: TrainingConfig = field(default_factory=TrainingConfig) 

62 

63 # parallelism configs 

64 accelerator: AcceleratorConfig = field(default_factory=AcceleratorConfig) 

65 fsdp_config: FSDP2Config = field(default_factory=FSDP2Config) 

66 plan_overrides: List[PlanOverride] = field(default_factory=list) 

67 mixed_precision: MixedPrecisionConfig = field( 

68 default_factory=MixedPrecisionConfig 

69 ) 

70 activation_checkpoint: ActivationCheckpointConfig = field( 

71 default_factory=ActivationCheckpointConfig 

72 ) 

73 activation_swap: Literal["none", "attention"] = "none" 

74 compile: CompileConfig = field(default_factory=CompileConfig) 

75 

76 # data 

77 dataset: Optional[DatasetConfig] = None 

78 dataloader: Optional[DataLoaderConfig] = None 

79 packed_sequence: Optional[Any] = None 

80 

81 checkpoint: CheckpointingConfig = field(default_factory=CheckpointingConfig) 

82 debug: DebugConfig = field(default_factory=DebugConfig) 

83 model_integration: ModelIntegrationConfig = field(default_factory=ModelIntegrationConfig) 

84 memory: MemoryConfig = field(default_factory=MemoryConfig) 

85 wandb: WandbConfig = field(default_factory=WandbConfig) 

86 profiler: ProfilingConfig = field(default_factory=ProfilingConfig) 

87 stress_detect: StressDetectConfig = field(default_factory=StressDetectConfig) 

88 magi: Optional[Any] = None 

89 peft: Optional[Any] = None 

90 

91 def __post_init__(self) -> None: 

92 """Validate combinations that span multiple config sections.""" 

93 if ( 

94 self.compile.enabled 

95 and not self.compile.selects_graph_compiler() 

96 and self.accelerator.pp_size > 1 

97 ): 

98 raise ValueError("compile is not supported together with pipeline parallelism") 

99 reduce_dtype = self.fsdp_config.mix_precision.reduce_dtype 

100 if self.optimizer.fp32_main_params and reduce_dtype != "float32": 

101 raise ValueError( 

102 "optimizer.fp32_main_params=true requires " 

103 "fsdp_config.mix_precision.reduce_dtype='float32'; " 

104 f"got {reduce_dtype!r}" 

105 ) 

106 

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

108 """Serialize the resolved trainer configuration for logging.""" 

109 return { 

110 config_field.name: _serialize_config_value( 

111 getattr(self, config_field.name) 

112 ) 

113 for config_field in fields(self) 

114 } 

115 

116 

117def save_configs(config: TrainerConfig, output_dir: str) -> None: 

118 """Accept trainer config persistence requests without writing files. 

119 

120 Args: 

121 config: Resolved trainer configuration. 

122 output_dir: Intended configuration output directory. 

123 """ 

124 del config, output_dir