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