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