| """ |
| Bits-for-LLM encoder (extended): token id -> fixed random n-bit code -> encoder -> hidden. |
| Backward compatible: defaults (enc_type='mlp', ffn_depth=2, out_scale=0) reproduce the |
| original 2-layer MLP exactly. New encoder-only tuning knobs for the 7B sweep: |
| - ffn_depth : number of Linear layers in the MLP (2 = original, 3 = deeper) |
| - out_scale : if >0, rescale the final projection so output std ~= out_scale (init-scale match) |
| - enc_type='tfenc' : small Transformer encoder over the bit-code (TF-Enc-H4-FFN512 replica) |
| """ |
| import torch, torch.nn as nn |
|
|
|
|
| class BitsEmbedding(nn.Module): |
| def __init__(self, vocab, hidden, nbits=64, W=256, seed=0, dtype=torch.bfloat16, |
| enc_type='mlp', ffn_depth=2, out_scale=0.0): |
| super().__init__() |
| g = torch.Generator().manual_seed(seed) |
| |
| self._codes = torch.randint(0, 2, (vocab, nbits), generator=g).to(dtype) |
| self.vocab, self.hidden, self.nbits, self.W = vocab, hidden, nbits, W |
| self.enc_type, self.ffn_depth = enc_type, ffn_depth |
| self.padding_idx = None |
| if enc_type == 'tfenc': |
| self.tok_bits = 16 |
| assert nbits % self.tok_bits == 0, "nbits must be divisible by 16 for tfenc" |
| self.seq = nbits // self.tok_bits |
| self.in_proj = nn.Linear(self.tok_bits, W) |
| self.pos = nn.Parameter(torch.zeros(self.seq, W)) |
| layer = nn.TransformerEncoderLayer(d_model=W, nhead=4, dim_feedforward=512, |
| activation='gelu', batch_first=True, norm_first=True) |
| self.tfenc = nn.TransformerEncoder(layer, num_layers=2) |
| self.out_proj = nn.Linear(W, hidden) |
| self.to(dtype) |
| else: |
| layers = [nn.Linear(nbits, W), nn.GELU()] |
| for _ in range(max(0, ffn_depth - 2)): |
| layers += [nn.Linear(W, W), nn.GELU()] |
| layers += [nn.Linear(W, hidden)] |
| self.ffn = nn.Sequential(*layers).to(dtype) |
| if out_scale and out_scale > 0: |
| self._match_out_scale(out_scale) |
|
|
| def _final(self): |
| return self.out_proj if self.enc_type == 'tfenc' else self.ffn[-1] |
|
|
| def _match_out_scale(self, target_std): |
| |
| |
| |
| with torch.no_grad(): |
| idx = torch.randint(0, self.vocab, (4096,)) |
| cur = self.forward(idx).float().std().item() |
| if cur > 1e-8: |
| k = target_std / cur |
| last = self._final() |
| last.weight.mul_(k) |
| if last.bias is not None: |
| last.bias.mul_(k) |
|
|
| def forward(self, ids): |
| if self._codes.device != ids.device: |
| self._codes = self._codes.to(ids.device) |
| codes = self._codes[ids] |
| if self.enc_type == 'tfenc': |
| *lead, n = codes.shape |
| x = codes.reshape(-1, self.seq, self.tok_bits) |
| x = self.in_proj(x) + self.pos |
| x = self.tfenc(x).mean(dim=1) |
| return self.out_proj(x).reshape(*lead, self.hidden) |
| return self.ffn(codes) |
|
|
| def n_params(self): |
| return sum(p.numel() for p in self.parameters()) |
|
|
|
|
| def swap_in_bits(model, nbits, W, seed, enc_type='mlp', ffn_depth=2, out_scale=0.0): |
| cfg = model.config |
| if getattr(cfg, "tie_word_embeddings", False): |
| out = model.get_output_embeddings() |
| if out is not None: |
| new = nn.Linear(out.weight.shape[1], out.weight.shape[0], bias=False) |
| new.weight = nn.Parameter(out.weight.detach().clone()) |
| model.set_output_embeddings(new); cfg.tie_word_embeddings = False |
| dtype = next(model.parameters()).dtype |
| bits = BitsEmbedding(cfg.vocab_size, cfg.hidden_size, nbits, W, seed, dtype, |
| enc_type=enc_type, ffn_depth=ffn_depth, out_scale=out_scale) |
| dev = model.get_input_embeddings().weight.device |
| model.set_input_embeddings(bits.to(dev)) |
| orig = cfg.vocab_size * cfg.hidden_size |
| print(f"[Bits enc={enc_type} depth={ffn_depth} W={W} nbits={nbits} scale={out_scale}] " |
| f"table {orig/1e9:.3f}B -> encoder {bits.n_params()/1e6:.2f}M " |
| f"= {orig/max(bits.n_params(),1):.0f}x compression", flush=True) |
| return bits |
|
|
|
|
| def get_dataset(tokenizer, max_len, n_tokens, split_stream): |
| from datasets import load_dataset |
| ds = load_dataset("allenai/c4", "en", split=split_stream, streaming=True) |
| def gen(): |
| buf = [] |
| while True: |
| for ex in ds: |
| ids = tokenizer(ex["text"], add_special_tokens=False)["input_ids"] |
| buf.extend(ids + [tokenizer.eos_token_id]) |
| while len(buf) >= max_len: |
| yield {"input_ids": buf[:max_len], "labels": buf[:max_len]} |
| buf = buf[max_len:] |
| from datasets import IterableDataset |
| return IterableDataset.from_generator(gen) |
|
|
