File size: 8,546 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 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 | """T2.2 student model on CPU, with a tiny random Qwen3-MoE (2 layers, 8 experts, group 128)."""
import copy
from pathlib import Path
import pytest
import torch
from lowbit.model import student as st
from lowbit.quant.ternary import QuantSpec, reference_codes_and_scale
REPO = Path(__file__).resolve().parent.parent
A0 = REPO / "configs" / "precision" / "A0.yaml"
def tiny(dtype=torch.float32, seed=0, impl="eager", hidden=128):
"""hidden=256 gives gate/up rows Q3_K can hold (D-51's int3)."""
from transformers import AutoModelForCausalLM, Qwen3MoeConfig
cfg = Qwen3MoeConfig(vocab_size=256, hidden_size=hidden, intermediate_size=256, moe_intermediate_size=128,
num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=2, head_dim=32,
num_experts=8, num_experts_per_tok=2, norm_topk_prob=True, max_position_embeddings=512,
tie_word_embeddings=True, decoder_sparse_step=1, mlp_only_layers=[])
torch.manual_seed(seed)
m = AutoModelForCausalLM.from_config(cfg, dtype=dtype, attn_implementation="sdpa", experts_implementation=impl)
with torch.no_grad(): # random init is near zero; give the experts realistic, distinct magnitudes
for layer in m.model.layers:
layer.mlp.experts.gate_up_proj.normal_(0, 0.05)
layer.mlp.experts.down_proj.normal_(0, 0.05)
return m.eval()
def ids(b=2, t=48, seed=1):
return torch.randint(0, 256, (b, t), generator=torch.Generator().manual_seed(seed))
def logits_of(model, x, **kw):
with torch.no_grad():
return st.logits(model, st.final_hidden(model, x, **kw)).float()
def test_quant_off_equals_teacher_bitwise():
teacher = tiny()
prec = st.Precision.load(A0)
prec.router_fp32 = True # the fp32-router path: with fp32 tiny weights the two computations coincide
student = st.build_student(copy.deepcopy(teacher), prec, quant=False)
x = ids()
assert torch.equal(logits_of(student, x), teacher(x).logits.float())
def test_quant_off_bf16_router_matches_hf():
teacher = tiny(torch.bfloat16)
prec = st.Precision.load(A0)
assert not prec.router_fp32 # D-28: A0's router is computed as stock Marco's
student = st.build_student(copy.deepcopy(teacher), prec, quant=False)
x = ids()
with torch.no_grad():
assert torch.equal(logits_of(student, x), teacher(x).logits.float())
def test_quant_on_matches_reference_path():
base = tiny()
student = st.build_student(copy.deepcopy(base), st.Precision.load(A0), quant=True)
ref = copy.deepcopy(base)
spec = QuantSpec()
with torch.no_grad():
for layer in ref.model.layers:
for name in ("gate_up_proj", "down_proj"):
w = getattr(layer.mlp.experts, name)
flat = w.reshape(-1, w.shape[-1])
codes, scales = reference_codes_and_scale(flat, spec)
w_hat = (codes.float().unflatten(-1, (-1, 128)) * scales.unsqueeze(-1)).flatten(-2)
w.copy_(w_hat.reshape(w.shape))
x = ids()
assert torch.equal(logits_of(student, x), ref(x).logits.float())
def test_packed_reload_reproduces_fake_quant(tmp_path):
prec = st.Precision.load(A0)
prec.router_fp32 = True # exercise the fp32-router export and reload path
student = st.build_student(tiny(torch.bfloat16), prec, quant=True)
man = st.export_packed_student(student, tmp_path / "packed")
assert man["router_fp32"] and len(man["expert_shapes"]) == 4
packed = st.load_packed_student(tmp_path / "packed")
x = ids()
with torch.no_grad():
assert torch.equal(logits_of(student, x), packed(x).logits.float())
def test_doc_masked_equals_separate_documents():
student = st.build_student(tiny(), st.Precision.load(A0), quant=True)
a, b = ids(1, 20, seed=2), ids(1, 28, seed=3)
packed = torch.cat([a, b], 1)
whole = logits_of(student, packed, mode="doc_masked", doc_starts=[[0, 20]])
sep = torch.cat([logits_of(student, a), logits_of(student, b)], 1)
assert torch.allclose(whole, sep, atol=1e-5)
causal = logits_of(student, packed, mode="causal_packed")
assert torch.allclose(causal[:, :20], sep[:, :20], atol=1e-5) # the first document is the same either way
assert not torch.allclose(causal[:, 20:], sep[:, 20:], atol=1e-3) # the second sees the first
def test_gradients_reach_trainable_classes_only():
student = st.build_student(tiny(), st.Precision.load(A0), quant=True)
classes = st.param_classes(student)
assert len(classes["experts"]) == 4 and not any(p.requires_grad for p in classes["router"] + classes["embedding"])
x = ids(1, 32)
h = st.final_hidden(student, x)
st.logits(student, h).float().logsumexp(-1).mean().backward()
for p in classes["experts"] + classes["attention"] + classes["norms"]:
assert p.grad is not None and p.grad.abs().sum() > 0
assert all(p.grad is None for p in classes["router"] + classes["embedding"])
ex = student.model.layers[0].mlp.experts
assert ex.gate_up_proj is None and ex.down_proj is None # no bf16 expert weights kept between forwards
@pytest.mark.parametrize("early_stop", [True, False])
def test_checkpointed_backward_keeps_no_expert_weights(early_stop):
student = st.build_student(tiny(torch.bfloat16), st.Precision.load(A0), quant=True)
if early_stop: # HF's default with non-reentrant checkpointing: the recompute stops before _post_hook
student.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
else:
st.enable_checkpointing(student)
student.train()
h = st.final_hidden(student, ids(1, 32))
st.logits(student, h).float().logsumexp(-1).mean().backward()
kept = [l.mlp.experts.gate_up_proj is not None for l in student.model.layers]
assert any(kept) if early_stop else not any(kept)
@pytest.mark.parametrize("method", ["absmean", "twn"])
def test_compiled_quantiser_gives_the_eager_codes(method):
from lowbit.quant.ternary import fake_quant
spec = QuantSpec(method=method)
w = tiny().model.layers[0].mlp.experts.gate_up_proj.detach().float()
w = torch.cat([w, w * 37.3, w * 1e-3]) # several scale ranges
def quant(x):
return fake_quant(x, spec)
with torch.no_grad(): # fp32 sums and division moved ~0.09% of codes under compile before the fix
assert torch.equal(quant(w), torch.compile(quant, dynamic=False)(w))
def test_compiled_layers_match_eager():
base = tiny()
plain = st.build_student(copy.deepcopy(base), st.Precision.load(A0), quant=True)
comp = st.build_student(copy.deepcopy(base), st.Precision.load(A0), quant=True)
for m in (plain, comp):
m.config._experts_implementation = "grouped_mm"
st.compile_layers(comp)
x = ids()
assert torch.allclose(logits_of(plain, x), logits_of(comp, x), atol=1e-5)
def test_grouped_mm_matches_eager():
base = tiny()
eager = st.build_student(copy.deepcopy(base), st.Precision.load(A0), quant=True)
grouped = st.build_student(copy.deepcopy(base), st.Precision.load(A0), quant=True)
grouped.config._experts_implementation = "grouped_mm"
x = ids()
assert torch.allclose(logits_of(eager, x), logits_of(grouped, x), atol=1e-4)
def test_a2_ternary_attention_and_a4_int8_layers(tmp_path):
for arm in ("A2", "A4"):
prec = st.Precision.load(REPO / "configs" / "precision" / f"{arm}.yaml")
if arm == "A4":
prec.int8_layers = (1,) # the tiny model has 2 layers
stu = st.build_student(tiny(torch.bfloat16), prec, quant=True)
x = ids()
if arm == "A2":
q = stu.model.layers[0].self_attn.q_proj
q(torch.zeros(1, q.in_features, dtype=torch.bfloat16))
assert q.weight is None # rebuilt per forward, not kept
assert len(st.param_classes(stu)["experts"]) == 4 + 4 # experts + q/o latents of 2 layers
h = st.final_hidden(stu, ids(1, 32))
st.logits(stu, h).float().logsumexp(-1).mean().backward()
lat = q.weight_latent if arm == "A2" else stu.model.layers[1].mlp.experts.gate_up_latent
assert lat.grad is not None and lat.grad.abs().sum() > 0 # the STE reaches the latents
man = st.export_packed_student(stu, tmp_path / arm)
assert man["linear_shapes" if arm == "A2" else "int8_shapes"]
packed = st.load_packed_student(tmp_path / arm)
with torch.no_grad():
assert torch.equal(logits_of(stu, x), packed(x).logits.float()), arm
|