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

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 

35from types import SimpleNamespace 

36 

37from hyper_parallel.components.modules.module_family import ( 

38 _family_mesh, _family_parameter_shapes, _record_family_layouts, 

39) 

40from hyper_parallel.distributed.module_binding import bind_parallel_modules 

41from hyper_parallel.distributed.plan import ShardingPlan 

42from hyper_parallel.distributed.recipe_spec import ( 

43 _normalize_out_fields, 

44) 

45from hyper_parallel.distributed._builder.applier import ( 

46 _apply_phase_c, 

47 _apply_plan_special_handlers, 

48 _get_active_mesh, 

49 _get_tp_submesh, 

50 _preflight_compute_injection, 

51) 

52from hyper_parallel.distributed._builder.parameter_sharding import ( 

53 _build_runtime_source_shard_info, 

54 _replicate_tied_weights, 

55 _resolve_parameter_source_meshes, 

56 _shard_planned_parameters, 

57 _set_param_by_path, 

58 detect_tied_weights, 

59) 

60 

61 

62def _prepare_sharding_meshes( 

63 model: Any, 

64 plan: ShardingPlan, 

65 mesh: Any, 

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

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

68 mesh_context = mesh if hasattr(mesh, "device_mesh") 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 original_mesh = mesh 

124 for part in model if isinstance(model, list) else [model]: 

125 original_mesh = _family_mesh(part, original_mesh) 

126 mesh = original_mesh 

127 shapes = {id(part): _family_parameter_shapes(part) 

128 for part in (model if isinstance(model, list) else [model])} 

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

130 _prepare_sharding_meshes(model, plan, mesh) 

131 ) 

132 

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

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

135 _normalize_out_fields(spec) 

136 

137 # ====== Phase A: parameter sharding ====== 

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

139 

140 # ====== Phase B: special handlers ====== 

141 _apply_plan_special_handlers(models, plan, mesh) 

142 

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

144 source_shard_info = _build_runtime_source_shard_info( 

145 models, 

146 plan, 

147 dense_source_mesh, 

148 expert_source_mesh, 

149 validate_mode, 

150 ) 

151 for part in models: 

152 for canonical, alias in plan.parameter_aliases: 

153 _set_param_by_path(part, alias, part.get_parameter(canonical)) 

154 

155 # ====== Phase C: wrap forward ====== 

156 for part in models: 

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

158 

159 # ====== Phase D: tied weights ====== 

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

161 for part in models: 

162 _replicate_tied_weights(part, tied_pairs) 

163 

164 runtime_mesh = original_mesh if hasattr(original_mesh, "device_mesh") else SimpleNamespace( 

165 device_mesh=original_mesh, fsdp_moe_mesh=expert_mesh, 

166 ) 

167 if plan.parallel_modules: 

168 for part in models: 

169 bind_parallel_modules(part, plan, runtime_mesh) 

170 

171 for part in models: 

172 _record_family_layouts(part, plan, runtime_mesh, shapes[id(part)], source_shard_info) 

173 

174 return model, source_shard_info