Hugging Face's logo

RegalInvader
/
Marco-Nano-Instruct-GGUF

Text Generation
GGUF
llama.cpp
pocketpal
on-device
Mixture of Experts
multilingual
quantization
distillation
1-bit
2-bit
3-bit
conversational
Model card Files
xet
 Community
Marco-Nano-Instruct-GGUF / code /tests /test_student.py
RegalInvader's picture
RegalInvader
Release: Marco-Nano-Instruct 1-, 2- and 3-bit GGUF for phones
e5235fb verified
Raw History Blame Contribute Delete
8.55 kB
"""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