初始化项目，由ModelHub XC社区提供模型

Model: jiamingshan/AHA-L2A-Qwen3-1.7B-repro
Source: Original Platform

recipe/router_training_utils.py

@@ -0,0 +1,86 @@
+"""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,
+    )