Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / dmodule / model_spec.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"""Per-model registration bundle for the training stack. 

16 

17Maps ``model.name`` in trainer YAML to construction and parallel hooks. 

18See :mod:`hyper_parallel.models.spec.registry` for ``register_spec`` / ``get_spec``. 

19""" 

20 

21from dataclasses import dataclass 

22from typing import Callable, Optional, Type 

23 

24from hyper_parallel.dmodule.model import BaseModel 

25 

26 

27@dataclass 

28class ModelSpec: 

29 """Per-model registration bundle. 

30 

31 Either ``build_model_fn`` (factory function) or ``model`` 

32 (:class:`~hyper_parallel.dmodule.model.BaseModel.Config`) must be set, 

33 but not both. Optional hooks customize parallelize, grad clip, pipeline, 

34 and checkpoint key mapping. 

35 

36 Args: 

37 name: Unique model id (matches trainer config ``model.name``). 

38 build_model_fn: ``(trainer_cfg) -> Module`` factory used by the trainer 

39 today in :mod:`hyper_parallel.models.spec`. 

40 model: Declarative :class:`BaseModel.Config` tree; built via 

41 ``spec.model.build()`` when the trainer supports this path. 

42 parallelize_fn: ``(model, mesh, trainer_cfg) -> model`` custom 

43 parallelization (required for current trainer). 

44 clip_grad_fn: Optional custom gradient clipping. 

45 pipelining_fn: Optional pipeline-parallel setup. 

46 state_dict_adapter: Optional checkpoint key translation class. 

47 

48 Example: 

49 Declarative config:: 

50 

51 model_cfg = MyModel.Config(num_layers=12) 

52 spec = ModelSpec( 

53 name="my_model", 

54 model=model_cfg, 

55 parallelize_fn=parallelize_my_model, 

56 ) 

57 

58 Factory function (current trainer default):: 

59 

60 def build_my_model(trainer_cfg): 

61 return MyModel(MyModel.Config(hidden=trainer_cfg.hidden_size)) 

62 

63 spec = ModelSpec( 

64 name="my_model", 

65 build_model_fn=build_my_model, 

66 parallelize_fn=parallelize_my_model, 

67 ) 

68 

69 Register before training:: 

70 

71 from hyper_parallel.models.spec import register_spec 

72 

73 register_spec("my_model", spec) 

74 """ 

75 

76 name: str 

77 build_model_fn: Optional[Callable] = None 

78 model: Optional[BaseModel.Config] = None 

79 parallelize_fn: Optional[Callable] = None 

80 clip_grad_fn: Optional[Callable] = None 

81 pipelining_fn: Optional[Callable] = None 

82 state_dict_adapter: Optional[Type] = None 

83 

84 def __post_init__(self) -> None: 

85 """Validate that exactly one model constructor is configured. 

86 

87 Raises: 

88 ValueError: If both or neither of ``build_model_fn`` and ``model`` 

89 are set. 

90 """ 

91 if self.build_model_fn is None and self.model is None: 

92 raise ValueError( 

93 f"ModelSpec '{self.name}' requires build_model_fn or model config" 

94 ) 

95 if self.build_model_fn is not None and self.model is not None: 

96 raise ValueError( 

97 f"ModelSpec '{self.name}' must not set both build_model_fn and model" 

98 ) 

99 

100 

101__all__ = ["ModelSpec"]