279 lines
12 KiB
Python
279 lines
12 KiB
Python
|
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()
|