Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / tools / model_integration / module_parity.py: 80%
212 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"""Generic final-module forward/intermediate/backward parity runner."""
16# pylint: disable=forbidden-backend-import
18from __future__ import annotations
20import fnmatch
21import traceback
22from dataclasses import dataclass
23from typing import Any, Mapping
25import torch
26from torch.utils._pytree import tree_flatten, tree_map
28from hyper_parallel.tools.model_integration.evidence_store import EvidenceStore
29from hyper_parallel.tools.model_integration.schemas import (
30 ComparisonMetric,
31 tolerance_value,
32)
33from hyper_parallel.models.validation_spec import ModelValidationSpec, ModuleParityCase
36@dataclass(frozen=True)
37class ParityContext:
38 """Resolved execution context supplied to model-owned parity builders."""
40 manifest: Any
41 dtype: str
42 device: str
43 evidence_store: EvidenceStore
44 options: dict[str, Any]
47@dataclass
48class ParityExecution:
49 """Values exposed to observation getters after one forward/backward."""
51 module: Any
52 output: Any
53 args: tuple[Any, ...]
54 kwargs: dict[str, Any]
57def _select_module(value: Any, selector: str) -> Any:
58 if selector in ("", "<root>"):
59 return value
60 module_by_fqn = dict(value.named_modules())
61 selected = module_by_fqn.get(selector)
62 if selected is None:
63 raise ValueError(f"candidate selector matched no final module: {selector!r}")
64 return selected
67def _clone_value(value: Any) -> Any:
68 if not torch.is_tensor(value):
69 return value
70 clone = value.detach().clone()
71 if value.requires_grad and (clone.is_floating_point() or clone.is_complex()):
72 clone.requires_grad_(True)
73 return clone
76def _clone_call_inputs(inputs: Any) -> tuple[tuple[Any, ...], dict[str, Any]]:
77 if isinstance(inputs, Mapping) and ("args" in inputs or "kwargs" in inputs):
78 args = tuple(inputs.get("args", ()))
79 kwargs = dict(inputs.get("kwargs", {}))
80 elif isinstance(inputs, Mapping):
81 args = ()
82 kwargs = dict(inputs)
83 elif isinstance(inputs, tuple):
84 args = inputs
85 kwargs = {}
86 else:
87 args = (inputs,)
88 kwargs = {}
89 return tree_map(_clone_value, args), tree_map(_clone_value, kwargs)
92def _floating_tensors(value: Any) -> list[torch.Tensor]:
93 return [
94 item for item in tree_flatten(value)[0]
95 if torch.is_tensor(item) and (item.is_floating_point() or item.is_complex())
96 ]
99def _default_objective(output: Any) -> torch.Tensor:
100 tensors = _floating_tensors(output)
101 if not tensors:
102 raise ValueError("parity output contains no differentiable tensor")
103 return sum(tensor.float().sum() for tensor in tensors)
106def _metric(
107 name: str,
108 reference: torch.Tensor,
109 candidate: torch.Tensor,
110 *,
111 exact: bool,
112 atol: float,
113 rtol: float,
114 relative_l2_limit: float | None = None,
115) -> ComparisonMetric:
116 reference_cpu = reference.detach().float().cpu()
117 candidate_cpu = candidate.detach().float().cpu()
118 same_shape = tuple(reference_cpu.shape) == tuple(candidate_cpu.shape)
119 finite = bool(torch.isfinite(reference_cpu).all() and torch.isfinite(candidate_cpu).all())
120 if not same_shape:
121 return ComparisonMetric(
122 name=name,
123 status="FAIL",
124 max_abs=None,
125 relative_l2=None,
126 finite=finite,
127 reference_shape=tuple(reference_cpu.shape),
128 candidate_shape=tuple(candidate_cpu.shape),
129 reason="shape mismatch",
130 )
131 equal = bool(torch.equal(reference_cpu, candidate_cpu))
132 delta = candidate_cpu.double() - reference_cpu.double()
133 max_abs = float(delta.abs().max()) if delta.numel() else 0.0
134 reference_norm = torch.linalg.vector_norm( # pylint: disable=not-callable
135 reference_cpu.double().reshape(-1)
136 )
137 delta_norm = torch.linalg.vector_norm(delta.reshape(-1)) # pylint: disable=not-callable
138 relative_l2 = float(delta_norm / reference_norm.clamp_min(torch.finfo(torch.float64).eps))
139 passed = equal if exact else bool(torch.allclose(reference_cpu, candidate_cpu, atol=atol, rtol=rtol))
140 if not exact and relative_l2_limit is not None:
141 passed = passed and relative_l2 <= relative_l2_limit
142 passed = passed and finite
143 return ComparisonMetric(
144 name=name,
145 status="PASS" if passed else "FAIL",
146 max_abs=max_abs,
147 relative_l2=relative_l2,
148 finite=finite,
149 reference_shape=tuple(reference_cpu.shape),
150 candidate_shape=tuple(candidate_cpu.shape),
151 )
154def _compare_tree(
155 prefix: str,
156 reference: Any,
157 candidate: Any,
158 *,
159 exact: bool,
160 atol: float,
161 rtol: float,
162 relative_l2_limit: float | None = None,
163) -> list[ComparisonMetric]:
164 reference_values, reference_spec = tree_flatten(reference)
165 candidate_values, candidate_spec = tree_flatten(candidate)
166 if reference_spec != candidate_spec:
167 return [
168 ComparisonMetric(
169 name=prefix,
170 status="FAIL",
171 max_abs=None,
172 relative_l2=None,
173 finite=True,
174 reference_shape=(),
175 candidate_shape=(),
176 reason=(
177 "output structure mismatch: "
178 f"reference={reference_spec!r}, candidate={candidate_spec!r}"
179 ),
180 )
181 ]
182 metrics = []
183 for index, (reference_value, candidate_value) in enumerate(zip(reference_values, candidate_values)):
184 name = f"{prefix}.{index}"
185 if torch.is_tensor(reference_value) and torch.is_tensor(candidate_value):
186 tensor_exact = exact or not (
187 reference_value.is_floating_point() or reference_value.is_complex()
188 )
189 metrics.append(
190 _metric(
191 name,
192 reference_value,
193 candidate_value,
194 exact=tensor_exact,
195 atol=atol,
196 rtol=rtol,
197 relative_l2_limit=relative_l2_limit,
198 )
199 )
200 elif reference_value != candidate_value:
201 metrics.append(
202 ComparisonMetric(
203 name=name,
204 status="FAIL",
205 max_abs=None,
206 relative_l2=None,
207 finite=True,
208 reference_shape=(),
209 candidate_shape=(),
210 reason=f"non-tensor values differ: {reference_value!r} != {candidate_value!r}",
211 )
212 )
213 return metrics
216def _parameter_mapping(
217 case: ModuleParityCase,
218 reference: Any,
219 candidate: Any,
220 context: ParityContext,
221) -> dict[str, str]:
222 if case.weight_adapter is None:
223 raise TypeError("in-process parity case requires weight_adapter")
224 mapping_result = case.weight_adapter(reference, candidate, context)
225 if mapping_result is None:
226 mapping = {
227 name: name for name, parameter in candidate.named_parameters()
228 if parameter.requires_grad
229 }
230 elif isinstance(mapping_result, Mapping):
231 mapping = {str(key): str(value) for key, value in mapping_result.items()}
232 else:
233 raise TypeError("parity weight_adapter must return a candidate-to-reference mapping or None")
234 candidate_trainable = {
235 name for name, parameter in candidate.named_parameters() if parameter.requires_grad
236 }
237 reference_names = dict(reference.named_parameters())
238 missing = sorted(candidate_trainable - set(mapping))
239 invalid = sorted(set(mapping.values()) - set(reference_names))
240 if missing or invalid:
241 raise ValueError(
242 "parity weight coverage is incomplete: "
243 f"missing_candidate={missing}, unknown_reference={invalid}"
244 )
245 for pattern in case.required_parameter_patterns:
246 if not any(fnmatch.fnmatchcase(name, pattern) for name in candidate_trainable):
247 raise ValueError(f"required parity parameter pattern matched nothing: {pattern!r}")
248 return mapping
251def _input_gradients(args: tuple[Any, ...], kwargs: dict[str, Any]) -> list[torch.Tensor | None]:
252 return [
253 value.grad if torch.is_tensor(value) and value.requires_grad else None
254 for value in tree_flatten((args, kwargs))[0]
255 ]
258def _build_parity_modules(case: ModuleParityCase, context: ParityContext) -> tuple[Any, Any]:
259 """Instantiate reference/candidate and reject aliased implementations."""
260 candidate_builder = case.candidate_builder
261 reference_builder = case.reference_builder
262 input_builder = case.input_builder
263 if not all(callable(value) for value in (candidate_builder, reference_builder, input_builder)):
264 raise TypeError("in-process parity case has an incomplete callable contract")
265 candidate_tree = candidate_builder(context)
266 candidate = _select_module(candidate_tree, case.candidate_selector)
267 reference = reference_builder(context)
268 if candidate is reference or type(candidate) is type(reference):
269 raise ValueError(
270 "reference and candidate must be independent implementations; "
271 f"got {type(candidate).__module__}.{type(candidate).__qualname__}"
272 )
273 return candidate, reference
276def _resolve_parity_tolerance(
277 context: ParityContext,
278) -> tuple[float, float, float | None]:
279 """Read the scalar tolerance contract from the parity context options."""
280 tolerance = context.options.get("tolerance", {})
281 atol = tolerance_value(tolerance, "atol", "max_abs")
282 rtol = tolerance_value(tolerance, "rtol")
283 relative_l2_limit = tolerance.get("relative_l2")
284 if relative_l2_limit is not None:
285 relative_l2_limit = float(relative_l2_limit)
286 return atol, rtol, relative_l2_limit
289def _observation_metrics(
290 case: ModuleParityCase,
291 reference_execution: ParityExecution,
292 candidate_execution: ParityExecution,
293 reference_output: Any,
294 candidate_output: Any,
295 atol: float,
296 rtol: float,
297 relative_l2_limit: float | None,
298) -> list[ComparisonMetric]:
299 """Compare every declared observation between reference and candidate."""
300 metrics = []
301 for observation in case.observations:
302 reference_value = (
303 observation.reference_getter(reference_execution)
304 if observation.reference_getter is not None
305 else reference_output
306 )
307 candidate_value = (
308 observation.candidate_getter(candidate_execution)
309 if observation.candidate_getter is not None
310 else candidate_output
311 )
312 if reference_value is None or candidate_value is None:
313 if observation.required:
314 raise ValueError(
315 f"required parity observation {observation.name!r} was not produced: "
316 f"reference={reference_value is not None}, "
317 f"candidate={candidate_value is not None}"
318 )
319 if reference_value is None and candidate_value is None:
320 continue
321 metrics.extend(
322 _compare_tree(
323 observation.name,
324 reference_value,
325 candidate_value,
326 exact=observation.comparison == "exact",
327 atol=observation.atol if observation.atol is not None else atol,
328 rtol=observation.rtol if observation.rtol is not None else rtol,
329 relative_l2_limit=relative_l2_limit,
330 )
331 )
332 return metrics
335def _input_gradient_metrics(
336 reference_args: tuple[Any, ...],
337 reference_kwargs: dict[str, Any],
338 candidate_args: tuple[Any, ...],
339 candidate_kwargs: dict[str, Any],
340 atol: float,
341 rtol: float,
342 relative_l2_limit: float | None,
343) -> list[ComparisonMetric]:
344 """Compare flattened input gradients pairwise after backward."""
345 metrics = []
346 reference_input_grads = _input_gradients(reference_args, reference_kwargs)
347 candidate_input_grads = _input_gradients(candidate_args, candidate_kwargs)
348 for index, (reference_grad, candidate_grad) in enumerate(
349 zip(reference_input_grads, candidate_input_grads)
350 ):
351 if reference_grad is None and candidate_grad is None:
352 continue
353 if reference_grad is None or candidate_grad is None:
354 raise ValueError(f"input gradient presence differs at flattened input {index}")
355 metrics.append(
356 _metric(
357 f"input_grad.{index}",
358 reference_grad,
359 candidate_grad,
360 exact=False,
361 atol=atol,
362 rtol=rtol,
363 relative_l2_limit=relative_l2_limit,
364 )
365 )
366 return metrics
369def _parameter_gradient_metrics(
370 mapping: dict[str, str],
371 reference: Any,
372 candidate: Any,
373 atol: float,
374 rtol: float,
375 relative_l2_limit: float | None,
376) -> list[ComparisonMetric]:
377 """Compare mapped parameter gradients pairwise after backward."""
378 metrics = []
379 reference_parameters = dict(reference.named_parameters())
380 candidate_parameters = dict(candidate.named_parameters())
381 for candidate_name, reference_name in sorted(mapping.items()):
382 reference_grad = reference_parameters[reference_name].grad
383 candidate_grad = candidate_parameters[candidate_name].grad
384 if reference_grad is None and candidate_grad is None:
385 continue
386 if reference_grad is None or candidate_grad is None:
387 raise ValueError(
388 f"parameter gradient presence differs: {candidate_name} -> {reference_name}"
389 )
390 metrics.append(
391 _metric(
392 f"parameter_grad.{candidate_name}",
393 reference_grad,
394 candidate_grad,
395 exact=False,
396 atol=atol,
397 rtol=rtol,
398 relative_l2_limit=relative_l2_limit,
399 )
400 )
401 return metrics
404def _run_in_process_case(case: ModuleParityCase, context: ParityContext) -> dict[str, Any]:
405 candidate, reference = _build_parity_modules(case, context)
406 mapping = _parameter_mapping(case, reference, candidate, context)
407 inputs = case.input_builder(context)
408 reference_args, reference_kwargs = _clone_call_inputs(inputs)
409 candidate_args, candidate_kwargs = _clone_call_inputs(inputs)
410 reference_output = reference(*reference_args, **reference_kwargs)
411 candidate_output = candidate(*candidate_args, **candidate_kwargs)
412 atol, rtol, relative_l2_limit = _resolve_parity_tolerance(context)
413 metrics = _compare_tree(
414 "output",
415 reference_output,
416 candidate_output,
417 exact=False,
418 atol=atol,
419 rtol=rtol,
420 relative_l2_limit=relative_l2_limit,
421 )
422 reference_execution = ParityExecution(reference, reference_output, reference_args, reference_kwargs)
423 candidate_execution = ParityExecution(candidate, candidate_output, candidate_args, candidate_kwargs)
424 metrics.extend(
425 _observation_metrics(
426 case,
427 reference_execution,
428 candidate_execution,
429 reference_output,
430 candidate_output,
431 atol,
432 rtol,
433 relative_l2_limit,
434 )
435 )
436 objective = case.objective or _default_objective
437 objective(reference_output).backward()
438 objective(candidate_output).backward()
439 metrics.extend(
440 _input_gradient_metrics(
441 reference_args,
442 reference_kwargs,
443 candidate_args,
444 candidate_kwargs,
445 atol,
446 rtol,
447 relative_l2_limit,
448 )
449 )
450 metrics.extend(
451 _parameter_gradient_metrics(
452 mapping,
453 reference,
454 candidate,
455 atol,
456 rtol,
457 relative_l2_limit,
458 )
459 )
460 status = "PASS" if all(metric.status == "PASS" for metric in metrics) else "FAIL"
461 return {
462 "case": case.name,
463 "status": status,
464 "candidate_type": f"{type(candidate).__module__}.{type(candidate).__qualname__}",
465 "reference_type": f"{type(reference).__module__}.{type(reference).__qualname__}",
466 "weight_mapping": mapping,
467 "metrics": [metric.to_dict() for metric in metrics],
468 }
471def run_module_parity(
472 validation_spec: ModelValidationSpec,
473 context: ParityContext,
474) -> dict[str, Any]:
475 """Execute every declared parity case and persist per-case evidence."""
476 reports = []
477 for case in validation_spec.module_cases:
478 if context.dtype not in case.required_dtypes:
479 continue
480 try:
481 if case.execution == "isolated_process":
482 report = case.isolated_runner(context) # type: ignore[misc]
483 if not isinstance(report, dict) or report.get("status") not in (
484 "PASS", "FAIL", "BLOCKED"
485 ):
486 raise ValueError("isolated parity runner returned an invalid report")
487 else:
488 report = _run_in_process_case(case, context)
489 except NotImplementedError as error:
490 report = {
491 "case": case.name,
492 "status": "BLOCKED",
493 "reason": "OPERATOR_UNSUPPORTED",
494 "error": str(error),
495 }
496 except Exception as error: # pylint: disable=broad-except
497 report = {
498 "case": case.name,
499 "status": "FAIL",
500 "error_type": type(error).__name__,
501 "error": str(error),
502 "traceback": traceback.format_exc(),
503 }
504 reports.append(report)
505 context.evidence_store.write_json(
506 f"module_parity/{case.name}/comparison.json",
507 report,
508 )
509 statuses = [report["status"] for report in reports]
510 status = (
511 "FAIL" if "FAIL" in statuses
512 else "BLOCKED" if "BLOCKED" in statuses or not statuses
513 else "PASS"
514 )
515 aggregate = {"status": status, "cases": reports}
516 if not reports:
517 aggregate["reason"] = (
518 "no module parity case supports the requested dtype, or the adapter "
519 "declared no authoritative case"
520 )
521 context.evidence_store.write_json("module_parity/comparison.json", aggregate)
522 return aggregate
525__all__ = ["ParityContext", "ParityExecution", "run_module_parity"]