| """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) |
|
|
