初始化项目，由ModelHub XC社区提供模型
Model: jiamingshan/AHA-L2A-Qwen3-1.7B-repro Source: Original Platform
This commit is contained in:
ModelHub XC
16
recipe/Dockerfile
Normal file
16
recipe/Dockerfile
Normal file
@@ -0,0 +1,16 @@
|
||||
FROM pytorch/pytorch:2.9.1-cuda12.8-cudnn9-runtime
|
||||
|
||||
ENV DEBIAN_FRONTEND=noninteractive \
|
||||
PYTHONUNBUFFERED=1 \
|
||||
TOKENIZERS_PARALLELISM=false \
|
||||
HF_HOME=/workspace/.cache/huggingface \
|
||||
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
|
||||
|
||||
RUN apt-get update \
|
||||
&& apt-get install -y --no-install-recommends git ca-certificates \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY requirements.txt /tmp/requirements.txt
|
||||
RUN python -m pip install --no-cache-dir -r /tmp/requirements.txt
|
||||
|
||||
WORKDIR /workspace/recipe
|
||||
18
recipe/data/am_distilled_long_mix/build_summary.json
Normal file
18
recipe/data/am_distilled_long_mix/build_summary.json
Normal file
@@ -0,0 +1,18 @@
|
||||
{
|
||||
"source_dataset": "/workspace/AHA/AHA-OLMO2/experiments/datasets/am-distilled-32768",
|
||||
"output_dataset": "/workspace/AHA/AHA-Qwen3/experiments/datasets/am-distilled-32768-mix-long12k512-rand512",
|
||||
"seed": 124,
|
||||
"long_n": 512,
|
||||
"rand_n": 512,
|
||||
"min_long": 12288,
|
||||
"rows": 1024,
|
||||
"min_tokens": 369,
|
||||
"p50": 12580,
|
||||
"p90": 21189,
|
||||
"p95": 24547,
|
||||
"p99": 30902,
|
||||
"max_tokens": 32618,
|
||||
"mean_tokens": 10600.341796875,
|
||||
"frac_ge_8192": 0.5517578125,
|
||||
"frac_ge_12288": 0.5302734375
|
||||
}
|
||||
1
recipe/data/am_distilled_long_mix/dataset_dict.json
Normal file
1
recipe/data/am_distilled_long_mix/dataset_dict.json
Normal file
@@ -0,0 +1 @@
|
||||
{"splits": ["train"]}
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:0135bf6e96f7a4f452894727b9be1ce92ff27cc65330b94faeab11e993f5aa3c
|
||||
size 61801216
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:27f158ba22512269d33e8fb14db97a84e541ca9b701fba73a98e70c73687f45c
|
||||
size 87171000
|
||||
152
recipe/data/am_distilled_long_mix/train/dataset_info.json
Normal file
152
recipe/data/am_distilled_long_mix/train/dataset_info.json
Normal file
@@ -0,0 +1,152 @@
|
||||
{
|
||||
"builder_name": "generator",
|
||||
"citation": "",
|
||||
"config_name": "default",
|
||||
"dataset_name": "generator",
|
||||
"dataset_size": 45895360821,
|
||||
"description": "",
|
||||
"download_checksums": {},
|
||||
"download_size": 0,
|
||||
"features": {
|
||||
"messages": {
|
||||
"feature": {
|
||||
"content": {
|
||||
"dtype": "string",
|
||||
"_type": "Value"
|
||||
},
|
||||
"role": {
|
||||
"dtype": "string",
|
||||
"_type": "Value"
|
||||
}
|
||||
},
|
||||
"_type": "List"
|
||||
},
|
||||
"input_ids": {
|
||||
"feature": {
|
||||
"dtype": "int32",
|
||||
"_type": "Value"
|
||||
},
|
||||
"_type": "List"
|
||||
},
|
||||
"attention_mask": {
|
||||
"feature": {
|
||||
"dtype": "int8",
|
||||
"_type": "Value"
|
||||
},
|
||||
"_type": "List"
|
||||
},
|
||||
"num_tokens": {
|
||||
"dtype": "int64",
|
||||
"_type": "Value"
|
||||
}
|
||||
},
|
||||
"homepage": "",
|
||||
"license": "",
|
||||
"size_in_bytes": 45895360821,
|
||||
"splits": {
|
||||
"train": {
|
||||
"name": "train",
|
||||
"num_bytes": 45895360821,
|
||||
"num_examples": 1890613,
|
||||
"shard_lengths": [
|
||||
15000,
|
||||
14000,
|
||||
14000,
|
||||
14000,
|
||||
14000,
|
||||
14000,
|
||||
14000,
|
||||
15000,
|
||||
14000,
|
||||
14000,
|
||||
14000,
|
||||
14000,
|
||||
14000,
|
||||
14000,
|
||||
14000,
|
||||
14000,
|
||||
14000,
|
||||
14000,
|
||||
14000,
|
||||
14000,
|
||||
14000,
|
||||
14000,
|
||||
14000,
|
||||
25000,
|
||||
26000,
|
||||
27000,
|
||||
26000,
|
||||
26000,
|
||||
25000,
|
||||
26000,
|
||||
25000,
|
||||
26000,
|
||||
25000,
|
||||
26000,
|
||||
25000,
|
||||
26000,
|
||||
26000,
|
||||
25000,
|
||||
25000,
|
||||
25000,
|
||||
24000,
|
||||
25000,
|
||||
24000,
|
||||
24000,
|
||||
24000,
|
||||
24000,
|
||||
23000,
|
||||
19000,
|
||||
18000,
|
||||
18000,
|
||||
18000,
|
||||
18000,
|
||||
18000,
|
||||
18000,
|
||||
18000,
|
||||
19000,
|
||||
22000,
|
||||
23000,
|
||||
23000,
|
||||
23000,
|
||||
23000,
|
||||
23000,
|
||||
23000,
|
||||
23000,
|
||||
23000,
|
||||
23000,
|
||||
23000,
|
||||
22000,
|
||||
22000,
|
||||
23000,
|
||||
22000,
|
||||
22000,
|
||||
22000,
|
||||
23000,
|
||||
22000,
|
||||
22000,
|
||||
22000,
|
||||
22000,
|
||||
23000,
|
||||
23000,
|
||||
23000,
|
||||
23000,
|
||||
22000,
|
||||
22000,
|
||||
28000,
|
||||
30000,
|
||||
30000,
|
||||
29000,
|
||||
30000,
|
||||
20613
|
||||
],
|
||||
"dataset_name": "generator"
|
||||
}
|
||||
},
|
||||
"version": {
|
||||
"version_str": "0.0.0",
|
||||
"major": 0,
|
||||
"minor": 0,
|
||||
"patch": 0
|
||||
}
|
||||
}
|
||||
18
recipe/data/am_distilled_long_mix/train/state.json
Normal file
18
recipe/data/am_distilled_long_mix/train/state.json
Normal file
@@ -0,0 +1,18 @@
|
||||
{
|
||||
"_data_files": [
|
||||
{
|
||||
"filename": "data-00000-of-00001.arrow"
|
||||
}
|
||||
],
|
||||
"_fingerprint": "8470c8174c1f6ae9",
|
||||
"_format_columns": [
|
||||
"attention_mask",
|
||||
"input_ids",
|
||||
"messages",
|
||||
"num_tokens"
|
||||
],
|
||||
"_format_kwargs": {},
|
||||
"_format_type": null,
|
||||
"_output_all_columns": false,
|
||||
"_split": null
|
||||
}
|
||||
BIN
recipe/data/eval_inputs/helmet_icl_8k_n50_per_config.jsonl
(Stored with Git LFS)
Normal file
BIN
recipe/data/eval_inputs/helmet_icl_8k_n50_per_config.jsonl
(Stored with Git LFS)
Normal file
Binary file not shown.
34
recipe/data/eval_inputs/input_manifest.json
Normal file
34
recipe/data/eval_inputs/input_manifest.json
Normal file
@@ -0,0 +1,34 @@
|
||||
{
|
||||
"tokenizer_path": "/data/sjm/AHA/models/Qwen3-1.7B",
|
||||
"thinking": false,
|
||||
"truncation": false,
|
||||
"target_context_k": 8,
|
||||
"helmet": {
|
||||
"path": "/data/sjm/AHA/AHA-Qwen3/experiments/qwen3_1p7b_fourbench_8k_broad_20260713/inputs/helmet_icl_8k_n50_per_config.jsonl",
|
||||
"source": "/data/sjm/AHA/AHA-Qwen3/experiments/qwen3_1p7b_fivebench_n100_20260712/helmet8k_rows_n100.jsonl",
|
||||
"configs": [
|
||||
"trec_coarse",
|
||||
"trec_fine",
|
||||
"banking77",
|
||||
"clinic150",
|
||||
"nlu"
|
||||
],
|
||||
"samples_per_config": 50,
|
||||
"total_samples": 250,
|
||||
"min_prompt_tokens": 6344,
|
||||
"max_prompt_tokens": 7413
|
||||
},
|
||||
"mrcr": {
|
||||
"path": "/data/sjm/AHA/AHA-Qwen3/experiments/qwen3_1p7b_fourbench_8k_broad_20260713/inputs/mrcr_8k_2_4_8needle_n10_per_config.jsonl",
|
||||
"source": "/data/sjm/AHA/AHA-Qwen3/experiments/qwen3_1p7b_fourbench_broad_20260713/inputs/mrcr_8k_16k_2_4_8needle_n10_per_config.jsonl",
|
||||
"configs": [
|
||||
"8k_2needle",
|
||||
"8k_4needle",
|
||||
"8k_8needle"
|
||||
],
|
||||
"samples_per_config": 10,
|
||||
"total_samples": 30,
|
||||
"min_prompt_tokens": 4312,
|
||||
"max_prompt_tokens": 8028
|
||||
}
|
||||
}
|
||||
BIN
recipe/data/eval_inputs/mrcr_8k_2_4_8needle_n10_per_config.jsonl
(Stored with Git LFS)
Normal file
BIN
recipe/data/eval_inputs/mrcr_8k_2_4_8needle_n10_per_config.jsonl
(Stored with Git LFS)
Normal file
Binary file not shown.
818
recipe/duo_train.py
Normal file
818
recipe/duo_train.py
Normal file
@@ -0,0 +1,818 @@
|
||||
"""Faithful DuoAttention-style training for Qwen3 in the AHA framework.
|
||||
|
||||
Reproduces the paper's training objective closely:
|
||||
L = distill + reg_weight * L1(alpha) + ce_weight * CE(student_logits, labels)
|
||||
distill = MSE(h_full, h_mix) on label positions
|
||||
h_full = forward with alpha = 1 everywhere (pure global attention)
|
||||
h_mix = forward with the currently learned alpha (blend global + streaming)
|
||||
|
||||
The CE anchor term defaults to 0 (paper-grade DuoAttention). It is required
|
||||
when --unfreeze_attn_proj is set, because once the backbone is trainable the
|
||||
distill objective degenerates: the teacher (alpha=1 forward) is rebuilt from
|
||||
the same drifting backbone, so distill becomes self-distillation against a
|
||||
moving target and admits collapse solutions where h_full ≡ h_mix but both
|
||||
generate garbled tokens. CE pins the backbone to "predicting labels under
|
||||
mix attention" and prevents that drift; see docs §9.4.4.
|
||||
|
||||
Data: synthetic multi-passkey retrieval on PaulGrahamEssays haystack,
|
||||
ported from duo-attention/duo_attn/data.py::MultiplePasskeyRetrievalDataset.
|
||||
|
||||
Trainable params:
|
||||
- default : 224 alpha scalars (28 layers * 8 kv_heads on Qwen3-0.6B).
|
||||
- --unfreeze_attn_proj : alphas + q/k/v/o_proj weights of every layer
|
||||
(Setting B; sink-ablation experiment).
|
||||
|
||||
Use `attn_implementation="eager"` or "sdpa" — no flash-attn dependency.
|
||||
|
||||
Usage (paper-grade, frozen backbone):
|
||||
python duo_train.py \\
|
||||
--model_path /workspace/AHA/models/Qwen3-0.6B \\
|
||||
--haystack_dir /workspace/AHA/third_party/duo-attention/eval/needle/PaulGrahamEssays \\
|
||||
--output_dir ckpts/duo_paper_s64_r256 \\
|
||||
--max_length 8192 --context_length_min 2000 --context_length_max 8000 \\
|
||||
--num_steps 800 --lr 0.02 --reg_weight 0.05 \\
|
||||
--sink_size 64 --recent_size 256
|
||||
|
||||
Usage (Setting B sink ablation, unfrozen backbone + CE anchor):
|
||||
python duo_train.py \\
|
||||
--model_path /workspace/AHA/models/Qwen3-0.6B \\
|
||||
--haystack_dir /workspace/AHA/third_party/duo-attention/eval/needle/PaulGrahamEssays \\
|
||||
--output_dir ckpts/duo_sinkabl_Bv3_ce \\
|
||||
--max_length 8192 --context_length_min 2000 --context_length_max 8000 \\
|
||||
--num_steps 400 --lr 0.02 --reg_weight 0.1 \\
|
||||
--sink_size 0 --recent_size 256 \\
|
||||
--unfreeze_attn_proj --backbone_lr 1e-5 --ce_weight 1.0
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import random
|
||||
import sys
|
||||
from typing import List
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from torch.nn.parallel import DistributedDataParallel as DDP
|
||||
from torch.utils.data import Dataset, DataLoader
|
||||
from torch.utils.data.distributed import DistributedSampler
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
HERE = os.path.dirname(os.path.abspath(__file__))
|
||||
if HERE not in sys.path:
|
||||
sys.path.insert(0, HERE)
|
||||
|
||||
from modeling_aha_qwen3 import AHAQwen3ForCausalLM # noqa: E402
|
||||
|
||||
|
||||
LONG_BENCH_PROMPT_TEMPLATES = {
|
||||
"qasper": (
|
||||
"You are given a scientific article and a question. "
|
||||
"Answer the question as concisely as you can, using a single phrase or sentence if possible. "
|
||||
"If the question cannot be answered based on the information in the article, write \"unanswerable\". "
|
||||
"If the question is a yes/no question, answer \"yes\", \"no\", or \"unanswerable\". "
|
||||
"Do not provide any explanation.\n\n"
|
||||
"Article: {context}\n\n"
|
||||
"Answer the question based on the above article as concisely as you can, using a single phrase or sentence if possible. "
|
||||
"If the question cannot be answered based on the information in the article, write \"unanswerable\". "
|
||||
"If the question is a yes/no question, answer \"yes\", \"no\", or \"unanswerable\". "
|
||||
"Do not provide any explanation.\n\nQuestion: {question}\nAnswer:"
|
||||
),
|
||||
"multifieldqa_en": (
|
||||
"Read the following text and answer briefly.\n\n"
|
||||
"{context}\n\n"
|
||||
"Now, answer the following question based on the above text, only give me the answer and do not output any other words.\n\n"
|
||||
"Question: {question}\nAnswer:"
|
||||
),
|
||||
"2wikimqa": (
|
||||
"Answer the question based on the given passages. Only give me the answer and do not output any other words.\n\n"
|
||||
"The following are given passages.\n{context}\n\n"
|
||||
"Answer the question based on the given passages. Only give me the answer and do not output any other words.\n\n"
|
||||
"Question: {question}\nAnswer:"
|
||||
),
|
||||
"passage_retrieval_en": (
|
||||
"Here are 30 paragraphs from Wikipedia, along with an abstract. "
|
||||
"Please determine which paragraph the abstract is from.\n\n"
|
||||
"{context}\n\n"
|
||||
"The following is an abstract.\n\n"
|
||||
"{question}\n\n"
|
||||
"Please enter the number of the paragraph that the abstract is from. "
|
||||
"The answer format must be like \"Paragraph 1\", \"Paragraph 2\", etc.\n\n"
|
||||
"The answer is: "
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Dataset (direct port of duo_attn/data.py::MultiplePasskeyRetrievalDataset)
|
||||
# -----------------------------------------------------------------------------
|
||||
PASSKEY_ALPHABET = [
|
||||
"alpha", "bravo", "charlie", "delta", "echo", "foxtrot", "golf", "hotel",
|
||||
"india", "juliett", "kilo", "lima", "mike", "november", "oscar", "papa",
|
||||
"quebec", "romeo", "sierra", "tango", "uniform", "victor", "whiskey",
|
||||
"xray", "yankee", "zulu",
|
||||
]
|
||||
ORDINAL_NUMBERS = [
|
||||
"first", "second", "third", "fourth", "fifth", "sixth", "seventh",
|
||||
"eighth", "ninth", "tenth", "eleventh", "twelfth", "thirteenth",
|
||||
"fourteenth", "fifteenth", "sixteenth", "seventeenth", "eighteenth",
|
||||
"nineteenth", "twentieth",
|
||||
]
|
||||
|
||||
|
||||
def _load_haystack_text(haystack_dir: str) -> str:
|
||||
parts = []
|
||||
for fname in sorted(os.listdir(haystack_dir)):
|
||||
if not fname.endswith(".txt"):
|
||||
continue
|
||||
with open(os.path.join(haystack_dir, fname), "r", encoding="utf-8", errors="ignore") as f:
|
||||
parts.append(f.read())
|
||||
return "\n\n".join(parts)
|
||||
|
||||
|
||||
class MultiPasskeyDataset(Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
tokenizer,
|
||||
haystack_text: str,
|
||||
context_length_min: int,
|
||||
context_length_max: int,
|
||||
context_lengths_num_intervals: int,
|
||||
depth_ratio_num_intervals: int,
|
||||
min_depth_ratio: float,
|
||||
max_depth_ratio: float,
|
||||
num_passkeys: int,
|
||||
passkey_length: int,
|
||||
pad_multiple: int = 16,
|
||||
buffer_size: int = 300,
|
||||
needle: str = "Remember this sequence of words, it's the {ordinal_number} passkey to the vault: ",
|
||||
retrieval_question: str = "Based on the content of the book, what is the {ordinal_number} passkey to the vault?\nPasskey: ",
|
||||
prompt1: str = "<|im_start|> This is a very long story book: <book> ",
|
||||
prompt2: str = " </book>.\n\n",
|
||||
seperator: str = "\n\n",
|
||||
):
|
||||
self.tokenizer = tokenizer
|
||||
self.num_passkeys = num_passkeys
|
||||
self.passkey_length = passkey_length
|
||||
self.pad_multiple = pad_multiple
|
||||
|
||||
self.context_length_intervals = torch.linspace(
|
||||
context_length_min, context_length_max,
|
||||
context_lengths_num_intervals, dtype=torch.int,
|
||||
).tolist()
|
||||
self.depth_ratio_intervals = torch.linspace(
|
||||
min_depth_ratio, max_depth_ratio, depth_ratio_num_intervals,
|
||||
).tolist()
|
||||
|
||||
self.needle_tokens_list = [
|
||||
tokenizer.encode(
|
||||
needle.format(ordinal_number=ord_), add_special_tokens=False
|
||||
) for ord_ in ORDINAL_NUMBERS[:num_passkeys]
|
||||
]
|
||||
self.retrieval_question_tokens_list = [
|
||||
tokenizer.encode(
|
||||
retrieval_question.format(ordinal_number=ord_), add_special_tokens=False
|
||||
) for ord_ in ORDINAL_NUMBERS[:num_passkeys]
|
||||
]
|
||||
|
||||
self.haystack_tokens = tokenizer.encode(haystack_text, add_special_tokens=False)
|
||||
if len(self.haystack_tokens) < context_length_max:
|
||||
# tile the corpus until long enough
|
||||
repeats = context_length_max // max(1, len(self.haystack_tokens)) + 2
|
||||
self.haystack_tokens = self.haystack_tokens * repeats
|
||||
self.haystack_tokens = self.haystack_tokens[: context_length_max + 200]
|
||||
|
||||
self.seperator_tokens = tokenizer.encode(seperator, add_special_tokens=False)
|
||||
self.prompt1_tokens = tokenizer.encode(prompt1, add_special_tokens=True)
|
||||
self.prompt2_tokens = tokenizer.encode(prompt2, add_special_tokens=False)
|
||||
self.buffer_size = buffer_size
|
||||
|
||||
def __len__(self):
|
||||
return 10 ** 9 # effectively infinite; trainer slices by num_steps
|
||||
|
||||
def _gen_passkey(self):
|
||||
seq = torch.randint(0, len(PASSKEY_ALPHABET), (self.passkey_length,))
|
||||
return " ".join(PASSKEY_ALPHABET[i] for i in seq)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
rng = random.Random(idx)
|
||||
context_length = int(rng.choice(self.context_length_intervals))
|
||||
depths = sorted(rng.sample(self.depth_ratio_intervals, self.num_passkeys))
|
||||
passkey_tokens_list = [
|
||||
self.tokenizer.encode(self._gen_passkey(), add_special_tokens=False)
|
||||
for _ in range(self.num_passkeys)
|
||||
]
|
||||
|
||||
haystack = self.haystack_tokens[:context_length]
|
||||
context = []
|
||||
last = 0
|
||||
for i, (d, pk) in enumerate(zip(depths, passkey_tokens_list)):
|
||||
ip = int(len(haystack) * d)
|
||||
needle = self.needle_tokens_list[i] + pk
|
||||
context += haystack[last:ip] + self.seperator_tokens + needle + self.seperator_tokens
|
||||
last = ip
|
||||
context += haystack[last:]
|
||||
|
||||
qa = []
|
||||
for i, pk in enumerate(passkey_tokens_list):
|
||||
qa += self.retrieval_question_tokens_list[i] + pk + self.seperator_tokens
|
||||
|
||||
ctx = self.prompt1_tokens + context + self.prompt2_tokens
|
||||
ids = ctx + qa
|
||||
# pad to multiple of 16
|
||||
pad = (-len(ids)) % self.pad_multiple
|
||||
if pad:
|
||||
ids = ids + self.haystack_tokens[-pad:]
|
||||
labels = [-100] * (len(ids) - len(qa)) + qa
|
||||
# clip pad-extension off labels
|
||||
labels = labels[: len(ids)]
|
||||
assert len(ids) == len(labels)
|
||||
return {"input_ids": torch.tensor(ids), "labels": torch.tensor(labels)}
|
||||
|
||||
|
||||
def collate(batch):
|
||||
return {
|
||||
"input_ids": torch.stack([b["input_ids"] for b in batch]),
|
||||
"labels": torch.stack([b["labels"] for b in batch]),
|
||||
}
|
||||
|
||||
|
||||
def _find_subsequence(haystack: List[int], needle: List[int]) -> int:
|
||||
if not needle or len(needle) > len(haystack):
|
||||
return -1
|
||||
last = len(haystack) - len(needle)
|
||||
for i in range(last + 1):
|
||||
if haystack[i:i + len(needle)] == needle:
|
||||
return i
|
||||
return -1
|
||||
|
||||
|
||||
class AmDistilledDataset(Dataset):
|
||||
"""Wraps a pre-tokenized HF dataset (e.g. /workspace/...am-distilled-8192).
|
||||
|
||||
Expects each sample to have an `input_ids` field already produced by
|
||||
`tokenize-am_distill.py`. By default labels = input_ids (full-token
|
||||
distill). With label_mode="answer_only", labels before the final answer
|
||||
span are masked to -100, matching DuoAttention's answer-only distill
|
||||
pressure more closely.
|
||||
|
||||
Used when `--data_source am_distilled` is set, as a drop-in replacement
|
||||
for the passkey synthetic dataset. Distill loss is then computed over
|
||||
real reasoning data instead of haystack passkey retrieval, which avoids
|
||||
the in-distribution overfitting documented in docs §9.4.7-8.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
ds_path: str,
|
||||
split: str,
|
||||
max_length: int,
|
||||
seed: int = 42,
|
||||
tokenizer=None,
|
||||
label_mode: str = "full",
|
||||
):
|
||||
import datasets as hf_datasets
|
||||
loaded = hf_datasets.load_from_disk(ds_path)
|
||||
if hasattr(loaded, "keys"):
|
||||
self.ds = loaded[split]
|
||||
else:
|
||||
self.ds = loaded
|
||||
self.max_length = max_length
|
||||
self.label_mode = label_mode
|
||||
self.answer_marker_ids = []
|
||||
self.think_end_ids = []
|
||||
self.assistant_marker_ids = []
|
||||
if tokenizer is not None:
|
||||
self.answer_marker_ids = tokenizer.encode("<answer>", add_special_tokens=False)
|
||||
self.think_end_ids = tokenizer.encode("</think>", add_special_tokens=False)
|
||||
self.assistant_marker_ids = tokenizer.encode("<|im_start|>assistant", add_special_tokens=False)
|
||||
self._order = list(range(len(self.ds)))
|
||||
random.Random(seed).shuffle(self._order)
|
||||
|
||||
def __len__(self):
|
||||
return len(self._order)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
real_idx = self._order[idx % len(self._order)]
|
||||
sample = self.ds[real_idx]
|
||||
ids = list(sample["input_ids"])[: self.max_length]
|
||||
labels = list(ids)
|
||||
if self.label_mode == "answer_only":
|
||||
labels = [-100] * len(ids)
|
||||
start = _find_subsequence(ids, self.answer_marker_ids)
|
||||
if start >= 0:
|
||||
start = start + len(self.answer_marker_ids)
|
||||
else:
|
||||
start = _find_subsequence(ids, self.think_end_ids)
|
||||
if start >= 0:
|
||||
start = start + len(self.think_end_ids)
|
||||
if start < 0:
|
||||
# Fallback for traces without explicit <answer> inside the
|
||||
# truncation window: supervise only assistant-side tokens.
|
||||
start = _find_subsequence(ids, self.assistant_marker_ids)
|
||||
if start >= 0:
|
||||
start = start + len(self.assistant_marker_ids)
|
||||
if 0 <= start < len(ids):
|
||||
labels[start:] = ids[start:]
|
||||
elif self.label_mode != "full":
|
||||
raise ValueError(f"unknown AM label_mode: {self.label_mode}")
|
||||
return {
|
||||
"input_ids": torch.tensor(ids, dtype=torch.long),
|
||||
"labels": torch.tensor(labels, dtype=torch.long),
|
||||
}
|
||||
|
||||
|
||||
class LongBenchLiteAnswerDataset(Dataset):
|
||||
"""Small target-distribution calibration set for alpha-only diagnostics.
|
||||
|
||||
Builds LongBench-lite prompts with the same templates as eval, appends one
|
||||
gold answer, and masks labels to answer tokens only. This is intentionally
|
||||
a diagnostic data source: it answers whether a high-sparsity static Duo mask
|
||||
exists on the target distribution.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tokenizer,
|
||||
tasks: List[str],
|
||||
samples_per_task: int,
|
||||
max_length: int,
|
||||
seed: int = 42,
|
||||
cache_dir: str = "/workspace/AHA/AHA-Qwen3/data/longbench_cache",
|
||||
pad_multiple: int = 16,
|
||||
):
|
||||
from datasets import load_dataset
|
||||
|
||||
self.tokenizer = tokenizer
|
||||
self.max_length = max_length
|
||||
self.pad_multiple = pad_multiple
|
||||
self.rows = []
|
||||
for task in tasks:
|
||||
template = LONG_BENCH_PROMPT_TEMPLATES[task]
|
||||
ds = load_dataset("Xnhyacinth/LongBench", task, split="test", cache_dir=cache_dir)
|
||||
n = min(samples_per_task, len(ds))
|
||||
for idx in range(n):
|
||||
sample = ds[idx]
|
||||
user_content = template.format(context=sample["context"], question=sample["question"])
|
||||
prompt = tokenizer.apply_chat_template(
|
||||
[{"role": "user", "content": user_content}],
|
||||
tokenize=False,
|
||||
add_generation_prompt=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
answers = sample["answers"] if isinstance(sample["answers"], list) else [sample["answers"]]
|
||||
answer = str(answers[0])
|
||||
prompt_ids = tokenizer(prompt, truncation=False, add_special_tokens=False)["input_ids"]
|
||||
answer_ids = tokenizer(answer, truncation=False, add_special_tokens=False)["input_ids"]
|
||||
budget = max_length - len(answer_ids) - 1
|
||||
if len(prompt_ids) > budget:
|
||||
half = max(1, budget // 2)
|
||||
prompt_ids = prompt_ids[:half] + prompt_ids[-(budget - half):]
|
||||
ids = prompt_ids + answer_ids
|
||||
pad = (-len(ids)) % pad_multiple
|
||||
if pad:
|
||||
ids = ids + [tokenizer.pad_token_id] * pad
|
||||
labels = [-100] * len(prompt_ids) + answer_ids + [-100] * pad
|
||||
self.rows.append({"input_ids": ids, "labels": labels, "task": task, "idx": idx})
|
||||
|
||||
random.Random(seed).shuffle(self.rows)
|
||||
|
||||
def __len__(self):
|
||||
return 10 ** 9
|
||||
|
||||
def __getitem__(self, idx):
|
||||
row = self.rows[idx % len(self.rows)]
|
||||
return {
|
||||
"input_ids": torch.tensor(row["input_ids"], dtype=torch.long),
|
||||
"labels": torch.tensor(row["labels"], dtype=torch.long),
|
||||
}
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Training
|
||||
# -----------------------------------------------------------------------------
|
||||
@torch.no_grad()
|
||||
def _set_alpha_full(model, value: float = 1.0):
|
||||
"""Temporarily overwrite `full_attention_heads` to a constant.
|
||||
|
||||
Used to compute the 'teacher' forward pass (full attention everywhere).
|
||||
Call `_restore_alpha` with the saved tensors afterwards.
|
||||
"""
|
||||
saved = []
|
||||
for layer in model.model.layers:
|
||||
p = layer.self_attn.full_attention_heads
|
||||
saved.append(p.data.clone())
|
||||
p.data.fill_(value)
|
||||
return saved
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def _restore_alpha(model, saved):
|
||||
for layer, s in zip(model.model.layers, saved):
|
||||
layer.self_attn.full_attention_heads.data.copy_(s)
|
||||
|
||||
|
||||
def log_alpha_stats(model) -> dict:
|
||||
with torch.no_grad():
|
||||
alphas = torch.stack([
|
||||
layer.self_attn.full_attention_heads.detach().float().clamp(0, 1)
|
||||
for layer in model.model.layers
|
||||
], dim=0)
|
||||
m = alphas.mean().item()
|
||||
return {
|
||||
"alpha_mean": m,
|
||||
"alpha_std": alphas.std().item(),
|
||||
"alpha_gt05": (alphas > 0.5).float().mean().item(),
|
||||
"alpha_min": alphas.min().item(),
|
||||
"alpha_max": alphas.max().item(),
|
||||
}
|
||||
|
||||
|
||||
def save_alpha_matrix(model, path: str):
|
||||
with torch.no_grad():
|
||||
alphas = torch.stack([
|
||||
layer.self_attn.full_attention_heads.detach().float().clamp(0, 1).cpu()
|
||||
for layer in model.model.layers
|
||||
], dim=0).numpy()
|
||||
np.savetxt(path, alphas, delimiter="\t")
|
||||
|
||||
|
||||
def _init_distributed():
|
||||
world_size = int(os.environ.get("WORLD_SIZE", "1"))
|
||||
if world_size <= 1:
|
||||
return False, 0, 0, 1, torch.device("cuda")
|
||||
local_rank = int(os.environ.get("LOCAL_RANK", "0"))
|
||||
rank = int(os.environ.get("RANK", "0"))
|
||||
torch.cuda.set_device(local_rank)
|
||||
dist.init_process_group(backend="nccl")
|
||||
return True, local_rank, rank, world_size, torch.device("cuda", local_rank)
|
||||
|
||||
|
||||
def main():
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--model_path", required=True)
|
||||
p.add_argument("--aha_checkpoint_path", default="",
|
||||
help="Optional AHA/Duo checkpoint to continue training from. "
|
||||
"When set, model weights and alpha scalars are loaded from "
|
||||
"this checkpoint instead of converting --model_path from base Qwen3.")
|
||||
p.add_argument("--haystack_dir", default="",
|
||||
help="Path to PaulGraham essays for the synthetic passkey dataset. "
|
||||
"Required when --data_source=passkey, ignored otherwise.")
|
||||
p.add_argument("--output_dir", required=True)
|
||||
p.add_argument("--max_length", type=int, default=8192)
|
||||
p.add_argument("--context_length_min", type=int, default=2000)
|
||||
p.add_argument("--context_length_max", type=int, default=8000)
|
||||
p.add_argument("--context_lengths_num_intervals", type=int, default=20)
|
||||
p.add_argument("--depth_ratio_num_intervals", type=int, default=1000)
|
||||
p.add_argument("--min_depth_ratio", type=float, default=0.05)
|
||||
p.add_argument("--max_depth_ratio", type=float, default=0.95)
|
||||
p.add_argument("--num_passkeys", type=int, default=10)
|
||||
p.add_argument("--passkey_length", type=int, default=32)
|
||||
p.add_argument("--num_steps", type=int, default=800)
|
||||
p.add_argument("--warmup_ratio", type=float, default=0.2)
|
||||
p.add_argument("--lr", type=float, default=0.02)
|
||||
p.add_argument("--reg_weight", type=float, default=0.05)
|
||||
p.add_argument("--sink_size", type=int, default=64)
|
||||
p.add_argument("--recent_size", type=int, default=256)
|
||||
p.add_argument("--batch_size", type=int, default=1)
|
||||
p.add_argument("--grad_accum", type=int, default=1)
|
||||
p.add_argument("--save_steps", type=int, default=200)
|
||||
p.add_argument("--log_steps", type=int, default=10)
|
||||
p.add_argument("--seed", type=int, default=42)
|
||||
p.add_argument("--dtype", default="bfloat16")
|
||||
p.add_argument("--attn_impl", default="sdpa", choices=["sdpa", "eager"])
|
||||
# Optional: unfreeze attention projection weights (q/k/v/o_proj) together
|
||||
# with the alpha scalars — this reproduces the senior-student experiment
|
||||
# where "retraining" the model removes the need for the attention sink.
|
||||
p.add_argument("--unfreeze_attn_proj", action="store_true",
|
||||
help="Also train q/k/v/o_proj weights alongside alpha scalars.")
|
||||
p.add_argument("--backbone_lr", type=float, default=1e-5,
|
||||
help="Learning rate for unfrozen backbone params (alpha keeps --lr).")
|
||||
# CE anchor: required to prevent self-distill collapse when backbone is unfrozen.
|
||||
# When backbone is frozen (paper-grade DuoAttention), distill alone is well-defined
|
||||
# because the teacher (alpha=1 forward) is a fixed pretrained reference; CE is
|
||||
# redundant. When backbone is trainable, the teacher itself drifts together with
|
||||
# the student, so distill becomes self-distillation against a moving target and
|
||||
# admits degenerate solutions (h_full ≡ h_mix but both wrong → garbled output).
|
||||
# CE on the student forward pins backbone to the "predicting labels correctly"
|
||||
# manifold, blocking that failure mode.
|
||||
p.add_argument("--ce_weight", type=float, default=0.0,
|
||||
help="Weight for cross-entropy anchor loss on labels (student/mix forward). "
|
||||
"0 disables (paper-grade DuoAttention). Recommended >0 when "
|
||||
"--unfreeze_attn_proj is set, to prevent backbone drift.")
|
||||
# Data source: passkey (DuoAttention legacy synthetic) or am_distilled
|
||||
# (real reasoning SFT data, see docs §9.4.7-8 for why we may want this).
|
||||
p.add_argument("--data_source", default="passkey",
|
||||
choices=["passkey", "am_distilled", "longbench_lite"],
|
||||
help="Training data: 'passkey' replicates DuoAttention's "
|
||||
"synthetic haystack retrieval; 'am_distilled' uses a "
|
||||
"pre-tokenized multi-task SFT dataset (e.g. AM-Thinking "
|
||||
"or AM-Qwen3-Distilled). 'longbench_lite' is a diagnostic "
|
||||
"target-distribution calibration source.")
|
||||
p.add_argument("--am_dataset_path", default="/workspace/Direct-Multitoken-Decoding/am-distilled-8192",
|
||||
help="Path to a `datasets.load_from_disk`-compatible dataset.")
|
||||
p.add_argument("--am_dataset_split", default="train")
|
||||
p.add_argument("--am_label_mode", default="full", choices=["full", "answer_only"],
|
||||
help="Label mask for --data_source=am_distilled. 'full' keeps the legacy "
|
||||
"all-token hidden-state distill; 'answer_only' masks tokens before "
|
||||
"the final <answer> span, closer to DuoAttention's QA-only objective.")
|
||||
p.add_argument("--longbench_tasks", nargs="+",
|
||||
default=["passage_retrieval_en", "multifieldqa_en", "qasper", "2wikimqa"])
|
||||
p.add_argument("--longbench_samples_per_task", type=int, default=30)
|
||||
p.add_argument("--longbench_cache_dir", default="/workspace/AHA/AHA-Qwen3/data/longbench_cache")
|
||||
args = p.parse_args()
|
||||
|
||||
distributed, local_rank, rank, world_size, device = _init_distributed()
|
||||
is_main = rank == 0
|
||||
|
||||
def log(*log_args, **log_kwargs):
|
||||
if is_main:
|
||||
print(*log_args, **log_kwargs)
|
||||
|
||||
torch.manual_seed(args.seed + rank)
|
||||
random.seed(args.seed + rank)
|
||||
np.random.seed(args.seed + rank)
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
if is_main:
|
||||
with open(os.path.join(args.output_dir, "duo_train_args.json"), "w") as f:
|
||||
saved_args = vars(args).copy()
|
||||
saved_args.update({"distributed": distributed, "world_size": world_size})
|
||||
json.dump(saved_args, f, indent=2)
|
||||
|
||||
log(f"[duo-train] model={args.model_path} ctx=[{args.context_length_min},{args.context_length_max}]")
|
||||
log(f"[duo-train] sink={args.sink_size} recent={args.recent_size} passkeys={args.num_passkeys}")
|
||||
log(f"[duo-train] lr={args.lr} reg_weight={args.reg_weight} num_steps={args.num_steps}")
|
||||
if distributed:
|
||||
log(f"[duo-train] distributed=torchrun world_size={world_size}")
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
|
||||
if tokenizer.pad_token_id is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
|
||||
if args.data_source == "passkey":
|
||||
if not args.haystack_dir:
|
||||
raise ValueError("--haystack_dir is required when --data_source=passkey")
|
||||
haystack_text = _load_haystack_text(args.haystack_dir)
|
||||
log(f"[duo-train] haystack char length: {len(haystack_text):,}")
|
||||
else:
|
||||
haystack_text = ""
|
||||
|
||||
# Build model in DUO mode. Default is initialising from base Qwen3 weights;
|
||||
# --aha_checkpoint_path is used for static-alpha continuation controls.
|
||||
dtype = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}[args.dtype]
|
||||
if args.aha_checkpoint_path:
|
||||
log(f"[duo-train] continuing from aha_checkpoint_path={args.aha_checkpoint_path}")
|
||||
model = AHAQwen3ForCausalLM.from_pretrained_aha(
|
||||
args.aha_checkpoint_path,
|
||||
torch_dtype=dtype,
|
||||
attn_implementation=args.attn_impl,
|
||||
).to(device)
|
||||
if getattr(model.config, "aha_mode", "") != "duo":
|
||||
raise ValueError("--aha_checkpoint_path for duo_train.py must have aha_mode='duo'")
|
||||
model.config.duo_sink_size = args.sink_size
|
||||
model.config.duo_recent_size = args.recent_size
|
||||
model.config.aha_distill_weight = 0.0
|
||||
model.config.aha_ce_weight = 0.0
|
||||
model.config.aha_lambda = 0.0
|
||||
model.config.aha_gate_target = 0.0
|
||||
model.config.aha_reg_weight = -1.0
|
||||
else:
|
||||
model = AHAQwen3ForCausalLM.from_pretrained_qwen3(
|
||||
args.model_path,
|
||||
aha_mode="duo",
|
||||
duo_sink_size=args.sink_size,
|
||||
duo_recent_size=args.recent_size,
|
||||
duo_alpha_init=1.0,
|
||||
aha_distill_weight=0.0, # we compute distill externally
|
||||
aha_ce_weight=0.0, # no CE in DuoAttention objective
|
||||
aha_lambda=0.0, # we compute L1 externally
|
||||
aha_gate_target=0.0,
|
||||
torch_dtype=dtype,
|
||||
attn_implementation=args.attn_impl,
|
||||
).to(device)
|
||||
|
||||
# Freeze everything except full_attention_heads
|
||||
for param in model.parameters():
|
||||
param.requires_grad = False
|
||||
alpha_params, backbone_params = [], []
|
||||
for layer in model.model.layers:
|
||||
layer.self_attn.full_attention_heads.requires_grad = True
|
||||
alpha_params.append(layer.self_attn.full_attention_heads)
|
||||
if args.unfreeze_attn_proj:
|
||||
# Sink-ablation setting B: also retrain attention projections so the
|
||||
# model can learn attention patterns that do not rely on sink tokens.
|
||||
for proj in ("q_proj", "k_proj", "v_proj", "o_proj"):
|
||||
mod = getattr(layer.self_attn, proj, None)
|
||||
if mod is None:
|
||||
continue
|
||||
for pname, param in mod.named_parameters():
|
||||
param.requires_grad = True
|
||||
backbone_params.append(param)
|
||||
|
||||
# Gradient checkpointing requires *some* input to require grad. In the
|
||||
# pure-alpha setting all backbone weights are frozen so we must manually
|
||||
# enable input grads; when `--unfreeze_attn_proj` is on the projections
|
||||
# themselves already require grad so this is still harmless but optional.
|
||||
model.enable_input_require_grads()
|
||||
model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
|
||||
core_model = model.model
|
||||
if distributed:
|
||||
core_model = DDP(core_model, device_ids=[local_rank], output_device=local_rank)
|
||||
n_alpha = sum(p.numel() for p in alpha_params)
|
||||
n_backbone = sum(p.numel() for p in backbone_params)
|
||||
log(f"[duo-train] trainable alpha scalars: {n_alpha}")
|
||||
if args.unfreeze_attn_proj:
|
||||
log(f"[duo-train] trainable backbone params (q/k/v/o_proj): {n_backbone:,}")
|
||||
else:
|
||||
log(f"[duo-train] backbone: frozen (paper-grade DuoAttention protocol)")
|
||||
|
||||
if args.data_source == "passkey":
|
||||
log(f"[duo-train] data_source=passkey, haystack_dir={args.haystack_dir}")
|
||||
dataset = MultiPasskeyDataset(
|
||||
tokenizer=tokenizer,
|
||||
haystack_text=haystack_text,
|
||||
context_length_min=args.context_length_min,
|
||||
context_length_max=args.context_length_max,
|
||||
context_lengths_num_intervals=args.context_lengths_num_intervals,
|
||||
depth_ratio_num_intervals=args.depth_ratio_num_intervals,
|
||||
min_depth_ratio=args.min_depth_ratio,
|
||||
max_depth_ratio=args.max_depth_ratio,
|
||||
num_passkeys=args.num_passkeys,
|
||||
passkey_length=args.passkey_length,
|
||||
)
|
||||
elif args.data_source == "am_distilled":
|
||||
log(f"[duo-train] data_source=am_distilled, path={args.am_dataset_path} split={args.am_dataset_split}")
|
||||
dataset = AmDistilledDataset(
|
||||
ds_path=args.am_dataset_path,
|
||||
split=args.am_dataset_split,
|
||||
max_length=args.max_length,
|
||||
seed=args.seed,
|
||||
tokenizer=tokenizer,
|
||||
label_mode=args.am_label_mode,
|
||||
)
|
||||
log(f"[duo-train] am_distilled dataset n={len(dataset):,}, max_length={args.max_length}, "
|
||||
f"label_mode={args.am_label_mode}")
|
||||
elif args.data_source == "longbench_lite":
|
||||
log(f"[duo-train] data_source=longbench_lite tasks={args.longbench_tasks} "
|
||||
f"samples_per_task={args.longbench_samples_per_task}")
|
||||
dataset = LongBenchLiteAnswerDataset(
|
||||
tokenizer=tokenizer,
|
||||
tasks=args.longbench_tasks,
|
||||
samples_per_task=args.longbench_samples_per_task,
|
||||
max_length=args.max_length,
|
||||
seed=args.seed,
|
||||
cache_dir=args.longbench_cache_dir,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"unknown data_source: {args.data_source}")
|
||||
sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=False) if distributed else None
|
||||
loader = DataLoader(
|
||||
dataset,
|
||||
batch_size=args.batch_size,
|
||||
shuffle=False,
|
||||
sampler=sampler,
|
||||
collate_fn=collate,
|
||||
num_workers=0,
|
||||
)
|
||||
data_iter = iter(loader)
|
||||
|
||||
# Two parameter groups with independent LR multipliers. Group 0 = alpha
|
||||
# scalars (base lr = args.lr, e.g. 0.02); group 1 = backbone (base lr =
|
||||
# args.backbone_lr, e.g. 1e-5). The trapezoidal schedule multiplies both.
|
||||
param_groups = [{"params": alpha_params, "base_lr": args.lr, "lr": args.lr}]
|
||||
if backbone_params:
|
||||
param_groups.append({"params": backbone_params, "base_lr": args.backbone_lr, "lr": args.backbone_lr})
|
||||
optim = torch.optim.AdamW(param_groups, weight_decay=0.0)
|
||||
warm = max(1, int(args.num_steps * args.warmup_ratio))
|
||||
|
||||
def lr_at(step):
|
||||
# trapezoidal schedule, same as DuoAttention's: ramp up over warm, hold, ramp down over warm
|
||||
if step < warm:
|
||||
return max(0.1, (step + 1) / warm)
|
||||
if step > args.num_steps - warm:
|
||||
return max(0.1, (args.num_steps - step) / warm)
|
||||
return 1.0
|
||||
|
||||
model.train()
|
||||
running_distill = running_reg = running_ce = 0.0
|
||||
steps_in_window = 0
|
||||
data_epoch = 0
|
||||
for step in range(args.num_steps):
|
||||
try:
|
||||
batch = next(data_iter)
|
||||
except StopIteration:
|
||||
data_epoch += 1
|
||||
if sampler is not None:
|
||||
sampler.set_epoch(data_epoch)
|
||||
data_iter = iter(loader)
|
||||
batch = next(data_iter)
|
||||
input_ids = batch["input_ids"].to(device)
|
||||
labels = batch["labels"].to(device)
|
||||
label_mask = labels != -100
|
||||
|
||||
# --- Teacher forward: force alpha = 1 everywhere (hidden_states) ----
|
||||
saved = _set_alpha_full(model, 1.0)
|
||||
old_teacher_fastpath = getattr(model.config, "_aha_teacher_full_fastpath", False)
|
||||
model.config._aha_teacher_full_fastpath = True
|
||||
try:
|
||||
with torch.no_grad():
|
||||
out_full = core_model(input_ids=input_ids, use_cache=False)
|
||||
h_full = out_full.last_hidden_state
|
||||
finally:
|
||||
model.config._aha_teacher_full_fastpath = old_teacher_fastpath
|
||||
_restore_alpha(model, saved)
|
||||
|
||||
# --- Student forward: current alpha --------------------------------
|
||||
out_mix = core_model(input_ids=input_ids, use_cache=False)
|
||||
h_mix = out_mix.last_hidden_state
|
||||
|
||||
# DuoAttention's exact distill: mean over hidden_dim, then mean over labelled tokens
|
||||
if label_mask.any():
|
||||
diff = (h_full.float() - h_mix.float())[label_mask] # [N_tok, d_model]
|
||||
distill = diff.pow(2).mean(dim=-1).mean()
|
||||
else:
|
||||
distill = (h_full.float() - h_mix.float()).pow(2).mean(dim=-1).mean()
|
||||
|
||||
# L1 on alpha (clamped)
|
||||
alpha_all = torch.cat([
|
||||
layer.self_attn.full_attention_heads.clamp(0.0, 1.0)
|
||||
for layer in model.model.layers
|
||||
])
|
||||
# DuoAttention uses sum/numel == mean; kept explicit for clarity.
|
||||
reg = alpha_all.abs().sum() / alpha_all.numel()
|
||||
|
||||
# CE anchor on the student (mix) forward. Only computed when ce_weight > 0
|
||||
# to keep paper-grade DuoAttention runs bit-identical to before.
|
||||
if args.ce_weight > 0.0:
|
||||
logits = model.lm_head(h_mix).float()
|
||||
shift_logits = logits[:, :-1, :].contiguous()
|
||||
shift_labels = labels[:, 1:].contiguous()
|
||||
ce = torch.nn.functional.cross_entropy(
|
||||
shift_logits.view(-1, shift_logits.size(-1)),
|
||||
shift_labels.view(-1),
|
||||
ignore_index=-100,
|
||||
)
|
||||
else:
|
||||
ce = h_mix.new_zeros((), dtype=torch.float32)
|
||||
|
||||
loss = distill + args.reg_weight * reg + args.ce_weight * ce
|
||||
(loss / args.grad_accum).backward()
|
||||
|
||||
if (step + 1) % args.grad_accum == 0:
|
||||
for g in optim.param_groups:
|
||||
g["lr"] = g["base_lr"] * lr_at(step)
|
||||
optim.step()
|
||||
optim.zero_grad()
|
||||
# hard clamp alpha into [0, 1]
|
||||
with torch.no_grad():
|
||||
for layer in model.model.layers:
|
||||
layer.self_attn.full_attention_heads.data.clamp_(0.0, 1.0)
|
||||
|
||||
running_distill += float(distill.detach())
|
||||
running_reg += float(reg.detach())
|
||||
running_ce += float(ce.detach())
|
||||
steps_in_window += 1
|
||||
if (step + 1) % args.log_steps == 0:
|
||||
stats = log_alpha_stats(model)
|
||||
lr_str = f"lr_alpha={optim.param_groups[0]['lr']:.4e}"
|
||||
if len(optim.param_groups) > 1:
|
||||
lr_str += f" lr_bb={optim.param_groups[1]['lr']:.2e}"
|
||||
ce_str = f"ce={running_ce/steps_in_window:.4f} " if args.ce_weight > 0.0 else ""
|
||||
log(
|
||||
f"[step {step+1:4d}/{args.num_steps}] "
|
||||
f"distill={running_distill/steps_in_window:.4f} "
|
||||
f"reg={running_reg/steps_in_window:.4f} "
|
||||
f"{ce_str}"
|
||||
f"alpha_mean={stats['alpha_mean']:.3f} "
|
||||
f"alpha_std={stats['alpha_std']:.3f} "
|
||||
f"alpha>0.5_frac={stats['alpha_gt05']:.3f} "
|
||||
f"{lr_str} "
|
||||
f"seq_len={input_ids.shape[1]}",
|
||||
flush=True,
|
||||
)
|
||||
running_distill = running_reg = running_ce = 0.0
|
||||
steps_in_window = 0
|
||||
|
||||
if (step + 1) % args.save_steps == 0 or (step + 1) == args.num_steps:
|
||||
sub = os.path.join(args.output_dir, f"checkpoint-{step+1}")
|
||||
if is_main:
|
||||
os.makedirs(sub, exist_ok=True)
|
||||
model.save_pretrained(sub, safe_serialization=True)
|
||||
tokenizer.save_pretrained(sub)
|
||||
save_alpha_matrix(model, os.path.join(sub, "full_attention_heads.tsv"))
|
||||
stats = log_alpha_stats(model)
|
||||
with open(os.path.join(sub, "duo_state.json"), "w") as f:
|
||||
json.dump({"step": step + 1, **stats}, f, indent=2)
|
||||
log(f"[duo-train] saved {sub}")
|
||||
if distributed:
|
||||
dist.barrier()
|
||||
|
||||
log(f"[duo-train] done. Final alpha stats: {log_alpha_stats(model)}")
|
||||
if distributed:
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
494
recipe/dynamic_duo_train.py
Normal file
494
recipe/dynamic_duo_train.py
Normal file
@@ -0,0 +1,494 @@
|
||||
"""Duo-style hidden-state distillation for AHA dynamic gates.
|
||||
|
||||
This is the clean static-to-dynamic continuation experiment:
|
||||
|
||||
loss = MSE(h_full, h_dynamic) on labelled tokens + reg_weight * mean(gate_soft)
|
||||
|
||||
where h_full is produced by the same frozen checkpoint with every dynamic gate
|
||||
forced to 1.0, and h_dynamic uses the learned per-token gate. For
|
||||
``aha_mode="duo_dynamic"``, "forced to 1.0" means "reproduce the locked static
|
||||
Duo full-head mask"; Duo streaming heads remain streaming. Backbone weights
|
||||
stay frozen; only q_proj gate rows are trainable.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import shutil
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
HERE = os.path.dirname(os.path.abspath(__file__))
|
||||
if HERE not in sys.path:
|
||||
sys.path.insert(0, HERE)
|
||||
|
||||
from duo_train import AmDistilledDataset, LongBenchLiteAnswerDataset, collate # noqa: E402
|
||||
from modeling_aha_qwen3 import ( # noqa: E402
|
||||
AHA_ROUTER_GRANULARITY,
|
||||
AHAQwen3Config,
|
||||
AHAQwen3ForCausalLM,
|
||||
aha_router_output_size,
|
||||
)
|
||||
from router_training_utils import configure_gate_only # noqa: E402
|
||||
|
||||
|
||||
def _set_force_gate(model, value):
|
||||
prev = getattr(model.config, "aha_force_gate_value", None)
|
||||
model.config.aha_force_gate_value = value
|
||||
return prev
|
||||
|
||||
|
||||
def _restore_force_gate(model, value):
|
||||
model.config.aha_force_gate_value = value
|
||||
|
||||
|
||||
def _gate_tensors(out):
|
||||
return [g.float() for g in out.all_gate_soft]
|
||||
|
||||
|
||||
def gate_stats_from_output(out) -> dict:
|
||||
with torch.no_grad():
|
||||
soft = torch.cat([g.reshape(-1) for g in _gate_tensors(out)])
|
||||
hard = torch.cat([g.float().reshape(-1) for g in out.all_gate_hard])
|
||||
return {
|
||||
"gate_soft_mean": soft.mean().item(),
|
||||
"gate_soft_std": soft.std().item(),
|
||||
"gate_hard_mean": hard.mean().item(),
|
||||
"gate_min": soft.min().item(),
|
||||
"gate_max": soft.max().item(),
|
||||
}
|
||||
|
||||
|
||||
def gate_param_stats(model) -> dict:
|
||||
weights = []
|
||||
biases = []
|
||||
num_heads = model.config.num_attention_heads
|
||||
head_dim = getattr(model.config, "head_dim", model.config.hidden_size // num_heads)
|
||||
q_rows = num_heads * head_dim
|
||||
with torch.no_grad():
|
||||
for layer in model.model.layers:
|
||||
q_proj = layer.self_attn.q_proj
|
||||
weights.append(q_proj.weight[q_rows:].detach().float().reshape(-1).cpu())
|
||||
if q_proj.bias is not None:
|
||||
biases.append(q_proj.bias[q_rows:].detach().float().cpu())
|
||||
w = torch.cat(weights)
|
||||
stats = {
|
||||
"gate_weight_l2": w.norm().item(),
|
||||
"gate_weight_abs_mean": w.abs().mean().item(),
|
||||
}
|
||||
if biases:
|
||||
b = torch.cat(biases)
|
||||
alpha = torch.sigmoid(b)
|
||||
stats.update({
|
||||
"gate_bias_mean": b.mean().item(),
|
||||
"gate_bias_std": b.std().item(),
|
||||
"gate_bias_alpha_mean": alpha.mean().item(),
|
||||
"gate_bias_alpha_gt05": (alpha > 0.5).float().mean().item(),
|
||||
})
|
||||
if getattr(model.config, "aha_mode", "dynamic") == "duo_dynamic":
|
||||
masks = []
|
||||
with torch.no_grad():
|
||||
for layer in model.model.layers:
|
||||
alpha_static = layer.self_attn.full_attention_heads.detach().float()
|
||||
masks.append((alpha_static > 0.5).reshape(-1).cpu())
|
||||
m = torch.cat(masks).float()
|
||||
stats.update({
|
||||
"duo_static_full_frac": m.mean().item(),
|
||||
"duo_static_streaming_frac": (1.0 - m).mean().item(),
|
||||
})
|
||||
return stats
|
||||
|
||||
|
||||
def build_reg_head_weights(model, mode: str, power: float):
|
||||
"""Optional per-head weights for the sparsity regularizer.
|
||||
|
||||
``duo_alpha_margin`` uses Duo's own static alpha as confidence: full heads
|
||||
barely above 0.5 receive little sparsity pressure, while high-alpha full
|
||||
heads receive normal pressure. This tests whether the alpha-low8 inference
|
||||
guard can be moved into training as a smooth objective.
|
||||
"""
|
||||
if mode == "uniform":
|
||||
return None
|
||||
if mode != "duo_alpha_margin":
|
||||
raise ValueError(f"unknown reg head weight mode: {mode}")
|
||||
if getattr(model.config, "aha_mode", "dynamic") != "duo_dynamic":
|
||||
raise ValueError(
|
||||
"--reg_head_weight_mode=duo_alpha_margin requires aha_mode=duo_dynamic"
|
||||
)
|
||||
if power <= 0.0:
|
||||
raise ValueError("--reg_head_weight_power must be positive")
|
||||
|
||||
weights = []
|
||||
with torch.no_grad():
|
||||
for layer in model.model.layers:
|
||||
alpha = layer.self_attn.full_attention_heads.detach().float().clamp(0.0, 1.0)
|
||||
duo_full = (alpha > 0.5).float()
|
||||
margin = ((alpha - 0.5) / 0.5).clamp(0.0, 1.0).pow(power)
|
||||
weights.append((margin * duo_full).view(1, 1, -1))
|
||||
flat = torch.cat([w.reshape(-1).cpu() for w in weights])
|
||||
print(
|
||||
"[dynamic-duo] reg_head_weight_mode=duo_alpha_margin "
|
||||
f"power={power:g} nonzero={int((flat > 0).sum().item())}/{flat.numel()} "
|
||||
f"mean={flat.mean().item():.4f}",
|
||||
flush=True,
|
||||
)
|
||||
return weights
|
||||
|
||||
|
||||
def sparsity_reg_from_gates(out, reg_head_weights):
|
||||
gate_layers = _gate_tensors(out)
|
||||
if reg_head_weights is None:
|
||||
gate_soft = torch.cat([g.reshape(-1) for g in gate_layers])
|
||||
return gate_soft.mean()
|
||||
|
||||
weighted = []
|
||||
for gate, weight in zip(gate_layers, reg_head_weights):
|
||||
w = weight.to(device=gate.device, dtype=gate.dtype)
|
||||
weighted.append((gate * w).reshape(-1))
|
||||
if not weighted:
|
||||
raise ValueError("empty weighted sparsity regularizer")
|
||||
return torch.cat(weighted).mean()
|
||||
|
||||
|
||||
def build_gate_optimizer(model, lr: float):
|
||||
setup = configure_gate_only(model)
|
||||
print(
|
||||
f"[dynamic-duo] router_granularity={model.config.aha_router_granularity} "
|
||||
f"gate_rows/layer={setup.gate_rows} effective trainable gate params: "
|
||||
f"{setup.effective_parameter_count:,}",
|
||||
flush=True,
|
||||
)
|
||||
return torch.optim.AdamW([{"params": setup.parameters, "lr": lr}], weight_decay=0.0)
|
||||
|
||||
|
||||
def save_checkpoint(model, tokenizer, output_dir: str, step: int, stats: dict):
|
||||
sub = os.path.join(output_dir, f"checkpoint-{step}")
|
||||
os.makedirs(sub, exist_ok=True)
|
||||
prev_force = getattr(model.config, "aha_force_gate_value", None)
|
||||
model.config.aha_force_gate_value = None
|
||||
model.save_pretrained(sub, safe_serialization=True)
|
||||
tokenizer.save_pretrained(sub)
|
||||
model.config.aha_force_gate_value = prev_force
|
||||
with open(os.path.join(sub, "dynamic_duo_state.json"), "w") as f:
|
||||
json.dump(
|
||||
{
|
||||
"step": step,
|
||||
"router_granularity": getattr(
|
||||
model.config,
|
||||
"aha_router_granularity",
|
||||
AHA_ROUTER_GRANULARITY,
|
||||
),
|
||||
"native_gate_rows_per_layer": aha_router_output_size(model.config),
|
||||
"effective_sparsity_denominator": "token x KV-head x layer",
|
||||
**stats,
|
||||
**gate_param_stats(model),
|
||||
},
|
||||
f,
|
||||
indent=2,
|
||||
)
|
||||
args_manifest = os.path.join(output_dir, "dynamic_duo_train_args.json")
|
||||
if os.path.exists(args_manifest):
|
||||
shutil.copyfile(
|
||||
args_manifest,
|
||||
os.path.join(sub, "dynamic_duo_train_args.json"),
|
||||
)
|
||||
print(f"[dynamic-duo] saved {sub}", flush=True)
|
||||
|
||||
|
||||
def main():
|
||||
AHAQwen3Config.register_for_auto_class()
|
||||
AHAQwen3ForCausalLM.register_for_auto_class("AutoModelForCausalLM")
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--aha_checkpoint", required=True)
|
||||
p.add_argument("--model_path", default="/workspace/AHA/models/Qwen3-0.6B")
|
||||
p.add_argument("--output_dir", required=True)
|
||||
p.add_argument("--am_dataset_path", default="/workspace/Direct-Multitoken-Decoding/am-distilled-8192")
|
||||
p.add_argument("--am_dataset_split", default="train")
|
||||
p.add_argument(
|
||||
"--data_source",
|
||||
default="am_distilled",
|
||||
choices=["am_distilled", "longbench_lite"],
|
||||
help=(
|
||||
"Training data for dynamic gate continuation. The default keeps the "
|
||||
"legacy AM-distilled behavior. longbench_lite is a diagnostic "
|
||||
"target-distribution calibration source, not a paper main protocol."
|
||||
),
|
||||
)
|
||||
p.add_argument(
|
||||
"--am_label_mode",
|
||||
default="full",
|
||||
choices=["full", "answer_only"],
|
||||
help=(
|
||||
"Label mask for --data_source=am_distilled. full keeps all-token "
|
||||
"hidden-state distill; answer_only distills only the final answer span."
|
||||
),
|
||||
)
|
||||
p.add_argument(
|
||||
"--longbench_tasks",
|
||||
nargs="+",
|
||||
default=["passage_retrieval_en", "multifieldqa_en", "qasper", "2wikimqa"],
|
||||
help="Tasks used when --data_source=longbench_lite.",
|
||||
)
|
||||
p.add_argument("--longbench_samples_per_task", type=int, default=30)
|
||||
p.add_argument("--longbench_cache_dir", default="/workspace/AHA/AHA-Qwen3/data/longbench_cache")
|
||||
p.add_argument("--max_length", type=int, default=8192)
|
||||
p.add_argument("--num_steps", type=int, default=400)
|
||||
p.add_argument("--warmup_ratio", type=float, default=0.2)
|
||||
p.add_argument("--lr", type=float, default=3e-5)
|
||||
p.add_argument("--reg_weight", type=float, default=0.05)
|
||||
p.add_argument(
|
||||
"--reg_head_weight_mode",
|
||||
default="uniform",
|
||||
choices=["uniform", "duo_alpha_margin"],
|
||||
help=(
|
||||
"Per-head weighting for the sparsity regularizer. uniform keeps "
|
||||
"the original mean(gate_soft). duo_alpha_margin downweights Duo "
|
||||
"static-full heads close to alpha=0.5 so fragile boundary full "
|
||||
"heads are not pushed local as strongly."
|
||||
),
|
||||
)
|
||||
p.add_argument(
|
||||
"--reg_head_weight_power",
|
||||
type=float,
|
||||
default=1.0,
|
||||
help=(
|
||||
"Power applied to the Duo alpha margin when "
|
||||
"--reg_head_weight_mode=duo_alpha_margin."
|
||||
),
|
||||
)
|
||||
p.add_argument("--ce_weight", type=float, default=0.0)
|
||||
p.add_argument(
|
||||
"--distill_tail_frac",
|
||||
type=float,
|
||||
default=0.0,
|
||||
help=(
|
||||
"If >0, add a tail-aware hidden-state distill term over the top "
|
||||
"fraction of labelled-token MSE values. This keeps the standard "
|
||||
"mean distill objective but prevents rare long-retrieval errors "
|
||||
"from being averaged away."
|
||||
),
|
||||
)
|
||||
p.add_argument(
|
||||
"--distill_tail_weight",
|
||||
type=float,
|
||||
default=0.0,
|
||||
help="Weight for the top-token MSE term enabled by --distill_tail_frac.",
|
||||
)
|
||||
p.add_argument("--batch_size", type=int, default=1)
|
||||
p.add_argument("--grad_accum", type=int, default=1)
|
||||
p.add_argument("--save_steps", type=int, default=100)
|
||||
p.add_argument("--log_steps", type=int, default=10)
|
||||
p.add_argument("--seed", type=int, default=42)
|
||||
p.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float16", "float32"])
|
||||
p.add_argument("--attn_impl", default="sdpa", choices=["sdpa", "eager"])
|
||||
p.add_argument("--aha_local_kind", default="sink_recent", choices=["sink_recent", "sliding_window"])
|
||||
p.add_argument(
|
||||
"--router_granularity",
|
||||
default=None,
|
||||
choices=["token", "token_kv_head"],
|
||||
help="Assert that the input checkpoint has this persisted router granularity.",
|
||||
)
|
||||
args = p.parse_args()
|
||||
|
||||
torch.manual_seed(args.seed)
|
||||
random.seed(args.seed)
|
||||
np.random.seed(args.seed)
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
with open(os.path.join(args.output_dir, "dynamic_duo_train_args.json"), "w") as f:
|
||||
json.dump(vars(args), f, indent=2)
|
||||
|
||||
dtype = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}[args.dtype]
|
||||
print(f"[dynamic-duo] checkpoint={args.aha_checkpoint}", flush=True)
|
||||
print(
|
||||
f"[dynamic-duo] loss=hidden_state_distill + {args.reg_weight} * mean(gate_soft), "
|
||||
f"ce_weight={args.ce_weight} tail_frac={args.distill_tail_frac} "
|
||||
f"tail_weight={args.distill_tail_weight}",
|
||||
flush=True,
|
||||
)
|
||||
if not (0.0 <= args.distill_tail_frac <= 1.0):
|
||||
raise ValueError("--distill_tail_frac must be in [0, 1]")
|
||||
if args.distill_tail_weight < 0.0:
|
||||
raise ValueError("--distill_tail_weight must be non-negative")
|
||||
|
||||
model = AHAQwen3ForCausalLM.from_pretrained_aha(
|
||||
args.aha_checkpoint,
|
||||
torch_dtype=dtype,
|
||||
attn_implementation=args.attn_impl,
|
||||
).cuda()
|
||||
if getattr(model.config, "aha_mode", "dynamic") not in ("dynamic", "duo_dynamic"):
|
||||
raise ValueError("dynamic_duo_train.py requires a dynamic or duo_dynamic AHA checkpoint")
|
||||
loaded_granularity = getattr(
|
||||
model.config, "aha_router_granularity", AHA_ROUTER_GRANULARITY
|
||||
)
|
||||
if args.router_granularity and args.router_granularity != loaded_granularity:
|
||||
raise ValueError(
|
||||
"--router_granularity does not match checkpoint architecture: "
|
||||
f"requested={args.router_granularity!r}, checkpoint={loaded_granularity!r}"
|
||||
)
|
||||
model.config.aha_local_kind = args.aha_local_kind
|
||||
model.config.aha_distill_weight = 0.0
|
||||
model.config.aha_ce_weight = 0.0
|
||||
model.config.aha_reg_weight = -1.0
|
||||
model.config.aha_force_gate_value = None
|
||||
print(
|
||||
f"[dynamic-duo] router_granularity={loaded_granularity} "
|
||||
f"native_gate_rows={aha_router_output_size(model.config)} "
|
||||
f"local_kind={model.config.aha_local_kind} "
|
||||
f"sink={getattr(model.config, 'duo_sink_size', None)} "
|
||||
f"recent={getattr(model.config, 'duo_recent_size', None)}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
|
||||
if tokenizer.pad_token_id is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
|
||||
optim = build_gate_optimizer(model, args.lr)
|
||||
reg_head_weights = build_reg_head_weights(
|
||||
model, args.reg_head_weight_mode, args.reg_head_weight_power
|
||||
)
|
||||
model.enable_input_require_grads()
|
||||
model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
|
||||
|
||||
if args.data_source == "am_distilled":
|
||||
dataset = AmDistilledDataset(
|
||||
ds_path=args.am_dataset_path,
|
||||
split=args.am_dataset_split,
|
||||
max_length=args.max_length,
|
||||
seed=args.seed,
|
||||
tokenizer=tokenizer,
|
||||
label_mode=args.am_label_mode,
|
||||
)
|
||||
print(
|
||||
f"[dynamic-duo] data_source=am_distilled path={args.am_dataset_path} "
|
||||
f"split={args.am_dataset_split} n={len(dataset):,} "
|
||||
f"label_mode={args.am_label_mode} max_length={args.max_length}",
|
||||
flush=True,
|
||||
)
|
||||
elif args.data_source == "longbench_lite":
|
||||
dataset = LongBenchLiteAnswerDataset(
|
||||
tokenizer=tokenizer,
|
||||
tasks=args.longbench_tasks,
|
||||
samples_per_task=args.longbench_samples_per_task,
|
||||
max_length=args.max_length,
|
||||
seed=args.seed,
|
||||
cache_dir=args.longbench_cache_dir,
|
||||
)
|
||||
print(
|
||||
f"[dynamic-duo] data_source=longbench_lite tasks={args.longbench_tasks} "
|
||||
f"samples_per_task={args.longbench_samples_per_task} "
|
||||
f"rows={len(dataset.rows):,} max_length={args.max_length}",
|
||||
flush=True,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"unknown data_source: {args.data_source}")
|
||||
loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False, collate_fn=collate, num_workers=0)
|
||||
data_iter = iter(loader)
|
||||
warm = max(1, int(args.num_steps * args.warmup_ratio))
|
||||
|
||||
def lr_at(step):
|
||||
if step < warm:
|
||||
return max(0.1, (step + 1) / warm)
|
||||
if step > args.num_steps - warm:
|
||||
return max(0.1, (args.num_steps - step) / warm)
|
||||
return 1.0
|
||||
|
||||
model.train()
|
||||
running = {"distill": 0.0, "reg": 0.0, "ce": 0.0, "loss": 0.0, "gate_soft": 0.0, "gate_hard": 0.0}
|
||||
steps_in_window = 0
|
||||
|
||||
for step in range(args.num_steps):
|
||||
try:
|
||||
batch = next(data_iter)
|
||||
except StopIteration:
|
||||
data_iter = iter(loader)
|
||||
batch = next(data_iter)
|
||||
|
||||
input_ids = batch["input_ids"].cuda()
|
||||
labels = batch["labels"].cuda()
|
||||
label_mask = labels != -100
|
||||
|
||||
prev_force = _set_force_gate(model, 1.0)
|
||||
with torch.no_grad():
|
||||
out_full = model.model(input_ids=input_ids, use_cache=False)
|
||||
h_full = out_full.last_hidden_state
|
||||
_restore_force_gate(model, prev_force)
|
||||
|
||||
out_mix = model.model(input_ids=input_ids, use_cache=False)
|
||||
h_mix = out_mix.last_hidden_state
|
||||
|
||||
if label_mask.any():
|
||||
diff = (h_full.float() - h_mix.float())[label_mask]
|
||||
else:
|
||||
diff = (h_full.float() - h_mix.float()).reshape(-1, h_mix.shape[-1])
|
||||
token_mse = diff.pow(2).mean(dim=-1)
|
||||
distill_mean = token_mse.mean()
|
||||
if args.distill_tail_frac > 0.0 and args.distill_tail_weight > 0.0:
|
||||
k = max(1, int(np.ceil(token_mse.numel() * args.distill_tail_frac)))
|
||||
distill_tail = torch.topk(token_mse, k=k, largest=True).values.mean()
|
||||
distill = distill_mean + args.distill_tail_weight * distill_tail
|
||||
else:
|
||||
distill_tail = h_mix.new_zeros((), dtype=torch.float32)
|
||||
distill = distill_mean
|
||||
|
||||
reg = sparsity_reg_from_gates(out_mix, reg_head_weights)
|
||||
|
||||
if args.ce_weight > 0.0:
|
||||
logits = model.lm_head(h_mix).float()
|
||||
shift_logits = logits[:, :-1, :].contiguous()
|
||||
shift_labels = labels[:, 1:].contiguous()
|
||||
ce = torch.nn.functional.cross_entropy(
|
||||
shift_logits.view(-1, shift_logits.size(-1)),
|
||||
shift_labels.view(-1),
|
||||
ignore_index=-100,
|
||||
)
|
||||
else:
|
||||
ce = h_mix.new_zeros((), dtype=torch.float32)
|
||||
|
||||
loss = distill + args.reg_weight * reg + args.ce_weight * ce
|
||||
(loss / args.grad_accum).backward()
|
||||
|
||||
if (step + 1) % args.grad_accum == 0:
|
||||
for group in optim.param_groups:
|
||||
group["lr"] = args.lr * lr_at(step)
|
||||
optim.step()
|
||||
optim.zero_grad()
|
||||
|
||||
stats = gate_stats_from_output(out_mix)
|
||||
running["distill"] += float(distill.detach())
|
||||
running["reg"] += float(reg.detach())
|
||||
running["ce"] += float(ce.detach())
|
||||
running["loss"] += float(loss.detach())
|
||||
running["gate_soft"] += stats["gate_soft_mean"]
|
||||
running["gate_hard"] += stats["gate_hard_mean"]
|
||||
steps_in_window += 1
|
||||
|
||||
if (step + 1) % args.log_steps == 0:
|
||||
denom = float(steps_in_window)
|
||||
print(
|
||||
f"[step {step + 1:4d}/{args.num_steps}] "
|
||||
f"distill={running['distill']/denom:.6f} "
|
||||
f"reg={running['reg']/denom:.6f} "
|
||||
f"loss={running['loss']/denom:.6f} "
|
||||
f"gate_soft={running['gate_soft']/denom:.4f} "
|
||||
f"gate_hard={running['gate_hard']/denom:.4f} "
|
||||
f"lr={optim.param_groups[0]['lr']:.4e} "
|
||||
f"seq_len={input_ids.shape[1]}",
|
||||
flush=True,
|
||||
)
|
||||
running = {k: 0.0 for k in running}
|
||||
steps_in_window = 0
|
||||
|
||||
if (step + 1) % args.save_steps == 0 or (step + 1) == args.num_steps:
|
||||
save_checkpoint(model, tokenizer, args.output_dir, step + 1, stats)
|
||||
|
||||
print(f"[dynamic-duo] done. Final gate param stats: {gate_param_stats(model)}", flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
13
recipe/manifest.json
Normal file
13
recipe/manifest.json
Normal file
@@ -0,0 +1,13 @@
|
||||
{
|
||||
"schema": "aha-l2a-qwen3-repro-v1",
|
||||
"source_commit": "b47a549",
|
||||
"files": {
|
||||
"../model.safetensors": "22c971a65f7f1835b0e2a38b45f2f92191f5d8b63b30969f2f61a69d30afbc05",
|
||||
"data/am_distilled_long_mix/train/data-00000-of-00001.arrow": "27f158ba22512269d33e8fb14db97a84e541ca9b701fba73a98e70c73687f45c",
|
||||
"data/eval_inputs/helmet_icl_8k_n50_per_config.jsonl": "76e4c8115014d11e2dbf47ca8a2e300c490fd8e454326d5e589b12682b7f140e",
|
||||
"data/eval_inputs/mrcr_8k_2_4_8needle_n10_per_config.jsonl": "7074281289fa051278e213ecea523cedf3f609e89b6ae13ab575c075f386ddc4",
|
||||
"modeling_aha_qwen3.py": "4b10c633ab62dbbacbd2a9ac0c69f8225876e97f70bd6ddc4b37367053e2fc1b",
|
||||
"dynamic_duo_train.py": "14188f72ae418831ccda997d0e4cae71c808a59b96aabd3be1e289536a1e073e",
|
||||
"sft.py": "d1f7ae5818639b25ed130433073f4717dcad310dc941bf4a13a9063730b4e555"
|
||||
}
|
||||
}
|
||||
1226
recipe/modeling_aha_qwen3.py
Normal file
1226
recipe/modeling_aha_qwen3.py
Normal file
File diff suppressed because it is too large
Load Diff
83
recipe/protocol.json
Normal file
83
recipe/protocol.json
Normal file
@@ -0,0 +1,83 @@
|
||||
{
|
||||
"release": "jiamingshan/AHA-L2A-Qwen3-1.7B-repro",
|
||||
"source_git": {
|
||||
"repository": "https://github.com/shanjiaming/AHA",
|
||||
"branch": "codex/l2a-router-granularity",
|
||||
"commit": "b47a549"
|
||||
},
|
||||
"base_checkpoint": {
|
||||
"location": "repository root",
|
||||
"description": "Qwen3-1.7B tuned vanilla, global step 6911",
|
||||
"sha256_model_safetensors": "22c971a65f7f1835b0e2a38b45f2f92191f5d8b63b30969f2f61a69d30afbc05"
|
||||
},
|
||||
"execution": {
|
||||
"distributed_training": false,
|
||||
"world_size_per_arm": 1,
|
||||
"gpu_count_per_arm": 1,
|
||||
"one_gpu_machine": "run AHA then L2A-style sequentially on GPU 0",
|
||||
"two_or_more_gpu_machine": "run AHA on GPU 0 and L2A-style on GPU 1 concurrently; remaining GPUs are unused",
|
||||
"arm_definition": "one independent experimental variant, not a processor architecture or GPU group"
|
||||
},
|
||||
"arms": {
|
||||
"aha": {
|
||||
"router_granularity": "token_kv_head",
|
||||
"description": "one native gate per token, KV head, and layer"
|
||||
},
|
||||
"l2a_style": {
|
||||
"router_granularity": "token",
|
||||
"description": "one native gate per token and layer, shared across all heads",
|
||||
"claim_boundary": "L2A-style shared-gate; not an official L2A reproduction"
|
||||
}
|
||||
},
|
||||
"common": {
|
||||
"local_attention": {
|
||||
"kind": "sink_recent",
|
||||
"sink_tokens": 64,
|
||||
"recent_tokens": 256
|
||||
},
|
||||
"gate_initialization": {
|
||||
"weight": "zeros",
|
||||
"full_probability": 0.9
|
||||
},
|
||||
"effective_sparsity_denominator": "token x KV-head x layer"
|
||||
},
|
||||
"stage_1": {
|
||||
"dataset": "recipe/data/am_distilled_long_mix",
|
||||
"rows": 1024,
|
||||
"max_length": 8192,
|
||||
"batch_size": 1,
|
||||
"optimizer_steps": 300,
|
||||
"epochs": 0.29296875,
|
||||
"trainable": "native gate rows only",
|
||||
"learning_rate": 0.00003,
|
||||
"warmup_ratio": 0.1,
|
||||
"hidden_state_distillation": 1.0,
|
||||
"regularizer": 0.1,
|
||||
"ce": 0.0,
|
||||
"train_threshold": 0.5,
|
||||
"seed": 42
|
||||
},
|
||||
"stage_2": {
|
||||
"dataset": "same 1024 rows",
|
||||
"max_length": 8192,
|
||||
"batch_size": 1,
|
||||
"optimizer_steps_run": 75,
|
||||
"selected_checkpoint_step": 25,
|
||||
"selected_epochs": 0.0244140625,
|
||||
"gate_learning_rate": 0.000003,
|
||||
"backbone_learning_rate": 0.0000003,
|
||||
"ce": 1.0,
|
||||
"attention_distillation": 0.5,
|
||||
"regularizer": 0.01,
|
||||
"train_threshold": 0.58,
|
||||
"freeze_embeddings_and_lm_head": true,
|
||||
"seed": 47
|
||||
},
|
||||
"evaluation": {
|
||||
"thresholds": [0.45, 0.5, 0.525, 0.55, 0.575, 0.6, 0.625, 0.65],
|
||||
"strict_prefill_and_decode_sparse_routing": true,
|
||||
"force_full_heads": false,
|
||||
"full_decode": false,
|
||||
"headline": "highest measured effective sparsity whose score is at least 95% of tuned vanilla on every suite"
|
||||
}
|
||||
}
|
||||
12
recipe/requirements.txt
Normal file
12
recipe/requirements.txt
Normal file
@@ -0,0 +1,12 @@
|
||||
transformers==4.54.0
|
||||
trl==0.19.1
|
||||
datasets==4.8.4
|
||||
accelerate==1.13.0
|
||||
safetensors==0.7.0
|
||||
huggingface-hub==0.36.2
|
||||
tokenizers==0.21.4
|
||||
pyarrow==24.0.0
|
||||
pandas==3.0.3
|
||||
numpy==2.3.4
|
||||
matplotlib==3.10.9
|
||||
lm_eval==0.4.11
|
||||
86
recipe/router_training_utils.py
Normal file
86
recipe/router_training_utils.py
Normal file
@@ -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,
|
||||
)
|
||||
82
recipe/scripts/eval_cell.sh
Normal file
82
recipe/scripts/eval_cell.sh
Normal file
@@ -0,0 +1,82 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
if [[ $# -ne 4 ]]; then
|
||||
echo "Usage: $0 METHOD MODEL_PATH GPU BENCHMARK" >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
METHOD="$1"
|
||||
MODEL_PATH="$2"
|
||||
GPU="$3"
|
||||
BENCHMARK="$4"
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
RECIPE="$(cd "$SCRIPT_DIR/.." && pwd)"
|
||||
REPO="$(cd "$RECIPE/.." && pwd)"
|
||||
EVAL_ROOT="${EVAL_ROOT:-$REPO/outputs/eval}"
|
||||
OUT="$EVAL_ROOT/$METHOD"
|
||||
LOG="$EVAL_ROOT/logs/$METHOD"
|
||||
TOKENIZER="${VANILLA_DIR:-$REPO}"
|
||||
HELMET_ROWS="$RECIPE/data/eval_inputs/helmet_icl_8k_n50_per_config.jsonl"
|
||||
MRCR_ROWS="$RECIPE/data/eval_inputs/mrcr_8k_2_4_8needle_n10_per_config.jsonl"
|
||||
RULER_LIMIT="${RULER_LIMIT:-20}"
|
||||
BABILONG_LIMIT="${BABILONG_LIMIT:-50}"
|
||||
AHA_GATE_HARD_THRESHOLD="${AHA_GATE_HARD_THRESHOLD:-0.5}"
|
||||
mkdir -p "$OUT" "$LOG"
|
||||
|
||||
if [[ -f "$OUT/${BENCHMARK}.DONE" ]]; then
|
||||
exit 0
|
||||
fi
|
||||
|
||||
export CUDA_VISIBLE_DEVICES="$GPU"
|
||||
export TOKENIZERS_PARALLELISM=false
|
||||
export PYTHONUNBUFFERED=1
|
||||
export AHA_SPARSITY_STATS_PATH="$OUT/sparsity_${BENCHMARK}.json"
|
||||
export AHA_GATE_HARD_THRESHOLD
|
||||
export AHA_FORCE_FULL_DECODE=0
|
||||
export AHA_FORCE_FULL_PREFILL_TAIL=0
|
||||
export AHA_FORCE_FULL_HEADS=
|
||||
export AHA_FORCE_LOW_ALPHA_FULL_HEADS=0
|
||||
export AHA_DUO_PREFILL_FULL=0
|
||||
|
||||
model_args="pretrained=${MODEL_PATH},trust_remote_code=True,dtype=bfloat16,max_length=16384,attn_implementation=sdpa"
|
||||
split_a="niah_single_1,niah_single_3,niah_multikey_2,niah_multiquery,ruler_vt,ruler_fwe,ruler_qa_hotpot"
|
||||
split_b="niah_single_2,niah_multikey_1,niah_multikey_3,niah_multivalue,ruler_cwe,ruler_qa_squad"
|
||||
|
||||
case "$BENCHMARK" in
|
||||
ruler_a|ruler_b)
|
||||
tasks="$split_a"
|
||||
[[ "$BENCHMARK" == ruler_b ]] && tasks="$split_b"
|
||||
python -m lm_eval \
|
||||
--model hf --model_args "$model_args" \
|
||||
--tasks "$tasks" --metadata '{"max_seq_lengths":[8192]}' \
|
||||
--batch_size 1 --limit "$RULER_LIMIT" --log_samples \
|
||||
--output_path "$OUT/ruler8k_splits/${BENCHMARK#ruler_}/lm_eval" \
|
||||
>"$LOG/${BENCHMARK}.log" 2>&1
|
||||
;;
|
||||
babilong)
|
||||
python -m lm_eval \
|
||||
--model hf --model_args "$model_args" \
|
||||
--tasks babilong_longctx --metadata '{"max_seq_lengths":"8k"}' \
|
||||
--num_fewshot 2 --batch_size 1 --limit "$BABILONG_LIMIT" --log_samples \
|
||||
--output_path "$OUT/babilong8k_qa1_qa5_n50" \
|
||||
>"$LOG/babilong.log" 2>&1
|
||||
;;
|
||||
helmet)
|
||||
python "$SCRIPT_DIR/eval_qwen_external_longctx_pilot.py" \
|
||||
--benchmark helmet_icl --model-path "$MODEL_PATH" --tokenizer-path "$TOKENIZER" \
|
||||
--attn-implementation sdpa --rows "$HELMET_ROWS" \
|
||||
--method "$METHOD" --output "$OUT/helmet_icl8k_n50.jsonl" \
|
||||
>"$LOG/helmet.log" 2>&1
|
||||
;;
|
||||
mrcr)
|
||||
python "$SCRIPT_DIR/eval_qwen_external_longctx_pilot.py" \
|
||||
--benchmark mrcr --model-path "$MODEL_PATH" --tokenizer-path "$TOKENIZER" \
|
||||
--attn-implementation sdpa --rows "$MRCR_ROWS" \
|
||||
--method "$METHOD" --output "$OUT/mrcr_8k_2_4_8needle_n10.jsonl" \
|
||||
>"$LOG/mrcr.log" 2>&1
|
||||
;;
|
||||
*) echo "Unknown benchmark: $BENCHMARK" >&2; exit 2 ;;
|
||||
esac
|
||||
|
||||
touch "$OUT/${BENCHMARK}.DONE"
|
||||
275
recipe/scripts/eval_qwen_external_longctx_pilot.py
Normal file
275
recipe/scripts/eval_qwen_external_longctx_pilot.py
Normal file
@@ -0,0 +1,275 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Aligned 8K HELMET-ICL and OpenAI-MRCR pilot for Qwen checkpoints.
|
||||
|
||||
This is deliberately a small diagnostic runner. HELMET prompt construction
|
||||
matches the ICL schedule used by the local OmniServe reproduction; MRCR uses
|
||||
the official OpenAI rows selected by that reproduction, but renders the
|
||||
messages with the Qwen chat template.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import random
|
||||
import re
|
||||
import string
|
||||
from collections import defaultdict
|
||||
from difflib import SequenceMatcher
|
||||
from pathlib import Path
|
||||
|
||||
import pandas as pd
|
||||
import torch
|
||||
from datasets import load_dataset
|
||||
from huggingface_hub import hf_hub_download
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
|
||||
HELMET_SPECS = {
|
||||
"trec_coarse": ("icl_trec_coarse_400shot_balance", 6),
|
||||
"trec_fine": ("icl_trec_fine_400shot_balance", 50),
|
||||
"banking77": ("icl_banking77_360shot_balance", 77),
|
||||
"clinic150": ("icl_clinic150_440shot_balance", 151),
|
||||
"nlu": ("icl_nlu_510shot_balance", 68),
|
||||
}
|
||||
|
||||
|
||||
def balanced(data, shots: int, label_field: str, seed: int):
|
||||
rng = random.Random(seed)
|
||||
by_label = defaultdict(list)
|
||||
for item in data:
|
||||
by_label[item[label_field]].append(item)
|
||||
rounds = math.ceil(shots / len(by_label))
|
||||
selected_rounds = [[] for _ in range(rounds)]
|
||||
for examples in by_label.values():
|
||||
indices = rng.sample(range(len(examples)), rounds % len(examples))
|
||||
while len(indices) < rounds:
|
||||
indices += rng.sample(range(len(examples)), min(rounds - len(indices), len(examples)))
|
||||
for index, example_index in enumerate(indices):
|
||||
selected_rounds[index].append(examples[example_index])
|
||||
for examples in selected_rounds:
|
||||
rng.shuffle(examples)
|
||||
return [item for group in selected_rounds for item in group][:shots]
|
||||
|
||||
|
||||
def helmet_source(family: str, seed: int):
|
||||
if family == "trec_coarse":
|
||||
source = load_dataset("CogComp/trec", trust_remote_code=True)
|
||||
return source["train"], source["test"], "text", "coarse_label"
|
||||
if family == "trec_fine":
|
||||
source = load_dataset("CogComp/trec", trust_remote_code=True)
|
||||
return source["train"], source["test"], "text", "fine_label"
|
||||
if family == "banking77":
|
||||
source = load_dataset("PolyAI/banking77", trust_remote_code=True)
|
||||
return source["train"], source["test"], "text", "label"
|
||||
if family == "clinic150":
|
||||
source = load_dataset("clinc/clinc_oos", "plus")
|
||||
return source["train"], source["validation"], "text", "intent"
|
||||
if family == "nlu":
|
||||
source = load_dataset("xingkunliuxtracta/nlu_evaluation_data", trust_remote_code=True)["train"]
|
||||
split = source.train_test_split(test_size=0.1, seed=seed)
|
||||
return split["train"], split["test"], "text", "label"
|
||||
raise ValueError(family)
|
||||
|
||||
|
||||
def build_helmet(seed: int, limit: int):
|
||||
rows = []
|
||||
user_template = (
|
||||
'Use the provided mapping from the text to label to assign a label to the text. '
|
||||
'Only output "label: {{label}}" and nothing else. \n\n{context}\n\n{question}'
|
||||
)
|
||||
for family, (dataset_name, num_labels) in HELMET_SPECS.items():
|
||||
shots = int(dataset_name.split("shot")[0].split("_")[-1])
|
||||
train, test, text_field, label_field = helmet_source(family, seed)
|
||||
samples = balanced(test, limit, label_field, seed)
|
||||
for sample_index, sample in enumerate(samples):
|
||||
local_seed = (int(hashlib.sha256(sample[text_field].encode()).hexdigest(), 16) + seed) % 2**31
|
||||
demos = balanced(train, shots, label_field, local_seed)
|
||||
mapping = list(range(num_labels))
|
||||
random.Random(local_seed).shuffle(mapping)
|
||||
context = "\n\n".join(
|
||||
f"{demo[text_field]}\nlabel: {mapping[int(demo[label_field])]}" for demo in demos
|
||||
)
|
||||
rows.append({
|
||||
"row_id": f"helmet_icl:8k:{family}:{sample_index}",
|
||||
"benchmark": "helmet_icl",
|
||||
"config": family,
|
||||
"prompt": user_template.format(context=context, question=sample[text_field]) + "\nlabel:",
|
||||
"answer": str(mapping[int(sample[label_field])]),
|
||||
"max_new_tokens": 20,
|
||||
"metadata": {"dataset": dataset_name, "shots": shots},
|
||||
})
|
||||
return rows
|
||||
|
||||
|
||||
def build_mrcr(tokenizer, selection_path: Path, limit: int):
|
||||
selection = [json.loads(line) for line in selection_path.read_text().splitlines() if line.strip()]
|
||||
chosen = sorted(
|
||||
(row for row in selection if int(row["context_k"]) == 8),
|
||||
key=lambda row: int(row["sample_index"]),
|
||||
)[:limit]
|
||||
if len(chosen) < limit:
|
||||
raise ValueError(f"MRCR selection contains {len(chosen)} 8K rows, requested {limit}")
|
||||
files = [
|
||||
hf_hub_download("openai/mrcr", filename=f"2needle/2needle_{shard}.parquet", repo_type="dataset")
|
||||
for shard in (0, 1)
|
||||
]
|
||||
frame = pd.concat([pd.read_parquet(path) for path in files], ignore_index=True)
|
||||
rows = []
|
||||
for selected in chosen:
|
||||
source = frame.iloc[int(selected["source_index"])]
|
||||
messages = json.loads(source["prompt"])
|
||||
prompt = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
answer = str(source["answer"])
|
||||
sample_index = int(selected["sample_index"])
|
||||
rows.append({
|
||||
"row_id": f"mrcr:8k:2needle:{sample_index}",
|
||||
"benchmark": "mrcr",
|
||||
"config": "2needle",
|
||||
"prompt": prompt,
|
||||
"answer": answer,
|
||||
"prefix": str(source["random_string_to_prepend"]),
|
||||
"max_new_tokens": min(768, len(tokenizer(answer, add_special_tokens=False).input_ids) + 64),
|
||||
"metadata": {"source_index": int(selected["source_index"]), "official_tokens": int(selected["official_tokens"])},
|
||||
})
|
||||
return rows
|
||||
|
||||
|
||||
def normalize(text: str) -> str:
|
||||
text = text.lower()
|
||||
text = re.sub(r"\b(a|an|the)\b", " ", text)
|
||||
text = "".join(ch for ch in text if ch not in string.punctuation)
|
||||
return " ".join(text.split())
|
||||
|
||||
|
||||
def score(row, prediction: str) -> float:
|
||||
if row["benchmark"] == "helmet_icl":
|
||||
first = prediction.strip().splitlines()[0] if prediction.strip() else ""
|
||||
parsed = re.sub(r"^label:", "", first, flags=re.I).strip()
|
||||
return float(normalize(parsed) == normalize(row["answer"]))
|
||||
prefix = row["prefix"]
|
||||
if not prediction.startswith(prefix):
|
||||
return 0.0
|
||||
response = prediction.removeprefix(prefix).strip()
|
||||
reference = row["answer"].removeprefix(prefix).strip()
|
||||
return float(SequenceMatcher(None, response, reference).ratio())
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--benchmark", choices=("helmet_icl", "mrcr"), required=True)
|
||||
parser.add_argument("--model-path", required=True)
|
||||
parser.add_argument("--tokenizer-path", required=True)
|
||||
parser.add_argument(
|
||||
"--attn-implementation",
|
||||
default="flash_attention_2",
|
||||
choices=("eager", "sdpa", "flash_attention_2"),
|
||||
help="Transformers attention backend; AHA 4-D local masks require SDPA or eager.",
|
||||
)
|
||||
parser.add_argument("--selection", type=Path)
|
||||
parser.add_argument("--rows", type=Path, help="Frozen aligned prompt rows (JSONL)")
|
||||
parser.add_argument("--method", required=True)
|
||||
parser.add_argument("--output", type=Path, required=True)
|
||||
parser.add_argument("--seed", type=int, default=20260710)
|
||||
parser.add_argument("--limit", type=int, default=1, help="Samples per benchmark config")
|
||||
parser.add_argument(
|
||||
"--stop-new-line",
|
||||
action="store_true",
|
||||
help="Use HELMET's newline stop-token policy during generation.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--duo-sink-size",
|
||||
type=int,
|
||||
help="Optional inference-only override for the AHA sink token count.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--duo-recent-size",
|
||||
type=int,
|
||||
help="Optional inference-only override for the AHA recent-window token count.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path, trust_remote_code=True, use_fast=False)
|
||||
if args.rows:
|
||||
rows = [json.loads(line) for line in args.rows.read_text().splitlines() if line.strip()]
|
||||
else:
|
||||
rows = (
|
||||
build_helmet(args.seed, args.limit)
|
||||
if args.benchmark == "helmet_icl"
|
||||
else build_mrcr(tokenizer, args.selection, args.limit)
|
||||
)
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
args.model_path,
|
||||
trust_remote_code=True,
|
||||
torch_dtype=torch.bfloat16,
|
||||
attn_implementation=args.attn_implementation,
|
||||
)
|
||||
for field, value in (
|
||||
("duo_sink_size", args.duo_sink_size),
|
||||
("duo_recent_size", args.duo_recent_size),
|
||||
):
|
||||
if value is not None:
|
||||
if value < 0:
|
||||
parser.error(f"--{field.replace('_', '-')} must be non-negative")
|
||||
previous = getattr(model.config, field, None)
|
||||
setattr(model.config, field, value)
|
||||
print(f"[eval] {field} override: {previous!r} -> {value!r}", flush=True)
|
||||
model = model.to("cuda").eval()
|
||||
stop_token_ids = model.generation_config.eos_token_id
|
||||
stop_token_ids = list(stop_token_ids) if isinstance(stop_token_ids, list) else [stop_token_ids]
|
||||
if args.stop_new_line:
|
||||
newline_tokens = ["\n", "Ċ", "ĊĊ", "<0x0A>"]
|
||||
stop_token_ids += [tokenizer.convert_tokens_to_ids(token) for token in newline_tokens]
|
||||
stop_token_ids = sorted(
|
||||
{
|
||||
token_id
|
||||
for token_id in stop_token_ids
|
||||
if token_id is not None and token_id != tokenizer.unk_token_id
|
||||
}
|
||||
)
|
||||
print(f"[eval] stop_token_ids={stop_token_ids}", flush=True)
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
outputs = []
|
||||
for ordinal, row in enumerate(rows, 1):
|
||||
encoded = tokenizer(row["prompt"], return_tensors="pt", add_special_tokens=False)
|
||||
prompt_tokens = int(encoded.input_ids.shape[1])
|
||||
print(f"[{args.benchmark}] {ordinal}/{len(rows)} {row['config']} tokens={prompt_tokens}", flush=True)
|
||||
encoded = {key: value.to("cuda") for key, value in encoded.items()}
|
||||
with torch.inference_mode():
|
||||
generated = model.generate(
|
||||
**encoded,
|
||||
do_sample=False,
|
||||
max_new_tokens=row["max_new_tokens"],
|
||||
use_cache=True,
|
||||
eos_token_id=stop_token_ids,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
)
|
||||
prediction = tokenizer.decode(generated[0, prompt_tokens:], skip_special_tokens=True).strip()
|
||||
outputs.append({
|
||||
**row,
|
||||
"method": args.method,
|
||||
"prompt_tokens": prompt_tokens,
|
||||
"prompt_sha256": hashlib.sha256(row["prompt"].encode()).hexdigest(),
|
||||
"prediction": prediction,
|
||||
"score": score(row, prediction),
|
||||
})
|
||||
args.output.write_text("".join(json.dumps(row, ensure_ascii=False) + "\n" for row in outputs))
|
||||
summary = args.output.with_suffix(".summary.csv")
|
||||
with summary.open("w", newline="") as handle:
|
||||
writer = csv.DictWriter(handle, fieldnames=["method", "benchmark", "config", "n", "score", "prompt_tokens"])
|
||||
writer.writeheader()
|
||||
for row in outputs:
|
||||
writer.writerow({"method": args.method, "benchmark": row["benchmark"], "config": row["config"], "n": 1, "score": row["score"], "prompt_tokens": row["prompt_tokens"]})
|
||||
print(summary)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
64
recipe/scripts/eval_sweep.sh
Normal file
64
recipe/scripts/eval_sweep.sh
Normal file
@@ -0,0 +1,64 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
RECIPE="$(cd "$SCRIPT_DIR/.." && pwd)"
|
||||
REPO="$(cd "$RECIPE/.." && pwd)"
|
||||
OUTPUT_ROOT="${OUTPUT_ROOT:-$REPO/outputs}"
|
||||
EVAL_ROOT="${EVAL_ROOT:-$OUTPUT_ROOT/eval}"
|
||||
VANILLA_MODEL="${VANILLA_DIR:-$REPO}"
|
||||
AHA_MODEL="${AHA_MODEL:-$OUTPUT_ROOT/aha/stage2/checkpoint-25}"
|
||||
L2A_MODEL="${L2A_MODEL:-$OUTPUT_ROOT/l2a_style/stage2/checkpoint-25}"
|
||||
GPU_LIST="${GPU_LIST:-0}"
|
||||
THRESHOLDS=(0.45 0.50 0.525 0.55 0.575 0.60 0.625 0.65)
|
||||
BENCHMARKS=(ruler_a ruler_b babilong helmet mrcr)
|
||||
IFS=',' read -r -a GPUS <<< "$GPU_LIST"
|
||||
declare -A SLOT_PIDS=()
|
||||
mkdir -p "$EVAL_ROOT"
|
||||
|
||||
for model in "$VANILLA_MODEL" "$AHA_MODEL" "$L2A_MODEL"; do
|
||||
[[ -f "$model/config.json" ]] || { echo "Missing model: $model" >&2; exit 2; }
|
||||
done
|
||||
|
||||
jobs=()
|
||||
for benchmark in "${BENCHMARKS[@]}"; do
|
||||
jobs+=("vanilla|$VANILLA_MODEL|0.5|$benchmark")
|
||||
done
|
||||
for threshold in "${THRESHOLDS[@]}"; do
|
||||
slug="${threshold/./}"
|
||||
for benchmark in "${BENCHMARKS[@]}"; do
|
||||
jobs+=("token_kv_head_t${slug}|$AHA_MODEL|$threshold|$benchmark")
|
||||
jobs+=("token_t${slug}|$L2A_MODEL|$threshold|$benchmark")
|
||||
done
|
||||
done
|
||||
|
||||
launch() {
|
||||
local gpu="$1" method="$2" model="$3" threshold="$4" benchmark="$5"
|
||||
EVAL_ROOT="$EVAL_ROOT" RULER_LIMIT=20 BABILONG_LIMIT=50 \
|
||||
AHA_GATE_HARD_THRESHOLD="$threshold" \
|
||||
"$SCRIPT_DIR/eval_cell.sh" "$method" "$model" "$gpu" "$benchmark" &
|
||||
LAST_PID="$!"
|
||||
}
|
||||
|
||||
for i in "${!jobs[@]}"; do
|
||||
slot=$((i % ${#GPUS[@]}))
|
||||
if [[ -n "${SLOT_PIDS[$slot]:-}" ]]; then
|
||||
wait "${SLOT_PIDS[$slot]}"
|
||||
fi
|
||||
IFS='|' read -r method model threshold benchmark <<< "${jobs[$i]}"
|
||||
launch "${GPUS[$slot]}" "$method" "$model" "$threshold" "$benchmark"
|
||||
SLOT_PIDS[$slot]="$LAST_PID"
|
||||
done
|
||||
|
||||
status=0
|
||||
for pid in "${SLOT_PIDS[@]}"; do
|
||||
wait "$pid" || status=1
|
||||
done
|
||||
[[ "$status" == 0 ]] || exit "$status"
|
||||
|
||||
python "$SCRIPT_DIR/summarize_qwen1p7b_router_granularity_20260714.py" \
|
||||
--repo "$REPO" \
|
||||
--eval-root "$EVAL_ROOT" \
|
||||
--baseline-root "$EVAL_ROOT/vanilla" \
|
||||
--input-dir "$RECIPE/data/eval_inputs" \
|
||||
--output-dir "$OUTPUT_ROOT/summary"
|
||||
20
recipe/scripts/evaluate.sh
Normal file
20
recipe/scripts/evaluate.sh
Normal file
@@ -0,0 +1,20 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
REPO="$(cd "$SCRIPT_DIR/../.." && pwd)"
|
||||
IMAGE="${IMAGE:-aha-l2a-qwen3-repro:torch2.9.1-cu128}"
|
||||
GPU_COUNT="$(nvidia-smi -L | wc -l)"
|
||||
if [[ "$GPU_COUNT" -lt 1 ]]; then
|
||||
echo "At least one NVIDIA GPU is required." >&2
|
||||
exit 2
|
||||
fi
|
||||
GPU_LIST="$(seq -s, 0 $((GPU_COUNT - 1)))"
|
||||
|
||||
docker build -t "$IMAGE" -f "$REPO/recipe/Dockerfile" "$REPO/recipe"
|
||||
docker run --rm --gpus all --ipc=host \
|
||||
--ulimit memlock=-1 --ulimit stack=67108864 \
|
||||
-e GPU_LIST="$GPU_LIST" \
|
||||
-v "$REPO:/workspace" \
|
||||
-w /workspace/recipe \
|
||||
"$IMAGE" bash scripts/eval_sweep.sh
|
||||
140
recipe/scripts/export_tuned_vanilla_aha_hotstart.py
Normal file
140
recipe/scripts/export_tuned_vanilla_aha_hotstart.py
Normal file
@@ -0,0 +1,140 @@
|
||||
#!/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()
|
||||
45
recipe/scripts/reproduce.sh
Normal file
45
recipe/scripts/reproduce.sh
Normal file
@@ -0,0 +1,45 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
REPO="$(cd "$SCRIPT_DIR/../.." && pwd)"
|
||||
IMAGE="${IMAGE:-aha-l2a-qwen3-repro:torch2.9.1-cu128}"
|
||||
GPU_COUNT="$(nvidia-smi -L | wc -l)"
|
||||
if [[ "$GPU_COUNT" -lt 1 ]]; then
|
||||
echo "At least one NVIDIA GPU is required." >&2
|
||||
exit 2
|
||||
fi
|
||||
GPU_AHA="${GPU_AHA:-0}"
|
||||
if [[ -z "${GPU_L2A+x}" ]]; then
|
||||
GPU_L2A=0
|
||||
if [[ "$GPU_COUNT" -ge 2 ]]; then
|
||||
GPU_L2A=1
|
||||
fi
|
||||
fi
|
||||
|
||||
for gpu in "$GPU_AHA" "$GPU_L2A"; do
|
||||
if [[ ! "$gpu" =~ ^[0-9]+$ ]] || (( gpu >= GPU_COUNT )); then
|
||||
echo "Invalid GPU index $gpu; nvidia-smi reports $GPU_COUNT visible GPU(s)." >&2
|
||||
exit 2
|
||||
fi
|
||||
done
|
||||
|
||||
echo "Training mode: one GPU per experimental variant (no DDP/FSDP/DeepSpeed)."
|
||||
echo "AHA variant GPU: $GPU_AHA"
|
||||
echo "L2A-style variant GPU: $GPU_L2A"
|
||||
if [[ "$GPU_AHA" == "$GPU_L2A" ]]; then
|
||||
echo "The two variants will run sequentially on the same GPU."
|
||||
else
|
||||
echo "The two variants will run concurrently on two independent GPUs."
|
||||
fi
|
||||
if [[ "$GPU_COUNT" -gt 2 ]]; then
|
||||
echo "Visible GPUs: $GPU_COUNT; this recipe intentionally does not use all GPUs."
|
||||
fi
|
||||
|
||||
docker build -t "$IMAGE" -f "$REPO/recipe/Dockerfile" "$REPO/recipe"
|
||||
docker run --rm --gpus all --ipc=host \
|
||||
--ulimit memlock=-1 --ulimit stack=67108864 \
|
||||
-e GPU_AHA="$GPU_AHA" -e GPU_L2A="$GPU_L2A" \
|
||||
-v "$REPO:/workspace" \
|
||||
-w /workspace/recipe \
|
||||
"$IMAGE" bash scripts/train_both.sh
|
||||
403
recipe/scripts/summarize_qwen1p7b_router_granularity_20260714.py
Normal file
403
recipe/scripts/summarize_qwen1p7b_router_granularity_20260714.py
Normal file
@@ -0,0 +1,403 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Audit and summarize the matched shared-gate vs KV-head-gate sweep."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import glob
|
||||
import hashlib
|
||||
import json
|
||||
import subprocess
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
|
||||
SUITES = ("RULER-local-full13", "BabiLong", "HELMET-ICL", "MRCR")
|
||||
THRESHOLDS = (0.45, 0.50, 0.525, 0.55, 0.575, 0.60, 0.625, 0.65)
|
||||
ARMS = ("token", "token_kv_head")
|
||||
SLUGS = {
|
||||
"RULER-local-full13": "ruler",
|
||||
"BabiLong": "babilong",
|
||||
"HELMET-ICL": "helmet",
|
||||
"MRCR": "mrcr",
|
||||
}
|
||||
SCORE_DEFINITIONS = {
|
||||
"RULER-local-full13": "unweighted macro of lm-eval 8192 string-match scores over the local 13-config set",
|
||||
"BabiLong": "unweighted macro exact-match accuracy over qa1-qa5",
|
||||
"HELMET-ICL": "unweighted macro label exact-match over five ICL configurations",
|
||||
"MRCR": "unweighted macro of prefix-check plus SequenceMatcher scores over 2/4/8-needle configurations",
|
||||
}
|
||||
|
||||
|
||||
def threshold_slug(value: float) -> str:
|
||||
return {
|
||||
0.45: "045",
|
||||
0.50: "050",
|
||||
0.525: "0525",
|
||||
0.55: "055",
|
||||
0.575: "0575",
|
||||
0.60: "060",
|
||||
0.625: "0625",
|
||||
0.65: "065",
|
||||
}[value]
|
||||
|
||||
|
||||
def one(pattern: str) -> Path:
|
||||
paths = [Path(path) for path in glob.glob(pattern, recursive=True)]
|
||||
if len(paths) != 1:
|
||||
raise RuntimeError(f"expected one path for {pattern}, found {paths}")
|
||||
return paths[0]
|
||||
|
||||
|
||||
def read_jsonl(path: Path) -> list[dict]:
|
||||
return [json.loads(line) for line in path.read_text().splitlines() if line.strip()]
|
||||
|
||||
|
||||
def collect_scores(root: Path) -> dict[tuple[str, str], dict]:
|
||||
rows: dict[tuple[str, str], dict] = {}
|
||||
for split in ("a", "b"):
|
||||
result = one(str(root / f"ruler8k_splits/{split}/lm_eval/**/results_*.json"))
|
||||
payload = json.loads(result.read_text())
|
||||
for config, metrics in payload["results"].items():
|
||||
if config in payload["n-samples"]:
|
||||
rows[("RULER-local-full13", config)] = {
|
||||
"score": float(metrics["8192,none"]),
|
||||
"n": int(payload["n-samples"][config]["effective"]),
|
||||
}
|
||||
result = one(str(root / "babilong8k_qa1_qa5_n50/**/results_*.json"))
|
||||
payload = json.loads(result.read_text())
|
||||
for config, metrics in payload["results"].items():
|
||||
if config in payload["n-samples"]:
|
||||
rows[("BabiLong", config)] = {
|
||||
"score": float(metrics["acc,none"]),
|
||||
"n": int(payload["n-samples"][config]["effective"]),
|
||||
}
|
||||
for suite, filename in (
|
||||
("HELMET-ICL", "helmet_icl8k_n50.jsonl"),
|
||||
("MRCR", "mrcr_8k_2_4_8needle_n10.jsonl"),
|
||||
):
|
||||
grouped: dict[str, list[dict]] = defaultdict(list)
|
||||
for row in read_jsonl(root / filename):
|
||||
grouped[row["config"]].append(row)
|
||||
for config, values in grouped.items():
|
||||
rows[(suite, config)] = {
|
||||
"score": sum(float(row["score"]) for row in values) / len(values),
|
||||
"n": len(values),
|
||||
}
|
||||
return rows
|
||||
|
||||
|
||||
def sample_signatures(root: Path, score_rows: dict[tuple[str, str], dict]) -> dict[str, str]:
|
||||
signatures: dict[str, str] = {}
|
||||
for suite, rel_dirs in (
|
||||
("RULER-local-full13", ("ruler8k_splits/a/lm_eval", "ruler8k_splits/b/lm_eval")),
|
||||
("BabiLong", ("babilong8k_qa1_qa5_n50",)),
|
||||
):
|
||||
configs = [config for row_suite, config in score_rows if row_suite == suite]
|
||||
for config in configs:
|
||||
matches = []
|
||||
for rel_dir in rel_dirs:
|
||||
matches.extend(root.glob(f"{rel_dir}/**/samples_{config}_*.jsonl"))
|
||||
if len(matches) != 1:
|
||||
raise RuntimeError(f"expected one sample log for {suite}/{config}: {matches}")
|
||||
values = [
|
||||
(
|
||||
row.get("doc_id"),
|
||||
row.get("doc_hash"),
|
||||
row.get("prompt_hash"),
|
||||
row.get("target_hash"),
|
||||
)
|
||||
for row in read_jsonl(matches[0])
|
||||
]
|
||||
signatures[f"{suite}/{config}"] = hashlib.sha256(
|
||||
json.dumps(values, sort_keys=True).encode()
|
||||
).hexdigest()
|
||||
for suite, filename in (
|
||||
("HELMET-ICL", "helmet_icl8k_n50.jsonl"),
|
||||
("MRCR", "mrcr_8k_2_4_8needle_n10.jsonl"),
|
||||
):
|
||||
grouped: dict[str, list[tuple]] = defaultdict(list)
|
||||
for row in read_jsonl(root / filename):
|
||||
grouped[row["config"]].append(
|
||||
(row["row_id"], row.get("prompt_sha256"), row.get("target"))
|
||||
)
|
||||
for config, values in grouped.items():
|
||||
signatures[f"{suite}/{config}"] = hashlib.sha256(
|
||||
json.dumps(values, sort_keys=True).encode()
|
||||
).hexdigest()
|
||||
return signatures
|
||||
|
||||
|
||||
def read_sparsity(root: Path, suite: str) -> dict:
|
||||
if suite == "RULER-local-full13":
|
||||
parts = [
|
||||
json.loads((root / f"sparsity_ruler_{split}.json").read_text())
|
||||
for split in ("a", "b")
|
||||
]
|
||||
sparse = sum(int(part["sparse_decisions"]) for part in parts)
|
||||
total = sum(int(part["total_decisions"]) for part in parts)
|
||||
granularities = {part.get("router_granularity") for part in parts}
|
||||
native = sum(int(part.get("native_router_decisions", 0)) for part in parts)
|
||||
phases: dict[str, dict[str, int]] = defaultdict(lambda: {"sparse": 0, "total": 0})
|
||||
for part in parts:
|
||||
for phase, values in part.get("by_phase", {}).items():
|
||||
phases[phase]["sparse"] += int(values["sparse_decisions"])
|
||||
phases[phase]["total"] += int(values["total_decisions"])
|
||||
else:
|
||||
part = json.loads((root / f"sparsity_{SLUGS[suite]}.json").read_text())
|
||||
sparse = int(part["sparse_decisions"])
|
||||
total = int(part["total_decisions"])
|
||||
granularities = {part.get("router_granularity")}
|
||||
native = int(part.get("native_router_decisions", 0))
|
||||
phases = {
|
||||
phase: {
|
||||
"sparse": int(values["sparse_decisions"]),
|
||||
"total": int(values["total_decisions"]),
|
||||
}
|
||||
for phase, values in part.get("by_phase", {}).items()
|
||||
}
|
||||
return {
|
||||
"sparse_decisions": sparse,
|
||||
"total_decisions": total,
|
||||
"sparsity": sparse / total,
|
||||
"full_attention_usage": 1.0 - sparse / total,
|
||||
"router_granularities": sorted(value for value in granularities if value),
|
||||
"native_router_decisions": native or None,
|
||||
"by_phase": phases,
|
||||
}
|
||||
|
||||
|
||||
def sha256(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as handle:
|
||||
for block in iter(lambda: handle.read(1024 * 1024), b""):
|
||||
digest.update(block)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def command_output(command: list[str], cwd: Path) -> str:
|
||||
try:
|
||||
return subprocess.run(
|
||||
command, cwd=cwd, check=True, text=True, capture_output=True
|
||||
).stdout.strip()
|
||||
except (OSError, subprocess.CalledProcessError):
|
||||
return "unavailable in downloaded HF snapshot"
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--repo", type=Path, default=Path("/data/sjm/AHA/AHA-Qwen3")
|
||||
)
|
||||
parser.add_argument(
|
||||
"--eval-root",
|
||||
type=Path,
|
||||
default=Path("/data/sjm/AHA/AHA-Qwen3/experiments/qwen3_1p7b_router_granularity_20260714/eval"),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--baseline-root",
|
||||
type=Path,
|
||||
default=Path("/data/sjm/AHA/AHA-Qwen3/experiments/qwen3_1p7b_fourbench_8k_broad_20260713/vanilla"),
|
||||
)
|
||||
parser.add_argument("--output-dir", type=Path)
|
||||
parser.add_argument("--input-dir", type=Path)
|
||||
args = parser.parse_args()
|
||||
output = args.output_dir or args.eval_root.parent / "summary"
|
||||
output.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
baseline = collect_scores(args.baseline_root)
|
||||
baseline_signatures = sample_signatures(args.baseline_root, baseline)
|
||||
curve_rows = []
|
||||
per_config = []
|
||||
points = []
|
||||
alignment = {}
|
||||
done_files = ("ruler_a.DONE", "ruler_b.DONE", "babilong.DONE", "helmet.DONE", "mrcr.DONE")
|
||||
|
||||
for arm in ARMS:
|
||||
for threshold in THRESHOLDS:
|
||||
method = f"{arm}_t{threshold_slug(threshold)}"
|
||||
root = args.eval_root / method
|
||||
missing = [name for name in done_files if not (root / name).exists()]
|
||||
if missing:
|
||||
raise RuntimeError(f"incomplete {method}: missing {missing}")
|
||||
scores = collect_scores(root)
|
||||
if set(scores) != set(baseline):
|
||||
raise RuntimeError(f"configuration mismatch for {method}")
|
||||
if any(scores[key]["n"] != baseline[key]["n"] for key in baseline):
|
||||
raise RuntimeError(f"sample-count mismatch for {method}")
|
||||
signatures = sample_signatures(root, scores)
|
||||
mismatches = sorted(
|
||||
key for key in baseline_signatures if signatures.get(key) != baseline_signatures[key]
|
||||
)
|
||||
alignment[method] = {"aligned": not mismatches, "mismatches": mismatches}
|
||||
if mismatches:
|
||||
raise RuntimeError(f"prompt/target hash mismatch for {method}: {mismatches}")
|
||||
|
||||
suite_rows = []
|
||||
sparse_sum = total_sum = 0
|
||||
for suite in SUITES:
|
||||
keys = sorted(key for key in baseline if key[0] == suite)
|
||||
baseline_score = sum(baseline[key]["score"] for key in keys) / len(keys)
|
||||
score = sum(scores[key]["score"] for key in keys) / len(keys)
|
||||
sparsity = read_sparsity(root, suite)
|
||||
expected_granularity = arm
|
||||
if sparsity["router_granularities"] != [expected_granularity]:
|
||||
raise RuntimeError(
|
||||
f"{method}/{suite} recorded {sparsity['router_granularities']}, "
|
||||
f"expected {[expected_granularity]}"
|
||||
)
|
||||
sparse_sum += sparsity["sparse_decisions"]
|
||||
total_sum += sparsity["total_decisions"]
|
||||
row = {
|
||||
"arm": arm,
|
||||
"router_granularity": arm,
|
||||
"threshold": threshold,
|
||||
"benchmark": suite,
|
||||
"config_count": len(keys),
|
||||
"samples_per_config": ",".join(
|
||||
str(value) for value in sorted({baseline[key]["n"] for key in keys})
|
||||
),
|
||||
"total_samples": sum(baseline[key]["n"] for key in keys),
|
||||
"tuned_vanilla_score": baseline_score,
|
||||
"quality_floor_95pct": 0.95 * baseline_score,
|
||||
"score": score,
|
||||
"retention": score / baseline_score if baseline_score else None,
|
||||
"quality_pass": score >= 0.95 * baseline_score,
|
||||
**{key: sparsity[key] for key in (
|
||||
"sparse_decisions", "total_decisions", "sparsity",
|
||||
"full_attention_usage", "native_router_decisions", "by_phase"
|
||||
)},
|
||||
}
|
||||
curve_rows.append(row)
|
||||
suite_rows.append(row)
|
||||
for key in keys:
|
||||
per_config.append({
|
||||
"arm": arm,
|
||||
"threshold": threshold,
|
||||
"benchmark": suite,
|
||||
"config": key[1],
|
||||
"n": baseline[key]["n"],
|
||||
"tuned_vanilla_score": baseline[key]["score"],
|
||||
"score": scores[key]["score"],
|
||||
})
|
||||
points.append({
|
||||
"arm": arm,
|
||||
"threshold": threshold,
|
||||
"all_suite_quality_pass": all(row["quality_pass"] for row in suite_rows),
|
||||
"decision_weighted_sparsity": sparse_sum / total_sum,
|
||||
"full_attention_usage": 1.0 - sparse_sum / total_sum,
|
||||
"sparse_decisions": sparse_sum,
|
||||
"total_decisions": total_sum,
|
||||
"suites": suite_rows,
|
||||
})
|
||||
|
||||
headline = {}
|
||||
for arm in ARMS:
|
||||
valid = [
|
||||
point for point in points
|
||||
if point["arm"] == arm and point["all_suite_quality_pass"]
|
||||
]
|
||||
if valid:
|
||||
best = max(valid, key=lambda point: point["decision_weighted_sparsity"])
|
||||
headline[arm] = {"found": True, **best}
|
||||
else:
|
||||
headline[arm] = {
|
||||
"found": False,
|
||||
"statement": "No measured threshold preserved at least 95% of tuned vanilla on every suite.",
|
||||
}
|
||||
|
||||
input_dir = args.input_dir or args.baseline_root.parent / "inputs"
|
||||
frozen_inputs = [
|
||||
input_dir / "helmet_icl_8k_n50_per_config.jsonl",
|
||||
input_dir / "mrcr_8k_2_4_8needle_n10_per_config.jsonl",
|
||||
]
|
||||
payload = {
|
||||
"protocol": {
|
||||
"model": "Qwen3-1.7B tuned vanilla",
|
||||
"context_length": 8192,
|
||||
"arms": {
|
||||
"token": "L2A-style shared-gate: one native gate per token/layer",
|
||||
"token_kv_head": "AHA: one native gate per token/KV-head/layer",
|
||||
},
|
||||
"local_attention": {"sink_tokens": 64, "recent_tokens": 256},
|
||||
"inference": "strict AHA routing in both prefill and decode; no force-full heads or full-decode fallback",
|
||||
"thresholds": list(THRESHOLDS),
|
||||
"quality_rule": "every suite macro score >= 95% of aligned tuned-vanilla macro score",
|
||||
"headline_rule": "highest decision-weighted measured sparsity among thresholds passing every suite",
|
||||
"sparsity_definition": "hard local routes / token x KV-head x layer effective decisions",
|
||||
"suite_scope": "RULER local full13 means this repository's fixed 13-config set, not complete upstream RULER",
|
||||
},
|
||||
"score_definitions": SCORE_DEFINITIONS,
|
||||
"headline": headline,
|
||||
"points": points,
|
||||
"curves": curve_rows,
|
||||
"per_config": per_config,
|
||||
"alignment": {"all_aligned": all(value["aligned"] for value in alignment.values()), "methods": alignment},
|
||||
"appendix_note": (
|
||||
"Upstream L2A results are not merged into this table because architecture, "
|
||||
"training data, objective, and evaluation protocols differ."
|
||||
),
|
||||
"reproducibility": {
|
||||
"git_commit": command_output(["git", "rev-parse", "HEAD"], args.repo),
|
||||
"git_status_short": command_output(["git", "status", "--short"], args.repo),
|
||||
"container": "dmtd-repro",
|
||||
"container_image": command_output(
|
||||
["docker", "inspect", "-f", "{{.Config.Image}}", "dmtd-repro"], args.repo
|
||||
),
|
||||
"baseline_root": str(args.baseline_root),
|
||||
"eval_root": str(args.eval_root),
|
||||
"frozen_input_sha256": {
|
||||
str(path): sha256(path) for path in frozen_inputs if path.exists()
|
||||
},
|
||||
},
|
||||
}
|
||||
(output / "router_granularity_results.json").write_text(
|
||||
json.dumps(payload, indent=2) + "\n"
|
||||
)
|
||||
|
||||
flat_curve_rows = [
|
||||
{key: value for key, value in row.items() if key != "by_phase"}
|
||||
for row in curve_rows
|
||||
]
|
||||
with (output / "quality_sparsity_curves.csv").open("w", newline="") as handle:
|
||||
writer = csv.DictWriter(handle, fieldnames=list(flat_curve_rows[0]))
|
||||
writer.writeheader()
|
||||
writer.writerows(flat_curve_rows)
|
||||
with (output / "per_config_scores.csv").open("w", newline="") as handle:
|
||||
writer = csv.DictWriter(handle, fieldnames=list(per_config[0]))
|
||||
writer.writeheader()
|
||||
writer.writerows(per_config)
|
||||
|
||||
figure, axes = plt.subplots(2, 2, figsize=(12, 9), constrained_layout=True)
|
||||
colors = {"token": "#d62728", "token_kv_head": "#1f77b4"}
|
||||
labels = {"token": "L2A-style shared-gate", "token_kv_head": "AHA KV-head gate"}
|
||||
for axis, suite in zip(axes.flat, SUITES):
|
||||
for arm in ARMS:
|
||||
values = sorted(
|
||||
(row for row in curve_rows if row["arm"] == arm and row["benchmark"] == suite),
|
||||
key=lambda row: row["sparsity"],
|
||||
)
|
||||
axis.plot(
|
||||
[100 * row["sparsity"] for row in values],
|
||||
[100 * row["retention"] for row in values],
|
||||
marker="o", color=colors[arm], label=labels[arm],
|
||||
)
|
||||
for row in values:
|
||||
axis.annotate(f"{row['threshold']:.3g}", (100 * row["sparsity"], 100 * row["retention"]), fontsize=7)
|
||||
axis.axhline(95, color="black", linestyle="--", linewidth=1)
|
||||
axis.set_title(suite)
|
||||
axis.set_xlabel("Measured effective sparsity (%)")
|
||||
axis.set_ylabel("Retention vs tuned vanilla (%)")
|
||||
axis.grid(alpha=0.25)
|
||||
axes.flat[0].legend()
|
||||
figure.suptitle("Qwen3-1.7B matched router-granularity quality–sparsity curves")
|
||||
figure.savefig(output / "quality_sparsity_curves.png", dpi=180)
|
||||
plt.close(figure)
|
||||
print(json.dumps(headline, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
4
recipe/scripts/train_aha.sh
Normal file
4
recipe/scripts/train_aha.sh
Normal file
@@ -0,0 +1,4 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
exec "$SCRIPT_DIR/train_one.sh" aha "${1:-0}"
|
||||
30
recipe/scripts/train_both.sh
Normal file
30
recipe/scripts/train_both.sh
Normal file
@@ -0,0 +1,30 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
RECIPE="$(cd "$SCRIPT_DIR/.." && pwd)"
|
||||
GPU_AHA="${GPU_AHA:-0}"
|
||||
GPU_L2A="${GPU_L2A:-1}"
|
||||
cd "$RECIPE"
|
||||
|
||||
echo "AHA = per-token, per-KV-head gate; assigned visible GPU $GPU_AHA."
|
||||
echo "L2A-style = per-token, head-shared gate; assigned visible GPU $GPU_L2A."
|
||||
echo "Each variant is a separate world-size-1 run with global batch size 1."
|
||||
|
||||
python "$SCRIPT_DIR/validate_release.py"
|
||||
python -m unittest -v tests.test_router_granularity
|
||||
|
||||
if [[ "$GPU_AHA" == "$GPU_L2A" ]]; then
|
||||
"$SCRIPT_DIR/train_one.sh" aha "$GPU_AHA"
|
||||
"$SCRIPT_DIR/train_one.sh" l2a_style "$GPU_L2A"
|
||||
else
|
||||
"$SCRIPT_DIR/train_one.sh" aha "$GPU_AHA" & aha_pid=$!
|
||||
"$SCRIPT_DIR/train_one.sh" l2a_style "$GPU_L2A" & l2a_pid=$!
|
||||
status=0
|
||||
wait "$aha_pid" || status=1
|
||||
wait "$l2a_pid" || status=1
|
||||
[[ "$status" == 0 ]] || exit "$status"
|
||||
fi
|
||||
|
||||
echo "AHA selected checkpoint: ${OUTPUT_ROOT:-$RECIPE/../outputs}/aha/stage2/checkpoint-25"
|
||||
echo "L2A-style selected checkpoint: ${OUTPUT_ROOT:-$RECIPE/../outputs}/l2a_style/stage2/checkpoint-25"
|
||||
4
recipe/scripts/train_l2a_style.sh
Normal file
4
recipe/scripts/train_l2a_style.sh
Normal file
@@ -0,0 +1,4 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
exec "$SCRIPT_DIR/train_one.sh" l2a_style "${1:-0}"
|
||||
127
recipe/scripts/train_one.sh
Normal file
127
recipe/scripts/train_one.sh
Normal file
@@ -0,0 +1,127 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
if [[ $# -lt 1 || $# -gt 2 ]]; then
|
||||
echo "Usage: $0 aha|l2a_style [GPU_ID]" >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
ARM="$1"
|
||||
GPU="${2:-0}"
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
RECIPE="$(cd "$SCRIPT_DIR/.." && pwd)"
|
||||
REPO="$(cd "$RECIPE/.." && pwd)"
|
||||
VANILLA_DIR="${VANILLA_DIR:-$REPO}"
|
||||
DATA_DIR="${DATA_DIR:-$RECIPE/data/am_distilled_long_mix}"
|
||||
OUTPUT_ROOT="${OUTPUT_ROOT:-$REPO/outputs}"
|
||||
|
||||
case "$ARM" in
|
||||
aha) GRANULARITY=token_kv_head ;;
|
||||
l2a_style) GRANULARITY=token ;;
|
||||
*) echo "Unknown arm: $ARM (expected aha or l2a_style)" >&2; exit 2 ;;
|
||||
esac
|
||||
|
||||
ARM_OUT="$OUTPUT_ROOT/$ARM"
|
||||
HOTSTART="$ARM_OUT/hotstart"
|
||||
STAGE1="$ARM_OUT/stage1"
|
||||
STAGE2="$ARM_OUT/stage2"
|
||||
LOG_DIR="$ARM_OUT/logs"
|
||||
mkdir -p "$LOG_DIR"
|
||||
|
||||
export CUDA_VISIBLE_DEVICES="$GPU"
|
||||
export PYTHONUNBUFFERED=1
|
||||
export TOKENIZERS_PARALLELISM=false
|
||||
export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"
|
||||
|
||||
if [[ ! -f "$VANILLA_DIR/model.safetensors" ]]; then
|
||||
echo "Missing tuned-vanilla checkpoint: $VANILLA_DIR/model.safetensors" >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
if [[ ! -f "$HOTSTART/config.json" ]]; then
|
||||
python "$SCRIPT_DIR/export_tuned_vanilla_aha_hotstart.py" \
|
||||
--vanilla-path "$VANILLA_DIR" \
|
||||
--output-path "$HOTSTART" \
|
||||
--window-size 128 \
|
||||
--local-kind sink_recent \
|
||||
--gate-init-full-prob 0.90 \
|
||||
--router-granularity "$GRANULARITY" \
|
||||
2>&1 | tee "$LOG_DIR/hotstart.log"
|
||||
fi
|
||||
|
||||
if [[ ! -f "$STAGE1/checkpoint-300/config.json" ]]; then
|
||||
AHA_TRAIN_GATE_HARD_THRESHOLD=0.50 \
|
||||
python "$RECIPE/dynamic_duo_train.py" \
|
||||
--aha_checkpoint "$HOTSTART" \
|
||||
--model_path "$VANILLA_DIR" \
|
||||
--output_dir "$STAGE1" \
|
||||
--data_source am_distilled \
|
||||
--am_dataset_path "$DATA_DIR" \
|
||||
--am_dataset_split train \
|
||||
--am_label_mode full \
|
||||
--max_length 8192 \
|
||||
--num_steps 300 \
|
||||
--warmup_ratio 0.10 \
|
||||
--lr 3e-5 \
|
||||
--reg_weight 0.1 \
|
||||
--ce_weight 0.0 \
|
||||
--batch_size 1 \
|
||||
--grad_accum 1 \
|
||||
--save_steps 100 \
|
||||
--log_steps 10 \
|
||||
--seed 42 \
|
||||
--dtype bfloat16 \
|
||||
--attn_impl sdpa \
|
||||
--aha_local_kind sink_recent \
|
||||
--router_granularity "$GRANULARITY" \
|
||||
2>&1 | tee "$LOG_DIR/stage1.log"
|
||||
fi
|
||||
|
||||
if [[ ! -f "$STAGE2/checkpoint-75/config.json" ]]; then
|
||||
MODEL_PATH="$VANILLA_DIR" \
|
||||
AHA_CHECKPOINT_PATH="$STAGE1/checkpoint-300" \
|
||||
DATASET_PATH="$DATA_DIR" \
|
||||
DATASET_SPLIT=train \
|
||||
OUTPUT_DIR="$STAGE2" \
|
||||
AHA_MODE=dynamic \
|
||||
AHA_ROUTER_GRANULARITY="$GRANULARITY" \
|
||||
AHA_LOCAL_KIND=sink_recent \
|
||||
AHA_CE_WEIGHT=1.0 \
|
||||
AHA_DISTILL_WEIGHT=0.5 \
|
||||
AHA_REG_WEIGHT=0.01 \
|
||||
AHA_TRAIN_GATE_HARD_THRESHOLD=0.58 \
|
||||
GROUPED_LR=1 \
|
||||
GATE_ONLY=0 \
|
||||
LEARNING_RATE=3e-6 \
|
||||
GATE_LEARNING_RATE=3e-6 \
|
||||
BACKBONE_LEARNING_RATE=3e-7 \
|
||||
LR_SCHEDULER_TYPE=constant_with_warmup \
|
||||
WARMUP_RATIO=0.10 \
|
||||
WEIGHT_DECAY=0.0 \
|
||||
MAX_SEQ_LENGTH=8192 \
|
||||
PER_DEVICE_TRAIN_BATCH_SIZE=1 \
|
||||
GRADIENT_ACCUMULATION_STEPS=1 \
|
||||
MAX_STEPS=75 \
|
||||
LOGGING_STEPS=5 \
|
||||
SAVE_STEPS=25 \
|
||||
SAVE_TOTAL_LIMIT=3 \
|
||||
FREEZE_EMBEDDINGS_LM_HEAD=1 \
|
||||
REPORT_TO=none \
|
||||
SEED=47 \
|
||||
python "$RECIPE/sft.py" \
|
||||
2>&1 | tee "$LOG_DIR/stage2.log"
|
||||
fi
|
||||
|
||||
python - "$STAGE2/checkpoint-25" "$GRANULARITY" <<'PY'
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
checkpoint = Path(sys.argv[1])
|
||||
expected = sys.argv[2]
|
||||
config = json.loads((checkpoint / "config.json").read_text())
|
||||
actual = config.get("aha_router_granularity", "token_kv_head")
|
||||
if actual != expected:
|
||||
raise RuntimeError(f"checkpoint granularity {actual!r} != {expected!r}")
|
||||
print(f"selected_checkpoint={checkpoint} router_granularity={actual}")
|
||||
PY
|
||||
72
recipe/scripts/validate_release.py
Normal file
72
recipe/scripts/validate_release.py
Normal file
@@ -0,0 +1,72 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Fail fast if a downloaded release is incomplete or has drifted."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
from collections import Counter
|
||||
from pathlib import Path
|
||||
|
||||
from datasets import load_from_disk
|
||||
|
||||
|
||||
RECIPE = Path(__file__).resolve().parents[1]
|
||||
REPO = RECIPE.parent
|
||||
|
||||
|
||||
def sha256(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as handle:
|
||||
for block in iter(lambda: handle.read(8 * 1024 * 1024), b""):
|
||||
digest.update(block)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def read_jsonl(path: Path) -> list[dict]:
|
||||
return [json.loads(line) for line in path.read_text().splitlines() if line.strip()]
|
||||
|
||||
|
||||
def main() -> None:
|
||||
manifest = json.loads((RECIPE / "manifest.json").read_text())
|
||||
checked = []
|
||||
for relative, expected in manifest["files"].items():
|
||||
path = (RECIPE / relative).resolve()
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(path)
|
||||
if path.name == "model.safetensors" and os.environ.get("SKIP_LARGE_HASH") == "1":
|
||||
continue
|
||||
actual = sha256(path)
|
||||
if actual != expected:
|
||||
raise RuntimeError(f"SHA256 mismatch for {path}: {actual} != {expected}")
|
||||
checked.append(str(path.relative_to(REPO)))
|
||||
|
||||
config = json.loads((REPO / "config.json").read_text())
|
||||
if config.get("model_type") != "qwen3":
|
||||
raise RuntimeError(f"unexpected tuned-vanilla model_type: {config.get('model_type')}")
|
||||
dataset = load_from_disk(str(RECIPE / "data/am_distilled_long_mix"))
|
||||
if len(dataset["train"]) != 1024:
|
||||
raise RuntimeError(f"expected 1024 training rows, found {len(dataset['train'])}")
|
||||
|
||||
helmet = read_jsonl(RECIPE / "data/eval_inputs/helmet_icl_8k_n50_per_config.jsonl")
|
||||
mrcr = read_jsonl(RECIPE / "data/eval_inputs/mrcr_8k_2_4_8needle_n10_per_config.jsonl")
|
||||
helmet_counts = Counter(row["config"] for row in helmet)
|
||||
mrcr_counts = Counter(row["config"] for row in mrcr)
|
||||
if sorted(helmet_counts.values()) != [50] * 5:
|
||||
raise RuntimeError(f"unexpected HELMET counts: {helmet_counts}")
|
||||
if sorted(mrcr_counts.values()) != [10] * 3:
|
||||
raise RuntimeError(f"unexpected MRCR counts: {mrcr_counts}")
|
||||
|
||||
print(json.dumps({
|
||||
"status": "ok",
|
||||
"checked_sha256": checked,
|
||||
"model_type": config["model_type"],
|
||||
"training_rows": len(dataset["train"]),
|
||||
"helmet_rows": len(helmet),
|
||||
"mrcr_rows": len(mrcr),
|
||||
}, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
642
recipe/sft.py
Normal file
642
recipe/sft.py
Normal file
@@ -0,0 +1,642 @@
|
||||
"""
|
||||
AHA-Qwen3 SFT training script.
|
||||
|
||||
Loads pretrained Qwen3-0.6B, converts to AHA-Qwen3 (adds gate to q_proj),
|
||||
and trains on am-distilled data. Only gate weights are randomly initialized;
|
||||
all other weights come from the pretrained model.
|
||||
"""
|
||||
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import shutil
|
||||
import socket
|
||||
import subprocess
|
||||
from typing import Optional
|
||||
|
||||
import datasets
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from trl import SFTConfig, SFTTrainer
|
||||
from transformers import AutoTokenizer, TrainerCallback
|
||||
from transformers.trainer_utils import get_last_checkpoint
|
||||
|
||||
from modeling_aha_qwen3 import (
|
||||
AHA_ROUTER_GRANULARITY,
|
||||
AHAQwen3Config,
|
||||
AHAQwen3ForCausalLM,
|
||||
aha_router_output_size,
|
||||
)
|
||||
from router_training_utils import RowWiseAdamW, configure_gate_only
|
||||
|
||||
|
||||
class AHASFTTrainer(SFTTrainer):
|
||||
"""Accumulate CE vs gate-aux vs distill breakdown and gate density; log with HF `loss` (total)."""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
# ``AHAQwen3ForCausalLM.forward`` accepts ``**kwargs`` for attention
|
||||
# backends, but its CE loss is an ordinary local token mean and does
|
||||
# not consume Trainer's ``num_items_in_batch``. Transformers 4.57
|
||||
# otherwise mistakes the variadic signature for a globally normalized
|
||||
# loss, multiplies it by world size, and skips the normal gradient-
|
||||
# accumulation division. Mark the actual loss contract explicitly so
|
||||
# lr=3e-5 has the same meaning at every world size.
|
||||
self.model_accepts_loss_kwargs = False
|
||||
self._aha_ce_sum = 0.0
|
||||
self._aha_gate_aux_sum = 0.0
|
||||
self._aha_distill_sum = 0.0
|
||||
self._aha_gate_soft_sum = 0.0
|
||||
self._aha_gate_hard_sum = 0.0
|
||||
self._aha_metric_count = 0
|
||||
# Cached for loss_total aggregation in log().
|
||||
self._aha_distill_weight = float(
|
||||
getattr(self.model.config, "aha_distill_weight", 0.0)
|
||||
)
|
||||
self._aha_ce_weight = float(
|
||||
getattr(self.model.config, "aha_ce_weight", 1.0)
|
||||
)
|
||||
|
||||
def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
|
||||
loss, outputs = super().compute_loss(
|
||||
model, inputs, return_outputs=True, num_items_in_batch=num_items_in_batch
|
||||
)
|
||||
if getattr(outputs, "ce_loss", None) is not None:
|
||||
self._aha_metric_count += 1
|
||||
self._aha_ce_sum += float(outputs.ce_loss.detach().float().mean().cpu())
|
||||
if getattr(outputs, "gate_aux_loss", None) is not None:
|
||||
self._aha_gate_aux_sum += float(outputs.gate_aux_loss.detach().float().mean().cpu())
|
||||
if getattr(outputs, "distill_loss", None) is not None:
|
||||
self._aha_distill_sum += float(outputs.distill_loss.detach().float().mean().cpu())
|
||||
if getattr(outputs, "gate_soft_mean", None) is not None:
|
||||
self._aha_gate_soft_sum += float(outputs.gate_soft_mean.detach().float().mean().cpu())
|
||||
if getattr(outputs, "gate_hard_mean", None) is not None:
|
||||
self._aha_gate_hard_sum += float(outputs.gate_hard_mean.detach().float().mean().cpu())
|
||||
if return_outputs:
|
||||
return (loss, outputs)
|
||||
return loss
|
||||
|
||||
def _distributed_mean_of_sums(self, sum_val: float, count: int) -> float:
|
||||
if count <= 0:
|
||||
return float("nan")
|
||||
device = self.accelerator.device
|
||||
if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1:
|
||||
t = torch.tensor([sum_val, float(count)], device=device, dtype=torch.float64)
|
||||
dist.all_reduce(t, op=dist.ReduceOp.SUM)
|
||||
return (t[0] / t[1]).item()
|
||||
return sum_val / float(count)
|
||||
|
||||
def log(self, logs: dict[str, float], start_time: Optional[float] = None) -> None:
|
||||
if self._aha_metric_count > 0:
|
||||
n = self._aha_metric_count
|
||||
ce_m = self._distributed_mean_of_sums(self._aha_ce_sum, n)
|
||||
aux_m = self._distributed_mean_of_sums(self._aha_gate_aux_sum, n)
|
||||
distill_m = self._distributed_mean_of_sums(self._aha_distill_sum, n)
|
||||
gs_m = self._distributed_mean_of_sums(self._aha_gate_soft_sum, n)
|
||||
gh_m = self._distributed_mean_of_sums(self._aha_gate_hard_sum, n)
|
||||
if not math.isnan(ce_m):
|
||||
logs["ce_loss"] = round(ce_m, 4)
|
||||
if not math.isnan(aux_m):
|
||||
logs["gate_aux_loss"] = round(aux_m, 6)
|
||||
if not math.isnan(distill_m) and distill_m != 0.0:
|
||||
logs["distill_loss"] = round(distill_m, 6)
|
||||
# loss_total = aha_ce_weight * ce + gate_aux
|
||||
# + aha_distill_weight * distill_loss
|
||||
# (distill_m / ce_m are unweighted; apply weights here so the
|
||||
# sum matches the scalar actually added to total loss.)
|
||||
if not math.isnan(ce_m) and not math.isnan(aux_m):
|
||||
lt = self._aha_ce_weight * ce_m + aux_m
|
||||
if not math.isnan(distill_m):
|
||||
lt = lt + self._aha_distill_weight * distill_m
|
||||
logs["loss_total"] = round(lt, 4)
|
||||
if "loss" in logs and lt > 1e-8:
|
||||
# With ``model_accepts_loss_kwargs=False``, HF's logging
|
||||
# window and the raw model loss should agree up to normal
|
||||
# batch-to-batch weighting differences. A large integer
|
||||
# ratio is a regression in Trainer loss normalization.
|
||||
logs["hf_loss_ratio"] = round(logs["loss"] / lt, 2)
|
||||
if not math.isnan(gs_m):
|
||||
logs["gate_soft_mean"] = round(gs_m, 4)
|
||||
if not math.isnan(gh_m):
|
||||
logs["gate_hard_mean"] = round(gh_m, 4)
|
||||
self._aha_ce_sum = 0.0
|
||||
self._aha_gate_aux_sum = 0.0
|
||||
self._aha_distill_sum = 0.0
|
||||
self._aha_gate_soft_sum = 0.0
|
||||
self._aha_gate_hard_sum = 0.0
|
||||
self._aha_metric_count = 0
|
||||
return super().log(logs, start_time)
|
||||
|
||||
|
||||
class AHAManifestCallback(TrainerCallback):
|
||||
"""Copy the immutable run manifest into each Trainer checkpoint."""
|
||||
|
||||
def __init__(self, manifest_path: str):
|
||||
self.manifest_path = manifest_path
|
||||
|
||||
def on_save(self, args, state, control, **kwargs):
|
||||
if state.is_world_process_zero and os.path.exists(self.manifest_path):
|
||||
checkpoint_dir = os.path.join(
|
||||
args.output_dir, f"checkpoint-{state.global_step}"
|
||||
)
|
||||
if os.path.isdir(checkpoint_dir):
|
||||
shutil.copyfile(
|
||||
self.manifest_path,
|
||||
os.path.join(checkpoint_dir, "aha_training_manifest.json"),
|
||||
)
|
||||
return control
|
||||
|
||||
|
||||
def get_optional_int(name: str) -> Optional[int]:
|
||||
value = os.environ.get(name)
|
||||
return None if value in {None, ""} else int(value)
|
||||
|
||||
|
||||
def resolve_resume_checkpoint(output_dir: str) -> Optional[str]:
|
||||
resume_from_checkpoint = os.environ.get("RESUME_FROM_CHECKPOINT")
|
||||
if resume_from_checkpoint in {None, ""}:
|
||||
return None
|
||||
if resume_from_checkpoint == "latest":
|
||||
if not os.path.isdir(output_dir):
|
||||
return None
|
||||
return get_last_checkpoint(output_dir)
|
||||
return resume_from_checkpoint
|
||||
|
||||
|
||||
def parse_fsdp_options(name: str) -> list[str]:
|
||||
value = os.environ.get(name, "").replace(",", " ").strip()
|
||||
return [item for item in value.split() if item]
|
||||
|
||||
|
||||
def parse_bool(name: str, default: bool = False) -> bool:
|
||||
value = os.environ.get(name)
|
||||
if value is None:
|
||||
return default
|
||||
return value.lower() in {"1", "true", "yes", "on"}
|
||||
|
||||
|
||||
def parse_report_to() -> list[str]:
|
||||
"""Comma-separated integrations, e.g. REPORT_TO=wandb or wandb,tensorboard. Empty / none / off -> []."""
|
||||
raw = os.environ.get("REPORT_TO", "").strip()
|
||||
if not raw or raw.lower() in ("none", "off"):
|
||||
return []
|
||||
return [x.strip() for x in raw.split(",") if x.strip()]
|
||||
|
||||
|
||||
def build_grouped_lr_optimizer(model: torch.nn.Module) -> torch.optim.Optimizer:
|
||||
"""Build an AdamW optimizer with separate *effective* LR for gate rows.
|
||||
|
||||
Dynamic-mode gate logits live in the final router rows of
|
||||
each attention ``q_proj`` instead of in a separate Parameter. PyTorch
|
||||
optimizer groups cannot split one Parameter by row, so q_proj tensors are
|
||||
placed in the gate-lr group and :class:`RowWiseAdamW` scales their realized
|
||||
Q-row updates by ``backbone_lr / gate_lr``. Other trainable parameters use
|
||||
the backbone-lr group directly.
|
||||
"""
|
||||
if getattr(model.config, "aha_mode", "dynamic") != "dynamic":
|
||||
raise ValueError("GROUPED_LR currently only supports AHA_MODE=dynamic")
|
||||
if gate_learning_rate <= 0.0 or backbone_learning_rate < 0.0:
|
||||
raise ValueError(
|
||||
"GATE_LEARNING_RATE must be positive and BACKBONE_LEARNING_RATE must be non-negative"
|
||||
)
|
||||
|
||||
num_heads = model.config.num_attention_heads
|
||||
head_dim = getattr(model.config, "head_dim", model.config.hidden_size // num_heads)
|
||||
q_rows = num_heads * head_dim
|
||||
q_row_scale = backbone_learning_rate / gate_learning_rate
|
||||
|
||||
gate_params = []
|
||||
gate_param_ids = set()
|
||||
|
||||
row_scales = []
|
||||
for layer in model.model.layers:
|
||||
q_proj = layer.self_attn.q_proj
|
||||
for p in (q_proj.weight, q_proj.bias):
|
||||
if p is None or not p.requires_grad:
|
||||
continue
|
||||
gate_params.append(p)
|
||||
gate_param_ids.add(id(p))
|
||||
row_scales.append((p, q_rows, q_row_scale))
|
||||
|
||||
backbone_params = [
|
||||
p for p in model.parameters()
|
||||
if p.requires_grad and id(p) not in gate_param_ids
|
||||
]
|
||||
if not gate_params:
|
||||
raise ValueError("GROUPED_LR found no trainable q_proj gate parameters")
|
||||
|
||||
gate_n = sum(p.numel() for p in gate_params)
|
||||
effective_gate_n = sum(
|
||||
layer.self_attn.aha_router_outputs
|
||||
* (layer.self_attn.q_proj.in_features + (layer.self_attn.q_proj.bias is not None))
|
||||
for layer in model.model.layers
|
||||
)
|
||||
backbone_n = sum(p.numel() for p in backbone_params)
|
||||
print(
|
||||
" GROUPED_LR optimizer: "
|
||||
f"gate_lr={gate_learning_rate:.3e}, backbone_lr={backbone_learning_rate:.3e}, "
|
||||
f"q_row_update_scale={q_row_scale:.3g}, "
|
||||
f"q_proj_params={gate_n:,}, effective_gate_params={effective_gate_n:,}, "
|
||||
f"other_trainable={backbone_n:,}"
|
||||
)
|
||||
|
||||
param_groups = [{"params": gate_params, "lr": gate_learning_rate}]
|
||||
if backbone_params:
|
||||
param_groups.append({"params": backbone_params, "lr": backbone_learning_rate})
|
||||
return RowWiseAdamW(
|
||||
param_groups,
|
||||
row_scales=row_scales,
|
||||
weight_decay=weight_decay,
|
||||
betas=(adam_beta1, adam_beta2),
|
||||
)
|
||||
|
||||
|
||||
# === Configuration (override via environment variables) ===
|
||||
model_path = os.environ.get("MODEL_PATH", "/workspace/AHA/models/Qwen3-0.6B")
|
||||
aha_checkpoint_path = os.environ.get("AHA_CHECKPOINT_PATH", "")
|
||||
dataset_path = os.environ.get("DATASET_PATH", "/workspace/Direct-Multitoken-Decoding/am-distilled-8192")
|
||||
dataset_split = os.environ.get("DATASET_SPLIT", "train")
|
||||
output_dir = os.environ.get("OUTPUT_DIR", "/workspace/AHA/AHA-Qwen3/ckpts/aha_qwen3_w1024")
|
||||
aha_window_size = int(os.environ.get("AHA_WINDOW_SIZE", "1024"))
|
||||
aha_lambda = float(os.environ.get("AHA_LAMBDA", "3e-4"))
|
||||
aha_distill_weight = float(os.environ.get("AHA_DISTILL_WEIGHT", "0.0"))
|
||||
# Language-modeling CE weight. Default 1.0 = standard SFT. Set to 0.0 to
|
||||
# drop CE from the training objective and shape the gate purely via
|
||||
# ``aux + distill`` (useful with GATE_ONLY + frozen backbone).
|
||||
aha_ce_weight = float(os.environ.get("AHA_CE_WEIGHT", "1.0"))
|
||||
# Hinge target on mean(gate_soft). Aux loss is ``λ · max(0, ḡ - τ)``.
|
||||
# Default 1.0 keeps the legacy unconditional aux (τ=1 makes the clamp
|
||||
# vacuous). Set e.g. 0.15 to cap global-attention density at 15% of
|
||||
# tokens in steady state -- this is the safety floor that prevents
|
||||
# gate collapse when training without CE. See AHAQwen3Config.
|
||||
aha_gate_target = float(os.environ.get("AHA_GATE_TARGET", "1.0"))
|
||||
# Direct Duo-style sparsity regularizer. When set to a non-negative value,
|
||||
# the AHA loss uses ``AHA_REG_WEIGHT * mean(gate_soft)`` instead of the legacy
|
||||
# hinge controlled by AHA_LAMBDA/AHA_GATE_TARGET.
|
||||
aha_reg_weight = float(os.environ.get("AHA_REG_WEIGHT", "-1.0"))
|
||||
# AHA mode: "dynamic" (per-(token, kv_head) MLP gate) or "duo" (per-(layer,
|
||||
# kv_head) static scalar, DuoAttention-style). DUO mode ignores GATE_ONLY
|
||||
# and instead uses its own ``DUO_ALPHA_ONLY`` freeze path.
|
||||
aha_mode = os.environ.get("AHA_MODE", "dynamic")
|
||||
aha_router_granularity_request = os.environ.get(
|
||||
"AHA_ROUTER_GRANULARITY", ""
|
||||
).strip()
|
||||
if aha_router_granularity_request and aha_router_granularity_request not in {
|
||||
"token",
|
||||
"token_kv_head",
|
||||
}:
|
||||
raise ValueError(
|
||||
"AHA_ROUTER_GRANULARITY must be 'token' or 'token_kv_head', got "
|
||||
f"{aha_router_granularity_request!r}"
|
||||
)
|
||||
# Override aha_local_kind on a loaded ckpt: "sink_recent" or
|
||||
# "sliding_window". Useful for ablating sink dependence in dynamic mode
|
||||
# from the same hot-start ckpt without retraining the source. Empty
|
||||
# string keeps the ckpt's existing config value.
|
||||
aha_local_kind_override = os.environ.get("AHA_LOCAL_KIND", "").strip()
|
||||
duo_sink_size = int(os.environ.get("DUO_SINK_SIZE", "64"))
|
||||
duo_recent_size = int(os.environ.get("DUO_RECENT_SIZE", "256"))
|
||||
duo_alpha_init = float(os.environ.get("DUO_ALPHA_INIT", "1.0"))
|
||||
duo_alpha_only = parse_bool("DUO_ALPHA_ONLY", default=(aha_mode == "duo"))
|
||||
gate_only = parse_bool("GATE_ONLY", default=False)
|
||||
grouped_lr = parse_bool("GROUPED_LR", default=False)
|
||||
max_seq_length = int(os.environ.get("MAX_SEQ_LENGTH", "8192"))
|
||||
min_train_tokens = get_optional_int("MIN_TRAIN_TOKENS")
|
||||
max_train_tokens = get_optional_int("MAX_TRAIN_TOKENS")
|
||||
max_train_samples = get_optional_int("MAX_TRAIN_SAMPLES")
|
||||
learning_rate = float(os.environ.get("LEARNING_RATE", "3e-5"))
|
||||
gate_learning_rate = float(os.environ.get("GATE_LEARNING_RATE", str(learning_rate)))
|
||||
backbone_learning_rate = float(os.environ.get("BACKBONE_LEARNING_RATE", str(learning_rate)))
|
||||
lr_scheduler_type = os.environ.get("LR_SCHEDULER_TYPE", "cosine")
|
||||
warmup_ratio = float(os.environ.get("WARMUP_RATIO", "0.03"))
|
||||
adam_beta1 = float(os.environ.get("ADAM_BETA1", "0.9"))
|
||||
adam_beta2 = float(os.environ.get("ADAM_BETA2", "0.999"))
|
||||
weight_decay = float(os.environ.get("WEIGHT_DECAY", "0.0"))
|
||||
max_grad_norm = float(os.environ.get("MAX_GRAD_NORM", "1.0"))
|
||||
per_device_batch_size = int(os.environ.get("PER_DEVICE_TRAIN_BATCH_SIZE", "2"))
|
||||
gradient_accumulation_steps = int(os.environ.get("GRADIENT_ACCUMULATION_STEPS", "16"))
|
||||
num_train_epochs = float(os.environ.get("NUM_TRAIN_EPOCHS", "1"))
|
||||
max_steps = int(os.environ.get("MAX_STEPS", "-1")) # -1 disables the cap; positive overrides num_train_epochs
|
||||
logging_steps = int(os.environ.get("LOGGING_STEPS", "10"))
|
||||
save_steps = int(os.environ.get("SAVE_STEPS", "500"))
|
||||
save_total_limit = int(os.environ.get("SAVE_TOTAL_LIMIT", "2"))
|
||||
save_strategy = os.environ.get("SAVE_STRATEGY", "steps")
|
||||
seed = int(os.environ.get("SEED", "42"))
|
||||
resume_from_checkpoint = resolve_resume_checkpoint(output_dir)
|
||||
report_to = parse_report_to()
|
||||
wandb_run_name = os.environ.get("WANDB_RUN_NAME") or os.environ.get("RUN_NAME") or os.path.basename(output_dir.rstrip("/"))
|
||||
fsdp_options = parse_fsdp_options("FSDP_OPTIONS")
|
||||
fsdp_transformer_layer_cls = os.environ.get("FSDP_TRANSFORMER_LAYER_CLS_TO_WRAP", "")
|
||||
fsdp_activation_checkpointing = parse_bool("FSDP_ACTIVATION_CHECKPOINTING", default=True)
|
||||
freeze_embeddings_lm_head = parse_bool("FREEZE_EMBEDDINGS_LM_HEAD", default=True)
|
||||
|
||||
|
||||
def main():
|
||||
AHAQwen3Config.register_for_auto_class()
|
||||
AHAQwen3ForCausalLM.register_for_auto_class("AutoModelForCausalLM")
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
print(f"Loading Qwen3 from {model_path} and converting to AHA-Qwen3...")
|
||||
if aha_checkpoint_path:
|
||||
print(f" aha_checkpoint_path={aha_checkpoint_path}")
|
||||
print(
|
||||
f" aha_mode={aha_mode}, "
|
||||
f"aha_router_granularity={aha_router_granularity_request or 'checkpoint/default'}, "
|
||||
f"aha_window_size={aha_window_size}, "
|
||||
f"aha_lambda={aha_lambda}, aha_distill_weight={aha_distill_weight}, "
|
||||
f"aha_ce_weight={aha_ce_weight}, aha_gate_target={aha_gate_target}, "
|
||||
f"aha_reg_weight={aha_reg_weight}, "
|
||||
f"gate_only={gate_only}, duo_alpha_only={duo_alpha_only}, grouped_lr={grouped_lr}"
|
||||
)
|
||||
if aha_mode == "duo":
|
||||
print(
|
||||
f" duo_sink_size={duo_sink_size}, duo_recent_size={duo_recent_size}, "
|
||||
f"duo_alpha_init={duo_alpha_init}"
|
||||
)
|
||||
print(f" dataset={dataset_path}[{dataset_split}]")
|
||||
print(f" output_dir={output_dir}")
|
||||
print(f" max_seq_length={max_seq_length}")
|
||||
print(
|
||||
" optimizer/schedule: "
|
||||
f"lr={learning_rate:.3e}, scheduler={lr_scheduler_type}, warmup_ratio={warmup_ratio}, "
|
||||
f"betas=({adam_beta1}, {adam_beta2}), weight_decay={weight_decay}, "
|
||||
f"max_grad_norm={max_grad_norm}"
|
||||
)
|
||||
world_size = int(os.environ.get("WORLD_SIZE", "1"))
|
||||
effective_batch_size = per_device_batch_size * gradient_accumulation_steps * world_size
|
||||
print(
|
||||
" batch: "
|
||||
f"per_device={per_device_batch_size}, grad_accum={gradient_accumulation_steps}, "
|
||||
f"world_size={world_size}, effective_global={effective_batch_size}"
|
||||
)
|
||||
if min_train_tokens is not None or max_train_tokens is not None:
|
||||
print(f" token_filter=[{min_train_tokens}, {max_train_tokens}]")
|
||||
if max_train_samples is not None:
|
||||
print(f" max_train_samples={max_train_samples}")
|
||||
if resume_from_checkpoint is not None:
|
||||
print(f" resume_from_checkpoint={resume_from_checkpoint}")
|
||||
if fsdp_options:
|
||||
print(f" fsdp_options={fsdp_options}")
|
||||
if fsdp_transformer_layer_cls:
|
||||
print(f" fsdp_transformer_layer_cls_to_wrap={fsdp_transformer_layer_cls}")
|
||||
print(f" report_to={report_to if report_to else 'none'}")
|
||||
if report_to and "wandb" in report_to:
|
||||
print(
|
||||
" wandb: set WANDB_PROJECT (required for cloud UI); optional WANDB_ENTITY, WANDB_RUN_GROUP, WANDB_TAGS. "
|
||||
"Run `wandb login` once. Metrics include train/ce_loss, train/gate_aux_loss, train/loss_total, train/loss (HF)."
|
||||
)
|
||||
|
||||
if aha_checkpoint_path:
|
||||
model = AHAQwen3ForCausalLM.from_pretrained_aha(
|
||||
aha_checkpoint_path,
|
||||
torch_dtype=torch.bfloat16,
|
||||
attn_implementation="sdpa",
|
||||
)
|
||||
# Config carried inside the checkpoint may predate aha_distill_weight /
|
||||
# aha_ce_weight / aha_gate_target; honor the env override either way.
|
||||
model.config.aha_distill_weight = aha_distill_weight
|
||||
model.config.aha_ce_weight = aha_ce_weight
|
||||
model.config.aha_lambda = aha_lambda
|
||||
model.config.aha_gate_target = aha_gate_target
|
||||
model.config.aha_reg_weight = aha_reg_weight
|
||||
loaded_granularity = getattr(
|
||||
model.config, "aha_router_granularity", AHA_ROUTER_GRANULARITY
|
||||
)
|
||||
if (
|
||||
aha_router_granularity_request
|
||||
and aha_router_granularity_request != loaded_granularity
|
||||
):
|
||||
raise ValueError(
|
||||
"AHA_ROUTER_GRANULARITY does not match checkpoint architecture: "
|
||||
f"requested={aha_router_granularity_request!r}, "
|
||||
f"checkpoint={loaded_granularity!r}"
|
||||
)
|
||||
if aha_local_kind_override:
|
||||
if aha_local_kind_override not in {"sink_recent", "sliding_window"}:
|
||||
raise ValueError(
|
||||
f"AHA_LOCAL_KIND must be 'sink_recent' or 'sliding_window', got {aha_local_kind_override!r}"
|
||||
)
|
||||
prev = getattr(model.config, "aha_local_kind", None)
|
||||
model.config.aha_local_kind = aha_local_kind_override
|
||||
print(f" AHA_LOCAL_KIND override: {prev!r} -> {aha_local_kind_override!r}")
|
||||
else:
|
||||
model = AHAQwen3ForCausalLM.from_pretrained_qwen3(
|
||||
model_path,
|
||||
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=aha_reg_weight,
|
||||
aha_mode=aha_mode,
|
||||
aha_router_granularity=(
|
||||
aha_router_granularity_request or AHA_ROUTER_GRANULARITY
|
||||
),
|
||||
duo_sink_size=duo_sink_size,
|
||||
duo_recent_size=duo_recent_size,
|
||||
duo_alpha_init=duo_alpha_init,
|
||||
torch_dtype=torch.bfloat16,
|
||||
attn_implementation="sdpa",
|
||||
)
|
||||
|
||||
# Historical Qwen experiments froze embeddings and lm_head. The rebuttal
|
||||
# scaling protocol disables this to match the OLMo-2 end-to-end recipe,
|
||||
# where the complete pretrained model and routers are fine-tuned jointly.
|
||||
if freeze_embeddings_lm_head:
|
||||
for param in model.model.embed_tokens.parameters():
|
||||
param.requires_grad = False
|
||||
for param in model.lm_head.parameters():
|
||||
param.requires_grad = False
|
||||
print(f" freeze_embeddings_lm_head={freeze_embeddings_lm_head}")
|
||||
|
||||
if duo_alpha_only:
|
||||
# DuoAttention-style: only the per-(layer, kv_head) scalar
|
||||
# ``full_attention_heads`` is trainable; everything else is
|
||||
# frozen. For Qwen3-0.6B this yields 224 trainable scalars.
|
||||
if getattr(model.config, "aha_mode", "dynamic") != "duo":
|
||||
raise ValueError(
|
||||
"DUO_ALPHA_ONLY=1 requires AHA_MODE=duo; otherwise there are no alpha params."
|
||||
)
|
||||
for param in model.parameters():
|
||||
param.requires_grad = False
|
||||
unfrozen = 0
|
||||
for layer in model.model.layers:
|
||||
p = layer.self_attn.full_attention_heads
|
||||
p.requires_grad = True
|
||||
unfrozen += p.numel()
|
||||
print(f" DUO_ALPHA_ONLY: trainable scalars {unfrozen:,}")
|
||||
elif gate_only:
|
||||
# Gate-only training: freeze everything, then unfreeze just the
|
||||
# gate rows of q_proj via gradient masks. q_proj.weight / bias
|
||||
# layout is [num_heads*head_dim Q rows, native router rows].
|
||||
setup = configure_gate_only(model)
|
||||
print(
|
||||
f" GATE_ONLY: granularity={model.config.aha_router_granularity}, "
|
||||
f"gate_rows/layer={setup.gate_rows}, "
|
||||
f"effective gate parameters {setup.effective_parameter_count:,} "
|
||||
"(Q rows grad-masked to 0)"
|
||||
)
|
||||
|
||||
n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||||
n_total = sum(p.numel() for p in model.parameters())
|
||||
print(f" Trainable: {n_trainable:,} / {n_total:,} ({n_trainable/n_total:.1%})")
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|
||||
|
||||
print(f"Loading dataset from {dataset_path}...")
|
||||
dataset_dict = datasets.load_from_disk(dataset_path)
|
||||
train_dataset = dataset_dict[dataset_split]
|
||||
print(f" loaded_rows={len(train_dataset):,}")
|
||||
|
||||
if min_train_tokens is not None:
|
||||
train_dataset = train_dataset.filter(lambda example: example["num_tokens"] >= min_train_tokens)
|
||||
print(f" rows_after_min_train_tokens={len(train_dataset):,}")
|
||||
if max_train_tokens is not None:
|
||||
train_dataset = train_dataset.filter(lambda example: example["num_tokens"] <= max_train_tokens)
|
||||
print(f" rows_after_max_train_tokens={len(train_dataset):,}")
|
||||
if max_train_samples is not None:
|
||||
keep = min(max_train_samples, len(train_dataset))
|
||||
train_dataset = train_dataset.shuffle(seed=seed).select(range(keep))
|
||||
print(f" rows_after_max_train_samples={len(train_dataset):,}")
|
||||
if len(train_dataset) == 0:
|
||||
raise ValueError("Training dataset is empty after filtering.")
|
||||
|
||||
if int(os.environ.get("RANK", "0")) == 0:
|
||||
try:
|
||||
git_commit = subprocess.check_output(
|
||||
["git", "rev-parse", "HEAD"], text=True
|
||||
).strip()
|
||||
git_status = subprocess.check_output(
|
||||
["git", "status", "--short"], text=True
|
||||
).splitlines()
|
||||
except (OSError, subprocess.CalledProcessError):
|
||||
git_commit, git_status = None, []
|
||||
manifest = {
|
||||
"schema": "aha-qwen3-training-manifest-v1",
|
||||
"git_commit": git_commit,
|
||||
"git_status": git_status,
|
||||
"host": socket.gethostname(),
|
||||
"model_path": str(model_path),
|
||||
"source_aha_checkpoint": str(aha_checkpoint_path) or None,
|
||||
"router_granularity": getattr(
|
||||
model.config, "aha_router_granularity", AHA_ROUTER_GRANULARITY
|
||||
),
|
||||
"native_gate_rows_per_layer": aha_router_output_size(model.config),
|
||||
"effective_gate_parameters": sum(
|
||||
layer.self_attn.aha_router_outputs
|
||||
* (
|
||||
layer.self_attn.q_proj.in_features
|
||||
+ (layer.self_attn.q_proj.bias is not None)
|
||||
)
|
||||
for layer in model.model.layers
|
||||
),
|
||||
"effective_sparsity_denominator": "token x KV-head x layer",
|
||||
"dataset": {
|
||||
"path": str(dataset_path),
|
||||
"split": dataset_split,
|
||||
"rows_after_filtering": len(train_dataset),
|
||||
"max_sequence_length": max_seq_length,
|
||||
},
|
||||
"training": {
|
||||
"gate_only": gate_only,
|
||||
"grouped_lr": grouped_lr,
|
||||
"learning_rate": learning_rate,
|
||||
"gate_learning_rate": gate_learning_rate,
|
||||
"backbone_learning_rate": backbone_learning_rate,
|
||||
"regularizer_weight": aha_reg_weight,
|
||||
"ce_weight": aha_ce_weight,
|
||||
"attention_distill_weight": aha_distill_weight,
|
||||
"train_gate_threshold": float(
|
||||
os.environ.get("AHA_TRAIN_GATE_HARD_THRESHOLD", "0.5")
|
||||
),
|
||||
"max_steps": max_steps,
|
||||
"seed": seed,
|
||||
"per_device_batch_size": per_device_batch_size,
|
||||
"gradient_accumulation_steps": gradient_accumulation_steps,
|
||||
"world_size": world_size,
|
||||
},
|
||||
}
|
||||
manifest_path = os.path.join(output_dir, "aha_training_manifest.json")
|
||||
with open(manifest_path, "w") as f:
|
||||
json.dump(manifest, f, indent=2)
|
||||
else:
|
||||
manifest_path = os.path.join(output_dir, "aha_training_manifest.json")
|
||||
|
||||
fsdp_config = None
|
||||
gradient_checkpointing = True
|
||||
gradient_checkpointing_kwargs = {"use_reentrant": False}
|
||||
if fsdp_options:
|
||||
fsdp_config = {}
|
||||
if fsdp_transformer_layer_cls:
|
||||
fsdp_config["transformer_layer_cls_to_wrap"] = [fsdp_transformer_layer_cls]
|
||||
if fsdp_activation_checkpointing:
|
||||
# Transformers warns that FSDP should use fsdp_config.activation_checkpointing
|
||||
# instead of Trainer-level gradient_checkpointing to avoid redundant all-gathers.
|
||||
fsdp_config["activation_checkpointing"] = True
|
||||
gradient_checkpointing = False
|
||||
gradient_checkpointing_kwargs = None
|
||||
|
||||
sft_config = SFTConfig(
|
||||
output_dir=output_dir,
|
||||
per_device_train_batch_size=per_device_batch_size,
|
||||
gradient_accumulation_steps=gradient_accumulation_steps,
|
||||
learning_rate=learning_rate,
|
||||
num_train_epochs=num_train_epochs,
|
||||
max_steps=max_steps,
|
||||
max_length=max_seq_length,
|
||||
lr_scheduler_type=lr_scheduler_type,
|
||||
warmup_ratio=warmup_ratio,
|
||||
adam_beta1=adam_beta1,
|
||||
adam_beta2=adam_beta2,
|
||||
weight_decay=weight_decay,
|
||||
max_grad_norm=max_grad_norm,
|
||||
logging_steps=logging_steps,
|
||||
save_steps=save_steps,
|
||||
save_total_limit=save_total_limit,
|
||||
save_strategy=save_strategy,
|
||||
bf16=True,
|
||||
tf32=True,
|
||||
gradient_checkpointing=gradient_checkpointing,
|
||||
gradient_checkpointing_kwargs=gradient_checkpointing_kwargs,
|
||||
dataloader_drop_last=True,
|
||||
remove_unused_columns=True,
|
||||
report_to=report_to if report_to else "none",
|
||||
run_name=wandb_run_name,
|
||||
seed=seed,
|
||||
# The custom model returns a standard local mean CE and does not use
|
||||
# ``num_items_in_batch``. Keep legacy OLMo-style per-microbatch means
|
||||
# instead of Transformers' global token-count scaling path.
|
||||
average_tokens_across_devices=False,
|
||||
fsdp=fsdp_options,
|
||||
fsdp_config=fsdp_config,
|
||||
)
|
||||
|
||||
trainer = AHASFTTrainer(
|
||||
model=model,
|
||||
args=sft_config,
|
||||
train_dataset=train_dataset,
|
||||
processing_class=tokenizer,
|
||||
optimizers=(build_grouped_lr_optimizer(model), None) if grouped_lr else (None, None),
|
||||
callbacks=[AHAManifestCallback(manifest_path)],
|
||||
)
|
||||
|
||||
print(
|
||||
"Starting training… Extra log fields: ce_loss, gate_aux_loss, loss_total (=ce+gate_aux), "
|
||||
"gate_soft_mean, gate_hard_mean, hf_loss_ratio (= key `loss` / loss_total). "
|
||||
"ce_loss / loss_total are raw model means (language quality + AHA aux). "
|
||||
"Key `loss` is HuggingFace Trainer’s gradient-accumulation-normalized aggregate; "
|
||||
"hf_loss_ratio should stay near 1.0."
|
||||
)
|
||||
print("Starting training...")
|
||||
train_result = trainer.train(resume_from_checkpoint=resume_from_checkpoint)
|
||||
trainer.save_model(output_dir)
|
||||
trainer.save_state()
|
||||
|
||||
metrics = train_result.metrics
|
||||
trainer.log_metrics("train", metrics)
|
||||
trainer.save_metrics("train", metrics)
|
||||
print(f"Training completed. Metrics: {metrics}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
9
recipe/source_training/all_results.json
Normal file
9
recipe/source_training/all_results.json
Normal file
@@ -0,0 +1,9 @@
|
||||
{
|
||||
"ce_loss": 0.8887,
|
||||
"total_flos": 6.189458802853544e+18,
|
||||
"train_loss": 0.9857371548739723,
|
||||
"train_ppl": 2.431889,
|
||||
"train_runtime": 23521.9665,
|
||||
"train_samples_per_second": 37.607,
|
||||
"train_steps_per_second": 0.294
|
||||
}
|
||||
9
recipe/source_training/train_results.json
Normal file
9
recipe/source_training/train_results.json
Normal file
@@ -0,0 +1,9 @@
|
||||
{
|
||||
"ce_loss": 0.8887,
|
||||
"total_flos": 6.189458802853544e+18,
|
||||
"train_loss": 0.9857371548739723,
|
||||
"train_ppl": 2.431889,
|
||||
"train_runtime": 23521.9665,
|
||||
"train_samples_per_second": 37.607,
|
||||
"train_steps_per_second": 0.294
|
||||
}
|
||||
8339
recipe/source_training/trainer_state.json
Normal file
8339
recipe/source_training/trainer_state.json
Normal file
File diff suppressed because it is too large
Load Diff
1
recipe/tests/__init__.py
Normal file
1
recipe/tests/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
"""Release-package tests."""
|
||||
278
recipe/tests/test_router_granularity.py
Normal file
278
recipe/tests/test_router_granularity.py
Normal file
@@ -0,0 +1,278 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
import torch
|
||||
from transformers import Qwen3Config, Qwen3ForCausalLM
|
||||
|
||||
import modeling_aha_qwen3 as aha_module
|
||||
from modeling_aha_qwen3 import (
|
||||
AHAQwen3Config,
|
||||
AHAQwen3ForCausalLM,
|
||||
aha_router_output_size,
|
||||
)
|
||||
from router_training_utils import RowWiseAdamW, configure_gate_only
|
||||
|
||||
|
||||
def base_config() -> Qwen3Config:
|
||||
config = Qwen3Config(
|
||||
vocab_size=97,
|
||||
hidden_size=32,
|
||||
intermediate_size=64,
|
||||
num_hidden_layers=2,
|
||||
num_attention_heads=4,
|
||||
num_key_value_heads=2,
|
||||
head_dim=8,
|
||||
max_position_embeddings=64,
|
||||
attention_dropout=0.0,
|
||||
attention_bias=True,
|
||||
tie_word_embeddings=False,
|
||||
)
|
||||
config._attn_implementation = "eager"
|
||||
return config
|
||||
|
||||
|
||||
def aha_config(granularity: str, *, force_gate_value=None) -> AHAQwen3Config:
|
||||
payload = base_config().to_dict()
|
||||
payload.update(
|
||||
aha_window_size=2,
|
||||
aha_local_kind="sliding_window",
|
||||
aha_mode="dynamic",
|
||||
aha_router_granularity=granularity,
|
||||
aha_force_gate_value=force_gate_value,
|
||||
aha_reg_weight=1.0,
|
||||
aha_ce_weight=0.0,
|
||||
model_type="aha_qwen3",
|
||||
)
|
||||
config = AHAQwen3Config(**payload)
|
||||
config._attn_implementation = "eager"
|
||||
return config
|
||||
|
||||
|
||||
def copy_base_weights(base: Qwen3ForCausalLM, target: AHAQwen3ForCausalLM) -> None:
|
||||
source = base.state_dict()
|
||||
destination = target.state_dict()
|
||||
q_rows = target.config.num_attention_heads * target.config.head_dim
|
||||
with torch.no_grad():
|
||||
for name, value in source.items():
|
||||
if name not in destination:
|
||||
continue
|
||||
if destination[name].shape == value.shape:
|
||||
destination[name].copy_(value)
|
||||
elif name.endswith("self_attn.q_proj.weight"):
|
||||
destination[name][:q_rows].copy_(value)
|
||||
destination[name][q_rows:].zero_()
|
||||
elif name.endswith("self_attn.q_proj.bias"):
|
||||
destination[name][:q_rows].copy_(value)
|
||||
destination[name][q_rows:].zero_()
|
||||
else:
|
||||
raise AssertionError(f"unexpected shape mismatch for {name}")
|
||||
target.load_state_dict(destination)
|
||||
|
||||
|
||||
def aligned_models():
|
||||
torch.manual_seed(7)
|
||||
base = Qwen3ForCausalLM(base_config()).eval()
|
||||
models = {}
|
||||
for granularity in ("token", "token_kv_head"):
|
||||
model = AHAQwen3ForCausalLM(
|
||||
aha_config(granularity, force_gate_value=1.0)
|
||||
).eval()
|
||||
copy_base_weights(base, model)
|
||||
models[granularity] = model
|
||||
return base, models
|
||||
|
||||
|
||||
class RouterGranularityTest(unittest.TestCase):
|
||||
def test_auto_class_checkpoint_embeds_modeling_source(self):
|
||||
AHAQwen3Config.register_for_auto_class()
|
||||
AHAQwen3ForCausalLM.register_for_auto_class("AutoModelForCausalLM")
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
AHAQwen3ForCausalLM(aha_config("token")).save_pretrained(
|
||||
tmp, safe_serialization=True
|
||||
)
|
||||
self.assertTrue((Path(tmp) / "modeling_aha_qwen3.py").exists())
|
||||
|
||||
def test_projection_and_effective_gate_shapes(self):
|
||||
_, models = aligned_models()
|
||||
input_ids = torch.tensor([[1, 2, 3, 4, 5]])
|
||||
expected_q_rows = 4 * 8
|
||||
for granularity, native_rows in (("token", 1), ("token_kv_head", 2)):
|
||||
with self.subTest(granularity=granularity):
|
||||
model = models[granularity]
|
||||
attention = model.model.layers[0].self_attn
|
||||
self.assertEqual(attention.q_proj.out_features, expected_q_rows + native_rows)
|
||||
self.assertEqual(attention.aha_router_outputs, native_rows)
|
||||
output = model.model(input_ids=input_ids, use_cache=False)
|
||||
self.assertEqual(len(output.all_gate_soft), 2)
|
||||
self.assertEqual(output.all_gate_soft[0].shape, (1, 5, 2))
|
||||
self.assertEqual(output.all_gate_hard[0].shape, (1, 5, 2))
|
||||
|
||||
def test_force_open_matches_full_attention_logits(self):
|
||||
base, models = aligned_models()
|
||||
input_ids = torch.tensor([[1, 4, 2, 8, 3, 9]])
|
||||
with torch.no_grad():
|
||||
expected = base(input_ids=input_ids, use_cache=False).logits
|
||||
for granularity, model in models.items():
|
||||
with self.subTest(granularity=granularity):
|
||||
actual = model(input_ids=input_ids, use_cache=False).logits
|
||||
torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-5)
|
||||
|
||||
def test_force_closed_uses_identical_local_branch(self):
|
||||
_, models = aligned_models()
|
||||
input_ids = torch.tensor([[1, 4, 2, 8, 3, 9]])
|
||||
for model in models.values():
|
||||
model.config.aha_force_gate_value = 0.0
|
||||
with torch.no_grad():
|
||||
token_logits = models["token"](input_ids=input_ids, use_cache=False).logits
|
||||
head_logits = models["token_kv_head"](
|
||||
input_ids=input_ids, use_cache=False
|
||||
).logits
|
||||
torch.testing.assert_close(token_logits, head_logits, atol=1e-6, rtol=1e-5)
|
||||
|
||||
def test_regularizer_uses_effective_kv_head_denominator(self):
|
||||
_, models = aligned_models()
|
||||
probability = 0.73
|
||||
bias = math.log(probability / (1.0 - probability))
|
||||
input_ids = torch.tensor([[1, 4, 2, 8, 3, 9]])
|
||||
outputs = {}
|
||||
for granularity, model in models.items():
|
||||
model.train()
|
||||
model.config.aha_force_gate_value = None
|
||||
q_rows = model.config.num_attention_heads * model.config.head_dim
|
||||
with torch.no_grad():
|
||||
for layer in model.model.layers:
|
||||
layer.self_attn.q_proj.weight[q_rows:].zero_()
|
||||
layer.self_attn.q_proj.bias[q_rows:].fill_(bias)
|
||||
outputs[granularity] = model(
|
||||
input_ids=input_ids, labels=input_ids.clone(), use_cache=False
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
outputs["token"].gate_soft_mean,
|
||||
outputs["token_kv_head"].gate_soft_mean,
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
outputs["token"].gate_aux_loss,
|
||||
outputs["token_kv_head"].gate_aux_loss,
|
||||
)
|
||||
self.assertAlmostEqual(outputs["token"].gate_soft_mean.item(), probability, places=6)
|
||||
|
||||
def test_gate_only_updates_only_appended_rows(self):
|
||||
for granularity, expected_rows in (("token", 1), ("token_kv_head", 2)):
|
||||
with self.subTest(granularity=granularity):
|
||||
model = AHAQwen3ForCausalLM(aha_config(granularity))
|
||||
setup = configure_gate_only(model)
|
||||
self.assertEqual(setup.gate_rows, expected_rows)
|
||||
q_proj = model.model.layers[0].self_attn.q_proj
|
||||
before = q_proj.weight.detach().clone()
|
||||
q_proj.weight.sum().backward()
|
||||
self.assertEqual(
|
||||
torch.count_nonzero(q_proj.weight.grad[: setup.q_rows]).item(), 0
|
||||
)
|
||||
self.assertGreater(
|
||||
torch.count_nonzero(q_proj.weight.grad[setup.q_rows :]).item(), 0
|
||||
)
|
||||
torch.optim.SGD(setup.parameters, lr=0.1).step()
|
||||
torch.testing.assert_close(
|
||||
q_proj.weight[: setup.q_rows], before[: setup.q_rows]
|
||||
)
|
||||
self.assertFalse(
|
||||
torch.equal(q_proj.weight[setup.q_rows :], before[setup.q_rows :])
|
||||
)
|
||||
|
||||
def test_rowwise_adamw_applies_real_ten_x_lr_ratio(self):
|
||||
parameter = torch.nn.Parameter(torch.zeros(3, 1))
|
||||
optimizer = RowWiseAdamW(
|
||||
[{"params": [parameter], "lr": 0.1}],
|
||||
row_scales=[(parameter, 2, 0.1)],
|
||||
weight_decay=0.0,
|
||||
betas=(0.9, 0.999),
|
||||
)
|
||||
parameter.grad = torch.ones_like(parameter)
|
||||
optimizer.step()
|
||||
backbone_update = parameter[:2].abs().mean().item()
|
||||
gate_update = parameter[2:].abs().mean().item()
|
||||
self.assertAlmostEqual(gate_update / backbone_update, 10.0, places=5)
|
||||
|
||||
def test_one_step_smoke_is_finite_and_reloadable(self):
|
||||
input_ids = torch.tensor([[1, 4, 2, 8, 3, 9]])
|
||||
for granularity in ("token", "token_kv_head"):
|
||||
with self.subTest(granularity=granularity), tempfile.TemporaryDirectory() as tmp:
|
||||
model = AHAQwen3ForCausalLM(aha_config(granularity)).train()
|
||||
setup = configure_gate_only(model)
|
||||
optimizer = torch.optim.AdamW(setup.parameters, lr=3e-5)
|
||||
output = model(
|
||||
input_ids=input_ids, labels=input_ids.clone(), use_cache=False
|
||||
)
|
||||
self.assertTrue(torch.isfinite(output.loss).item())
|
||||
output.loss.backward()
|
||||
self.assertTrue(
|
||||
all(
|
||||
parameter.grad is None
|
||||
or torch.isfinite(parameter.grad).all().item()
|
||||
for parameter in setup.parameters
|
||||
)
|
||||
)
|
||||
optimizer.step()
|
||||
model.save_pretrained(tmp, safe_serialization=True)
|
||||
reloaded = AHAQwen3ForCausalLM.from_pretrained_aha(
|
||||
tmp, torch_dtype=torch.float32, attn_implementation="eager"
|
||||
)
|
||||
self.assertEqual(
|
||||
reloaded.config.aha_router_granularity, granularity
|
||||
)
|
||||
|
||||
def test_legacy_default_and_save_load_roundtrip(self):
|
||||
legacy = AHAQwen3Config(**base_config().to_dict())
|
||||
self.assertEqual(legacy.aha_router_granularity, "token_kv_head")
|
||||
self.assertEqual(aha_router_output_size(legacy), legacy.num_key_value_heads)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
model = AHAQwen3ForCausalLM(aha_config("token"))
|
||||
model.save_pretrained(tmp, safe_serialization=True)
|
||||
loaded = AHAQwen3ForCausalLM.from_pretrained_aha(
|
||||
tmp, torch_dtype=torch.float32, attn_implementation="eager"
|
||||
)
|
||||
self.assertEqual(loaded.config.aha_router_granularity, "token")
|
||||
self.assertEqual(loaded.model.layers[0].self_attn.aha_router_outputs, 1)
|
||||
|
||||
def test_sparsity_tracker_records_native_and_effective_counts(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
output = os.path.join(tmp, "sparsity.json")
|
||||
with mock.patch.dict(os.environ, {"AHA_SPARSITY_STATS_PATH": output}):
|
||||
tracker = aha_module._AHAInferenceSparsityTracker()
|
||||
gate_hard = torch.tensor([[[1.0, 1.0], [0.0, 0.0]]])
|
||||
gate_soft = torch.tensor([[[0.8, 0.8], [0.2, 0.2]]])
|
||||
tracker.update(
|
||||
gate_hard,
|
||||
gate_soft,
|
||||
layer_idx=0,
|
||||
phase="prefill",
|
||||
router_granularity="token",
|
||||
native_router_width=1,
|
||||
)
|
||||
tracker.update(
|
||||
gate_hard[:, :1],
|
||||
gate_soft[:, :1],
|
||||
layer_idx=0,
|
||||
phase="decode",
|
||||
router_granularity="token",
|
||||
native_router_width=1,
|
||||
)
|
||||
tracker.write_stats()
|
||||
payload = json.loads(Path(output).read_text())
|
||||
self.assertEqual(payload["router_granularity"], "token")
|
||||
self.assertEqual(payload["native_router_decisions"], 3)
|
||||
self.assertEqual(payload["effective_router_decisions"], 6)
|
||||
self.assertAlmostEqual(payload["sparsity"], 1 / 3)
|
||||
self.assertEqual(payload["by_phase"]["decode"]["total_decisions"], 2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user