Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / dmodule / __init__.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"""Declarative distributed modules for training (M1, ``dmodule``). 

16 

17Public surface for building models from nested configs, applying 

18:class:`~hyper_parallel.dmodule.sharding.ShardingConfig`, and registering 

19models for :class:`~hyper_parallel.trainer.base.BaseTrainer`. 

20 

21Typical flow:: 

22 

23 from hyper_parallel.dmodule import BaseModel, Module, ShardingConfig 

24 from hyper_parallel.dmodule.types import MeshAxisName 

25 

26 class Linear(Module): 

27 class Config(Module.Config): 

28 ... 

29 

30 class MyModel(BaseModel): 

31 class Config(BaseModel.Config): 

32 layers: list[Linear.Config] = ... 

33 

34 cfg = MyModel.Config(...) 

35 cfg.build().init_states() 

36 cfg.build().parallelize(tp_mesh) 

37""" 

38 

39from hyper_parallel.config import Configurable 

40from hyper_parallel.dmodule.model import BaseModel, ModelConfigConverter 

41from hyper_parallel.dmodule.model_spec import ModelSpec 

42from hyper_parallel.dmodule.module import Module 

43from hyper_parallel.dmodule.sharding import ( 

44 LocalMapConfig, 

45 ShardingConfig, 

46 resolve_placements, 

47) 

48from hyper_parallel.dmodule.types import MeshAxisName, NamedPlacement 

49 

50__all__ = [ 

51 "BaseModel", 

52 "Configurable", 

53 "LocalMapConfig", 

54 "MeshAxisName", 

55 "ModelConfigConverter", 

56 "ModelSpec", 

57 "Module", 

58 "NamedPlacement", 

59 "ShardingConfig", 

60 "resolve_placements", 

61]