819 lines
36 KiB
Python
819 lines
36 KiB
Python
"""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()
|