Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / compile / graph_parallel_plan.py: 94%
85 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
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"""
16Graph Parallel Plan - declarative module selection for graph-mode passes.
18Declares which modules the graph-mode ``FSDPPass`` should shard and which
19modules each pipeline stage owns. ``FSDPPass`` itself owns all the actual
20sharding logic (all_gather on parameter placeholders, reduce_scatter on
21gradient outputs, in-place live-model sharding), so this is purely a *which
22modules* lookup.
24Note:
25 Earlier revisions stored an ``FSDPModuleConfig`` dataclass per entry and
26 exposed ``get_fsdp_config`` / ``merge``. The dataclass carried no field
27 beyond the FQN that already keys it, and neither helper had a production
28 consumer. They were removed as dead surface; when a real per-module
29 option (reshard, CPU offload) lands it can re-add a value type with a
30 real consumer.
31"""
33__all__ = [
34 "GraphParallelPlan",
35 "create_plan_from_yaml",
36 "create_all_fsdp_plan",
37]
39import fnmatch
40from dataclasses import dataclass, field
41from pathlib import Path
42from typing import List, Optional, Set
44import yaml
46DEFAULT_CONFIG_DIR = Path(__file__).parent / "examples"
49@dataclass
50class GraphParallelPlan:
51 """Declare the graph-mode parallel assignment of modules.
53 ``fsdp_modules`` holds exact module FQNs to FSDP-shard and
54 ``fsdp_patterns`` holds ``fnmatch`` wildcards; :meth:`is_marked_for_fsdp`
55 checks the exact set first, then the patterns.
57 Pipeline-parallel stage assignment lives in
58 ``pp_module_fqns_per_stage``: a list whose i-th entry is the list of
59 module FQNs (exact, no wildcards) assigned to stage ``i``. When it is
60 ``None`` the ``PpPass`` falls back to an automatic even split of the
61 model's layer-like children (see ``pp_pass._auto_stage_split``).
63 Example:
64 plan = GraphParallelPlan()
65 plan.fsdp_mark("tok_embeddings")
66 plan.fsdp_mark_pattern("layers.*")
67 plan.pp_stage(0, ["tok_embeddings", "layers.0"])
68 plan.pp_stage(1, ["layers.1", "norm", "lm_head"])
69 """
71 fsdp_modules: Set[str] = field(default_factory=set)
72 fsdp_patterns: Set[str] = field(default_factory=set)
73 pp_module_fqns_per_stage: Optional[List[List[str]]] = None
75 def fsdp_mark(self, module_fqn: str) -> "GraphParallelPlan":
76 """Mark a specific module for FSDP sharding (exact match).
78 Args:
79 module_fqn: Module fully qualified name.
81 Returns:
82 ``self`` (chainable).
84 Example:
85 plan.fsdp_mark("tok_embeddings")
86 """
87 self.fsdp_modules.add(module_fqn)
88 return self
90 def fsdp_mark_pattern(self, pattern: str) -> "GraphParallelPlan":
91 """Mark modules for FSDP sharding (wildcard match).
93 Args:
94 pattern: Module FQN pattern (``fnmatch`` wildcards, e.g. ``*``,
95 ``layers.*``).
97 Returns:
98 ``self`` (chainable).
100 Example:
101 plan.fsdp_mark_pattern("layers.*")
102 """
103 self.fsdp_patterns.add(pattern)
104 return self
106 def is_marked_for_fsdp(self, module_fqn: str) -> bool:
107 """Whether ``module_fqn`` is marked for FSDP sharding.
109 Pure exact/pattern lookup: an exact mark or any matching pattern
110 returns ``True``. Ancestor matching (a mark on ``layers.0`` covering
111 ``layers.0.attention.weight``) is NOT done here — callers such as
112 ``FSDPPass`` walk a parameter's ancestors and query each level.
113 """
114 if module_fqn in self.fsdp_modules:
115 return True
117 return any(
118 fnmatch.fnmatch(module_fqn, pattern) for pattern in self.fsdp_patterns
119 )
121 def pp_stage(self, stage_idx: int, module_fqns: List[str]) -> "GraphParallelPlan":
122 """Declare the module FQNs of one pipeline stage (exact match).
124 Args:
125 stage_idx: Zero-based stage index, ``stage_idx >= 0``. Stages
126 may be declared in any order / sparsely; ``PpPass``
127 validates completeness (every stage ``0..pp_degree-1``
128 declared, no FQN assigned twice) before splitting.
129 module_fqns: Module FQNs owned by this stage, in model order.
130 Exact FQNs only — wildcards are rejected because a stage
131 cut must be unambiguous.
133 Returns:
134 ``self`` (chainable).
136 Raises:
137 ValueError: If ``stage_idx`` is negative or a module FQN
138 contains wildcard characters.
140 Example:
141 plan.pp_stage(0, ["tok_embeddings", "layers.0"])
142 plan.pp_stage(1, ["layers.1", "norm", "lm_head"])
143 """
144 if stage_idx < 0:
145 raise ValueError(f"stage_idx must be >= 0, got {stage_idx}")
146 for fqn in module_fqns:
147 if any(ch in fqn for ch in "*?["):
148 raise ValueError(
149 f"pp_stage() takes exact module FQNs, got wildcard pattern "
150 f"'{fqn}' — a stage cut must be unambiguous"
151 )
152 if self.pp_module_fqns_per_stage is None:
153 self.pp_module_fqns_per_stage = []
154 while len(self.pp_module_fqns_per_stage) <= stage_idx:
155 self.pp_module_fqns_per_stage.append([])
156 self.pp_module_fqns_per_stage[stage_idx] = list(module_fqns)
157 return self
160def create_plan_from_yaml(
161 config_path: Optional[str] = None,
162 model_name: Optional[str] = None,
163) -> GraphParallelPlan:
164 """Create a ``GraphParallelPlan`` from a YAML configuration file.
166 Args:
167 config_path: Path to YAML config file.
168 model_name: Model name (looks up in ``examples/{model_name}/config.yaml``).
170 Returns:
171 GraphParallelPlan object.
173 Raises:
174 ValueError: When neither argument is given, ``model_name`` is empty
175 or contains path separators, or the YAML is not a mapping.
176 FileNotFoundError: When the resolved config file does not exist.
178 Example:
179 plan = create_plan_from_yaml(model_name="llama3")
180 plan = create_plan_from_yaml(config_path="path/to/config.yaml")
181 """
182 if config_path is None and model_name is None:
183 raise ValueError("Must provide either config_path or model_name")
185 if config_path is None:
186 if not model_name or not isinstance(model_name, str):
187 raise ValueError("model_name must be a non-empty string")
188 if ".." in model_name or "/" in model_name or "\\" in model_name:
189 raise ValueError(
190 f"Invalid model_name '{model_name}': must not contain path "
191 "separators or parent directory references"
192 )
193 config_path = DEFAULT_CONFIG_DIR / model_name / "config.yaml"
194 else:
195 config_path = Path(config_path)
197 if not config_path.exists():
198 raise FileNotFoundError(f"Config file not found: {config_path}")
200 with open(config_path, "r", encoding="utf-8") as f:
201 config = yaml.safe_load(f)
203 if config is None:
204 raise ValueError(f"YAML config file is empty: {config_path}")
205 if not isinstance(config, dict):
206 raise ValueError(
207 f"YAML config must be a mapping (dict), "
208 f"got {type(config).__name__}: {config_path}"
209 )
211 plan = GraphParallelPlan()
213 fsdp_config = _yaml_section(config, "fsdp", config_path)
214 _maybe_process_fsdp(plan, fsdp_config)
216 pp_config = _yaml_section(config, "pp", config_path)
217 if pp_config.get("stages"):
218 _process_pp(plan, pp_config)
220 return plan
223def _yaml_section(config: dict, key: str, config_path: Path) -> dict:
224 """Return YAML ``key`` as a mapping; empty/``None`` becomes ``{}``.
226 ``or {}``: a YAML key present but empty (``fsdp:`` with only comments
227 under it) parses to None, and ``dict.get(key, {})`` then returns
228 None instead of the default.
229 """
230 section = config.get(key) or {}
231 if not isinstance(section, dict):
232 # A present-but-empty section parses to None and is normalized to {}
233 # above, so anything landing here is a real scalar/sequence typo.
234 raise ValueError(
235 f"YAML '{key}' section must be a mapping (e.g. nested keys or an "
236 f"empty section); got {type(section).__name__} in {config_path}"
237 )
238 return section
241def _maybe_process_fsdp(plan: GraphParallelPlan, fsdp_config: dict) -> None:
242 """Run the FSDP processor when the section is enabled.
244 Enabled defaults to True when explicit modules/patterns are declared.
245 """
246 has_explicit = bool(fsdp_config.get("modules")) or bool(fsdp_config.get("patterns"))
247 if fsdp_config.get("enabled", has_explicit):
248 _process_fsdp(plan, fsdp_config)
251def _process_fsdp(plan: GraphParallelPlan, fsdp_config: dict) -> None:
252 """Process FSDP configuration (modules + patterns) into the plan."""
253 for module_config in fsdp_config.get("modules", []):
254 plan.fsdp_mark(module_config["name"])
256 for pattern_config in fsdp_config.get("patterns", []):
257 plan.fsdp_mark_pattern(pattern_config["pattern"])
260def _process_pp(plan: GraphParallelPlan, pp_config: dict) -> None:
261 """Process PP configuration (explicit per-stage FQN lists) into the plan.
263 YAML shape::
265 pp:
266 stages:
267 - [tok_embeddings, layers.0]
268 - [layers.1, norm, lm_head]
270 An entry may also be a mapping with ``stage`` / ``modules`` keys for
271 readability::
273 pp:
274 stages:
275 - stage: 0
276 modules: [tok_embeddings, layers.0]
277 """
278 for idx, stage in enumerate(pp_config["stages"]):
279 if isinstance(stage, dict):
280 stage_idx = stage.get("stage", idx)
281 module_fqns = list(stage.get("modules", []))
282 else:
283 stage_idx = idx
284 module_fqns = list(stage)
285 plan.pp_stage(stage_idx, module_fqns)
288def create_all_fsdp_plan() -> GraphParallelPlan:
289 """Create a plan that FSDP-marks every module (``*`` pattern).
291 Convenience for tests / quick demos.
292 """
293 plan = GraphParallelPlan()
294 plan.fsdp_mark_pattern("*")
295 return plan