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(``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 inner_wrapper, 

44 local_compute, 

45 parse_named_placement, 

46 parse_placement, 

47 resolve_placements, 

48) 

49from hyper_parallel.distributed._builder.default_templates import ( 

50 ShardingTemplate, 

51) 

52from hyper_parallel.distributed.activation_checkpoint import ( 

53 _apply_activation_checkpointing, 

54 compile_selective_checkpoint_policy, 

55 make_selective_checkpoint_context_fn, 

56) 

57from hyper_parallel.distributed.attention_swap import ( 

58 apply_attention_swap, 

59 attention_swap_policy, 

60 validate_attention_swap, 

61) 

62from hyper_parallel.distributed.compile import ( 

63 apply_compile, 

64 get_compile_layers, 

65 resolve_compile_kwargs, 

66) 

67 

68__all__ = [ 

69 "CP", 

70 "DP", 

71 "EP", 

72 "TP", 

73 "MeshAxisName", 

74 "ModuleShardingSpec", 

75 "NamedPlacement", 

76 "PlacementMismatchError", 

77 "ShardingPlan", 

78 "ShardingPlanner", 

79 "ShardingTemplate", 

80 "inner_wrapper", 

81 "local_compute", 

82 "parse_named_placement", 

83 "parse_placement", 

84 "resolve_placements", 

85 "apply_sharding_plan", 

86 "validate_model_compatibility", 

87 "_apply_activation_checkpointing", 

88 "compile_selective_checkpoint_policy", 

89 "make_selective_checkpoint_context_fn", 

90 "apply_attention_swap", 

91 "attention_swap_policy", 

92 "validate_attention_swap", 

93 "apply_compile", 

94 "get_compile_layers", 

95 "resolve_compile_kwargs", 

96]