lilo_common.py

"""LILO — Lightweight Local Intelligence & Optimizer

Helpers for loading and running LILO-100M on CPU. Variants:
  - lilo-100m (fp16) : fp16 weights (221MB on disk) — fastest, quality gate
                       PPL delta 0.08% vs fp32, greedy outputs 5/5 identical
  - lilo-100m (int8) : INT8 group-wise weights (group=64, 119MB on disk) —
                       smallest/RAM-light, PPL delta 0.32%, greedy 5/5 identical

Both ship patched modeling code (freqs_cis fix + working KV cache) so single-token
decode is cheap. Generate with `generate_lilo()` (explicit KV cache).

CPU speed (measured on 2 vCPU AMD EPYC 7K62 Zen2, AVX2-only, under load):
  fp16  14-22 tok/s at --threads 1   (recommended default)
  int8  9-19 tok/s at --threads 1
More threads HURT on tiny models (sync overhead); scale up only on 8+ cores.
"""

import json
import os
from typing import List, Optional, Sequence

import torch
from safetensors import safe_open
from safetensors.torch import load_file

from transformers import AutoConfig, AutoTokenizer

from configuration_deepseek_v4 import DeepseekV4Config
from modeling_deepseek_v4 import (
    DeepseekV4ForCausalLM,
    QuantizedLinear,
    QuantizedEmbedding,
)


# submodule leaf names that get INT8-quantized in the int8 variant
_LINEAR_TARGETS = (
    "wq_a", "wq_b", "wkv", "wo_a", "wo_b",   # attention projections
    "w1", "w2", "w3",                        # experts (gate/up/down)
    "lm_head",
)


def _swap_quantized(model, quantize_embed=True):
    """Swap nn.Linear -> QuantizedLinear and nn.Embedding -> QuantizedEmbedding.
    Returns number of modules swapped."""
    swapped = 0
    for parent in list(dict(model.named_modules()).values()):
        for name, child in list(parent.named_children()):
            target = None
            if name in _LINEAR_TARGETS and child.__class__.__name__ == "Linear":
                target = QuantizedLinear(child.in_features, child.out_features)
            elif quantize_embed and name == "embed_tokens" and child.__class__.__name__ == "Embedding":
                target = QuantizedEmbedding(child.num_embeddings, child.embedding_dim)
            if target is not None:
                setattr(parent, name, target)
                swapped += 1
    return swapped


def load_lilo(path_or_repo, variant="int8", device="cpu", threads=2, hf_token=None):
    """Load a LILO model. variant: "int8" | "fp16" | "fp32"."""
    path = path_or_repo
    if not os.path.isdir(path):
        from huggingface_hub import snapshot_download
        path = snapshot_download(repo_id=path_or_repo, token=hf_token)

    torch.set_num_threads(threads)

    config = DeepseekV4Config.from_pretrained(path, trust_remote_code=True)
    tok = AutoTokenizer.from_pretrained(path, trust_remote_code=True)

    if variant == "int8":
        model = DeepseekV4ForCausalLM(config)
        from safetensors.torch import load_file
        sd = load_file(os.path.join(path, "model.int8.safetensors"))
        with safe_open(os.path.join(path, "model.int8.safetensors"), framework="pt") as _sf:
            emb_dtype = _sf.get_slice("model.embed_tokens.weight").get_dtype()
        n = _swap_quantized(model, quantize_embed=(emb_dtype == "I8"))
        missing, unexpected = model.load_state_dict(sd, strict=False)
        bad_missing = [k for k in missing if not k.startswith("freqs_cis")]
        assert not bad_missing and not unexpected, f"missing={bad_missing[:5]} unexpected={unexpected[:5]}"
        model.eval()
        model.model.get_freqs_cis()
        print(f"[lilo] loaded INT8 variant from {path} ({n} modules swapped, "
              f"{sum(t.numel() for t in sd.values())/1e6:.0f}M params)")
    elif variant == "fp16":
        model = DeepseekV4ForCausalLM(config)
        from safetensors.torch import load_file
        sd = load_file(os.path.join(path, "model.fp16.safetensors"))
        missing, unexpected = model.load_state_dict(sd, strict=False)
        assert not missing and not unexpected, f"missing={missing[:5]} unexpected={unexpected[:5]}"
        model.eval()
        model = model.to(torch.float16)
        model.model.get_freqs_cis()
        print(f"[lilo] loaded FP16 variant from {path}")
    else:
        model = DeepseekV4ForCausalLM.from_pretrained(
            path, trust_remote_code=True, dtype=torch.float32, attn_implementation="eager"
        )
        model.eval()
        model.model.get_freqs_cis()
        print(f"[lilo] loaded FP32 variant from {path}")

    model.to(device)
    return tok, model


@torch.inference_mode()
def generate_lilo(
    model,
    tok,
    prompt_ids: Sequence[int],
    max_new_tokens: int = 200,
    temperature: float = 0.7,
    top_p: float = 0.9,
    repetition_penalty: float = 1.0,
    eos_token_id: int = 1,
    stream: bool = False,
):
    """Greedy/sampling generation with an explicit KV cache.

    Returns (generated_ids list). Repetition penalty is applied on the growing
    sequence (standard formulation)."""
    ids = torch.tensor([list(prompt_ids)], dtype=torch.long)
    generated: List[int] = []
    past = None

    logits = None
    for step in range(max_new_tokens):
        if step == 0:
            # chunked prefill to bound peak memory
            seq = ids
            chunk = 256
            for s in range(0, seq.shape[1], chunk):
                piece = seq[:, s : s + chunk]
                out = model(input_ids=piece, past_key_values=past, use_cache=True)
                past = out.past_key_values
            logits = out.logits[0, -1].float()
        else:
            out = model(
                input_ids=torch.tensor([[generated[-1]]], dtype=torch.long),
                past_key_values=past,
                use_cache=True,
            )
            past = out.past_key_values
            logits = out.logits[0, -1].float()

        if repetition_penalty != 1.0:
            for t in set(generated):
                logits[t] /= repetition_penalty if logits[t] > 0 else 1.0 / repetition_penalty

        if temperature <= 0:
            nxt = int(torch.argmax(logits))
        else:
            probs = torch.softmax(logits / max(temperature, 1e-5), dim=-1)
            if top_p < 1.0:
                sp, si = torch.sort(probs, descending=True)
                cum = torch.cumsum(sp, dim=-1)
                mask = cum - sp > top_p
                sp[mask] = 0.0
                probs = torch.zeros_like(probs).scatter_(0, si, sp)
                probs = probs / probs.sum()
            nxt = int(torch.multinomial(probs, 1))
        generated.append(nxt)
        if stream:
            print(tok.decode([nxt]), end="", flush=True)
        if nxt == eos_token_id:
            break
    return generated


def chat(tok, model, messages, **kw):
    prompt = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    ids = tok.encode(prompt, add_special_tokens=False)
    gen = generate_lilo(model, tok, ids, **kw)
    return tok.decode(gen, skip_special_tokens=True)