| import math |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| class DraftBlock(nn.Module): |
| def __init__(self, D, H): |
| super().__init__() |
| self.fc1 = nn.Linear(D, H, bias=False) |
| self.fc2 = nn.Linear(H, D, bias=False) |
| self.alpha = nn.Parameter(torch.zeros(())) |
| self.sqrtD = math.sqrt(D) |
|
|
| def forward(self, s): |
| ns = s * (self.sqrtD / (s.norm(dim=-1, keepdim=True) + 1e-12)) |
| return ns + self.fc2(F.silu(self.fc1(ns))) * self.alpha |
|
|
|
|
| class DraftDecoderLayer(nn.Module): |
| def __init__(self, D, H): |
| super().__init__() |
| self.block = DraftBlock(D, H) |
|
|
| def forward(self, hidden_states, *a, **k): |
| return self.block(hidden_states) |
|
|
|
|
| def load_merged(model_dir, dtype=torch.bfloat16, device="cuda", chain_hidden=2048): |
| from transformers import AutoConfig, AutoModelForCausalLM |
| from safetensors.torch import load_file |
| cfg = AutoConfig.from_pretrained(model_dir) |
| m = AutoModelForCausalLM.from_config(cfg).to(device).to(dtype).eval() |
| n = cfg.num_hidden_layers |
| for j in range(4): |
| m.model.layers[n - 4 + j] = DraftDecoderLayer(cfg.hidden_size, chain_hidden).to(device, dtype=dtype) |
| m.model.norm = nn.Identity() |
| sd = load_file(f"{model_dir}/model.safetensors", device=device) |
| sd = {k.replace("model.language_model.", "model.", 1) if k.startswith("model.language_model.") else k: v |
| for k, v in sd.items()} |
| if "lm_head.weight" not in sd: |
| sd["lm_head.weight"] = sd["model.embed_tokens.weight"] |
| m.load_state_dict(sd, strict=True) |
| return m |
|
|
