87 lines
2.6 KiB
Python
87 lines
2.6 KiB
Python
|
"""Shared training helpers for AHA q_proj router rows."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from dataclasses import dataclass
|
||
|
|
|
||
|
|
import torch
|
||
|
|
|
||
|
|
from modeling_aha_qwen3 import aha_router_output_size
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class GateOnlySetup:
|
||
|
|
parameters: list[torch.nn.Parameter]
|
||
|
|
effective_parameter_count: int
|
||
|
|
q_rows: int
|
||
|
|
gate_rows: int
|
||
|
|
|
||
|
|
|
||
|
|
class RowWiseAdamW(torch.optim.AdamW):
|
||
|
|
"""AdamW with an exact lower LR on prefixes of selected tensors."""
|
||
|
|
|
||
|
|
def __init__(self, params, *, row_scales, **kwargs):
|
||
|
|
super().__init__(params, **kwargs)
|
||
|
|
self._row_scales = row_scales
|
||
|
|
|
||
|
|
@torch.no_grad()
|
||
|
|
def step(self, closure=None):
|
||
|
|
before = [p[:n_rows].detach().clone() for p, n_rows, _ in self._row_scales]
|
||
|
|
loss = super().step(closure=closure)
|
||
|
|
for (parameter, n_rows, scale), old in zip(self._row_scales, before):
|
||
|
|
if scale != 1.0:
|
||
|
|
new = parameter[:n_rows]
|
||
|
|
new.copy_(old + scale * (new - old))
|
||
|
|
return loss
|
||
|
|
|
||
|
|
|
||
|
|
def q_projection_rows(config) -> int:
|
||
|
|
head_dim = getattr(
|
||
|
|
config,
|
||
|
|
"head_dim",
|
||
|
|
config.hidden_size // config.num_attention_heads,
|
||
|
|
)
|
||
|
|
return int(config.num_attention_heads * head_dim)
|
||
|
|
|
||
|
|
|
||
|
|
def configure_gate_only(model: torch.nn.Module) -> GateOnlySetup:
|
||
|
|
"""Freeze a model and expose only the appended q_proj gate rows.
|
||
|
|
|
||
|
|
PyTorch cannot mark only a slice of a Parameter trainable, so each q_proj
|
||
|
|
tensor remains trainable while a hook zeros the ordinary Q-row gradient.
|
||
|
|
The returned parameter count is the effective native gate parameter count,
|
||
|
|
not the full q_proj tensor size seen by the optimizer.
|
||
|
|
"""
|
||
|
|
|
||
|
|
for parameter in model.parameters():
|
||
|
|
parameter.requires_grad = False
|
||
|
|
|
||
|
|
q_rows = q_projection_rows(model.config)
|
||
|
|
gate_rows = aha_router_output_size(model.config)
|
||
|
|
|
||
|
|
def mask_q_rows(gradient: torch.Tensor) -> torch.Tensor:
|
||
|
|
masked = gradient.clone()
|
||
|
|
masked[:q_rows] = 0.0
|
||
|
|
return masked
|
||
|
|
|
||
|
|
parameters: list[torch.nn.Parameter] = []
|
||
|
|
effective = 0
|
||
|
|
for layer in model.model.layers:
|
||
|
|
q_proj = layer.self_attn.q_proj
|
||
|
|
q_proj.weight.requires_grad = True
|
||
|
|
q_proj.weight.register_hook(mask_q_rows)
|
||
|
|
parameters.append(q_proj.weight)
|
||
|
|
effective += gate_rows * q_proj.in_features
|
||
|
|
if q_proj.bias is not None:
|
||
|
|
q_proj.bias.requires_grad = True
|
||
|
|
q_proj.bias.register_hook(mask_q_rows)
|
||
|
|
parameters.append(q_proj.bias)
|
||
|
|
effective += gate_rows
|
||
|
|
|
||
|
|
return GateOnlySetup(
|
||
|
|
parameters=parameters,
|
||
|
|
effective_parameter_count=effective,
|
||
|
|
q_rows=q_rows,
|
||
|
|
gate_rows=gate_rows,
|
||
|
|
)
|