Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / expert_parallel / structure.py: 0%

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"""Conservative generation-time recognition of supported MoE structures.""" 

16 

17import ast 

18from dataclasses import dataclass 

19import inspect 

20import textwrap 

21from typing import Any 

22 

23from torch import nn 

24 

25 

26class UnsupportedModuleStructure(ValueError): 

27 """A matched module has no proven parallel recipe.""" 

28 

29 

30@dataclass(frozen=True) 

31class MoeStructure: 

32 """Serializable structural decisions frozen before distributed execution.""" 

33 

34 router: str 

35 shared_experts: str 

36 expert_storage: str 

37 jitter: bool = False 

38 

39 

40def _unsupported(module: Any, detail: str) -> UnsupportedModuleStructure: 

41 return UnsupportedModuleStructure( 

42 f"{type(module).__name__}: {detail}. Declare an explicit " 

43 "local_compute_fn._target_ to opt out of structure detection." 

44 ) 

45 

46 

47def _forward_tree(module: Any) -> ast.FunctionDef: 

48 try: 

49 source = textwrap.dedent(inspect.getsource(type(module).forward)) 

50 tree = ast.parse(source) 

51 except (AttributeError, OSError, TypeError, SyntaxError) as exc: 

52 raise _unsupported(module, "forward source is unavailable") from exc 

53 functions = [node for node in tree.body if isinstance(node, ast.FunctionDef)] 

54 if len(functions) != 1: 

55 raise _unsupported(module, "forward source is ambiguous") 

56 return functions[0] 

57 

58 

59def detect_router_kind(module: Any) -> str: 

60 """Recognize a router without invoking it or consulting a model identity. 

61 

62 Tuple routers keep their own scoring algorithm, including grouped sigmoid 

63 routing. Recomputing scores from their first return value would lose gate 

64 configuration and risks applying scaling twice. 

65 """ 

66 gate = getattr(module, "gate", None) 

67 if gate is None: 

68 gate = getattr(module, "router", None) 

69 if gate is None: 

70 raise _unsupported(module, "missing gate/router") 

71 if type(gate) is nn.Linear: # pylint: disable=C0123 # exact match: subclasses take the AST path 

72 config = getattr(module, "config", None) 

73 if not any(getattr(owner, name, None) is not None 

74 for owner in (module, config) for name in ("top_k", "num_experts_per_tok")): 

75 raise _unsupported(module, "linear router has no explicit top-k setting") 

76 # A bare-linear gate feeds either softmax top-k routing or 

77 # group-limited sigmoid routing; both gates are literally nn.Linear. 

78 # The merge contract is told apart by the sigmoid/group markers the 

79 # MOE declares on itself or its config (DeepSeek-V3 / GLM-4-MoE keep 

80 # routed_scaling_factor + n_group + topk_group; plain softmax routers 

81 # carry none). This mirrors the runtime router adapters, which key on 

82 # the same attributes. 

83 if any(getattr(owner, name, None) is not None 

84 for owner in (module, config) for name in ("routed_scaling_factor", "n_group", "topk_group")): 

85 return "sigmoid_group" 

86 return "softmax_topk" 

87 forward = _forward_tree(gate) 

88 returns = [node.value for node in ast.walk(forward) if isinstance(node, ast.Return)] 

89 calls = {node.func.attr for node in ast.walk(forward) 

90 if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute)} 

91 index_names = _topk_index_names(forward) 

92 if (returns and all(isinstance(value, ast.Tuple) and len(value.elts) == 3 

93 and isinstance(value.elts[2], ast.Name) 

94 and value.elts[2].id in index_names for value in returns) 

95 and "topk" in calls and "linear" in calls): 

96 return "topk_router_module" 

97 raise _unsupported(module, "router does not expose a supported logits/scores/indices contract") 

98 

99 

100def _topk_index_names(forward: ast.FunctionDef) -> set[str]: 

101 """Prove the final tuple item comes from a top-k index result.""" 

102 names = set() 

103 for node in ast.walk(forward): 

104 if not isinstance(node, ast.Assign) or len(node.targets) != 1: 

105 continue 

106 target, value = node.targets[0], node.value 

107 if (isinstance(target, (ast.Tuple, ast.List)) and len(target.elts) == 2 

108 and isinstance(target.elts[1], ast.Name) and _is_topk_call(value)): 

109 names.add(target.elts[1].id) 

110 elif (isinstance(target, ast.Name) and isinstance(value, ast.Subscript) 

111 and isinstance(value.slice, ast.Constant) and value.slice.value == 1 

112 and _is_topk_call(value.value)): 

113 names.add(target.id) 

114 return names 

115 

116 

117def _is_topk_call(value: ast.AST) -> bool: 

118 return (isinstance(value, ast.Call) and isinstance(value.func, ast.Attribute) 

119 and value.func.attr == "topk") 

120 

121 

122def detect_moe_structure(module: Any) -> MoeStructure: 

123 """Freeze the router, shared branch and expert parameter storage contract.""" 

124 experts = getattr(module, "experts", None) 

125 if experts is None: 

126 raise _unsupported(module, "missing experts") 

127 if isinstance(experts, nn.ModuleList): 

128 if not experts or not all(all(hasattr(expert, name) for name in 

129 ("gate_proj", "up_proj", "down_proj")) for expert in experts): 

130 raise _unsupported(module, "expert list is not a SwiGLU projection list") 

131 storage = "module_list" 

132 elif all(isinstance(getattr(experts, name, None), nn.Parameter) 

133 and getattr(experts, name).ndim == 3 for name in ("gate_up_proj", "down_proj")): 

134 storage = "batched_parameters" 

135 else: 

136 raise _unsupported(module, "unsupported expert parameter storage") 

137 shared_names = {name for name in ("shared_expert", "shared_expert_gate", "shared_experts") 

138 if getattr(module, name, None) is not None} 

139 if not shared_names: 

140 shared = "none" 

141 elif shared_names == {"shared_expert", "shared_expert_gate"}: 

142 shared = "gated" 

143 elif shared_names == {"shared_experts"}: 

144 shared = "additive" 

145 else: 

146 raise _unsupported(module, f"ambiguous shared expert branches {sorted(shared_names)}") 

147 return MoeStructure(detect_router_kind(module), shared, storage, hasattr(module, "jitter_noise"))