File size: 2,685 Bytes
46b9eea | 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 | """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,
)
|