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

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"""apply: public entry point for applying a ShardingPlan (05 §4 canonical). 

16 

17apply_sharding_plan: Phase 0 normalization -> A parameter sharding -> B special 

18handlers -> C entry unpack + source_shard_info -> C forward wrapping 

19(production/validate/moe/cp/vocab_embed, five paths) -> D tied weights. 

20 

21Dual-mode architecture constraint (05 §1.4): production has zero DTensor 

22dispatch (build-time unpack + PrecompiledBoundary); the only difference between 

23validate and production is the boundary stitching method -- for any module whose 

24DTensor dispatch hides data-dependent logic (embedding mask / attention K/V 

25gather / MoE all-to-all), both modes explicitly reconstruct it with the same 

26local-region wrapper (D-01''/D-02/D-03'). 

27 

28This module is a thin public facade: the implementation lives in 

29``_builder/applier.py`` (preflight/mesh/phase orchestration), 

30``_builder/parameter_sharding.py`` (parameter localize/stack/placement) and 

31``_builder/forward_rewriter.py`` (forward mutation). 

32""" 

33 

34from typing import Any, Optional, Tuple 

35 

36from hyper_parallel.distributed.plan import ShardingPlan 

37from hyper_parallel.distributed.recipe_spec import ( 

38 _normalize_out_fields, 

39) 

40from hyper_parallel.distributed._builder.applier import ( 

41 _apply_phase_c, 

42 _apply_plan_special_handlers, 

43 _get_active_mesh, 

44 _get_tp_submesh, 

45 _preflight_compute_injection, 

46) 

47from hyper_parallel.distributed._builder.parameter_sharding import ( 

48 _build_runtime_source_shard_info, 

49 _replicate_tied_weights, 

50 _resolve_parameter_source_meshes, 

51 _shard_planned_parameters, 

52 detect_tied_weights, 

53) 

54 

55 

56def _prepare_sharding_meshes( 

57 model: Any, 

58 plan: ShardingPlan, 

59 mesh: Any, 

60) -> Tuple[Any, Any, Any, Any, Any]: 

61 """Resolve active and parameter source meshes before plan application.""" 

62 # Lazy import: components.distributed.__init__ re-exports this module, so 

63 # a module-level import of the mesh module would cycle through the 

64 # package __init__. 

65 from hyper_parallel.distributed.mesh import ( # pylint: disable=C0415 

66 MeshContext, 

67 ) 

68 mesh_context = mesh if isinstance(mesh, MeshContext) else None 

69 if mesh_context is None: 

70 device_mesh = mesh 

71 else: 

72 device_mesh = mesh_context.device_mesh 

73 if device_mesh is None: 

74 raise ValueError("apply_sharding_plan requires a DeviceMesh") 

75 

76 mesh_dim_names = plan.mesh_dim_names 

77 # Active sub-mesh: the planner strips size=1 axes (plan.mesh_dim_names), but the 

78 # passed-in mesh may still contain those axes -- placements are resolved against 

79 # plan.mesh_dim_names, so the dimensionality must align with the mesh, otherwise 

80 # distribute_tensor will silently shard along the wrong axis. 

81 full_mesh = device_mesh 

82 active_mesh = _get_active_mesh(device_mesh, mesh_dim_names) 

83 tp_mesh = _get_tp_submesh(active_mesh, mesh_dim_names) 

84 models = model if isinstance(model, list) else [model] 

85 

86 # Explicit-injection guard: CP/EP sharding without an explicit compute 

87 # injection fails fast here, BEFORE any parameter is touched 

88 _preflight_compute_injection(plan, active_mesh, model=models[0]) 

89 

90 expert_mesh, dense_source_mesh, expert_source_mesh = ( 

91 _resolve_parameter_source_meshes(plan, mesh_context, full_mesh, tp_mesh) 

92 ) 

93 return active_mesh, models, expert_mesh, dense_source_mesh, expert_source_mesh 

94 

95 

96def apply_sharding_plan( 

97 model: Any, 

98 plan: ShardingPlan, 

99 mesh: Any, 

100 *, 

101 validate_mode: bool = False, 

102) -> Tuple[Any, Optional[Any]]: 

103 """Apply a ShardingPlan using a DeviceMesh or MeshContext. 

104 

105 Args: 

106 model: The model to shard (an HF-style ``nn.Module``, or a list of 

107 per-part models in PP scenarios). 

108 plan: The :class:`ShardingPlan` produced by the planner. 

109 mesh: A ``DeviceMesh`` or a :class:`MeshContext` carrying one. 

110 validate_mode: When ``True``, keep parameters as DTensors and wrap 

111 forwards for placement-propagation validation instead of the 

112 production local-tensor path. 

113 

114 Returns (model, source_shard_info): 

115 - production: at the Phase C entry, a one-shot `_local_params_context` permanently 

116 unwraps DTensor parameters into plain local tensors, and builds source_shard_info 

117 for fully_shard to use. Entries record the complete source layout: 

118 ``{param_fqn: (placements_tuple, source_sub_mesh)}`` with one placement per 

119 non-FSDP source axis (dense entries derive from the dense-FSDP mesh, routed 

120 experts from the expert mesh); 

121 - validate: no unwrap (parameters remain DTensors); source_shard_info is None. 

122 """ 

123 mesh, models, expert_mesh, dense_source_mesh, expert_source_mesh = ( 

124 _prepare_sharding_meshes(model, plan, mesh) 

125 ) 

126 

127 # ====== Phase 0: normalize out_src/out_dst scalar shorthand (idempotent, covers user-injected paths) ====== 

128 for spec in plan.modules.values(): 

129 _normalize_out_fields(spec) 

130 

131 # ====== Phase A: parameter sharding ====== 

132 _shard_planned_parameters(models, plan, mesh, expert_mesh, validate_mode) 

133 

134 # ====== Phase B: special handlers ====== 

135 _apply_plan_special_handlers(models, plan, mesh) 

136 

137 # ====== Phase C entry: one-shot unpack at build time (production only) ====== 

138 source_shard_info = _build_runtime_source_shard_info( 

139 models, 

140 plan, 

141 dense_source_mesh, 

142 expert_source_mesh, 

143 validate_mode, 

144 ) 

145 

146 # ====== Phase C: wrap forward ====== 

147 for part in models: 

148 _apply_phase_c(part, plan, mesh, validate_mode, expert_mesh=expert_mesh) 

149 

150 # ====== Phase D: tied weights ====== 

151 tied_pairs = list(plan.tied_pairs) or detect_tied_weights(models[0]) 

152 for part in models: 

153 _replicate_tied_weights(part, tied_pairs) 

154 

155 return model, source_shard_info