File size: 3,909 Bytes
3e363e2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | """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
#: Section 6 budget ranges, as fractions of total trainable parameters.
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)
)
|