Spaces:
Running on Zero
Running on Zero
| """Load Whittle-Qwen-3.8 checkpoints (35B, 45B) with stock transformers (>= 5.18, Qwen4Exp). | |
| Whittle keeps the Qwen4Exp architecture and its sharded n-gram table format, with two differences from the stock code: | |
| * geometry: 8 hash heads of exactly 4,880,000 rows each (the geometry its GGUFs carry for llama.cpp), where the stock | |
| code sizes each head as a distinct prime above `ngram_vocab_size_base`; | |
| * key prefix: the table pieces are saved as `model.ngram_embedding.shard_{i}.weight` and the hash buffers as | |
| `model.ple_embedding.*`, where the stock code uses `model.layers.{L}.ple.ple_embedding.*`. | |
| This loader sizes the PLE n-gram table from the checkpoint's own hash buffers, maps the two prefixes onto the stock path | |
| and lets transformers do the rest (it concatenates the pieces on CPU and keeps the ~20 GB lookup table off the GPU). | |
| """ | |
| import json | |
| import mmap | |
| import os | |
| import re | |
| import struct | |
| import torch | |
| from huggingface_hub import hf_hub_download | |
| from safetensors import safe_open | |
| from transformers import AutoConfig, AutoModelForCausalLM | |
| from transformers.models.qwen4_exp import modeling_qwen4_exp as _qwen4exp | |
| def _local(repo: str, filename: str) -> str: | |
| return os.path.join(repo, filename) if os.path.isdir(repo) else hf_hub_download(repo, filename) | |
| FULL_INDEX = "whittle_full_index.json" # optional: full index kept beside a load index that omits the table pieces | |
| def _weight_map(repo: str) -> dict: | |
| if os.path.isdir(repo) and os.path.isfile(os.path.join(repo, FULL_INDEX)): | |
| return json.load(open(os.path.join(repo, FULL_INDEX)))["weight_map"] | |
| return json.load(open(_local(repo, "model.safetensors.index.json")))["weight_map"] | |
| def _hash_buffers(repo: str) -> tuple[dict, int]: | |
| """Whittle's hash buffers, and the number of table rows its checkpoint pieces hold.""" | |
| weight_map = _weight_map(repo) | |
| out = {} | |
| for name in ("ngram_heads_vocab_sizes", "ngram_heads_offsets", "layer_multipliers"): | |
| key = f"model.ple_embedding.{name}" | |
| with safe_open(_local(repo, weight_map[key]), "pt") as f: | |
| out[name] = f.get_tensor(key).tolist() | |
| rows = 0 | |
| for key in (k for k in weight_map if k.startswith("model.ngram_embedding.shard_")): | |
| with safe_open(_local(repo, weight_map[key]), "pt") as f: | |
| rows += f.get_slice(key).get_shape()[0] | |
| return out, rows | |
| class MmapNGramTable(torch.nn.Module): | |
| """The n-gram table read in place from the checkpoint's row pieces (memory-mapped): only looked-up rows are read, | |
| so the ~20 GB table costs almost no RAM - the same idea as llama.cpp's mmap. Behaves like nn.Embedding for lookups.""" | |
| def __init__(self, pieces, dim: int, dtype): | |
| super().__init__() | |
| self._pieces, self._parts, starts, row = list(pieces), [], [], 0 | |
| for path, offset, rows, stored in pieces: | |
| with open(path, "rb") as fh: | |
| buf = mmap.mmap(fh.fileno(), 0, access=mmap.ACCESS_READ) | |
| self._parts.append(torch.frombuffer(buf, dtype=stored, count=rows * dim, offset=offset).view(rows, dim)) | |
| starts.append(row); row += rows | |
| self.num_embeddings, self.embedding_dim, self._starts = row, dim, starts | |
| self.register_buffer("_bounds", torch.tensor(starts[1:], dtype=torch.long), persistent=False) | |
| self.register_buffer("weight", torch.empty(0, dim, dtype=dtype), persistent=False) # device/dtype marker only | |
| def to_ram(self, chunk: int = 64 << 20) -> torch.nn.Embedding: | |
| """Copy the whole table into one in-memory nn.Embedding with large sequential file reads (much faster than | |
| faulting the memory map in page by page, especially on a network mount).""" | |
| stored = self._pieces[0][3] | |
| assert all(p[3] == stored for p in self._pieces), "pieces with mixed dtypes" | |
| table = torch.empty((self.num_embeddings, self.embedding_dim), dtype=stored) | |
| raw = memoryview(table.view(torch.uint8).numpy().reshape(-1)) | |
| pos = 0 | |
| for path, offset, rows, _ in self._pieces: | |
| nbytes = rows * self.embedding_dim * table.element_size() | |
| with open(path, "rb", buffering=0) as fh: | |
| fh.seek(offset) | |
| done = 0 | |
| while done < nbytes: | |
| n = fh.readinto(raw[pos + done: pos + min(nbytes, done + chunk)]) | |
| if not n: | |
| raise IOError(f"short read in {path} at {offset + done}") | |
| done += n | |
| pos += nbytes | |
| return torch.nn.Embedding.from_pretrained(table.to(self.weight.dtype), freeze=True) | |
| def forward(self, ids: torch.Tensor) -> torch.Tensor: | |
| flat = ids.reshape(-1).cpu() | |
| which = torch.bucketize(flat, self._bounds, right=True) | |
| out = torch.empty(flat.numel(), self.embedding_dim, dtype=self.weight.dtype) | |
| for i, part in enumerate(self._parts): | |
| sel = which == i | |
| if sel.any(): | |
| out[sel] = part[flat[sel] - self._starts[i]].to(out.dtype) | |
| return out.view(*ids.shape, self.embedding_dim) | |
| def _table_pieces(repo: str): | |
| """(path, absolute byte offset, rows, stored dtype) of each n-gram table piece, in row order.""" | |
| weight_map = _weight_map(repo) | |
| keys = sorted((k for k in weight_map if re.fullmatch(r"model\.ngram_embedding\.shard_\d+\.weight", k)), | |
| key=lambda k: int(k.split(".")[2].split("_")[1])) | |
| out = [] | |
| for key in keys: | |
| path = _local(repo, weight_map[key]) | |
| with open(path, "rb") as fh: | |
| n = struct.unpack("<Q", fh.read(8))[0] | |
| header = json.loads(fh.read(n)) | |
| start, _ = header[key]["data_offsets"] | |
| stored = {"BF16": torch.bfloat16, "F16": torch.float16, "F32": torch.float32}[header[key]["dtype"]] | |
| out.append((path, 8 + n + start, header[key]["shape"][0], stored)) | |
| return out | |
| def load_whittle(repo: str, dtype=torch.bfloat16, device_map="cuda", table: str = "load", **kwargs): | |
| """table=\"load\" puts the n-gram table in CPU memory (~20 GB, plus a transient copy while loading); | |
| table=\"mmap\" reads it in place from the checkpoint files (needs the shards on a local or mounted path).""" | |
| assert table in ("load", "mmap"), table | |
| cfg = AutoConfig.from_pretrained(repo) | |
| prefix = f"model.layers.{cfg.ple_layer_ids[0] - 1}.ple.ple_embedding" | |
| hb, rows = _hash_buffers(repo) | |
| sizes, offsets = hb["ngram_heads_vocab_sizes"], hb["ngram_heads_offsets"] | |
| total = offsets[-1] + sizes[-1] | |
| assert rows >= total, f"table pieces hold {rows:,} rows, the hash addresses {total:,}" | |
| original_init = _qwen4exp.Qwen4ExpTextNGramEmbedding.__init__ | |
| def whittle_init(self, config, embedding_dim, layer_idx, ple_layer_index=0): | |
| original_init(self, config, embedding_dim, layer_idx, ple_layer_index) | |
| assert len(sizes) == self.ngram_heads, (len(sizes), self.ngram_heads) | |
| self.head_vocab_sizes, self.head_offsets, self.total_vocab_size = list(sizes), list(offsets), total | |
| self.ngram_heads_vocab_sizes = torch.nn.Buffer(torch.tensor(sizes, dtype=torch.long)) | |
| self.ngram_heads_offsets = torch.nn.Buffer(torch.tensor(offsets, dtype=torch.long)) | |
| dim = self.ngram_embedding.embedding_dim | |
| # load: exactly the checkpoint's rows; mmap: a one-row placeholder, replaced by the memory-mapped table after loading | |
| self.ngram_embedding = torch.nn.Embedding(rows if table == "load" else 1, dim) | |
| _qwen4exp.Qwen4ExpTextNGramEmbedding.__init__ = whittle_init | |
| try: | |
| model = AutoModelForCausalLM.from_pretrained( | |
| repo, dtype=dtype, device_map=device_map, | |
| key_mapping={r"^model\.ple_embedding\.": f"{prefix}.", | |
| # mmap: rename the pieces so transformers neither concatenates nor loads them | |
| r"^model\.ngram_embedding\.": f"{prefix}.ngram_embedding." if table == "load" else "whittle_unloaded_table."}, | |
| **kwargs, | |
| ) | |
| finally: | |
| _qwen4exp.Qwen4ExpTextNGramEmbedding.__init__ = original_init | |
| emb = model.get_submodule(prefix) | |
| if table == "mmap": | |
| emb.ngram_embedding = MmapNGramTable(_table_pieces(repo), emb.ngram_embedding.embedding_dim, dtype) | |
| for name, want in hb.items(): # the forward pass hashes with these buffers: they must be Whittle's | |
| got = getattr(emb, name).tolist() | |
| assert got == want, f"{name}: {got[:3]}... != checkpoint {want[:3]}..." | |
| assert getattr(emb.ngram_embedding, "num_embeddings", emb.ngram_embedding.weight.shape[0]) >= total | |
| return model | |