Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / planner_overlay.py: 71%

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"""Codegen-only additions to native sharding-plan discovery.""" 

16 

17from contextlib import contextmanager 

18from fnmatch import fnmatchcase 

19from typing import Any, Iterator, Optional 

20 

21from hyper_parallel.core.dtensor.placement_types import Replicate 

22from hyper_parallel.distributed._builder.default_templates import _multi_dim 

23from hyper_parallel.distributed._builder.planner import ShardingPlanner 

24from hyper_parallel.distributed.recipe_spec import ModuleShardingSpec, NamedPlacement 

25from hyper_parallel.distributed.tensor_parallel.param_role import ParamRole 

26from hyper_parallel.models._transformers import auto_model, model_builder 

27 

28 

29_CODEGEN_LEAF_SEGMENTS = frozenset({"linear_fc1", "linear_fc2"}) 

30 

31 

32def _override_declared_params(plan_overrides: dict[str, Any], model: Any) -> set[str]: 

33 """Return parameter FQNs already owned by exact or glob overrides.""" 

34 module_names = tuple(name for name, _ in model.named_modules()) 

35 declared = set() 

36 for key, spec in plan_overrides.items(): 

37 matches = ( 

38 (name for name in module_names if fnmatchcase(name, key)) 

39 if any(char in key for char in "*?[") 

40 else (key,) 

41 ) 

42 for module_fqn in matches: 

43 prefix = f"{module_fqn}." if module_fqn else "" 

44 declared.update(prefix + name for name in (spec.params or {})) 

45 return declared 

46 

47 

48class CodegenShardingPlanner(ShardingPlanner): 

49 """Add audit-only coverage without changing the Trainer's native planner.""" 

50 

51 def _infer_boundary_type(self, fqn: str, group: Any) -> str: 

52 """Let Megatron-style projection leaves merge into their MLP owner.""" 

53 if fqn.rsplit(".", 1)[-1].lower() in _CODEGEN_LEAF_SEGMENTS: 

54 return "unknown" 

55 return super()._infer_boundary_type(fqn, group) 

56 

57 def _group_by_boundary(self, param_roles: dict[str, ParamRole]) -> dict[str, list[tuple[str, ParamRole]]]: 

58 """Keep unknown root leftovers attached to each real parameter owner.""" 

59 groups = super()._group_by_boundary(param_roles) 

60 normalized: dict[str, list[tuple[str, ParamRole]]] = {} 

61 for module_fqn, group in groups.items(): 

62 if self._infer_boundary_type(module_fqn, group) != "unknown": 

63 normalized[module_fqn] = group 

64 continue 

65 for param_fqn, role in group: 

66 owner = param_fqn.rsplit(".", 1)[0] 

67 normalized.setdefault(owner, []).append((param_fqn, role)) 

68 return normalized 

69 

70 def _derive_boundary_specs(self, plan: Any, boundary_groups: Any, **kwargs: Any) -> None: 

71 """Append explicit replicated audit entries after native derivation.""" 

72 super()._derive_boundary_specs(plan, boundary_groups, **kwargs) 

73 if not self._derive: 

74 return 

75 model = kwargs["model"] 

76 parameters_by_name = dict(model.named_parameters()) 

77 override_declared = _override_declared_params(self._plan_overrides, model) 

78 for boundary_fqn, group in boundary_groups.items(): 

79 if boundary_fqn in plan.modules: 

80 continue 

81 fallback = self._build_replicated_fallback_spec( 

82 boundary_fqn, 

83 group, 

84 parameters_by_name, 

85 override_declared, 

86 ) 

87 if fallback is not None: 

88 plan.modules[boundary_fqn] = fallback 

89 

90 @staticmethod 

91 def _build_replicated_fallback_spec( 

92 boundary_fqn: str, 

93 group: list[tuple[str, ParamRole]], 

94 parameters_by_name: dict[str, Any], 

95 override_declared: set[str], 

96 ) -> Optional[ModuleShardingSpec]: 

97 """Declare communication-free metadata for replicated leftovers.""" 

98 replicable = (ParamRole.REPLICATED, ParamRole.NORM, ParamRole.BIAS) 

99 if any(role not in replicable for _, role in group): 

100 return None 

101 prefix = f"{boundary_fqn}." if boundary_fqn else "" 

102 params: dict[str, NamedPlacement] = {} 

103 for param_fqn, _ in group: 

104 parameter = parameters_by_name.get(param_fqn) 

105 if parameter is None or not parameter.requires_grad or param_fqn in override_declared: 

106 continue 

107 relative = param_fqn[len(prefix):] if prefix else param_fqn 

108 params[relative] = _multi_dim(tp=Replicate(), cp=Replicate(), ep=Replicate()) 

109 if not params: 

110 return None 

111 return ModuleShardingSpec(params=params, in_src={}, in_dst={}, is_boundary=False) 

112 

113 

114@contextmanager 

115def use_codegen_planner_overlay() -> Iterator[None]: 

116 """Scope planner observation and suppress FSDP runtime construction.""" 

117 original = auto_model.instantiate_infrastructure 

118 original_fsdp = model_builder._instantiate_fsdp2 # pylint: disable=protected-access 

119 

120 def skip_fsdp_runtime(*_args: Any, **_kwargs: Any) -> None: 

121 """Keep strategy configuration while suppressing device execution.""" 

122 return None 

123 

124 def instantiate_infrastructure(*args: Any, **kwargs: Any) -> Any: 

125 """Upgrade only an exact native planner created inside this context.""" 

126 planner, fsdp_manager = original(*args, **kwargs) 

127 if type(planner) is ShardingPlanner: # pylint: disable=unidiomatic-typecheck 

128 planner.__class__ = CodegenShardingPlanner 

129 return planner, fsdp_manager 

130 

131 auto_model.instantiate_infrastructure = instantiate_infrastructure 

132 model_builder._instantiate_fsdp2 = skip_fsdp_runtime # pylint: disable=protected-access 

133 try: 

134 yield 

135 finally: 

136 auto_model.instantiate_infrastructure = original 

137 model_builder._instantiate_fsdp2 = original_fsdp # pylint: disable=protected-access 

138 

139 

140__all__ = ["CodegenShardingPlanner", "use_codegen_planner_overlay"]