Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / dmodule / model.py: 80%

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"""Whole-model protocol: :class:`BaseModel` and :class:`ModelConfigConverter`.""" 

16 

17from abc import abstractmethod 

18from dataclasses import dataclass 

19 

20from hyper_parallel.config import Configurable 

21from hyper_parallel.dmodule.module import Module 

22 

23 

24class ModelConfigConverter(Configurable): 

25 """Transform a model config tree before :meth:`BaseModel.Config.build`. 

26 

27 Subclasses implement :meth:`convert` to walk or rewrite nested 

28 :class:`~hyper_parallel.config.configurable.Configurable.Config` nodes 

29 (typically via :meth:`~hyper_parallel.config.configurable.Configurable.Config.traverse`). 

30 

31 Example: 

32 Run ``convert`` on the model config, then ``build`` the model:: 

33 

34 model_cfg = MyModel.Config(num_layers=4) 

35 MyConverter.Config(enabled=True).build().convert(model_cfg) 

36 model = model_cfg.build() 

37 

38 Minimal ``convert`` implementation:: 

39 

40 class MyConverter(ModelConfigConverter): 

41 @dataclass(kw_only=True, slots=True) 

42 class Config(ModelConfigConverter.Config): 

43 enabled: bool = True 

44 

45 def convert(self, model_config: Configurable.Config) -> None: 

46 for _fqn, cfg, parent, attr in model_config.traverse(Module.Config): 

47 new_cfg = cfg # replace with another Config type as needed 

48 setattr(parent, attr, new_cfg) 

49 """ 

50 

51 @dataclass(kw_only=True, slots=True) 

52 class Config(Configurable.Config): 

53 """Converter hyperparameters (subclass adds fields such as ``rank``).""" 

54 

55 @abstractmethod 

56 def convert(self, model_config: Configurable.Config) -> None: 

57 """Rewrite *model_config* in place. 

58 

59 Args: 

60 model_config: Root :class:`BaseModel.Config` (or compatible tree) 

61 to modify before ``build()``. 

62 

63 Example:: 

64 

65 LoRAConverter.Config(rank=16).build().convert(llama_cfg) 

66 """ 

67 raise NotImplementedError 

68 

69 

70class BaseModel(Module): 

71 """Base class for declarative whole-model components. 

72 

73 Subclasses define ``Config(BaseModel.Config)`` and assemble child 

74 :class:`~hyper_parallel.dmodule.module.Module` layers. Register with 

75 :class:`~hyper_parallel.dmodule.model_spec.ModelSpec` via ``model=...Config``. 

76 

77 Example:: 

78 

79 @dataclass(kw_only=True, slots=True) 

80 class MyModelConfig(BaseModel.Config): 

81 hidden: int = 128 

82 

83 def get_nparams_and_flops(self, model: Module, seq_len: int) -> tuple[int, int]: 

84 return sum(p.numel() for p in model.parameters()), 0 

85 

86 class MyModel(BaseModel): 

87 def __init__(self, config: MyModelConfig): 

88 super().__init__() 

89 self.config = config 

90 

91 cfg = MyModelConfig(hidden=256) 

92 cfg.update_from_config(trainer_config=trainer_cfg) 

93 model = cfg.build() 

94 model.init_states() 

95 model.verify_module_protocol() 

96 """ 

97 

98 def init_weights(self, **kwargs) -> None: 

99 """Initialize parameters and buffers (alias for :meth:`init_states`). 

100 

101 Args: 

102 **kwargs: Forwarded to :meth:`init_states` (for example 

103 ``buffer_device`` for RoPE-style buffers). 

104 

105 Example:: 

106 

107 model.init_weights(buffer_device=torch.device("npu", 0)) 

108 """ 

109 buffer_device = kwargs.get("buffer_device") 

110 self.init_states(buffer_device=buffer_device) 

111 

112 def verify_module_protocol(self) -> None: 

113 """Ensure every submodule is a :class:`~hyper_parallel.dmodule.module.Module`. 

114 

115 Raises: 

116 RuntimeError: If any ``named_modules()`` entry is not a 

117 :class:`Module` instance. 

118 

119 Example:: 

120 

121 model = cfg.build() 

122 model.verify_module_protocol() # before parallelize / training 

123 """ 

124 failures: list[tuple[str, str]] = [] 

125 for fqn, mod in self.named_modules(): 

126 if not isinstance(mod, Module): 

127 failures.append((fqn, type(mod).__name__)) 

128 if failures: 

129 details = ", ".join(f"'{fqn}' ({cls})" for fqn, cls in failures) 

130 raise RuntimeError( 

131 f"The following modules do not satisfy the Module protocol: {details}" 

132 ) 

133 

134 @dataclass(kw_only=True, slots=True) 

135 class Config(Module.Config): 

136 """Base config for whole models (extends :class:`Module.Config`).""" 

137 

138 def update_from_config(self, *, trainer_config, **kwargs) -> None: 

139 """Inject trainer-derived fields into this config before ``build``. 

140 

141 Override in model configs to copy parallelism flags, sharding specs, 

142 or other values from the trainer configuration. 

143 

144 Args: 

145 trainer_config: Trainer or job configuration object. 

146 **kwargs: Optional extra inputs for specific models. 

147 

148 Example:: 

149 

150 cfg = LlamaModel.Config(num_layers=32) 

151 cfg.update_from_config(trainer_config=trainer_cfg) 

152 model = cfg.build() 

153 """ 

154 del trainer_config, kwargs 

155 

156 def get_nparams_and_flops(self, model: Module, seq_len: int) -> tuple[int, int]: 

157 """Return parameter count and estimated FLOPs for logging. 

158 

159 Args: 

160 model: Built model instance. 

161 seq_len: Sequence length used for FLOPs estimation. 

162 

163 Returns: 

164 ``(num_params, num_flops)``. 

165 

166 Raises: 

167 NotImplementedError: If the concrete model config does not 

168 override this method. 

169 

170 Example:: 

171 

172 nparams, nflops = cfg.get_nparams_and_flops(model, seq_len=4096) 

173 """ 

174 del model, seq_len 

175 raise NotImplementedError( 

176 f"{type(self).__name__} must implement get_nparams_and_flops" 

177 ) 

178 

179 

180__all__ = ["BaseModel", "ModelConfigConverter"]