File size: 5,287 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 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | #!/usr/bin/env python3
"""Convert a tuned vanilla Qwen3 checkpoint into a quality-safe AHA hot start.
All vanilla weights are preserved. Dynamic gate weights start at zero and a
trainable bias initializes every gate to the requested full-attention
probability, so hard routing is initially exactly vanilla full attention.
"""
from __future__ import annotations
import argparse
import json
import math
import sys
from pathlib import Path
import torch
from torch import nn
from transformers import AutoTokenizer
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from modeling_aha_qwen3 import AHAQwen3Config, AHAQwen3ForCausalLM
def add_zero_bias(linear: nn.Linear) -> nn.Linear:
new = nn.Linear(
linear.in_features,
linear.out_features,
bias=True,
device=linear.weight.device,
dtype=linear.weight.dtype,
)
with torch.no_grad():
new.weight.copy_(linear.weight)
new.bias.zero_()
return new
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--vanilla-path", required=True)
parser.add_argument("--output-path", required=True)
parser.add_argument("--window-size", type=int, default=128)
parser.add_argument("--gate-init-full-prob", type=float, default=0.90)
parser.add_argument(
"--local-kind", choices=("sliding_window", "sink_recent"), default="sliding_window"
)
parser.add_argument(
"--router-granularity",
choices=("token", "token_kv_head"),
default="token_kv_head",
help="Native dynamic gate shape. token is one shared gate per layer/token.",
)
args = parser.parse_args()
if not 0.5 < args.gate_init_full_prob < 1.0:
raise ValueError("--gate-init-full-prob must be strictly between 0.5 and 1")
output = Path(args.output_path)
output.mkdir(parents=True, exist_ok=True)
AHAQwen3Config.register_for_auto_class()
AHAQwen3ForCausalLM.register_for_auto_class("AutoModelForCausalLM")
model = AHAQwen3ForCausalLM.from_pretrained_qwen3(
args.vanilla_path,
aha_window_size=args.window_size,
aha_local_kind=args.local_kind,
aha_router_granularity=args.router_granularity,
aha_mode="dynamic",
aha_gate_target=0.70,
aha_reg_weight=-1.0,
torch_dtype=torch.bfloat16,
attn_implementation="sdpa",
)
# Qwen3 normally has bias-free projections. A gate bias gives a stable,
# token-independent full-attention hot start while zero biases on q/k/v/o
# keep the original attention computation bit-for-bit unchanged.
if not model.config.attention_bias:
for layer in model.model.layers:
attn = layer.self_attn
attn.q_proj = add_zero_bias(attn.q_proj)
attn.k_proj = add_zero_bias(attn.k_proj)
attn.v_proj = add_zero_bias(attn.v_proj)
attn.o_proj = add_zero_bias(attn.o_proj)
model.config.attention_bias = True
q_rows = model.config.num_attention_heads * model.config.head_dim
gate_logit = math.log(args.gate_init_full_prob / (1.0 - args.gate_init_full_prob))
with torch.no_grad():
for layer in model.model.layers:
q_proj = layer.self_attn.q_proj
q_proj.weight[q_rows:].zero_()
q_proj.bias[q_rows:].fill_(gate_logit)
# Keep a distinct LM head so the custom checkpoint has an explicit,
# self-contained state dict under safetensors.
model.config.tie_word_embeddings = False
if model.lm_head.weight.data_ptr() == model.model.embed_tokens.weight.data_ptr():
model.lm_head.weight = nn.Parameter(model.lm_head.weight.detach().clone())
model.config.aha_hotstart_source = str(Path(args.vanilla_path).resolve())
model.config.aha_gate_init_full_prob = args.gate_init_full_prob
tokenizer = AutoTokenizer.from_pretrained(args.vanilla_path, trust_remote_code=True)
model.save_pretrained(output, safe_serialization=True)
tokenizer.save_pretrained(output)
gate_rows = model.model.layers[0].self_attn.aha_router_outputs
gate_parameters = len(model.model.layers) * gate_rows * (
model.config.hidden_size + 1
)
(output / "aha_hotstart_manifest.json").write_text(
json.dumps(
{
"source_checkpoint": str(Path(args.vanilla_path).resolve()),
"router_granularity": args.router_granularity,
"native_gate_rows_per_layer": gate_rows,
"effective_gate_parameters": gate_parameters,
"gate_init_full_probability": args.gate_init_full_prob,
"gate_weight_init": "zeros",
"gate_bias_logit": gate_logit,
"local_attention": {
"kind": args.local_kind,
"sink_size": model.config.duo_sink_size,
"recent_size": model.config.duo_recent_size,
},
},
indent=2,
)
+ "\n"
)
print(
f"saved={output} init_full_prob={args.gate_init_full_prob:.4f} "
f"gate_logit={gate_logit:.6f} hard_sparsity=0.0 "
f"router_granularity={args.router_granularity} gate_params={gate_parameters:,}"
)
if __name__ == "__main__":
main()
|