| """Programmatic parameter accounting and budget validation (PLAN section 6). |
| |
| Section 6 asks for two things that the previous implementation did not do: a |
| *concrete* allocation across buckets, and programmatic validation of it. Only |
| the model total was checked before, so an allocation that put 46% of the budget |
| in PAT (against a 15-25% target) and 1.4% in the heads (against 10-15%) passed. |
| |
| :func:`validate_budget` therefore checks every bucket fraction as well as the |
| total, and reports all violations at once rather than failing on the first. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from collections.abc import Mapping |
| from dataclasses import dataclass |
| from typing import Protocol, runtime_checkable |
|
|
| from torch import nn |
|
|
| |
| BUDGET: Mapping[str, tuple[float, float]] = { |
| "graph_encoder": (0.10, 0.20), |
| "pat": (0.15, 0.25), |
| "reasoning_transformer": (0.45, 0.60), |
| "heads": (0.10, 0.15), |
| "embeddings": (0.05, 0.10), |
| } |
|
|
|
|
| @runtime_checkable |
| class BucketedModule(Protocol): |
| """A module that can partition its own parameters into budget buckets.""" |
|
|
| def parameter_buckets(self) -> Mapping[str, int]: ... |
|
|
|
|
| @dataclass(frozen=True) |
| class ParameterReport: |
| total: int |
| trainable: int |
| buckets: Mapping[str, int] |
|
|
| @property |
| def fractions(self) -> Mapping[str, float]: |
| return { |
| name: (count / self.total if self.total else 0.0) |
| for name, count in self.buckets.items() |
| } |
|
|
| def to_dict(self) -> dict[str, object]: |
| return { |
| "total_parameters": self.total, |
| "trainable_parameters": self.trainable, |
| **dict(self.buckets), |
| "fractions": { |
| name: round(value, 5) for name, value in self.fractions.items() |
| }, |
| } |
|
|
|
|
| def count_parameters(model: nn.Module) -> ParameterReport: |
| total = sum(parameter.numel() for parameter in model.parameters()) |
| trainable = sum( |
| parameter.numel() for parameter in model.parameters() if parameter.requires_grad |
| ) |
| buckets = ( |
| dict(model.parameter_buckets()) if isinstance(model, BucketedModule) else {} |
| ) |
| return ParameterReport(total, trainable, buckets) |
|
|
|
|
| def validate_target( |
| report: ParameterReport, target: int, tolerance: float = 0.20 |
| ) -> None: |
| """Check the trainable total lies within ``tolerance`` of ``target``.""" |
| if target <= 0: |
| raise ValueError("target must be positive") |
| lower, upper = int(target * (1 - tolerance)), int(target * (1 + tolerance)) |
| if not lower <= report.trainable <= upper: |
| raise ValueError( |
| f"trainable parameter count {report.trainable:,} is outside " |
| f"[{lower:,}, {upper:,}] for target {target:,}" |
| ) |
|
|
|
|
| def budget_violations( |
| report: ParameterReport, budget: Mapping[str, tuple[float, float]] | None = None |
| ) -> tuple[str, ...]: |
| """Return a human-readable violation per bucket outside its budget range.""" |
| ranges = budget or BUDGET |
| fractions = report.fractions |
| violations: list[str] = [] |
| for name, (low, high) in ranges.items(): |
| value = fractions.get(name) |
| if value is None: |
| violations.append(f"{name}: bucket missing from parameter report") |
| elif not low <= value <= high: |
| violations.append( |
| f"{name}: {value * 100:.1f}% outside target {low * 100:.0f}-{high * 100:.0f}%" |
| ) |
| return tuple(violations) |
|
|
|
|
| def validate_budget( |
| report: ParameterReport, budget: Mapping[str, tuple[float, float]] | None = None |
| ) -> None: |
| """Fail with every budget violation listed, not just the first.""" |
| violations = budget_violations(report, budget) |
| if violations: |
| raise ValueError( |
| "parameter allocation violates the PLAN section 6 budget:\n " |
| + "\n ".join(violations) |
| ) |
|
|
