|
|
|
|
|
import torch |
|
|
import torch.nn as nn |
|
|
import math |
|
|
from transformers import PreTrainedModel, PreTrainedConfig |
|
|
from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions |
|
|
|
|
|
class NanoGPTCompressedConfig(PreTrainedConfig): |
|
|
model_type = "nanogpt_compressed" |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
vocab_size=6060, |
|
|
block_size=1024, |
|
|
n_layer=8, |
|
|
n_head=8, |
|
|
n_embd=512, |
|
|
dropout=0.0, |
|
|
bias=True, |
|
|
compression_method="fixed_low_rank_mlp", |
|
|
compression_rank=128, |
|
|
compressed_layers=[1], |
|
|
**kwargs |
|
|
): |
|
|
self.vocab_size = vocab_size |
|
|
self.block_size = block_size |
|
|
self.n_layer = n_layer |
|
|
self.n_head = n_head |
|
|
self.n_embd = n_embd |
|
|
self.dropout = dropout |
|
|
self.bias = bias |
|
|
self.compression_method = compression_method |
|
|
self.compression_rank = compression_rank |
|
|
self.compressed_layers = compressed_layers |
|
|
super().__init__(**kwargs) |
|
|
|
|
|
class LowRankLinear(nn.Module): |
|
|
def __init__(self, input_dim, output_dim, rank=16, bias=True): |
|
|
super().__init__() |
|
|
self.rank = rank |
|
|
self.input_dim = input_dim |
|
|
self.output_dim = output_dim |
|
|
|
|
|
self.U = nn.Parameter(torch.randn(input_dim, rank) * 0.02) |
|
|
self.V = nn.Parameter(torch.randn(rank, output_dim) * 0.02) |
|
|
|
|
|
if bias: |
|
|
self.bias = nn.Parameter(torch.zeros(output_dim)) |
|
|
else: |
|
|
self.register_parameter('bias', None) |
|
|
|
|
|
def forward(self, x): |
|
|
result = (x @ self.U) @ self.V |
|
|
if self.bias is not None: |
|
|
result = result + self.bias |
|
|
return result |
|
|
|
|
|
class LayerNorm(nn.Module): |
|
|
def __init__(self, ndim, bias): |
|
|
super().__init__() |
|
|
self.weight = nn.Parameter(torch.ones(ndim)) |
|
|
self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None |
|
|
|
|
|
def forward(self, input): |
|
|
return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5) |
|
|
|
|
|
class CausalSelfAttention(nn.Module): |
|
|
def __init__(self, config): |
|
|
super().__init__() |
|
|
assert config.n_embd % config.n_head == 0 |
|
|
self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias) |
|
|
self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias) |
|
|
self.attn_dropout = nn.Dropout(config.dropout) |
|
|
self.resid_dropout = nn.Dropout(config.dropout) |
|
|
self.n_head = config.n_head |
|
|
self.n_embd = config.n_embd |
|
|
self.dropout = config.dropout |
|
|
self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size)) |
|
|
.view(1, 1, config.block_size, config.block_size)) |
|
|
|
|
|
def forward(self, x): |
|
|
B, T, C = x.size() |
|
|
q, k, v = self.c_attn(x).split(self.n_embd, dim=2) |
|
|
k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) |
|
|
q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) |
|
|
v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) |
|
|
att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) |
|
|
att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf')) |
|
|
att = F.softmax(att, dim=-1) |
|
|
att = self.attn_dropout(att) |
|
|
y = att @ v |
|
|
y = y.transpose(1, 2).contiguous().view(B, T, C) |
|
|
y = self.resid_dropout(self.c_proj(y)) |
|
|
return y |
|
|
|
|
|
class MLP(nn.Module): |
|
|
def __init__(self, config, layer_idx=None): |
|
|
super().__init__() |
|
|
self.layer_idx = layer_idx |
|
|
|
|
|
|
|
|
if (hasattr(config, 'compressed_layers') and |
|
|
layer_idx is not None and |
|
|
layer_idx in config.compressed_layers): |
|
|
|
|
|
print(f"Creating compressed MLP for layer {layer_idx}") |
|
|
rank = getattr(config, 'compression_rank', 128) |
|
|
self.c_fc = LowRankLinear(config.n_embd, 4 * config.n_embd, rank, bias=config.bias) |
|
|
self.c_proj = LowRankLinear(4 * config.n_embd, config.n_embd, rank, bias=config.bias) |
|
|
else: |
|
|
self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias) |
|
|
self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias) |
|
|
|
|
|
self.dropout = nn.Dropout(config.dropout) |
|
|
|
|
|
def forward(self, x): |
|
|
x = self.c_fc(x) |
|
|
x = F.gelu(x) |
|
|
x = self.c_proj(x) |
|
|
x = self.dropout(x) |
|
|
return x |
|
|
|
|
|
class Block(nn.Module): |
|
|
def __init__(self, config, layer_idx=None): |
|
|
super().__init__() |
|
|
self.ln_1 = LayerNorm(config.n_embd, bias=config.bias) |
|
|
self.attn = CausalSelfAttention(config) |
|
|
self.ln_2 = LayerNorm(config.n_embd, bias=config.bias) |
|
|
self.mlp = MLP(config, layer_idx=layer_idx) |
|
|
|
|
|
def forward(self, x): |
|
|
x = x + self.attn(self.ln_1(x)) |
|
|
x = x + self.mlp(self.ln_2(x)) |
|
|
return x |
|
|
|
|
|
class NanoGPTCompressedModel(PreTrainedModel): |
|
|
config_class = NanoGPTCompressedConfig |
|
|
|
|
|
def __init__(self, config): |
|
|
super().__init__(config) |
|
|
self.config = config |
|
|
|
|
|
self.transformer = nn.ModuleDict(dict( |
|
|
wte = nn.Embedding(config.vocab_size, config.n_embd), |
|
|
wpe = nn.Embedding(config.block_size, config.n_embd), |
|
|
drop = nn.Dropout(config.dropout), |
|
|
h = nn.ModuleList([Block(config, layer_idx=i) for i in range(config.n_layer)]), |
|
|
ln_f = LayerNorm(config.n_embd, bias=config.bias), |
|
|
)) |
|
|
self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) |
|
|
|
|
|
|
|
|
self.transformer.wte.weight = self.lm_head.weight |
|
|
|
|
|
|
|
|
self.apply(self._init_weights) |
|
|
for pn, p in self.named_parameters(): |
|
|
if pn.endswith('c_proj.weight') or pn.endswith('c_proj.V'): |
|
|
torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer)) |
|
|
|
|
|
def _init_weights(self, module): |
|
|
if isinstance(module, nn.Linear): |
|
|
torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) |
|
|
if module.bias is not None: |
|
|
torch.nn.init.zeros_(module.bias) |
|
|
elif isinstance(module, nn.Embedding): |
|
|
torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) |
|
|
|
|
|
def forward(self, idx, targets=None): |
|
|
device = idx.device |
|
|
b, t = idx.size() |
|
|
assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}" |
|
|
pos = torch.arange(0, t, dtype=torch.long, device=device) |
|
|
|
|
|
tok_emb = self.transformer.wte(idx) |
|
|
pos_emb = self.transformer.wpe(pos) |
|
|
x = self.transformer.drop(tok_emb + pos_emb) |
|
|
for block in self.transformer.h: |
|
|
x = block(x) |
|
|
x = self.transformer.ln_f(x) |
|
|
|
|
|
if targets is not None: |
|
|
logits = self.lm_head(x) |
|
|
loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1) |
|
|
else: |
|
|
logits = self.lm_head(x[:, [-1], :]) |
|
|
loss = None |
|
|
|
|
|
return CausalLMOutputWithCrossAttentions( |
|
|
loss=loss, |
|
|
logits=logits, |
|
|
past_key_values=None, |
|
|
hidden_states=None, |
|
|
attentions=None, |
|
|
cross_attentions=None, |
|
|
) |
|
|
|
|
|
@torch.no_grad() |
|
|
def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None): |
|
|
for _ in range(max_new_tokens): |
|
|
idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:] |
|
|
logits = self(idx_cond).logits |
|
|
logits = logits[:, -1, :] / temperature |
|
|
if top_k is not None: |
|
|
v, _ = torch.topk(logits, min(top_k, logits.size(-1))) |
|
|
logits[logits < v[:, [-1]]] = -float('Inf') |
|
|
probs = F.softmax(logits, dim=-1) |
|
|
idx_next = torch.multinomial(probs, num_samples=1) |
|
|
idx = torch.cat((idx, idx_next), dim=1) |
|
|
return idx |
|
|
|
