Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / __init__.py: 100%

1# Copyright 2025-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"""distributed: model-level DTensor sharding plan and recipe contract. 

16 

17Public surface (migration plan §7.1): the model-level ``ShardingPlan``, the 

18recipe input contract (``ModuleShardingSpec`` / ``NamedPlacement`` / 

19``MeshAxisName`` / placement DSL parsers) and the public template type 

20``ShardingTemplate``, plan construction (``ShardingPlanner`` / 

21``validate_model_compatibility``) and the injection decorators 

22(``forward_replacement`` / ``local_compute`` / ``inner_wrapper``) and plan application 

23(``apply_sharding_plan``). 

24""" 

25 

26from hyper_parallel.distributed.apply import ( 

27 apply_sharding_plan, 

28) 

29from hyper_parallel.distributed.plan import ShardingPlan 

30from hyper_parallel.distributed._builder.planner import ( 

31 ShardingPlanner, 

32 validate_model_compatibility, 

33) 

34from hyper_parallel.distributed.recipe_spec import ( 

35 CP, 

36 DP, 

37 EP, 

38 TP, 

39 MeshAxisName, 

40 ModuleShardingSpec, 

41 NamedPlacement, 

42 PlacementMismatchError, 

43 forward_replacement, 

44 inner_wrapper, 

45 local_compute, 

46 parse_named_placement, 

47 parse_placement, 

48 resolve_placements, 

49) 

50from hyper_parallel.distributed._builder.default_templates import ( 

51 ShardingTemplate, 

52) 

53from hyper_parallel.distributed.activation_checkpoint import ( 

54 _apply_activation_checkpointing, 

55 compile_selective_checkpoint_policy, 

56 make_selective_checkpoint_context_fn, 

57) 

58from hyper_parallel.distributed.attention_swap import ( 

59 apply_attention_swap, 

60 attention_swap_policy, 

61 validate_attention_swap, 

62) 

63from hyper_parallel.distributed.compile import ( 

64 apply_compile, 

65 get_compile_layers, 

66 resolve_compile_kwargs, 

67) 

68 

69__all__ = [ 

70 "CP", 

71 "DP", 

72 "EP", 

73 "TP", 

74 "MeshAxisName", 

75 "ModuleShardingSpec", 

76 "NamedPlacement", 

77 "PlacementMismatchError", 

78 "ShardingPlan", 

79 "ShardingPlanner", 

80 "ShardingTemplate", 

81 "forward_replacement", 

82 "inner_wrapper", 

83 "local_compute", 

84 "parse_named_placement", 

85 "parse_placement", 

86 "resolve_placements", 

87 "apply_sharding_plan", 

88 "validate_model_compatibility", 

89 "_apply_activation_checkpointing", 

90 "compile_selective_checkpoint_policy", 

91 "make_selective_checkpoint_context_fn", 

92 "apply_attention_swap", 

93 "attention_swap_policy", 

94 "validate_attention_swap", 

95 "apply_compile", 

96 "get_compile_layers", 

97 "resolve_compile_kwargs", 

98]