Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / tools / model_integration / integration_validator.py: 74%
194 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"""Build/runtime model integration validators backed by final model and FSDP state facts."""
17from __future__ import annotations
19import fnmatch
20from collections import defaultdict
21from dataclasses import dataclass
22from typing import Any, Iterable, Mapping, Optional
24from hyper_parallel.tools.model_integration.structure_inventory import build_model_inventory
25from hyper_parallel.tools.model_integration.checkpoint_coverage import build_checkpoint_coverage
26from hyper_parallel.tools.model_integration.schemas import (
27 ModelIntegrationReport,
28 FindingSeverity,
29 ModelIntegrationFinding,
30)
31from hyper_parallel.models.registry import get_model_adapter
32from hyper_parallel.models.validation_spec import ModelValidationSpec
35@dataclass(frozen=True)
36class ParameterOwnership:
37 """One final parameter's inferred FSDP owner and gradient domain."""
39 parameter_fqn: str
40 owner_fqn: str
41 gradient_domain: str
42 source_mesh: tuple[str, ...]
43 source_placements: tuple[str, ...]
44 global_shape: tuple[int, ...]
45 local_shape: tuple[int, ...]
46 parameter_identity: int = 0
48 def to_dict(self) -> dict[str, Any]:
49 """Return a flat representation suitable for CSV or JSON."""
50 return {
51 "parameter_fqn": self.parameter_fqn,
52 "owner_fqn": self.owner_fqn,
53 "gradient_domain": self.gradient_domain,
54 "source_mesh": ",".join(self.source_mesh),
55 "source_placements": ",".join(self.source_placements),
56 "global_shape": repr(self.global_shape),
57 "local_shape": repr(self.local_shape),
58 }
61def resolve_model_validation_spec(model: Any) -> Optional[ModelValidationSpec]:
62 """Resolve and type-check the final model family's validation provider."""
63 config = getattr(model, "config", None)
64 identities = [getattr(config, "model_type", None)]
65 identities.extend(getattr(config, "architectures", None) or ())
66 adapter_spec = None
67 for identity in identities:
68 if not identity:
69 continue
70 adapter_spec = get_model_adapter(identity)
71 if adapter_spec is not None:
72 break
73 provider = getattr(adapter_spec, "validation", None)
74 if provider is None:
75 return None
76 if not callable(provider):
77 raise TypeError("ModelAdapterSpec.validation must be callable")
78 validation_spec = provider() # pylint: disable=not-callable
79 if not isinstance(validation_spec, ModelValidationSpec):
80 raise TypeError(
81 "ModelAdapterSpec.validation must return ModelValidationSpec, "
82 f"got {type(validation_spec).__name__}"
83 )
84 return validation_spec
87def _tensor_local_shape(tensor: Any) -> tuple[int, ...]:
88 to_local = getattr(tensor, "to_local", None)
89 local = to_local() if callable(to_local) else tensor
90 return tuple(int(size) for size in getattr(local, "shape", ()))
93def _source_layout(hsdp_param: Any) -> tuple[tuple[str, ...], tuple[str, ...]]:
94 source_info = getattr(hsdp_param, "source_shard_info", None)
95 source_mesh = getattr(source_info, "mesh", None)
96 names = tuple(getattr(source_mesh, "mesh_dim_names", None) or ())
97 placements = tuple(repr(value) for value in getattr(source_info, "placements", ()) or ())
98 return names, placements
101def _gradient_domain(hsdp_param: Any) -> str:
102 names, placements = _source_layout(hsdp_param)
103 expert_axis = any(name in ("ep", "expert") or name.startswith("ep_") for name in names)
104 expert_shard = any("Shard" in placement for placement in placements)
105 if expert_axis and expert_shard:
106 return "EXPERT"
107 return "DENSE"
110def collect_parameter_ownership(model: Any) -> tuple[list[ParameterOwnership], dict[str, Any]]:
111 """Collect parameter-to-FSDP-unit ownership from final runtime state."""
112 ownership = []
113 units = {}
114 seen_parameters = set()
115 for module_fqn, module in model.named_modules():
116 scheduler = getattr(module, "hsdp_scheduler", None)
117 state = getattr(scheduler, "hsdp_state", None)
118 hsdp_params = tuple(getattr(state, "hsdp_params", ()) or ())
119 if scheduler is None:
120 continue
121 owner_fqn = module_fqn or "<root>"
122 unit_domains = defaultdict(list)
123 for parameter_index, hsdp_param in enumerate(hsdp_params):
124 parameter = getattr(hsdp_param, "sharded_param", None)
125 parameter_fqn = getattr(hsdp_param, "_param_fqn", None) or (
126 f"{owner_fqn}.param[{parameter_index}]"
127 )
128 names, placements = _source_layout(hsdp_param)
129 domain = _gradient_domain(hsdp_param)
130 unit_domains[domain].append(parameter_fqn)
131 global_shape = tuple(
132 int(size) for size in getattr(hsdp_param, "_orig_size", getattr(parameter, "shape", ()))
133 )
134 ownership.append(
135 ParameterOwnership(
136 parameter_fqn=parameter_fqn,
137 owner_fqn=owner_fqn,
138 gradient_domain=domain,
139 source_mesh=names,
140 source_placements=placements,
141 global_shape=global_shape,
142 local_shape=_tensor_local_shape(parameter),
143 parameter_identity=id(parameter) if parameter is not None else 0,
144 )
145 )
146 if parameter is not None:
147 seen_parameters.add(id(parameter))
148 units[owner_fqn] = {
149 "domains": {name: sorted(fqns) for name, fqns in sorted(unit_domains.items())},
150 "parameter_count": len(hsdp_params),
151 "forward_prefetch_targets": [
152 _module_fqn(model, target)
153 for target in getattr(scheduler, "forward_prefetch_cells", ())
154 ],
155 "backward_prefetch_targets": [
156 _module_fqn(model, target)
157 for target in getattr(scheduler, "backward_prefetch_cells", ())
158 ],
159 }
160 for parameter_fqn, parameter in model.named_parameters(remove_duplicate=False):
161 if id(parameter) in seen_parameters:
162 continue
163 ownership.append(
164 ParameterOwnership(
165 parameter_fqn=parameter_fqn,
166 owner_fqn="<unmanaged>",
167 gradient_domain="LOCAL" if not parameter.requires_grad else "DENSE",
168 source_mesh=(),
169 source_placements=(),
170 global_shape=tuple(int(size) for size in parameter.shape),
171 local_shape=_tensor_local_shape(parameter),
172 parameter_identity=id(parameter),
173 )
174 )
175 return ownership, units
178def _module_fqn(model: Any, target: Any) -> str:
179 for module_fqn, module in model.named_modules():
180 if module is target:
181 return module_fqn or "<root>"
182 return f"<untracked:{type(target).__name__}>"
185def _validate_gradient_domains(
186 ownership: Iterable[ParameterOwnership],
187) -> list[ModelIntegrationFinding]:
188 """Emit findings for conflicting ownership or gradient domains."""
189 grouped = defaultdict(lambda: defaultdict(list))
190 for entry in ownership:
191 if entry.owner_fqn == "<unmanaged>" or entry.gradient_domain == "LOCAL":
192 continue
193 grouped[entry.owner_fqn][entry.gradient_domain].append(entry)
194 findings = []
195 aliases = defaultdict(list)
196 for entry in ownership:
197 if entry.parameter_identity:
198 aliases[entry.parameter_identity].append(entry)
199 for entries in aliases.values():
200 contracts = {
201 (
202 entry.owner_fqn,
203 entry.gradient_domain,
204 entry.source_mesh,
205 entry.source_placements,
206 )
207 for entry in entries
208 }
209 if len(contracts) <= 1:
210 continue
211 findings.append(
212 ModelIntegrationFinding(
213 phase="D3",
214 owner_fqn=entries[0].parameter_fqn,
215 severity=FindingSeverity.ERROR,
216 message="aliased parameter has conflicting ownership or gradient domains",
217 facts={
218 "aliases": [entry.parameter_fqn for entry in entries],
219 "contracts": [repr(contract) for contract in sorted(contracts, key=repr)],
220 },
221 why_unsafe="one storage cannot be reduced or sharded under two incompatible contracts",
222 remediation="assign every alias of the tied parameter to one owner and source layout",
223 related_config=("plan_overrides", "model adapter fsdp_wrap_modules"),
224 )
225 )
226 for owner_fqn, domain_entries in sorted(grouped.items()):
227 domains = set(domain_entries)
228 if len(domains) > 1:
229 facts = {
230 domain: [entry.parameter_fqn for entry in entries]
231 for domain, entries in sorted(domain_entries.items())
232 }
233 findings.append(
234 ModelIntegrationFinding(
235 phase="D3",
236 owner_fqn=owner_fqn,
237 severity=FindingSeverity.ERROR,
238 message="FSDP unit has incompatible gradient domains",
239 facts=facts,
240 why_unsafe=(
241 "one FSDP unit selects one collective mesh, so dense and expert "
242 "parameters cannot both receive their required gradient reduction"
243 ),
244 remediation=(
245 "declare homogeneous child FSDP units and keep dense gates/norms "
246 "in a dense owner"
247 ),
248 related_config=("model adapter fsdp_wrap_modules", "plan_overrides"),
249 )
250 )
251 if owner_fqn == "<root>" and "EXPERT" in domains:
252 findings.append(
253 ModelIntegrationFinding(
254 phase="D3",
255 owner_fqn=owner_fqn,
256 severity=FindingSeverity.ERROR,
257 message="expert-domain parameter remains in the dense root FSDP unit",
258 facts={
259 "expert_parameters": [
260 entry.parameter_fqn for entry in domain_entries["EXPERT"]
261 ]
262 },
263 why_unsafe="the dense root does not own the expert reduction domain",
264 remediation="declare a homogeneous nested expert FSDP unit",
265 related_config=("model adapter fsdp_wrap_modules",),
266 )
267 )
268 return findings
271def _validate_materialization(model: Any) -> list[ModelIntegrationFinding]:
272 """Emit findings for meta tensors and incomplete buffer rebuilds."""
273 meta_state = []
274 for name, value in (
275 *model.named_parameters(remove_duplicate=False),
276 *model.named_buffers(remove_duplicate=False),
277 ):
278 if bool(getattr(value, "is_meta", False)):
279 meta_state.append(name)
280 findings = []
281 if meta_state:
282 findings.append(
283 ModelIntegrationFinding(
284 phase="D4",
285 owner_fqn="<root>",
286 severity=FindingSeverity.ERROR,
287 message="materialized model still contains meta tensors",
288 facts={"state_fqns": sorted(meta_state)},
289 why_unsafe="forward or checkpoint access would read storage that does not exist",
290 remediation="materialize storage and run the adapter rebuild lifecycle before validation",
291 related_config=("training.init_device", "model adapter materialization"),
292 )
293 )
294 invalid_rebuildable_buffers = []
295 for module_fqn, module in model.named_modules(remove_duplicate=False):
296 specs = module.__dict__.get("_hp_rebuildable_buffer_specs", {})
297 for buffer_name, spec in specs.items():
298 buffer = module._buffers.get(buffer_name) # pylint: disable=protected-access
299 reason = None
300 if buffer is None:
301 reason = "buffer is missing"
302 elif buffer_name not in module._non_persistent_buffers_set: # pylint: disable=protected-access
303 reason = "buffer became persistent"
304 elif bool(getattr(buffer, "is_meta", False)):
305 reason = "buffer still has meta storage"
306 elif getattr(spec, "source", None) is None and getattr(spec, "factory", None) is None:
307 reason = "rebuild recipe has neither source nor factory"
308 if reason is not None:
309 invalid_rebuildable_buffers.append(
310 {
311 "fqn": f"{module_fqn}.{buffer_name}" if module_fqn else buffer_name,
312 "reason": reason,
313 }
314 )
315 if invalid_rebuildable_buffers:
316 findings.append(
317 ModelIntegrationFinding(
318 phase="D4",
319 owner_fqn="<root>",
320 severity=FindingSeverity.ERROR,
321 message="rebuildable derived-buffer lifecycle is incomplete",
322 facts={"buffers": invalid_rebuildable_buffers},
323 why_unsafe="to_empty() may leave model-derived non-persistent state undefined",
324 remediation="register and execute a deterministic rebuild recipe after materialization",
325 related_config=("training.init_device", "model adapter materialization"),
326 )
327 )
328 return findings
331def _validate_parameter_probes(
332 model: Any,
333 validation_spec: Optional[ModelValidationSpec],
334) -> list[ModelIntegrationFinding]:
335 """Emit findings for required parameter probes that miss."""
336 if validation_spec is None:
337 return []
338 parameter_names = tuple(name for name, _ in model.named_parameters(remove_duplicate=False))
339 findings = []
340 for probe in validation_spec.parameter_probes:
341 matches = [name for name in parameter_names if fnmatch.fnmatchcase(name, probe.match)]
342 if probe.required and not matches:
343 findings.append(
344 ModelIntegrationFinding(
345 phase="V0",
346 owner_fqn=probe.match,
347 severity=FindingSeverity.ERROR,
348 message="required parameter probe matched no final parameter",
349 facts={"pattern": probe.match},
350 why_unsafe="a renamed or missing critical parameter would be silently unobserved",
351 remediation="update the model validation provider to match final parameter FQNs",
352 related_config=("ModelAdapterSpec.validation.parameter_probes",),
353 )
354 )
355 return findings
358def validate_state_invariants(
359 model: Any,
360 validation_spec: Optional[ModelValidationSpec],
361 context: Any,
362 phases: tuple[str, ...] = ("structure", "materialization", "checkpoint"),
363) -> list[ModelIntegrationFinding]:
364 """Evaluate model-owned invariants at their declared lifecycle phase."""
365 findings = []
366 if validation_spec is None:
367 return findings
368 phase_by_name = {
369 "structure": "D2",
370 "materialization": "D4",
371 "checkpoint": "D5",
372 "runtime": "D7",
373 }
374 for invariant in validation_spec.state_invariants:
375 if invariant.phase not in phase_by_name:
376 findings.append(
377 ModelIntegrationFinding(
378 phase="D1",
379 owner_fqn=invariant.name,
380 severity=FindingSeverity.ERROR,
381 message="model state invariant declares an unsupported lifecycle phase",
382 facts={"phase": invariant.phase, "supported": tuple(phase_by_name)},
383 why_unsafe="an invariant with an unknown phase would never be evaluated",
384 remediation="use structure, materialization, checkpoint, or runtime",
385 related_config=("ModelAdapterSpec.validation.state_invariants",),
386 )
387 )
388 continue
389 if invariant.phase not in phases:
390 continue
391 result = invariant.checker(model, context)
392 if result is True or result is None:
393 continue
394 facts = result if isinstance(result, Mapping) else {"result": result}
395 findings.append(
396 ModelIntegrationFinding(
397 phase=phase_by_name[invariant.phase],
398 owner_fqn=invariant.name,
399 severity=FindingSeverity.ERROR,
400 message=f"model state invariant failed: {invariant.name}",
401 facts=dict(facts),
402 why_unsafe="the final model state does not satisfy its authoritative lifecycle contract",
403 remediation="fix model initialization/checkpoint mapping or the adapter invariant declaration",
404 related_config=("model adapter validation",),
405 )
406 )
407 return findings
410def validate_final_model(
411 model: Any,
412 *,
413 context: Any = None,
414) -> tuple[ModelIntegrationReport, list[ParameterOwnership], dict[str, Any]]:
415 """Run model/FSDP-state model integration on the final built model.
417 This function is observational. It does not change module structure,
418 parameters, communication state, or optimizer state.
419 """
420 validation_spec = resolve_model_validation_spec(model)
421 ownership, fsdp_units = collect_parameter_ownership(model)
422 report = ModelIntegrationReport(
423 metadata={
424 "model_type": getattr(getattr(model, "config", None), "model_type", None),
425 "inventory_counts": {
426 "modules": len(list(model.modules())),
427 "parameters": len(list(model.parameters())),
428 "buffers": len(list(model.buffers())),
429 },
430 }
431 )
432 report.extend(_validate_materialization(model))
433 report.extend(_validate_gradient_domains(ownership))
434 checkpoint_spec = validation_spec.checkpoint if validation_spec is not None else None
435 checkpoint_coverage, checkpoint_findings = build_checkpoint_coverage(model, checkpoint_spec)
436 report.metadata["checkpoint_coverage"] = checkpoint_coverage
437 report.extend(checkpoint_findings)
438 report.extend(_validate_parameter_probes(model, validation_spec))
439 report.extend(validate_state_invariants(model, validation_spec, context))
440 return report, ownership, fsdp_units
443def inventory_final_model(model: Any, source: Optional[dict[str, Any]] = None) -> dict[str, Any]:
444 """Build a serializable final-tree inventory for integration evidence."""
445 return build_model_inventory(model, source=source).to_dict()
448def collect_activation_checkpoint_wrappers(model: Any) -> dict[str, Any]:
449 """List final activation-checkpoint wrappers without backend imports."""
450 wrappers = []
451 for module_fqn, module in model.named_modules(remove_duplicate=False):
452 if type(module).__name__ != "CheckpointWrapper":
453 continue
454 wrapped = getattr(module, "_wrapped_module", None)
455 wrappers.append(
456 {
457 "fqn": module_fqn or "<root>",
458 "wrapped_type": (
459 None
460 if wrapped is None
461 else f"{type(wrapped).__module__}.{type(wrapped).__qualname__}"
462 ),
463 }
464 )
465 return {"count": len(wrappers), "wrappers": wrappers}
468__all__ = [
469 "ParameterOwnership",
470 "collect_activation_checkpoint_wrappers",
471 "collect_parameter_ownership",
472 "inventory_final_model",
473 "resolve_model_validation_spec",
474 "validate_final_model",
475 "validate_state_invariants",
476]