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

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.training import ( 

39 DebugConfig, 

40 MemoryConfig, 

41 ModelIntegrationConfig, 

42 ProfilingConfig, 

43 TrainingConfig, 

44 WandbConfig, 

45) 

46 

47 

48@dataclass 

49class TrainerConfig: 

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

51 

52 model: Target[Any] 

53 optimizer: OptimizerConfig 

54 

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

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

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

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

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

60 training: TrainingConfig = field(default_factory=TrainingConfig) 

61 

62 # parallelism configs 

63 accelerator: AcceleratorConfig = field(default_factory=AcceleratorConfig) 

64 fsdp_config: FSDP2Config = field(default_factory=FSDP2Config) 

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

66 mixed_precision: MixedPrecisionConfig = field( 

67 default_factory=MixedPrecisionConfig 

68 ) 

69 activation_checkpoint: ActivationCheckpointConfig = field( 

70 default_factory=ActivationCheckpointConfig 

71 ) 

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

73 compile: CompileConfig = field(default_factory=CompileConfig) 

74 

75 # codegen 

76 codegen: bool = False 

77 

78 # data 

79 dataset: Optional[DatasetConfig] = None 

80 dataloader: Optional[DataLoaderConfig] = None 

81 packed_sequence: Optional[Any] = None 

82 

83 checkpoint: CheckpointingConfig = field(default_factory=CheckpointingConfig) 

84 debug: DebugConfig = field(default_factory=DebugConfig) 

85 model_integration: ModelIntegrationConfig = field(default_factory=ModelIntegrationConfig) 

86 memory: MemoryConfig = field(default_factory=MemoryConfig) 

87 wandb: WandbConfig = field(default_factory=WandbConfig) 

88 profiler: ProfilingConfig = field(default_factory=ProfilingConfig) 

89 magi: Optional[Any] = None 

90 peft: Optional[Any] = None 

91 

92 def __post_init__(self) -> None: 

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

94 if ( 

95 self.compile.enabled 

96 and not self.compile.selects_graph_compiler() 

97 and self.accelerator.pp_size > 1 

98 ): 

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

100 reduce_dtype = self.fsdp_config.mix_precision.reduce_dtype 

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

102 raise ValueError( 

103 "optimizer.fp32_main_params=true requires " 

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

105 f"got {reduce_dtype!r}" 

106 ) 

107 

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

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

110 return { 

111 config_field.name: _serialize_config_value( 

112 getattr(self, config_field.name) 

113 ) 

114 for config_field in fields(self) 

115 } 

116 

117 

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

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

120 

121 Args: 

122 config: Resolved trainer configuration. 

123 output_dir: Intended configuration output directory. 

124 """ 

125 del config, output_dir