1227 lines
57 KiB
Python
1227 lines
57 KiB
Python
"""
|
|
AHA-Qwen3: All-or-Here Attention applied to Qwen3.
|
|
|
|
Minimal modification to standard Qwen3 attention: adds a binary gate that
|
|
dynamically toggles between full attention and local attention. The persisted
|
|
``aha_router_granularity`` config selects either one gate per (token, KV-head)
|
|
or one gate per token shared by every head in the layer.
|
|
|
|
Based on: transformers.models.qwen3.modeling_qwen3 (Qwen3Attention)
|
|
Changes:
|
|
1. q_proj outputs either one extra gate value (``token``) or
|
|
``num_key_value_heads`` values (``token_kv_head``). The shared-token
|
|
value is broadcast to all KV/query heads before attention mixing.
|
|
2. Gate: sigmoid → hard threshold → STE for training
|
|
3. Double-mix: global attention + local window attention, blended by
|
|
gate (broadcast across the GQA group)
|
|
4. Auxiliary loss: λ * mean(gate_soft) to encourage sparsity
|
|
"""
|
|
|
|
import atexit
|
|
import json
|
|
import os
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Callable, Optional, Union
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from safetensors.torch import load_file
|
|
|
|
from transformers import Qwen3Config, Qwen3ForCausalLM
|
|
from transformers.modeling_layers import GradientCheckpointingLayer
|
|
from transformers.models.qwen3.modeling_qwen3 import (
|
|
Qwen3Attention,
|
|
Qwen3DecoderLayer,
|
|
Qwen3Model,
|
|
Qwen3RMSNorm,
|
|
apply_rotary_pos_emb,
|
|
)
|
|
from transformers.cache_utils import Cache, DynamicCache
|
|
from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
|
|
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
|
|
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
|
|
from transformers.processing_utils import Unpack
|
|
from transformers.utils.generic import TransformersKwargs
|
|
|
|
# AHA constants
|
|
AHA_WINDOW_SIZE = 1024
|
|
AHA_LAMBDA = 3e-4
|
|
AHA_DISTILL_WEIGHT = 0.0
|
|
AHA_CE_WEIGHT = 1.0
|
|
# Gate-density ceiling (hinge target). The aux loss only fires when
|
|
# ``mean(gate_soft) > AHA_GATE_TARGET``, so sparsity cannot drop
|
|
# below ``1 - AHA_GATE_TARGET`` in steady state regardless of how
|
|
# large ``aha_lambda`` is. Setting this to 1.0 recovers the legacy
|
|
# behaviour (unconditional downward pressure).
|
|
AHA_GATE_TARGET = 1.0
|
|
|
|
# DuoAttention baseline knobs. When ``aha_mode == "duo"`` the per-(token,
|
|
# kv_head) gate is replaced by a per-(layer, kv_head) learnable scalar
|
|
# (``duo_alpha``, 224 scalars total for Qwen3-0.6B), and the "local"
|
|
# branch becomes a streaming mask with ``duo_sink_size`` attention sinks
|
|
# plus a recent window of ``duo_recent_size``. This is a strict reduction
|
|
# of the dynamic AHA gate and reproduces the DuoAttention paper's head
|
|
# classifier (Xiao et al. 2024) on the Qwen3 backbone.
|
|
AHA_MODE = "dynamic"
|
|
DUO_SINK_SIZE = 64
|
|
DUO_RECENT_SIZE = 256
|
|
DUO_ALPHA_INIT = 1.0
|
|
# Local-branch mask kind for the dual-branch attention (global vs local).
|
|
# - "sliding_window": last ``aha_window_size`` tokens, no attention sinks.
|
|
# Historical default for ``aha_mode="dynamic"`` (Bv3 was trained with this).
|
|
# - "sink_recent": first ``duo_sink_size`` tokens + last ``duo_recent_size``
|
|
# tokens, matching the DuoAttention paper. Historical hardcoded default for
|
|
# ``aha_mode="duo"`` (still forced for duo regardless of this field).
|
|
# § 9.4.13.3 found that hot-starting a dynamic-mode ckpt from a duo ckpt with
|
|
# sliding_window local loses ~20pp LongBench retention vs the duo source,
|
|
# almost entirely from passage_retrieval_en (sinks are critical for long
|
|
# passage retrieval). Setting this to "sink_recent" for hot-started dynamic
|
|
# ckpts closes that architectural gap.
|
|
AHA_LOCAL_KIND = "sliding_window"
|
|
AHA_ROUTER_GRANULARITY = "token_kv_head"
|
|
|
|
|
|
def aha_router_output_size(config: Qwen3Config) -> int:
|
|
"""Return the number of native dynamic-router logits per token."""
|
|
|
|
granularity = getattr(config, "aha_router_granularity", AHA_ROUTER_GRANULARITY)
|
|
if granularity == "token":
|
|
return 1
|
|
if granularity == "token_kv_head":
|
|
return int(config.num_key_value_heads)
|
|
raise ValueError(f"Unsupported aha_router_granularity={granularity!r}")
|
|
|
|
|
|
class _AHAInferenceSparsityTracker:
|
|
"""Accumulate exact inference hard-route counts without changing outputs."""
|
|
|
|
def __init__(self) -> None:
|
|
# Keep counters as device scalars until process exit. Calling .item()
|
|
# once per layer and decode token serialized CUDA and made the
|
|
# diagnostic tracker materially slow generation.
|
|
self.sparse_decisions: dict[int, torch.Tensor] = {}
|
|
self.total_decisions: dict[int, int] = {}
|
|
self.phase_sparse_decisions: dict[str, torch.Tensor] = {}
|
|
self.phase_total_decisions: dict[str, int] = {}
|
|
raw_thresholds = os.environ.get("AHA_COUNTERFACTUAL_THRESHOLDS", "")
|
|
self.counterfactual_thresholds = tuple(
|
|
sorted({float(value) for value in raw_thresholds.split(",") if value.strip()})
|
|
)
|
|
if any(value < 0.0 or value > 1.0 for value in self.counterfactual_thresholds):
|
|
raise ValueError("AHA_COUNTERFACTUAL_THRESHOLDS values must be in [0, 1]")
|
|
self.counterfactual_sparse: dict[float, torch.Tensor] = {}
|
|
self.router_granularity: Optional[str] = None
|
|
self.native_router_decisions = 0
|
|
atexit.register(self.write_stats)
|
|
|
|
def update(
|
|
self,
|
|
gate_hard: torch.Tensor,
|
|
gate_soft: torch.Tensor,
|
|
layer_idx: int,
|
|
phase: str,
|
|
router_granularity: str,
|
|
native_router_width: int,
|
|
) -> None:
|
|
if self.router_granularity is None:
|
|
self.router_granularity = router_granularity
|
|
elif self.router_granularity != router_granularity:
|
|
raise RuntimeError(
|
|
"AHA sparsity tracker received mixed router granularities: "
|
|
f"{self.router_granularity!r} and {router_granularity!r}"
|
|
)
|
|
sparse = torch.count_nonzero(gate_hard == 0).to(torch.int64)
|
|
total = int(gate_hard.numel())
|
|
effective_width = int(gate_hard.shape[-1])
|
|
self.native_router_decisions += total // effective_width * native_router_width
|
|
if layer_idx not in self.sparse_decisions:
|
|
self.sparse_decisions[layer_idx] = sparse
|
|
else:
|
|
self.sparse_decisions[layer_idx].add_(sparse)
|
|
self.total_decisions[layer_idx] = self.total_decisions.get(layer_idx, 0) + total
|
|
if phase not in self.phase_sparse_decisions:
|
|
self.phase_sparse_decisions[phase] = sparse.clone()
|
|
else:
|
|
self.phase_sparse_decisions[phase].add_(sparse)
|
|
self.phase_total_decisions[phase] = self.phase_total_decisions.get(phase, 0) + total
|
|
if self.counterfactual_thresholds:
|
|
flat = gate_soft.detach().reshape(-1, 1)
|
|
thresholds = flat.new_tensor(self.counterfactual_thresholds).reshape(1, -1)
|
|
counts = (flat <= thresholds).sum(dim=0, dtype=torch.int64)
|
|
for threshold, count in zip(self.counterfactual_thresholds, counts.unbind()):
|
|
if threshold not in self.counterfactual_sparse:
|
|
self.counterfactual_sparse[threshold] = count
|
|
else:
|
|
self.counterfactual_sparse[threshold].add_(count)
|
|
|
|
def write_stats(self) -> None:
|
|
output = os.environ.get("AHA_SPARSITY_STATS_PATH")
|
|
if not output:
|
|
return
|
|
sparse_by_layer = {
|
|
layer_idx: int(value.item())
|
|
for layer_idx, value in self.sparse_decisions.items()
|
|
}
|
|
sparse = sum(sparse_by_layer.values())
|
|
total = sum(self.total_decisions.values())
|
|
per_layer = {}
|
|
for layer_idx in sorted(self.total_decisions):
|
|
layer_sparse = sparse_by_layer.get(layer_idx, 0)
|
|
layer_total = self.total_decisions[layer_idx]
|
|
per_layer[str(layer_idx)] = {
|
|
"sparse_decisions": layer_sparse,
|
|
"total_decisions": layer_total,
|
|
"sparsity": layer_sparse / layer_total if layer_total else None,
|
|
}
|
|
payload = {
|
|
"definition": (
|
|
"hard AHA gate zeros / all token x KV-head x layer decisions, "
|
|
"token-weighted over prefill and decode"
|
|
),
|
|
"router_granularity": self.router_granularity,
|
|
"native_router_decisions": self.native_router_decisions,
|
|
"effective_router_decisions": total,
|
|
"native_router_note": (
|
|
"Native decisions count learned gate logits before broadcast; effective "
|
|
"decisions always use token x KV-head x layer for matched sparsity."
|
|
),
|
|
"sparse_decisions": sparse,
|
|
"total_decisions": total,
|
|
"sparsity": sparse / total if total else None,
|
|
"full_attention_usage": 1.0 - sparse / total if total else None,
|
|
"per_layer": per_layer,
|
|
"by_phase": {
|
|
phase: {
|
|
"sparse_decisions": int(self.phase_sparse_decisions[phase].item()),
|
|
"total_decisions": phase_total,
|
|
"sparsity": int(self.phase_sparse_decisions[phase].item()) / phase_total
|
|
if phase_total else None,
|
|
"full_attention_usage": 1.0
|
|
- int(self.phase_sparse_decisions[phase].item()) / phase_total
|
|
if phase_total else None,
|
|
}
|
|
for phase, phase_total in sorted(self.phase_total_decisions.items())
|
|
},
|
|
}
|
|
if self.counterfactual_thresholds:
|
|
payload["counterfactual_sparsity_by_threshold"] = {
|
|
f"{threshold:.6g}": {
|
|
"sparse_decisions": int(self.counterfactual_sparse[threshold].item()),
|
|
"total_decisions": total,
|
|
"sparsity": int(self.counterfactual_sparse[threshold].item()) / total
|
|
if total else None,
|
|
}
|
|
for threshold in self.counterfactual_thresholds
|
|
}
|
|
payload["counterfactual_note"] = (
|
|
"Routing counts only; benchmark quality must be measured in a separate "
|
|
"run with AHA_GATE_HARD_THRESHOLD set to the selected threshold."
|
|
)
|
|
path = Path(output)
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text(json.dumps(payload, indent=2) + "\n")
|
|
|
|
|
|
_AHA_INFERENCE_SPARSITY_TRACKER: Optional[_AHAInferenceSparsityTracker] = None
|
|
|
|
|
|
def _track_aha_inference_sparsity(
|
|
gate_hard: torch.Tensor,
|
|
gate_soft: torch.Tensor,
|
|
layer_idx: int,
|
|
phase: str,
|
|
router_granularity: str,
|
|
native_router_width: int,
|
|
) -> None:
|
|
global _AHA_INFERENCE_SPARSITY_TRACKER
|
|
if not os.environ.get("AHA_SPARSITY_STATS_PATH"):
|
|
return
|
|
if _AHA_INFERENCE_SPARSITY_TRACKER is None:
|
|
_AHA_INFERENCE_SPARSITY_TRACKER = _AHAInferenceSparsityTracker()
|
|
_AHA_INFERENCE_SPARSITY_TRACKER.update(
|
|
gate_hard.detach(), gate_soft.detach(), layer_idx, phase,
|
|
router_granularity, native_router_width,
|
|
)
|
|
|
|
|
|
def _get_optional_nonnegative_int_env(name: str) -> Optional[int]:
|
|
raw = os.environ.get(name, "").strip()
|
|
if not raw:
|
|
return None
|
|
value = int(raw)
|
|
if value < 0:
|
|
raise ValueError(f"{name} must be non-negative, got {value}")
|
|
return value
|
|
|
|
|
|
def _apply_low_alpha_force_full_heads_from_env(model: nn.Module) -> None:
|
|
"""Auto-protect the weakest static-full heads in duo_dynamic mode.
|
|
|
|
``AHA_FORCE_LOW_ALPHA_FULL_HEADS=K`` selects the K smallest
|
|
``full_attention_heads`` values among heads that static Duo still kept
|
|
full (alpha > 0.5). These are the fragile boundary full heads: raising the
|
|
dynamic threshold can close them, but static Duo's alpha says they should
|
|
not be fully discarded. The selected heads are OR-ed with any manually
|
|
supplied ``AHA_FORCE_FULL_HEADS`` mask.
|
|
"""
|
|
k = _get_optional_nonnegative_int_env("AHA_FORCE_LOW_ALPHA_FULL_HEADS")
|
|
if not k:
|
|
return
|
|
if getattr(model.config, "aha_mode", "dynamic") != "duo_dynamic":
|
|
return
|
|
|
|
candidates: list[tuple[float, int, int]] = []
|
|
with torch.no_grad():
|
|
for layer_idx, layer in enumerate(model.model.layers):
|
|
attn = layer.self_attn
|
|
alpha = getattr(attn, "full_attention_heads", None)
|
|
if alpha is None:
|
|
continue
|
|
for head_idx, value in enumerate(alpha.detach().float().cpu().tolist()):
|
|
if value > 0.5:
|
|
candidates.append((float(value), layer_idx, head_idx))
|
|
selected = sorted(candidates)[: min(k, len(candidates))]
|
|
if not selected:
|
|
return
|
|
|
|
selected_by_layer: dict[int, list[int]] = {}
|
|
for _, layer_idx, head_idx in selected:
|
|
selected_by_layer.setdefault(layer_idx, []).append(head_idx)
|
|
with torch.no_grad():
|
|
for layer_idx, heads in selected_by_layer.items():
|
|
mask = model.model.layers[layer_idx].self_attn._aha_force_full_heads_mask
|
|
for head_idx in heads:
|
|
mask[head_idx] = True
|
|
model.config.aha_force_low_alpha_full_heads = int(k)
|
|
formatted = ",".join(f"{layer}:{head}" for _, layer, head in selected)
|
|
print(
|
|
f"[AHA] AHA_FORCE_LOW_ALPHA_FULL_HEADS={k} selected {len(selected)} "
|
|
f"duo_dynamic heads: {formatted}",
|
|
flush=True,
|
|
)
|
|
|
|
|
|
class AHAQwen3Config(Qwen3Config):
|
|
"""Qwen3Config with AHA-specific parameters."""
|
|
model_type = "aha_qwen3"
|
|
|
|
def __init__(
|
|
self,
|
|
aha_window_size=AHA_WINDOW_SIZE,
|
|
aha_lambda=AHA_LAMBDA,
|
|
aha_distill_weight=AHA_DISTILL_WEIGHT,
|
|
aha_ce_weight=AHA_CE_WEIGHT,
|
|
aha_gate_target=AHA_GATE_TARGET,
|
|
aha_reg_weight=-1.0,
|
|
aha_force_gate_value=None,
|
|
aha_mode=AHA_MODE,
|
|
duo_sink_size=DUO_SINK_SIZE,
|
|
duo_recent_size=DUO_RECENT_SIZE,
|
|
duo_alpha_init=DUO_ALPHA_INIT,
|
|
aha_local_kind=AHA_LOCAL_KIND,
|
|
aha_router_granularity=AHA_ROUTER_GRANULARITY,
|
|
**kwargs,
|
|
):
|
|
super().__init__(**kwargs)
|
|
self.aha_window_size = aha_window_size
|
|
self.aha_lambda = aha_lambda
|
|
# Per-layer attention distillation weight. When > 0 and model is
|
|
# training, each layer contributes ``mean(((1-g)·(global-local))**2)``
|
|
# to the total loss. This is a dense per-(token, head) signal that
|
|
# pushes the gate toward 1 where SWA is insufficient to approximate
|
|
# full attention. Cost: ~0 extra forward compute (global/local are
|
|
# already computed for the mix).
|
|
self.aha_distill_weight = aha_distill_weight
|
|
# Language-modeling CE weight for the training objective. Default 1.0
|
|
# keeps the standard SFT behaviour. Set to 0.0 to drop CE entirely
|
|
# and train the gate purely from ``aux + distill`` (useful when the
|
|
# backbone is frozen and we only want to shape the gate). Evaluation
|
|
# always reports unweighted ``ce_loss`` regardless of this setting.
|
|
self.aha_ce_weight = aha_ce_weight
|
|
# Direct Duo-style sparsity regularizer. When >= 0, this replaces the
|
|
# legacy hinge term and uses ``aha_reg_weight * mean(gate_soft)``.
|
|
# Default -1 keeps older checkpoints on their original hinge objective.
|
|
self.aha_reg_weight = float(aha_reg_weight)
|
|
# Optional override used by Duo-style dynamic distillation. When set
|
|
# to 1.0, dynamic mode behaves as a full-attention teacher for one
|
|
# forward; None keeps the learned gate.
|
|
self.aha_force_gate_value = aha_force_gate_value
|
|
# Target ceiling on mean(gate_soft). Aux loss is hinge-shaped:
|
|
# ``λ · max(0, mean(gate_soft) - τ)``. Setting τ < 1.0 guarantees
|
|
# the gate cannot drop below ``1 - τ`` local-attention fraction
|
|
# in steady state -- the failure mode we observed with the legacy
|
|
# unconditional aux, where ``λ = 1.0`` drove mean(gate_hard) to
|
|
# 0.005 (99.5% local) and killed long-context accuracy
|
|
# (gsm8k 74% -> 19%). Default 1.0 is a no-op for backward
|
|
# compatibility with checkpoints trained on the old loss.
|
|
self.aha_gate_target = aha_gate_target
|
|
# DuoAttention-style baseline. ``dynamic`` keeps the original per-
|
|
# (token, kv_head) MLP gate. ``duo`` replaces it with a per-(layer,
|
|
# kv_head) learnable scalar and swaps the local branch's mask for a
|
|
# streaming (sink + recent) pattern. ``duo_dynamic`` keeps the Duo
|
|
# static mask for KV-capacity planning, but applies a dynamic AHA gate
|
|
# only on the static-full heads: effective_full = duo_full AND dyn_gate.
|
|
# Sink / recent sizes are fixed at train time; set ``duo_sink_size=0``
|
|
# for the sink-ablation run.
|
|
if aha_mode not in ("dynamic", "duo", "duo_dynamic"):
|
|
raise ValueError(
|
|
f"aha_mode must be 'dynamic', 'duo', or 'duo_dynamic', got {aha_mode!r}"
|
|
)
|
|
self.aha_mode = aha_mode
|
|
if aha_router_granularity not in ("token", "token_kv_head"):
|
|
raise ValueError(
|
|
"aha_router_granularity must be 'token' or 'token_kv_head', "
|
|
f"got {aha_router_granularity!r}"
|
|
)
|
|
# Persist this even for static Duo checkpoints. Dynamic AHA uses it to
|
|
# choose the native q_proj gate-row count; old configs omit the field
|
|
# and therefore retain the historical token_kv_head behavior.
|
|
self.aha_router_granularity = aha_router_granularity
|
|
self.duo_sink_size = int(duo_sink_size)
|
|
self.duo_recent_size = int(duo_recent_size)
|
|
self.duo_alpha_init = float(duo_alpha_init)
|
|
if aha_local_kind not in ("sliding_window", "sink_recent"):
|
|
raise ValueError(
|
|
f"aha_local_kind must be 'sliding_window' or 'sink_recent', "
|
|
f"got {aha_local_kind!r}"
|
|
)
|
|
# When aha_mode="duo" the local branch is always sink+recent regardless
|
|
# of this field (forced to the historical DuoAttention behaviour). The
|
|
# field only controls the dynamic mode's local mask.
|
|
self.aha_local_kind = aha_local_kind
|
|
|
|
|
|
def _build_streaming_causal_mask(
|
|
seq_len_q: int,
|
|
seq_len_kv: int,
|
|
sink_size: int,
|
|
recent_size: int,
|
|
device: torch.device,
|
|
dtype: torch.dtype,
|
|
kv_offset: int = 0,
|
|
) -> torch.Tensor:
|
|
"""Construct a streaming attention mask (sink + recent, causal).
|
|
|
|
Returns an additive mask of shape ``[1, 1, Q, K]`` where
|
|
disallowed positions are ``finfo(dtype).min`` and allowed positions
|
|
are ``0``. A query at position ``q`` (absolute in the sequence,
|
|
with ``kv_offset`` tokens already in the KV cache) attends to a key
|
|
at position ``k`` iff ``k <= q_abs`` AND (``k < sink_size`` OR
|
|
``q_abs - k < recent_size``).
|
|
"""
|
|
q_abs = torch.arange(seq_len_q, device=device).unsqueeze(-1) + kv_offset # [Q, 1]
|
|
k_abs = torch.arange(seq_len_kv, device=device).unsqueeze(0) # [1, K]
|
|
causal = k_abs <= q_abs
|
|
sink = k_abs < sink_size if sink_size > 0 else torch.zeros_like(causal)
|
|
recent = (q_abs - k_abs) < recent_size
|
|
allowed = causal & (sink | recent)
|
|
mask = torch.zeros(seq_len_q, seq_len_kv, dtype=dtype, device=device)
|
|
mask.masked_fill_(~allowed, torch.finfo(dtype).min)
|
|
return mask.unsqueeze(0).unsqueeze(0)
|
|
|
|
|
|
class AHAQwen3Attention(nn.Module):
|
|
"""Qwen3 attention with AHA gate for per-head local/global routing."""
|
|
|
|
def __init__(self, config: AHAQwen3Config, layer_idx: int):
|
|
super().__init__()
|
|
self.config = config
|
|
self.layer_idx = layer_idx
|
|
self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
|
|
self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
|
|
self.scaling = self.head_dim ** -0.5
|
|
self.attention_dropout = config.attention_dropout
|
|
self.is_causal = True
|
|
self.aha_window_size = config.aha_window_size
|
|
|
|
# q_proj: extra router outputs for either one shared per-token gate or
|
|
# one per-(token, KV-head) gate.
|
|
# DUO mode uses a per-(layer, kv_head) static scalar instead, so
|
|
# q_proj keeps its vanilla Qwen3 shape and is binary-compatible
|
|
# with the base model checkpoint.
|
|
self._has_dyn_gate_logits = getattr(config, "aha_mode", "dynamic") in (
|
|
"dynamic",
|
|
"duo_dynamic",
|
|
)
|
|
self.aha_router_outputs = aha_router_output_size(config)
|
|
q_out_dim = config.num_attention_heads * self.head_dim
|
|
if self._has_dyn_gate_logits:
|
|
q_out_dim += self.aha_router_outputs
|
|
self.q_proj = nn.Linear(
|
|
config.hidden_size,
|
|
q_out_dim,
|
|
bias=config.attention_bias,
|
|
)
|
|
self.k_proj = nn.Linear(
|
|
config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
|
|
)
|
|
self.v_proj = nn.Linear(
|
|
config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
|
|
)
|
|
self.o_proj = nn.Linear(
|
|
config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
|
|
)
|
|
# Qwen3 applies q_norm/k_norm per head_dim (not full projection)
|
|
self.q_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps)
|
|
self.k_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps)
|
|
|
|
# DuoAttention baseline: per-(layer, kv_head) learnable scalar.
|
|
# Populated in ``duo`` and ``duo_dynamic``. In ``duo_dynamic`` this
|
|
# scalar is the frozen Duo capacity mask; the dynamic gate is only
|
|
# allowed to close heads where this scalar says "full". Registered
|
|
# as ``full_attention_heads`` to match the name used in the
|
|
# official DuoAttention repo (``duo_attn.patch.llama``), so
|
|
# per-head statistics (e.g. retrieval heads vs streaming heads)
|
|
# can be read back with the same tooling.
|
|
if getattr(config, "aha_mode", "dynamic") in ("duo", "duo_dynamic"):
|
|
init = float(getattr(config, "duo_alpha_init", 1.0))
|
|
self.full_attention_heads = nn.Parameter(
|
|
torch.full((config.num_key_value_heads,), init, dtype=torch.float32)
|
|
)
|
|
else:
|
|
self.full_attention_heads = None
|
|
force_full_mask = torch.zeros(config.num_key_value_heads, dtype=torch.bool)
|
|
for raw in os.environ.get("AHA_FORCE_FULL_HEADS", "").split(","):
|
|
raw = raw.strip()
|
|
if not raw:
|
|
continue
|
|
layer_s, head_s = raw.split(":", 1)
|
|
if int(layer_s) == int(layer_idx):
|
|
force_full_mask[int(head_s)] = True
|
|
self.register_buffer(
|
|
"_aha_force_full_heads_mask",
|
|
force_full_mask,
|
|
persistent=False,
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
position_embeddings: tuple[torch.Tensor, torch.Tensor],
|
|
attention_mask: Optional[dict], # dict with "global" and "local" keys
|
|
past_key_values: Optional[Cache] = None,
|
|
cache_position: Optional[torch.LongTensor] = None,
|
|
**kwargs,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
|
|
input_shape = hidden_states.shape[:-1]
|
|
hidden_shape = (*input_shape, -1, self.head_dim)
|
|
|
|
# === Q projection + (optional) gate extraction ===
|
|
q_proj_out = self.q_proj(hidden_states)
|
|
if self._has_dyn_gate_logits:
|
|
query_states, gate = torch.split(
|
|
q_proj_out,
|
|
[
|
|
self.config.num_attention_heads * self.head_dim,
|
|
self.aha_router_outputs,
|
|
],
|
|
dim=-1,
|
|
)
|
|
else:
|
|
query_states = q_proj_out
|
|
# DUO mode: use a zero placeholder so downstream dtype/device
|
|
# checks stay unchanged. The actual gate comes from the
|
|
# per-kv-head scalar ``full_attention_heads`` below.
|
|
gate = torch.zeros(
|
|
*input_shape, self.config.num_key_value_heads,
|
|
dtype=query_states.dtype, device=query_states.device,
|
|
)
|
|
|
|
# Qwen3-style: norm then reshape then transpose
|
|
query_states = self.q_norm(query_states.view(hidden_shape)).transpose(1, 2)
|
|
key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
|
|
value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
|
|
|
|
cos, sin = position_embeddings
|
|
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
|
|
|
|
if past_key_values is not None:
|
|
cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
|
|
key_states, value_states = past_key_values.update(
|
|
key_states, value_states, self.layer_idx, cache_kwargs
|
|
)
|
|
|
|
# === Gate computation ===
|
|
# gate shape: (B, Q, H_kv) — per-(token, kv_head) in dynamic mode,
|
|
# or per-kv_head broadcast to (B, Q, H_kv) in DUO mode.
|
|
# At inference, threshold is read from env var AHA_GATE_HARD_THRESHOLD
|
|
# so users can trade sparsity for quality without re-training.
|
|
# Default 0.5 keeps training-time behaviour. Lowering to e.g. 0.35
|
|
# lets more (token, kv_head) pairs route through global attention,
|
|
# which recovers long-CoT quality at the cost of sparsity.
|
|
aha_mode = getattr(self.config, "aha_mode", "dynamic")
|
|
if aha_mode == "duo":
|
|
# DuoAttention: alpha is per-(layer, kv_head) learnable scalar.
|
|
# gate_soft = clamp(alpha, 0, 1); gate_hard = (alpha > 0.5).
|
|
# Both broadcast to (B, Q, H_kv) for a uniform downstream API
|
|
# (so the distill / aux / blend paths do not need to branch).
|
|
qhead_alpha = getattr(self, "_aha_qhead_full_attention_heads", None)
|
|
if qhead_alpha is not None:
|
|
# Eval-only oracle diagnostic: route at Q-head granularity
|
|
# instead of KV-head/GQA-group granularity. This does not
|
|
# represent a KV-capacity-saving deployment point; it isolates
|
|
# whether GQA grouping itself is the quality bottleneck.
|
|
alpha = qhead_alpha.to(gate.dtype).clamp(0.0, 1.0)
|
|
gate_soft = alpha.view(1, 1, -1).expand(*input_shape, -1)
|
|
import os as _os
|
|
_thr = float(_os.environ.get("AHA_GATE_HARD_THRESHOLD", "0.5"))
|
|
gate_hard = (alpha > _thr).to(gate.dtype).view(1, 1, -1).expand(*input_shape, -1)
|
|
gate_ste = gate_hard
|
|
|
|
else:
|
|
alpha = self.full_attention_heads.to(gate.dtype).clamp(0.0, 1.0)
|
|
gate_soft = alpha.view(1, 1, -1).expand(*input_shape, -1)
|
|
if self.training:
|
|
gate_hard = (alpha > 0.5).to(gate.dtype).view(1, 1, -1).expand(*input_shape, -1)
|
|
# Straight-through estimator through the clamped alpha so
|
|
# L1 / distill gradients flow into the scalar.
|
|
gate_ste = gate_soft
|
|
else:
|
|
import os as _os
|
|
_thr = float(_os.environ.get("AHA_GATE_HARD_THRESHOLD", "0.5"))
|
|
gate_hard = (alpha > _thr).to(gate.dtype).view(1, 1, -1).expand(*input_shape, -1)
|
|
gate_ste = gate_hard
|
|
elif aha_mode == "duo_dynamic":
|
|
# Duo + AHA-on-full-heads:
|
|
# duo_full_mask = 1 for static Duo full/retrieval heads
|
|
# dyn_gate = per-token AHA decision inside those heads
|
|
# effective_full = duo_full_mask AND dyn_gate
|
|
#
|
|
# Duo streaming heads are locked to local/streaming forever. This
|
|
# preserves Duo's KV-capacity saving: only static-full heads need a
|
|
# full KV cache. ``aha_force_gate_value=1`` therefore reproduces
|
|
# the static Duo deployment point exactly, not all-full attention.
|
|
alpha = self.full_attention_heads.to(gate.dtype).clamp(0.0, 1.0)
|
|
duo_full = (alpha > 0.5).to(gate.dtype).view(1, 1, -1)
|
|
dyn_soft = torch.sigmoid(gate)
|
|
force_gate_value = getattr(self.config, "aha_force_gate_value", None)
|
|
if force_gate_value is not None:
|
|
dyn_soft = torch.full_like(dyn_soft, float(force_gate_value))
|
|
gate_soft = dyn_soft * duo_full
|
|
gate_hard = gate_soft
|
|
gate_ste = gate_soft
|
|
elif self.training:
|
|
train_thr = float(os.environ.get("AHA_TRAIN_GATE_HARD_THRESHOLD", "0.5"))
|
|
gate_soft = dyn_soft * duo_full
|
|
dyn_hard = (dyn_soft > train_thr).to(gate.dtype)
|
|
gate_hard = dyn_hard * duo_full
|
|
gate_ste = gate_hard + (gate_soft - gate_soft.detach())
|
|
else:
|
|
import os as _os
|
|
_thr = float(_os.environ.get("AHA_GATE_HARD_THRESHOLD", "0.5"))
|
|
gate_soft = dyn_soft * duo_full
|
|
dyn_hard = (dyn_soft > _thr).to(gate.dtype)
|
|
gate_hard = dyn_hard * duo_full
|
|
gate_ste = gate_hard
|
|
force_full = self._aha_force_full_heads_mask.to(gate.dtype).view(1, 1, -1)
|
|
if bool(force_full.any().item()):
|
|
force_full = force_full * duo_full
|
|
gate_hard = torch.maximum(gate_hard, force_full)
|
|
gate_ste = torch.maximum(gate_ste, force_full)
|
|
else:
|
|
gate_soft = torch.sigmoid(gate)
|
|
force_gate_value = getattr(self.config, "aha_force_gate_value", None)
|
|
if force_gate_value is not None:
|
|
force = float(force_gate_value)
|
|
gate_soft = torch.full_like(gate_soft, force)
|
|
gate_hard = torch.full_like(gate_soft, force)
|
|
gate_ste = gate_soft
|
|
elif self.training:
|
|
train_thr = float(os.environ.get("AHA_TRAIN_GATE_HARD_THRESHOLD", "0.5"))
|
|
gate_hard = (gate_soft > train_thr).to(gate.dtype)
|
|
gate_ste = gate_hard + (gate_soft - gate_soft.detach())
|
|
else:
|
|
import os as _os
|
|
_thr = float(_os.environ.get("AHA_GATE_HARD_THRESHOLD", "0.5"))
|
|
gate_hard = (gate_soft > _thr).to(gate.dtype)
|
|
gate_ste = gate_hard
|
|
|
|
# Normalize the downstream contract to effective
|
|
# token x KV-head decisions. In shared-token mode there is only one
|
|
# learned logit per token, but regularization, sparsity accounting, and
|
|
# attention mixing all see the same KV-head denominator as the
|
|
# head-granular arm.
|
|
if gate_soft.shape[-1] == 1 and self.config.num_key_value_heads != 1:
|
|
effective_shape = (*gate_soft.shape[:-1], self.config.num_key_value_heads)
|
|
gate_soft = gate_soft.expand(effective_shape)
|
|
gate_hard = gate_hard.expand(effective_shape)
|
|
gate_ste = gate_ste.expand(effective_shape)
|
|
|
|
if (
|
|
aha_mode in ("duo", "duo_dynamic")
|
|
and not self.training
|
|
and os.environ.get("AHA_DUO_PREFILL_FULL", "0") == "1"
|
|
and query_states.shape[-2] == key_states.shape[-2]
|
|
):
|
|
# Diagnostic parity with the official DuoAttention HF eval path:
|
|
# initial prefill is full attention; streaming heads are only
|
|
# localized during subsequent decode steps.
|
|
gate_hard = torch.ones_like(gate_hard)
|
|
gate_ste = torch.ones_like(gate_ste)
|
|
|
|
if not self.training:
|
|
phase = (
|
|
"decode"
|
|
if hidden_states.shape[1] == 1 and key_states.shape[-2] > 1
|
|
else "prefill"
|
|
)
|
|
_track_aha_inference_sparsity(
|
|
gate_hard, gate_soft, self.layer_idx, phase,
|
|
getattr(
|
|
self.config,
|
|
"aha_router_granularity",
|
|
AHA_ROUTER_GRANULARITY,
|
|
),
|
|
self.aha_router_outputs if self._has_dyn_gate_logits
|
|
else self.config.num_key_value_heads,
|
|
)
|
|
|
|
# === Attention: double_mix (global + local, blended by gate) ===
|
|
from transformers.models.qwen3.modeling_qwen3 import eager_attention_forward
|
|
|
|
attention_interface: Callable = eager_attention_forward
|
|
if self.config._attn_implementation != "eager":
|
|
attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
|
|
|
|
# Global attention (full context)
|
|
global_attn_output, _ = attention_interface(
|
|
self, query_states, key_states, value_states,
|
|
attention_mask["global"],
|
|
dropout=0.0 if not self.training else self.attention_dropout,
|
|
scaling=self.scaling,
|
|
**kwargs,
|
|
)
|
|
|
|
if getattr(self.config, "_aha_teacher_full_fastpath", False):
|
|
attn_output = global_attn_output.reshape(*input_shape, -1).contiguous()
|
|
attn_output = self.o_proj(attn_output)
|
|
return attn_output, gate_soft, gate_hard, None
|
|
|
|
# Local / streaming branch. The local mask was pre-built by the
|
|
# model's ``forward`` either as sliding-window (legacy dynamic mode
|
|
# default) or sink+recent (duo mode, and dynamic mode when
|
|
# ``aha_local_kind="sink_recent"``). Only pass ``sliding_window`` to
|
|
# the SDPA backend when the mask is the sliding-window kind, otherwise
|
|
# the backend will double-mask and wipe out the sink positions.
|
|
local_attn_kwargs = dict(
|
|
dropout=0.0 if not self.training else self.attention_dropout,
|
|
scaling=self.scaling,
|
|
)
|
|
local_kind = "sink_recent" if aha_mode in ("duo", "duo_dynamic") else getattr(
|
|
self.config, "aha_local_kind", "sliding_window"
|
|
)
|
|
if local_kind == "sliding_window":
|
|
local_attn_kwargs["sliding_window"] = self.aha_window_size
|
|
local_attn_output, _ = attention_interface(
|
|
self, query_states, key_states, value_states,
|
|
attention_mask["local"],
|
|
**local_attn_kwargs,
|
|
**kwargs,
|
|
)
|
|
|
|
# Blend: normal gate_ste is (B, Q, H_kv); broadcast to (B, Q, H, 1)
|
|
# by repeating each kv_head decision. The optional qhead oracle above
|
|
# already emits (B, Q, H), so it bypasses this broadcast.
|
|
if gate_ste.shape[-1] == self.config.num_attention_heads:
|
|
g_bqh = gate_ste
|
|
else:
|
|
g_bqh = gate_ste.repeat_interleave(self.num_key_value_groups, dim=-1)
|
|
attn_output = (
|
|
global_attn_output * g_bqh.unsqueeze(-1)
|
|
+ local_attn_output * (1 - g_bqh.unsqueeze(-1))
|
|
)
|
|
|
|
# Per-layer attention distillation signal. Returns a scalar so the
|
|
# outer model can aggregate across layers. When gate → 1 the term
|
|
# vanishes; when gate → 0 on positions where global/local disagree
|
|
# strongly, the gradient pushes gate back toward 1.
|
|
#
|
|
# ``diff`` is ``detach``-ed so distill only trains the gate; it does
|
|
# not reshape attention outputs through the backbone. This also
|
|
# prevents the squared-difference graph from retaining activations
|
|
# for all 28 layers' global/local branches, which was the source of
|
|
# OOM when gradient checkpointing is disabled on the custom forward.
|
|
distill_per_layer: Optional[torch.Tensor] = None
|
|
if self.training and getattr(self.config, "aha_distill_weight", 0.0) > 0.0:
|
|
with torch.no_grad():
|
|
diff_sq = (global_attn_output - local_attn_output).float() ** 2
|
|
one_minus_g_sq = ((1.0 - g_bqh.float()) ** 2).unsqueeze(-1)
|
|
distill_per_layer = (one_minus_g_sq * diff_sq).mean()
|
|
|
|
attn_output = attn_output.reshape(*input_shape, -1).contiguous()
|
|
attn_output = self.o_proj(attn_output)
|
|
return attn_output, gate_soft, gate_hard, distill_per_layer
|
|
|
|
|
|
class AHAQwen3DecoderLayer(GradientCheckpointingLayer):
|
|
"""Decoder layer with gradient-checkpointing support.
|
|
|
|
Inheriting from ``GradientCheckpointingLayer`` lets HuggingFace Trainer's
|
|
``gradient_checkpointing_enable`` take effect on this layer without any
|
|
explicit ``_gradient_checkpointing_func`` call in the outer forward.
|
|
"""
|
|
|
|
def __init__(self, config: AHAQwen3Config, layer_idx: int):
|
|
super().__init__()
|
|
self.hidden_size = config.hidden_size
|
|
self.self_attn = AHAQwen3Attention(config=config, layer_idx=layer_idx)
|
|
from transformers.models.qwen3.modeling_qwen3 import Qwen3MLP
|
|
self.mlp = Qwen3MLP(config)
|
|
self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
|
self.post_attention_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
attention_mask: Optional[dict] = None,
|
|
position_ids: Optional[torch.LongTensor] = None,
|
|
past_key_values: Optional[Cache] = None,
|
|
use_cache: Optional[bool] = False,
|
|
cache_position: Optional[torch.LongTensor] = None,
|
|
position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
|
|
**kwargs,
|
|
):
|
|
residual = hidden_states
|
|
hidden_states = self.input_layernorm(hidden_states)
|
|
hidden_states, gate_soft, gate_hard, distill_per_layer = self.self_attn(
|
|
hidden_states=hidden_states,
|
|
attention_mask=attention_mask,
|
|
position_ids=position_ids,
|
|
past_key_values=past_key_values,
|
|
use_cache=use_cache,
|
|
cache_position=cache_position,
|
|
position_embeddings=position_embeddings,
|
|
**kwargs,
|
|
)
|
|
hidden_states = residual + hidden_states
|
|
|
|
residual = hidden_states
|
|
hidden_states = self.post_attention_layernorm(hidden_states)
|
|
hidden_states = self.mlp(hidden_states)
|
|
hidden_states = residual + hidden_states
|
|
return hidden_states, gate_soft, gate_hard, distill_per_layer
|
|
|
|
|
|
@dataclass
|
|
class AHACausalLMOutputWithPast(CausalLMOutputWithPast):
|
|
"""Causal LM output with AHA loss breakdown and gate density stats (training / eval with labels)."""
|
|
|
|
ce_loss: Optional[torch.FloatTensor] = None
|
|
gate_aux_loss: Optional[torch.FloatTensor] = None
|
|
distill_loss: Optional[torch.FloatTensor] = None
|
|
gate_soft_mean: Optional[torch.FloatTensor] = None
|
|
gate_hard_mean: Optional[torch.FloatTensor] = None
|
|
|
|
|
|
class AHAModelOutputWithPast(BaseModelOutputWithPast):
|
|
def __init__(self, last_hidden_state, past_key_values=None, hidden_states=None, attentions=None,
|
|
all_gate_soft=None, all_gate_hard=None, all_distill_per_layer=None):
|
|
super().__init__(last_hidden_state=last_hidden_state, past_key_values=past_key_values,
|
|
hidden_states=hidden_states, attentions=attentions)
|
|
self.all_gate_soft = all_gate_soft
|
|
self.all_gate_hard = all_gate_hard
|
|
# Tuple of per-layer distill scalars, or empty tuple when distillation is off.
|
|
self.all_distill_per_layer = all_distill_per_layer if all_distill_per_layer is not None else ()
|
|
|
|
|
|
class AHAQwen3Model(Qwen3Model):
|
|
config_class = AHAQwen3Config
|
|
|
|
def __init__(self, config: AHAQwen3Config):
|
|
from transformers import PreTrainedModel
|
|
PreTrainedModel.__init__(self, config)
|
|
self.padding_idx = config.pad_token_id
|
|
self.vocab_size = config.vocab_size
|
|
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
|
|
self.layers = nn.ModuleList(
|
|
[AHAQwen3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
|
|
)
|
|
self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
|
from transformers.models.qwen3.modeling_qwen3 import Qwen3RotaryEmbedding
|
|
self.rotary_emb = Qwen3RotaryEmbedding(config=config)
|
|
self.gradient_checkpointing = False
|
|
self.post_init()
|
|
|
|
def forward(
|
|
self,
|
|
input_ids: Optional[torch.LongTensor] = None,
|
|
attention_mask: Optional[torch.Tensor] = None,
|
|
position_ids: Optional[torch.LongTensor] = None,
|
|
past_key_values: Optional[Cache] = None,
|
|
inputs_embeds: Optional[torch.FloatTensor] = None,
|
|
use_cache: Optional[bool] = None,
|
|
cache_position: Optional[torch.LongTensor] = None,
|
|
**kwargs,
|
|
) -> AHAModelOutputWithPast:
|
|
if (input_ids is None) ^ (inputs_embeds is not None):
|
|
raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
|
|
if inputs_embeds is None:
|
|
inputs_embeds = self.embed_tokens(input_ids)
|
|
if use_cache and past_key_values is None:
|
|
past_key_values = DynamicCache(config=self.config)
|
|
if cache_position is None:
|
|
past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
|
|
cache_position = torch.arange(
|
|
past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
|
|
)
|
|
if position_ids is None:
|
|
position_ids = cache_position.unsqueeze(0)
|
|
|
|
# Build two masks: global (full causal) and local (sliding window
|
|
# for dynamic mode, streaming sink+recent for DUO mode).
|
|
orig_sliding_window = getattr(self.config, "sliding_window", None)
|
|
self.config.sliding_window = self.config.aha_window_size
|
|
mask_kwargs = {
|
|
"config": self.config,
|
|
"input_embeds": inputs_embeds,
|
|
"attention_mask": attention_mask,
|
|
"cache_position": cache_position,
|
|
"past_key_values": past_key_values,
|
|
"position_ids": position_ids,
|
|
}
|
|
global_mask = create_causal_mask(**mask_kwargs)
|
|
aha_mode = getattr(self.config, "aha_mode", "dynamic")
|
|
# duo / duo_dynamic force sink_recent for backward compat. dynamic mode
|
|
# reads the local kind from config (default sliding_window for legacy
|
|
# Bv3 ckpts; hot-started ckpts from duo set sink_recent — see
|
|
# §9.4.13.3).
|
|
local_kind = "sink_recent" if aha_mode in ("duo", "duo_dynamic") else getattr(
|
|
self.config, "aha_local_kind", "sliding_window"
|
|
)
|
|
if local_kind == "sink_recent":
|
|
seq_len_q = inputs_embeds.shape[1]
|
|
past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0
|
|
seq_len_kv = seq_len_q + past_seen
|
|
local_mask = _build_streaming_causal_mask(
|
|
seq_len_q=seq_len_q,
|
|
seq_len_kv=seq_len_kv,
|
|
sink_size=int(getattr(self.config, "duo_sink_size", 64)),
|
|
recent_size=int(getattr(self.config, "duo_recent_size", 256)),
|
|
device=inputs_embeds.device,
|
|
dtype=inputs_embeds.dtype,
|
|
kv_offset=past_seen,
|
|
)
|
|
else:
|
|
local_mask = create_sliding_window_causal_mask(**mask_kwargs)
|
|
causal_mask_dict = {"global": global_mask, "local": local_mask}
|
|
self.config.sliding_window = orig_sliding_window
|
|
|
|
hidden_states = inputs_embeds
|
|
position_embeddings = self.rotary_emb(hidden_states, position_ids)
|
|
|
|
all_gate_soft = ()
|
|
all_gate_hard = ()
|
|
all_distill_per_layer: tuple[torch.Tensor, ...] = ()
|
|
for decoder_layer in self.layers:
|
|
hidden_states, gate_soft, gate_hard, distill_per_layer = decoder_layer(
|
|
hidden_states,
|
|
attention_mask=causal_mask_dict,
|
|
position_ids=position_ids,
|
|
past_key_values=past_key_values,
|
|
use_cache=use_cache,
|
|
cache_position=cache_position,
|
|
position_embeddings=position_embeddings,
|
|
**kwargs,
|
|
)
|
|
all_gate_soft += (gate_soft,)
|
|
all_gate_hard += (gate_hard,)
|
|
if distill_per_layer is not None:
|
|
all_distill_per_layer += (distill_per_layer,)
|
|
|
|
hidden_states = self.norm(hidden_states)
|
|
return AHAModelOutputWithPast(
|
|
last_hidden_state=hidden_states,
|
|
past_key_values=past_key_values if use_cache else None,
|
|
all_gate_soft=all_gate_soft,
|
|
all_gate_hard=all_gate_hard,
|
|
all_distill_per_layer=all_distill_per_layer,
|
|
)
|
|
|
|
|
|
class AHAQwen3ForCausalLM(Qwen3ForCausalLM):
|
|
config_class = AHAQwen3Config
|
|
_tied_weights_keys = ["lm_head.weight"]
|
|
|
|
def __init__(self, config: AHAQwen3Config):
|
|
# Skip Qwen3ForCausalLM.__init__ which creates Qwen3Model
|
|
# Go to PreTrainedModel.__init__ directly
|
|
from transformers import PreTrainedModel
|
|
PreTrainedModel.__init__(self, config)
|
|
self.model = AHAQwen3Model(config)
|
|
self.vocab_size = config.vocab_size
|
|
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
|
self.post_init()
|
|
|
|
def forward(
|
|
self,
|
|
input_ids: Optional[torch.LongTensor] = None,
|
|
attention_mask: Optional[torch.Tensor] = None,
|
|
position_ids: Optional[torch.LongTensor] = None,
|
|
past_key_values: Optional[Cache] = None,
|
|
inputs_embeds: Optional[torch.FloatTensor] = None,
|
|
labels: Optional[torch.LongTensor] = None,
|
|
use_cache: Optional[bool] = None,
|
|
cache_position: Optional[torch.LongTensor] = None,
|
|
logits_to_keep: Union[int, torch.Tensor] = 0,
|
|
**kwargs,
|
|
) -> AHACausalLMOutputWithPast:
|
|
outputs = self.model(
|
|
input_ids=input_ids,
|
|
attention_mask=attention_mask,
|
|
position_ids=position_ids,
|
|
past_key_values=past_key_values,
|
|
inputs_embeds=inputs_embeds,
|
|
use_cache=use_cache,
|
|
cache_position=cache_position,
|
|
**kwargs,
|
|
)
|
|
hidden_states = outputs.last_hidden_state
|
|
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
|
|
logits = self.lm_head(hidden_states[:, slice_indices, :])
|
|
|
|
loss = None
|
|
ce_loss = None
|
|
gate_aux_loss = None
|
|
distill_loss = None
|
|
gate_soft_mean = None
|
|
gate_hard_mean = None
|
|
|
|
if labels is not None:
|
|
from torch.nn import CrossEntropyLoss
|
|
shift_logits = logits[..., :-1, :].contiguous()
|
|
shift_labels = labels[..., 1:].contiguous()
|
|
loss_fct = CrossEntropyLoss()
|
|
ce_loss = loss_fct(shift_logits.view(-1, self.vocab_size), shift_labels.view(-1))
|
|
|
|
all_gate_soft = outputs.all_gate_soft
|
|
all_gate_hard = outputs.all_gate_hard
|
|
if labels is not None and len(all_gate_soft) > 0 and len(all_gate_hard) > 0:
|
|
gate_soft_stack = torch.stack(
|
|
[g.to(hidden_states.device) for g in all_gate_soft], dim=-1
|
|
).float()
|
|
gate_hard_stack = torch.stack(
|
|
[g.to(hidden_states.device).float() for g in all_gate_hard], dim=-1
|
|
)
|
|
shift_gate_soft = gate_soft_stack[:, :-1, :].contiguous()
|
|
shift_gate_hard = gate_hard_stack[:, :-1, :].contiguous()
|
|
# Match the CE/PPL population exactly. The old unmasked mean
|
|
# counted padded sequence positions as real gate decisions, which
|
|
# made the training sparsity curve disagree with inference-time
|
|
# sparsity on variable-length batches.
|
|
valid_gate = shift_labels.ne(-100)
|
|
if attention_mask is not None:
|
|
valid_gate = valid_gate & attention_mask[:, 1:].to(torch.bool)
|
|
if bool(valid_gate.any()):
|
|
valid_gate = valid_gate[..., None, None].to(shift_gate_soft.dtype)
|
|
normalizer = (
|
|
valid_gate.sum()
|
|
* shift_gate_soft.shape[-2]
|
|
* shift_gate_soft.shape[-1]
|
|
)
|
|
gate_soft_mean = (shift_gate_soft * valid_gate).sum() / normalizer
|
|
gate_hard_mean = (shift_gate_hard * valid_gate).sum() / normalizer
|
|
else:
|
|
gate_soft_mean = shift_gate_soft.mean()
|
|
gate_hard_mean = shift_gate_hard.mean()
|
|
|
|
# Per-layer distillation loss: mean over layers of
|
|
# ``mean(((1-g)·(global-local))**2)``. Populated only when training
|
|
# with ``config.aha_distill_weight > 0`` (see AHAQwen3Attention).
|
|
all_distill = outputs.all_distill_per_layer
|
|
if self.training and len(all_distill) > 0:
|
|
distill_loss = torch.stack(list(all_distill)).mean()
|
|
|
|
# === Loss assembly ===
|
|
# Eval path: standard LM loss for compatibility with HF metrics.
|
|
# Train path: weighted sum of (ce, gate_aux, distill) so callers can
|
|
# zero any term by setting its weight to 0 (e.g. ``aha_ce_weight=0``
|
|
# to train the gate purely from ``aux + distill``).
|
|
if not self.training:
|
|
loss = ce_loss
|
|
else:
|
|
ce_weight = float(getattr(self.config, "aha_ce_weight", 1.0))
|
|
distill_weight = float(getattr(self.config, "aha_distill_weight", 0.0))
|
|
terms: list[torch.Tensor] = []
|
|
if ce_loss is not None and ce_weight > 0.0:
|
|
terms.append(ce_weight * ce_loss)
|
|
if gate_soft_mean is not None:
|
|
reg_weight = float(getattr(self.config, "aha_reg_weight", -1.0))
|
|
if reg_weight >= 0.0:
|
|
# Duo-style direct sparsity term:
|
|
# alpha/gate near 1 means full/global attention, so penalize
|
|
# mean gate usage directly.
|
|
gate_aux_loss = reg_weight * gate_soft_mean
|
|
else:
|
|
# Legacy hinge aux: penalise only the excess above the target
|
|
# ceiling. Below ``τ`` the gradient is zero, so CE and
|
|
# distill freely shape where the gate opens without an
|
|
# unbounded downward push from aux. See AHAQwen3Config
|
|
# for the rationale and the collapse that motivated it.
|
|
gate_target = float(getattr(self.config, "aha_gate_target", 1.0))
|
|
excess = torch.clamp(gate_soft_mean - gate_target, min=0.0)
|
|
gate_aux_loss = self.config.aha_lambda * excess
|
|
terms.append(gate_aux_loss)
|
|
if distill_loss is not None and distill_weight > 0.0:
|
|
terms.append(distill_weight * distill_loss)
|
|
if terms:
|
|
loss = terms[0]
|
|
for t in terms[1:]:
|
|
loss = loss + t
|
|
|
|
return AHACausalLMOutputWithPast(
|
|
loss=loss,
|
|
logits=logits,
|
|
past_key_values=outputs.past_key_values,
|
|
ce_loss=ce_loss,
|
|
gate_aux_loss=gate_aux_loss,
|
|
distill_loss=distill_loss,
|
|
gate_soft_mean=gate_soft_mean,
|
|
gate_hard_mean=gate_hard_mean,
|
|
)
|
|
|
|
@classmethod
|
|
def from_pretrained_qwen3(cls, qwen3_path: str, aha_window_size: int = AHA_WINDOW_SIZE,
|
|
aha_lambda: float = AHA_LAMBDA,
|
|
aha_distill_weight: float = AHA_DISTILL_WEIGHT,
|
|
aha_ce_weight: float = AHA_CE_WEIGHT,
|
|
aha_gate_target: float = AHA_GATE_TARGET,
|
|
aha_reg_weight: float = -1.0,
|
|
aha_force_gate_value=None,
|
|
aha_mode: str = AHA_MODE,
|
|
duo_sink_size: int = DUO_SINK_SIZE,
|
|
duo_recent_size: int = DUO_RECENT_SIZE,
|
|
duo_alpha_init: float = DUO_ALPHA_INIT,
|
|
aha_local_kind: str = AHA_LOCAL_KIND,
|
|
aha_router_granularity: str = AHA_ROUTER_GRANULARITY,
|
|
**kwargs):
|
|
"""Load a pretrained Qwen3 model and convert to AHA-Qwen3.
|
|
|
|
Initializes the gate weights in q_proj randomly (the extra num_heads outputs).
|
|
All other weights are copied from the pretrained model.
|
|
"""
|
|
# Load original Qwen3
|
|
original = Qwen3ForCausalLM.from_pretrained(qwen3_path, **kwargs)
|
|
original_config = original.config
|
|
|
|
# Create AHA config
|
|
config_dict = original_config.to_dict()
|
|
config_dict["aha_window_size"] = aha_window_size
|
|
config_dict["aha_lambda"] = aha_lambda
|
|
config_dict["aha_distill_weight"] = aha_distill_weight
|
|
config_dict["aha_ce_weight"] = aha_ce_weight
|
|
config_dict["aha_gate_target"] = aha_gate_target
|
|
config_dict["aha_reg_weight"] = aha_reg_weight
|
|
config_dict["aha_force_gate_value"] = aha_force_gate_value
|
|
config_dict["aha_mode"] = aha_mode
|
|
config_dict["duo_sink_size"] = duo_sink_size
|
|
config_dict["duo_recent_size"] = duo_recent_size
|
|
config_dict["duo_alpha_init"] = duo_alpha_init
|
|
config_dict["aha_local_kind"] = aha_local_kind
|
|
config_dict["aha_router_granularity"] = aha_router_granularity
|
|
config_dict["model_type"] = "aha_qwen3"
|
|
aha_config = AHAQwen3Config(**config_dict)
|
|
aha_config._attn_implementation = kwargs.get("attn_implementation", "sdpa")
|
|
|
|
# Create AHA model
|
|
aha_model = cls(aha_config)
|
|
|
|
# Copy weights
|
|
# embed_tokens, lm_head, norm
|
|
aha_model.model.embed_tokens.load_state_dict(original.model.embed_tokens.state_dict())
|
|
aha_model.lm_head.load_state_dict(original.lm_head.state_dict())
|
|
aha_model.model.norm.load_state_dict(original.model.norm.state_dict())
|
|
aha_model.model.rotary_emb.load_state_dict(original.model.rotary_emb.state_dict())
|
|
|
|
# Copy per-layer weights
|
|
num_heads = original_config.num_attention_heads
|
|
head_dim = aha_config.head_dim
|
|
|
|
for i in range(original_config.num_hidden_layers):
|
|
orig_layer = original.model.layers[i]
|
|
aha_layer = aha_model.model.layers[i]
|
|
|
|
# MLP + norms: direct copy
|
|
aha_layer.mlp.load_state_dict(orig_layer.mlp.state_dict())
|
|
aha_layer.input_layernorm.load_state_dict(orig_layer.input_layernorm.state_dict())
|
|
aha_layer.post_attention_layernorm.load_state_dict(orig_layer.post_attention_layernorm.state_dict())
|
|
|
|
# Attention: k_proj, v_proj, o_proj, q_norm, k_norm — direct copy
|
|
aha_attn = aha_layer.self_attn
|
|
orig_attn = orig_layer.self_attn
|
|
aha_attn.k_proj.load_state_dict(orig_attn.k_proj.state_dict())
|
|
aha_attn.v_proj.load_state_dict(orig_attn.v_proj.state_dict())
|
|
aha_attn.o_proj.load_state_dict(orig_attn.o_proj.state_dict())
|
|
aha_attn.q_norm.load_state_dict(orig_attn.q_norm.state_dict())
|
|
aha_attn.k_norm.load_state_dict(orig_attn.k_norm.state_dict())
|
|
|
|
# q_proj:
|
|
# - dynamic / duo_dynamic mode: copy original Q rows, init extra
|
|
# gate rows random (~0.5).
|
|
# - duo mode: q_proj is identical shape to base Qwen3 (no gate
|
|
# rows), so direct load_state_dict works.
|
|
orig_q_weight = orig_attn.q_proj.weight.data # (num_heads*head_dim, hidden_size)
|
|
aha_q_weight = aha_attn.q_proj.weight.data
|
|
if aha_attn._has_dyn_gate_logits:
|
|
aha_q_weight[:num_heads * head_dim, :] = orig_q_weight
|
|
nn.init.normal_(aha_q_weight[num_heads * head_dim:, :], mean=0.0, std=0.01)
|
|
if orig_attn.q_proj.bias is not None and aha_attn.q_proj.bias is not None:
|
|
aha_attn.q_proj.bias.data[:num_heads * head_dim] = orig_attn.q_proj.bias.data
|
|
nn.init.zeros_(aha_attn.q_proj.bias.data[num_heads * head_dim:])
|
|
else:
|
|
aha_attn.q_proj.load_state_dict(orig_attn.q_proj.state_dict())
|
|
|
|
del original
|
|
return aha_model
|
|
|
|
@classmethod
|
|
def from_pretrained_aha(cls, checkpoint_path: str, **kwargs):
|
|
"""Load an AHA-Qwen3 checkpoint saved by `trainer.save_model`.
|
|
|
|
Supports both single-file (``model.safetensors``) and multi-shard
|
|
(``model.safetensors.index.json`` + ``model-*-of-*.safetensors``)
|
|
layouts. The latter is automatically used by HF when a checkpoint
|
|
exceeds the single-file size threshold (e.g. Qwen3-1.7B+).
|
|
"""
|
|
config = AHAQwen3Config.from_pretrained(checkpoint_path)
|
|
config._attn_implementation = kwargs.get(
|
|
"attn_implementation",
|
|
getattr(config, "_attn_implementation", "sdpa"),
|
|
)
|
|
|
|
model = cls(config)
|
|
|
|
single_path = os.path.join(checkpoint_path, "model.safetensors")
|
|
index_path = os.path.join(checkpoint_path, "model.safetensors.index.json")
|
|
if os.path.exists(single_path):
|
|
state_dict = load_file(single_path)
|
|
elif os.path.exists(index_path):
|
|
with open(index_path) as f:
|
|
shard_index = json.load(f)
|
|
shard_files = sorted(set(shard_index["weight_map"].values()))
|
|
state_dict = {}
|
|
for shard in shard_files:
|
|
state_dict.update(load_file(os.path.join(checkpoint_path, shard)))
|
|
else:
|
|
raise FileNotFoundError(
|
|
f"No safetensors file or index found in {checkpoint_path}. "
|
|
f"Expected either model.safetensors or model.safetensors.index.json."
|
|
)
|
|
missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
|
|
allowed_missing = {"lm_head.weight"} if getattr(config, "tie_word_embeddings", False) else set()
|
|
extra_missing = set(missing_keys) - allowed_missing
|
|
if extra_missing or unexpected_keys:
|
|
raise RuntimeError(
|
|
f"Unexpected checkpoint mismatch. missing={sorted(extra_missing)}, "
|
|
f"unexpected={sorted(unexpected_keys)}"
|
|
)
|
|
if "lm_head.weight" in missing_keys:
|
|
model.tie_weights()
|
|
|
|
_apply_low_alpha_force_full_heads_from_env(model)
|
|
|
|
torch_dtype = kwargs.get("torch_dtype")
|
|
if torch_dtype is not None:
|
|
model = model.to(dtype=torch_dtype)
|
|
return model
|