Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / plan / freeze.py: 23%
189 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# Licensed under the Apache License, Version 2.0
3# ============================================================================
4"""Freeze a ``ShardingPlan`` into JSON-safe structures for meta / generated code.
6The planner's plan objects are live (``Placement`` instances, ``Target``
7references, callables, internal ``init=False`` fields). The generated
8modeling file and ``codegen_meta.json`` must consume the SAME plan without
9re-running the planner — and meta must stay JSON-serializable. The freeze
10layer is that projection: placements become stable strings / dicts, targets
11become ``to_dict()``, internal EP metadata is made explicit, and the full
12sharded-parameter tree is expanded to concrete FQNs.
14Everything here is a pure function over planner output; nothing imports
15torch or the training stack.
16"""
17from __future__ import annotations
19import re
20from dataclasses import dataclass, field
21from typing import Any
23from hyper_parallel.codegen.plan.spec_fields import (
24 INJECTION_SPEC_FIELDS,
25 PARAM_PLAN_SPEC_FIELDS,
26)
28# ---------------------------------------------------------------------------
29# Placement serialization (placement_types.py)
30# ---------------------------------------------------------------------------
32# Registry of known placement types (imported lazily — freeze.py must stay
33# importable without torch) mapped to their string-form serializers. The
34# registry stays empty until the first serialization; ``_register_placement_types``
35# is the ONLY writer. The serializer classes are stateless holders for the
36# per-kind string logic — they must NOT self-register (a decorated registry
37# keyed by the serializer classes would poison the isinstance dispatch below).
38_PLACEMENT_SERIALIZERS: dict[type, type] = {}
41class _StridedShardSerializer:
42 @staticmethod
43 def to_string(p) -> str:
44 dim = p.dim if not p.uneven_shard else -p.dim - 1
45 return f"SS({dim},{p.split_factor})"
48class _ShardSerializer:
49 @staticmethod
50 def to_string(p) -> str:
51 dim = p.dim if not p.uneven_shard else -p.dim - 1
52 return f"S({dim})"
55class _ReplicateSerializer:
56 @staticmethod
57 def to_string(p) -> str:
58 return "R"
61class _PartialSerializer:
62 @staticmethod
63 def to_string(p) -> str:
64 return f"P({p.reduce_op})"
67def _register_placement_types() -> None:
68 """Register the real placement classes against their serializers.
70 ``placement_types`` is imported lazily (torch dependency) — the registry
71 stays empty until the first actual serialization, so an environment
72 without torch can still import ``freeze.py``.
74 Order matters: ``StridedShard`` subclasses ``Shard``, so it must be
75 registered FIRST or ``isinstance`` dispatch would serialize every
76 strided shard as a plain ``Shard``.
77 """
78 from hyper_parallel.core.dtensor.placement_types import (
79 Partial,
80 Replicate,
81 Shard,
82 StridedShard,
83 )
85 for cls, serializer in (
86 (StridedShard, _StridedShardSerializer),
87 (Shard, _ShardSerializer),
88 (Replicate, _ReplicateSerializer),
89 (Partial, _PartialSerializer),
90 ):
91 _PLACEMENT_SERIALIZERS[cls] = serializer
94def placement_to_string(placement: Any) -> str:
95 """Serialize a ``Placement`` instance to its canonical string form.
97 Shard(dim) -> "S(dim)"; Replicate() -> "R"; Partial(reduce_op) -> "P";
98 StridedShard(dim, split_factor) -> "SS(dim,factor)". ``uneven_shard`` is
99 encoded as a negative dim (``Shard(-dim-1)``) so it survives the
100 round-trip without an extra field. Unknown placement types fail fast —
101 a silently dropped placement would corrupt the frozen plan.
102 """
103 if not _PLACEMENT_SERIALIZERS:
104 _register_placement_types()
105 for cls, serializer in _PLACEMENT_SERIALIZERS.items():
106 if isinstance(placement, cls):
107 return serializer.to_string(placement)
108 raise TypeError(
109 f"cannot serialize placement {placement!r} of unknown type "
110 f"{type(placement).__name__}"
111 )
114# ---------------------------------------------------------------------------
115# Named-placement serialization
116# ---------------------------------------------------------------------------
118def named_placement_to_dict(named: Any) -> dict[str, str]:
119 """Serialize a ``NamedPlacement`` / output-contract dict to JSON.
121 ``in_*`` entries are keyed by input name (``{"hidden_states": ...}``);
122 ``out_*`` entries are normalized by the planner into
123 ``{"output": ...}`` — both wrap the same axis -> ``Placement``
124 dictionary, so the value at every level may itself be a dict and is
125 recursed until a leaf ``Placement`` (or a plain scalar) is reached.
126 Axis keys are normalized to strings and sorted for stability; a
127 ``None`` value (unset) serializes to an empty dict.
128 """
129 if not named:
130 return {}
131 return {
132 str(getattr(axis, "value", axis)): _value_to_dict(placement)
133 for axis, placement in named.items()
134 }
137# ---------------------------------------------------------------------------
138# Placement deserialization (inverse of the serializers above)
139# ---------------------------------------------------------------------------
141def parse_placement(text: str) -> Any:
142 """Inverse of :func:`placement_to_string`: canonical string -> Placement.
144 Compact forms: ``S(dim)`` / ``R`` / ``P(reduce_op)`` / ``SS(dim,factor)``;
145 ``uneven_shard`` is recovered from the negative-dim encoding
146 (``S(-1)`` -> ``Shard(0, uneven_shard=True)``). Unknown or malformed
147 forms fail fast — a silently downgraded placement would corrupt the
148 frozen plan (e.g. quietly becoming ``Replicate()``).
149 """
150 from hyper_parallel.core.dtensor.placement_types import (
151 Partial,
152 Replicate,
153 Shard,
154 StridedShard,
155 )
157 if not isinstance(text, str) or not text:
158 raise ValueError(f"cannot parse placement from {text!r}")
160 if text == "R":
161 return Replicate()
163 match = re.fullmatch(r"([SP]|SS)\((.*)\)", text)
164 if match is None:
165 raise ValueError(
166 f"unknown placement form {text!r} (expected S(dim) / R / "
167 f"P(reduce_op) / SS(dim,factor))"
168 )
169 kind, args = match.groups()
171 if kind == "P":
172 reduce_op = args
173 if not reduce_op:
174 raise ValueError(f"malformed partial placement {text!r}: missing reduce_op")
175 return Partial(reduce_op)
177 parts = [part.strip() for part in args.split(",")]
178 try:
179 dim = int(parts[0])
180 except (IndexError, ValueError) as exc:
181 raise ValueError(f"malformed placement {text!r}: invalid dim") from exc
182 uneven = dim < 0
183 dim = -dim - 1 if uneven else dim
185 if kind == "S":
186 if len(parts) != 1:
187 raise ValueError(f"malformed shard placement {text!r}: too many args")
188 return Shard(dim, uneven_shard=uneven)
189 if kind == "SS":
190 if len(parts) != 2:
191 raise ValueError(f"malformed strided shard {text!r}: expected SS(dim,factor)")
192 try:
193 split_factor = int(parts[1])
194 except ValueError as exc:
195 raise ValueError(
196 f"malformed strided shard {text!r}: invalid split_factor"
197 ) from exc
198 return StridedShard(dim, split_factor, uneven_shard=uneven)
199 raise ValueError(f"unknown placement kind {kind!r} in {text!r}")
202def _value_from_dict(value: Any) -> Any:
203 """Inverse of :func:`_value_to_dict`: JSON-safe value -> planner-ish value.
205 Placement strings become ``Placement`` objects; nested dicts / lists
206 recurse; plain scalars pass through unchanged. ``Target`` objects
207 serialized via ``to_dict()`` are NOT reconstructed — the frozen boundary
208 contract is data, and the generated module resolves runtime targets
209 itself.
210 """
211 if isinstance(value, str):
212 return parse_placement(value)
213 if isinstance(value, dict):
214 return {str(getattr(k, "value", k)): _value_from_dict(v) for k, v in value.items()}
215 if isinstance(value, list):
216 return [_value_from_dict(v) for v in value]
217 return value
220def parse_named_placement(data: dict[str, Any]) -> Any:
221 """Inverse of :func:`named_placement_to_dict`: JSON dict -> named placement.
223 Round-trips ``named_placement_to_dict()`` output back to a
224 ``NamedPlacement`` / output-contract dict with real ``Placement`` leaves.
225 """
226 if not data:
227 return {}
228 return _value_from_dict(data)
231# ---------------------------------------------------------------------------
232# FrozenPlan
233# ---------------------------------------------------------------------------
235@dataclass
236class FrozenPlan:
237 """JSON-safe projection of a ``ShardingPlan`` for meta / generated code.
239 - ``param_plan``: per-boundary parameter placements (the sharding
240 contract the generated file executes).
241 - ``injections``: per-boundary injection declarations (CP wrappers /
242 local compute / inner targets), already serialized to strings/dicts.
243 - ``tied_pairs``: tied-weight FQN pairs.
244 - ``special_handlers``: special-handler name per parameter.
245 - ``frozen_sharded_params``: the full FQN list of every parameter that
246 takes a non-``Replicate`` placement anywhere in the plan — what the
247 generated file must create as DTensors at load time.
248 - ``boundary_classes``: {边界FQN: 类名} — the class that owns each
249 boundary's ``forward``. The frozen plan is keyed by
250 boundary FQN, but the generated file rewrites ``class.forward``, which
251 one class shares across many instances; the emitter groups FQNs by this
252 map (failed to resolve a class = a boundary whose forward is never
253 rewritten = silent drop).
254 - ``mesh_dim_names``: the plan's active axes (``ShardingPlan.
255 mesh_dim_names`` verbatim — ``tp``/``cp``/``ep`` with size-1 axes and
256 dp* axes already stripped). The runtime slices the live mesh down to
257 these axes (mirroring ``_get_active_mesh``) so a placement never lands
258 on a dp axis the plan never shards on.
259 """
261 param_plan: dict[str, dict[str, Any]] = field(default_factory=dict)
262 injections: list[dict[str, Any]] = field(default_factory=list)
263 tied_pairs: list[list[str]] = field(default_factory=list)
264 special_handlers: dict[str, str] = field(default_factory=dict)
265 frozen_sharded_params: list[str] = field(default_factory=list)
266 boundary_classes: dict[str, str] = field(default_factory=dict)
267 mesh_dim_names: tuple[str, ...] = ()
270def freeze_plan(plan: Any, model: Any) -> FrozenPlan:
271 """Freeze a planner ``ShardingPlan`` into a ``FrozenPlan``.
273 ``model`` is the (meta-device) model the plan was derived from — used to
274 expand the sharded-parameter tree to concrete FQNs and to record each
275 boundary FQN's owning class (``boundary_classes``). ``model``
276 may be ``None`` when only the plan shape is frozen (class names are then
277 left unresolvable — the manager passes a live model).
278 """
279 if plan is None:
280 return FrozenPlan()
281 return FrozenPlan(
282 param_plan=freeze_param_plan(plan),
283 injections=freeze_injections(plan),
284 tied_pairs=freeze_tied_pairs(plan),
285 special_handlers=dict(getattr(plan, "special_handlers", {}) or {}),
286 frozen_sharded_params=expand_frozen_sharded_params(plan),
287 boundary_classes=_freeze_boundary_classes(plan, model),
288 mesh_dim_names=tuple(getattr(plan, "mesh_dim_names", ()) or ()),
289 )
292def _freeze_boundary_classes(plan: Any, model: Any) -> dict[str, str]:
293 """Map each boundary FQN to the class that owns its ``forward``.
295 The generated file rewrites ``class.forward``, which one class shares
296 across many instances — but the frozen plan is keyed by boundary FQN. For
297 every ``plan.modules`` key (``model.layers.0.self_attn``), record the type
298 of the submodule at that FQN. ``get_submodule("")`` returns the root
299 module and yields the model class itself.
300 """
301 if model is None:
302 return {}
303 modules = getattr(plan, "modules", {}) or {}
304 result: dict[str, str] = {}
305 for fqn in modules:
306 cls = type(model.get_submodule(fqn)).__name__
307 result[fqn] = cls
308 return result
311def freeze_param_plan(plan: Any) -> dict[str, dict[str, Any]]:
312 """Freeze the per-boundary parameter sharding plan.
314 Output shape (one entry per boundary module)::
316 {
317 "model.layers.0.self_attn": {
318 "params": { # per-parameter placement
319 "q_proj.weight": {"tp": "S(0)"},
320 "k_proj.weight": {"tp": "S(0)"},
321 ...
322 },
323 "in_src": {"hidden_states": {"tp": "S(-1)"}}, # optional
324 ...
325 },
326 ...
327 }
329 Non-``Replicate`` placements are the sharding contract; ``Replicate``
330 entries are kept too (an explicit "this stays replicated") so the
331 generated file never has to guess. Placement values are always strings
332 via :func:`placement_to_string`; parameter names are kept as declared
333 (with the ``.weight`` / ``.bias`` suffix).
335 Which fields are frozen (and how) is declared in
336 :mod:`hyper_parallel.codegen.plan.spec_fields` — the table both sides
337 of the meta boundary walk, so freeze and rebuild cannot drift apart.
338 """
339 modules = getattr(plan, "modules", {}) or {}
340 result: dict[str, dict[str, Any]] = {}
341 for fqn, spec in modules.items():
342 entry: dict[str, Any] = {}
343 for spec_field in PARAM_PLAN_SPEC_FIELDS:
344 frozen = _freeze_spec_field(spec_field, getattr(spec, spec_field.name, None))
345 if frozen is not _MISSING:
346 entry[spec_field.key] = frozen
347 result[fqn] = entry
348 return result
351#: Sentinel distinguishing "field absent from the frozen entry" from a
352#: legitimately frozen ``None``/empty value (``{"params": {}}``).
353_MISSING = object()
356def _freeze_spec_field(spec_field: Any, value: Any) -> Any:
357 """Serialize one spec field per its :class:`SpecField` kind contract.
359 Returns ``_MISSING`` when the field must not appear in the frozen
360 output. The write condition is part of the kind (see spec_fields.py):
361 ``params`` uses ``is not None`` so an explicit empty dict ("shards
362 nothing") survives — a truthy check here used to drop it and leave the
363 rebuilt spec with ``params=None``, which Phase A cannot iterate.
364 """
365 kind = spec_field.freeze
366 if kind in ("placement_map", "named_placements", "scalar",
367 "value_to_dict", "plain", "attr_list"):
368 if value is None:
369 return _MISSING
370 elif not value: # name_list / dict_copy / int_truthy: truthy to write
371 return _MISSING
372 if kind == "placement_map":
373 return {
374 name: {
375 str(getattr(axis, "value", axis)): placement_to_string(placement)
376 for axis, placement in named.items()
377 }
378 for name, named in value.items()
379 }
380 if kind == "named_placements":
381 return named_placement_to_dict(value)
382 if kind in ("name_list", "attr_list"):
383 return list(value)
384 if kind == "value_to_dict":
385 return _value_to_dict(value)
386 if kind == "dict_copy":
387 return dict(value)
388 # scalar / plain / int_truthy: the value is already JSON-safe.
389 return value
392def freeze_injections(plan: Any) -> list[dict[str, Any]]:
393 """Freeze per-boundary injection declarations.
395 CP attention wrappers (``inner_wrapper``), local compute regions
396 (``local_compute_fn``), and inner targets — everything the generated
397 file must attach at apply time. ``Target`` values are resolved to their
398 ``to_dict()`` form; callables are recorded by name (they cannot cross
399 the meta boundary — the generated file re-resolves them from the same
400 runtime registry). Field set and write conditions come from
401 :mod:`hyper_parallel.codegen.plan.spec_fields`.
402 """
403 modules = getattr(plan, "modules", {}) or {}
404 injections: list[dict[str, Any]] = []
405 for fqn, spec in modules.items():
406 entry: dict[str, Any] = {}
407 for spec_field in INJECTION_SPEC_FIELDS:
408 frozen = _freeze_spec_field(spec_field, getattr(spec, spec_field.name, None))
409 if frozen is not _MISSING:
410 entry[spec_field.key] = frozen
411 if entry:
412 injections.append({"match": fqn, **entry})
413 return injections
416def freeze_tied_pairs(plan: Any) -> list[list[str]]:
417 """Freeze tied-weight pairs as stable string lists."""
418 pairs = getattr(plan, "tied_pairs", None) or []
419 return [list(pair) for pair in pairs]
422def expand_frozen_sharded_params(plan: Any) -> list[str]:
423 """Expand every sharded parameter in the plan to a concrete FQN.
425 A parameter is "sharded" when any axis of its placement is a non-
426 ``Replicate`` placement (``Shard`` / ``StridedShard`` / ``Partial``).
427 The FQN is the module FQN + ``.weight`` / ``.bias`` suffix — the same
428 key the generated runtime uses to look up the parameter.
429 """
430 modules = getattr(plan, "modules", {}) or {}
431 fqns: list[str] = []
432 seen: set[str] = set()
433 for fqn, spec in modules.items():
434 for name, named in (getattr(spec, "params", {}) or {}).items():
435 if _is_sharded(named):
436 fqn_key = f"{fqn}.{name}" if name else fqn
437 if fqn_key not in seen:
438 seen.add(fqn_key)
439 fqns.append(fqn_key)
440 return sorted(fqns)
443# ---------------------------------------------------------------------------
444# helpers
445# ---------------------------------------------------------------------------
447def _is_sharded(named: Any) -> bool:
448 """Whether any axis of a named placement is non-``Replicate``."""
449 if not named:
450 return False
451 for placement in named.values():
452 if placement is None:
453 continue
454 kind = type(placement).__name__
455 if kind in ("Shard", "StridedShard", "Partial"):
456 return True
457 return False
460def _value_to_dict(value: Any) -> Any:
461 """Serialize a spec value: Target -> to_dict(), Placement -> string,
462 named placement -> dict, callable -> name."""
463 from hyper_parallel.core.dtensor.placement_types import Placement
465 if isinstance(value, Placement):
466 return placement_to_string(value)
467 if hasattr(value, "to_dict"):
468 return value.to_dict()
469 if isinstance(value, dict):
470 return {
471 str(getattr(k, "value", k)): _value_to_dict(v)
472 for k, v in value.items()
473 }
474 if isinstance(value, (list, tuple)):
475 return [_value_to_dict(v) for v in value]
476 if callable(value):
477 return getattr(value, "__name__", str(value))
478 return value
481__all__ = [
482 "FrozenPlan",
483 "expand_frozen_sharded_params",
484 "freeze_injections",
485 "freeze_param_plan",
486 "freeze_plan",
487 "freeze_tied_pairs",
488 "named_placement_to_dict",
489 "parse_named_placement",
490 "parse_placement",
491 "placement_to_string",
492]