File size: 3,272 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 | """The one model wrapper every eval goes through (spec §7, T0.3).
lm-evaluation-harness's HF model, registered as `marco`, with three additions:
- a quantisation switch: `quant=off` (the model as loaded) or `quant=ternary` (routed experts replaced
in place by their ternary fake-quant, spec §4: Q-A absmean, g128 by default) — the PTQ view of a
checkpoint without exporting it;
- a loader hook for packed students, `packed=<dir>` (filled in T2.6: codes + scales → expert weights);
- stop tokens: generation stops on both `<|im_end|>` and `<|endoftext|>`, because Marco's
generation_config.json lists only the latter while every chat turn ends with the former (D-12).
Use: `lm_eval --model marco --model_args pretrained=...,quant=ternary --include_path configs/eval/lm_eval_tasks`.
"""
from __future__ import annotations
import torch
from lm_eval.api.registry import register_model
from lm_eval.models.huggingface import HFLM
from lowbit.quant.ternary import QuantSpec, fake_quant
STOP_TOKENS = ("<|im_end|>", "<|endoftext|>")
def routed_expert_params(model: torch.nn.Module):
"""(name, parameter) for every routed-expert weight: the fused [E, out, in] tensors of transformers 5
(`experts.gate_up_proj`, `experts.down_proj`) or per-expert Linear weights in older layouts."""
for name, p in model.named_parameters():
if ".mlp.experts." in name and name.endswith(("gate_up_proj", "down_proj", "gate_proj.weight",
"up_proj.weight", "down_proj.weight")):
yield name, p
@torch.no_grad()
def ternarise_experts(model: torch.nn.Module, spec: QuantSpec) -> int:
"""Replace routed-expert weights by s * q in place (groups along the input dim); returns #weights."""
n = 0
for _, p in routed_expert_params(model):
p.copy_(fake_quant(p.data.float(), spec).to(p.dtype))
n += p.numel()
return n
def load_packed(model: torch.nn.Module, path: str) -> None:
raise NotImplementedError("packed-student loading is T2.6")
@register_model("marco")
class MarcoLM(HFLM):
def __init__(self, pretrained: str, quant: str = "off", packed: str | None = None,
quant_method: str = "absmean", group_size: int = 128, **kwargs):
super().__init__(pretrained=pretrained, **kwargs)
if quant not in ("off", "ternary"):
raise ValueError(f"quant must be off or ternary, got {quant!r}")
self.quant, self.packed = quant, packed
if packed:
load_packed(self.model, packed)
if quant == "ternary":
n = ternarise_experts(self.model, QuantSpec(method=quant_method, group_size=int(group_size)))
if n == 0:
raise ValueError("quant=ternary but the model has no routed-expert weights")
stop_ids = [self.tokenizer.convert_tokens_to_ids(t) for t in STOP_TOKENS]
self.stop_token_ids = [i for i in stop_ids if isinstance(i, int) and i >= 0]
def _model_generate(self, context, max_length, stop, **generation_kwargs):
generation_kwargs.setdefault("eos_token_id", self.stop_token_ids)
stop = list(dict.fromkeys([*stop, *STOP_TOKENS]))
return super()._model_generate(context, max_length, stop, **generation_kwargs)
|