| import torch | |
| import torch.nn as nn | |
| from model.config import ModelConfig | |
| from model.transformer import GPT, Block, RMSNorm | |
| class DraftModel(nn.Module): | |
| """ | |
| A very small version of the main GPT model used for Speculative Decoding. | |
| Usually 2-4 layers with smaller embedding dimension. | |
| """ | |
| def __init__(self, config: ModelConfig, draft_layers=2, draft_embd=128): | |
| super().__init__() | |
| # Create a modified config for the draft model | |
| self.config = config | |
| self.token_embedding_table = nn.Embedding(config.vocab_size, draft_embd) | |
| # We reuse the Block architecture but with smaller dims | |
| # Note: In a real scenario, DraftModel would have its own specific smaller Block | |
| # For simplicity, we implement a simple one here | |
| self.layers = nn.ModuleList([ | |
| nn.TransformerEncoderLayer( | |
| d_model=draft_embd, | |
| nhead=4, | |
| dim_feedforward=draft_embd * 4, | |
| dropout=0.0, | |
| activation='gelu', | |
| batch_first=True, | |
| norm_first=True | |
| ) for _ in range(draft_layers) | |
| ]) | |
| self.ln_f = nn.LayerNorm(draft_embd) | |
| self.lm_head = nn.Linear(draft_embd, config.vocab_size, bias=False) | |
| def forward(self, idx): | |
| B, T = idx.shape | |
| x = self.token_embedding_table(idx) | |
| for layer in self.layers: | |
| x = layer(x) | |
| x = self.ln_f(x) | |
| logits = self.lm_head(x) | |
| return logits, None # Matches GPT interface for sampling | |