File size: 1,506 Bytes
9756f8b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 | import torch
import torch.nn as nn
from trainDraft import DraftChain
class DraftDecoderLayer(nn.Module):
def __init__(self, block):
super().__init__()
self.block = block
def forward(self, hidden_states, position_embeddings=None, attention_mask=None,
position_ids=None, past_key_values=None, use_cache=None, **kwargs):
return self.block(hidden_states)
class _Identity(nn.Module):
def forward(self, x):
return x
def freeze_chain(chain):
for b in chain.blocks:
for name in ("fc1", "fc2"):
wn = getattr(b, name)
w = wn.weight().detach()
lin = nn.Linear(w.shape[1], w.shape[0], bias=False, dtype=w.dtype)
lin.weight.data.copy_(w)
setattr(b, name, lin)
return chain
def load_chain(path, dtype=None):
ckpt = torch.load(path, weights_only=True)
chain = DraftChain(ckpt["D"], ckpt["hidden"], ckpt["k"])
chain.load_state_dict(ckpt["chain"])
if dtype is not None:
chain = chain.to(dtype)
chain.eval()
return freeze_chain(chain)
def apply_merge(model, chain):
n = model.config.num_hidden_layers
k = len(chain.blocks)
for _ in range(k):
del model.model.layers[-1]
for j in range(k):
layer = DraftDecoderLayer(chain.blocks[j])
layer.to(model.device)
model.model.layers.append(layer)
model.model.norm = _Identity()
return model
|