Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / config / training.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"""Training-loop, debug, profiler, and remote-logging configuration sections.""" 

16 

17from dataclasses import dataclass, field 

18from typing import Literal, Optional 

19 

20from hyper_parallel.components.quantization.config import LowPrecisionConfig 

21 

22 

23@dataclass 

24class TrainingConfig: 

25 """Training-loop parameters exposed by the initial YAML schema.""" 

26 

27 train_iters: Optional[int] = None 

28 train_samples: Optional[int] = None 

29 lr_scheduler_iters: Optional[int] = None 

30 eval_iters: int = 0 

31 

32 global_batch_size: int = 8 

33 micro_batch_size: int = 1 

34 

35 backend: Literal["nccl", "hccl", "gloo"] = "nccl" 

36 max_grad_norm: float = 1.0 

37 init_device: Literal["meta", "cpu", "cuda", "npu"] = "meta" 

38 loss_aggregation: Literal["token_weighted", "rank_average"] = "token_weighted" 

39 seed: Optional[int] = None 

40 enable_full_determinism: bool = False 

41 gc_steps: int = 0 

42 empty_cache_steps: int = 0 

43 empty_cache_before_backward: bool = False 

44 eval_steps: int = 0 

45 eval_epochs: int = 0 

46 logging_steps: int = 1 

47 log_tflops: bool = False 

48 """Publish analytical per-device model FLOPs throughput. 

49 

50 When enabled, the environment meter adds 

51 ``performance/tflops_per_device`` (TFLOP/s, cluster average per device) 

52 to each step's metrics. The estimate charges every model-input token, 

53 including prompt tokens masked out of the loss, and accumulates causal 

54 attention over the real sequence boundaries. 

55 """ 

56 low_precision: LowPrecisionConfig = field(default_factory=LowPrecisionConfig) 

57 

58 

59@dataclass 

60class DebugConfig: 

61 """Debug parameters exposed by the initial YAML schema.""" 

62 

63 check_dataset: Optional[Literal["debug", "info", "warn"]] = None 

64 log_first_rendered_template: bool = False 

65 check_nan_inf: bool = False 

66 check_fsdp_runtime: bool = False 

67 

68 

69@dataclass 

70class ModelIntegrationConfig: 

71 """Enable model-integration checks at one well-defined lifecycle scope. 

72 

73 ``build`` validates the final module tree, sharding ownership, data 

74 contract, optimizer layout, and checkpoint coverage. ``runtime`` includes 

75 those checks and records one-step input, gradient, shared-state, and layout 

76 evidence. Output location and probe selection are framework policy rather 

77 than user-tunable training semantics. 

78 """ 

79 

80 mode: Literal["off", "build", "runtime"] = "off" 

81 

82 def __post_init__(self) -> None: 

83 """Reject ambiguous validation scopes at the configuration boundary.""" 

84 if self.mode not in ("off", "build", "runtime"): 

85 raise ValueError( 

86 "model_integration.mode must be one of: off, build, runtime" 

87 ) 

88 

89 

90@dataclass 

91class MemoryConfig: 

92 """Configure allocator memory-history snapshots during training. 

93 

94 Args: 

95 enable: Whether allocator history recording is enabled. 

96 start_step: First session-relative step included in the snapshot; zero includes model construction. 

97 end_step: Last session-relative step included in the snapshot. 

98 save_path: Directory used for snapshot files. 

99 dump_ranks: Global ranks that write snapshot files. 

100 stacks: Stack information recorded by the allocator. 

101 mem_info: Whether peak allocator statistics are logged before each step. 

102 """ 

103 

104 enable: bool = False 

105 start_step: int = 1 

106 end_step: int = 2 

107 save_path: str = "./memory_snapshot" 

108 dump_ranks: list[int] = field(default_factory=lambda: [0]) 

109 stacks: Literal["python", "all"] = "all" 

110 mem_info: bool = False 

111 

112 

113@dataclass 

114class WandbConfig: 

115 """WandB remote-logging parameters (03 §4.2.5: read by build_callback_manager).""" 

116 

117 enabled: bool = False 

118 project: str = "" 

119 entity: Optional[str] = None 

120 

121 

122@dataclass 

123class ProfilingAnalysisConfig: 

124 """Configure Ascend NPU trace analysis after collection. 

125 

126 Args: 

127 enable: Whether collected data is automatically analyzed. 

128 async_mode: Whether automatic analysis runs asynchronously; ignored when disabled. 

129 """ 

130 

131 enable: bool = True 

132 async_mode: bool = False 

133 

134 

135@dataclass 

136class ProfilingConfig: 

137 """Configure bounded PyTorch profiler collection during training. 

138 

139 HyperParallel groups these options under ``TrainerConfig.profiler``, so 

140 their names describe only the option itself rather than repeating a 

141 ``profile_`` prefix. 

142 

143 Args: 

144 enable: Whether profiler collection is enabled. 

145 start_step: First training step included in the profile. 

146 stop_step: Last training step included in the profile. 

147 warmup_steps: Number of warmup steps immediately before ``start_step``; must be smaller than it. 

148 memory: Whether tensor memory events are collected. 

149 rank_ids: Explicit ranks that collect a profile. ``None`` or an empty list selects all ranks. 

150 output_path: Output root. Each rank writes under ``profile/rank_<rank>``. 

151 level: Ascend collection level, from 0 through 2. ``None`` selects level 0. 

152 with_stack: Whether Python stack traces are collected. 

153 data_simplification: Whether torch-npu removes selected auxiliary CANN data after export. 

154 analysis: Ascend NPU automatic trace analysis options; ignored on CPU/CUDA. 

155 """ 

156 

157 enable: bool = False 

158 start_step: int = 1 

159 stop_step: int = 10 

160 warmup_steps: int = 0 

161 memory: bool = True 

162 rank_ids: Optional[list[int]] = None 

163 output_path: Optional[str] = None 

164 level: Optional[int] = 1 

165 with_stack: bool = False 

166 data_simplification: bool = False 

167 analysis: ProfilingAnalysisConfig = field(default_factory=ProfilingAnalysisConfig)