File size: 1,549 Bytes
e5235fb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 | """T0.3 wrapper: the ternary switch replaces routed experts by their fake-quant, and nothing else."""
import torch
from transformers import Qwen3MoeConfig, Qwen3MoeForCausalLM
from lowbit.eval.wrapper import routed_expert_params, ternarise_experts
from lowbit.quant.ternary import QuantSpec, fake_quant
def tiny_moe():
torch.manual_seed(0)
cfg = Qwen3MoeConfig(vocab_size=256, hidden_size=128, intermediate_size=256, moe_intermediate_size=128,
num_hidden_layers=2, num_attention_heads=2, num_key_value_heads=1, head_dim=64,
num_experts=4, num_experts_per_tok=2, max_position_embeddings=64)
m = Qwen3MoeForCausalLM(cfg)
with torch.no_grad():
for p in m.parameters():
p.normal_(0, 0.02)
return m
def test_ternarise_experts_only_touches_experts():
m = tiny_moe()
before = {n: p.detach().clone() for n, p in m.named_parameters()}
expert_names = {n for n, _ in routed_expert_params(m)}
assert expert_names, "no routed-expert weights found"
spec = QuantSpec()
n = ternarise_experts(m, spec)
assert n == sum(before[k].numel() for k in expert_names)
for name, p in m.named_parameters():
if name in expert_names:
assert torch.equal(p, fake_quant(before[name].float(), spec).to(p.dtype))
g = p.reshape(-1, 128) # each g128 group holds at most {-s, 0, +s}
assert all(len(torch.unique(row.abs())) <= 2 for row in g[:16])
else:
assert torch.equal(p, before[name]), name
|