| from __future__ import annotations |
|
|
| import math |
| from dataclasses import asdict, dataclass |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
| from .quantum_residual import QuantumResidualMergerWrapper |
|
|
|
|
| @dataclass(frozen=True) |
| class ClassicalControlConfig: |
| hidden_size: int = 5120 |
| bottleneck_size: int = 8 |
| depth: int = 2 |
| scale_init: float = 0.3 |
| scale_max: float = 1.0 |
| up_init_std: float = 0.01 |
|
|
|
|
| class MatchedClassicalResidualAdapter(nn.Module): |
| """Exact trainable-parameter-count control for QH-001b (81,961 parameters).""" |
|
|
| def __init__(self, config: ClassicalControlConfig | None = None) -> None: |
| super().__init__() |
| self.config = config or ClassicalControlConfig() |
| self.down = nn.Linear(self.config.hidden_size, self.config.bottleneck_size, bias=True) |
| self.core_affine = nn.Parameter( |
| torch.zeros(self.config.depth, self.config.bottleneck_size, 2, dtype=torch.float32) |
| ) |
| self.up = nn.Linear(self.config.bottleneck_size, self.config.hidden_size, bias=False) |
| probability = self.config.scale_init / self.config.scale_max |
| self.raw_scale = nn.Parameter(torch.tensor(math.log(probability / (1.0 - probability)))) |
| nn.init.xavier_uniform_(self.down.weight) |
| nn.init.zeros_(self.down.bias) |
| nn.init.normal_(self.up.weight, mean=0.0, std=self.config.up_init_std) |
|
|
| @property |
| def scale(self) -> torch.Tensor: |
| return self.config.scale_max * torch.sigmoid(self.raw_scale) |
|
|
| def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| original_shape = hidden_states.shape |
| flat = hidden_states.reshape(-1, original_shape[-1]) |
| normalized = F.layer_norm(flat.float(), (self.config.hidden_size,)).to(flat.dtype) |
| features = torch.tanh(self.down(normalized).float()) |
| for layer in range(self.config.depth): |
| gain = 1.0 + 0.05 * self.core_affine[layer, :, 0] |
| bias = 0.05 * self.core_affine[layer, :, 1] |
| features = torch.tanh(features * gain + bias) |
| residual = self.up(features.to(self.up.weight.dtype)).to(flat.dtype) |
| output = flat + residual * self.scale.to(residual.dtype) |
| return output.reshape(original_shape) |
|
|
| def audit(self) -> dict: |
| return { |
| "candidate": "CC-001", |
| "config": asdict(self.config), |
| "current_scale": float(self.scale.detach().cpu()), |
| "trainable_parameters": sum(p.numel() for p in self.parameters() if p.requires_grad), |
| } |
|
|
|
|
| def freeze_and_inject_cc001(model: nn.Module, config: ClassicalControlConfig | None = None) -> dict: |
| for parameter in model.parameters(): |
| parameter.requires_grad_(False) |
| visual = model.model.visual |
| old_merger = visual.merger |
| reference = next(old_merger.parameters()) |
| adapter = MatchedClassicalResidualAdapter(config).to(device=reference.device) |
| adapter.down.to(dtype=reference.dtype) |
| adapter.up.to(dtype=reference.dtype) |
| adapter.core_affine.data = adapter.core_affine.data.float() |
| visual.merger = QuantumResidualMergerWrapper(old_merger, adapter) |
| trainable = [(name, p.numel()) for name, p in model.named_parameters() if p.requires_grad] |
| return { |
| "candidate": "CC-001", |
| "injection_path": "model.model.visual.merger", |
| "adapter": adapter.audit(), |
| "trainable_parameter_count": sum(count for _, count in trainable), |
| "trainable_tensors": trainable, |
| } |
|
|
|
|
