recipe/router_training_utils.py

"""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,
    )