|
|
from torch import nn as nn |
|
|
import torch |
|
|
from torch.nn import functional as F |
|
|
from transformers import PretrainedConfig, PreTrainedModel |
|
|
|
|
|
|
|
|
class MiniLLMConfig(PretrainedConfig): |
|
|
model_type = "mini-llm" |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
vocab_size=32000, |
|
|
num_layers=12, |
|
|
dim=1024, |
|
|
rope_base=10000, |
|
|
num_attention_q_heads=64, |
|
|
num_attention_kv_heads=32, |
|
|
qkv_bias=True, |
|
|
drop_rate=0.0, |
|
|
**kwargs, |
|
|
): |
|
|
super().__init__(**kwargs) |
|
|
self.vocab_size = vocab_size |
|
|
self.num_layers = num_layers |
|
|
self.dim = dim |
|
|
self.rope_base = rope_base |
|
|
self.num_attention_q_heads = num_attention_q_heads |
|
|
self.num_attention_kv_heads = num_attention_kv_heads |
|
|
self.qkv_bias = qkv_bias |
|
|
self.drop_rate = drop_rate |
|
|
self.auto_map = { |
|
|
"AutoConfig": "mini_llm.MiniLLMConfig", |
|
|
"AutoModelForCausalLM": "mini_llm.MiniLLM", |
|
|
} |
|
|
|
|
|
|
|
|
class RMSNorm(nn.Module): |
|
|
def __init__(self, dim): |
|
|
super().__init__() |
|
|
self.weight = nn.Parameter(torch.ones(dim)) |
|
|
|
|
|
def forward(self, x: torch.Tensor): |
|
|
norm_x = x / torch.sqrt(x.pow(2).mean(dim=-1, keepdim=True) + 1e-8) |
|
|
output = self.weight * norm_x |
|
|
return output |
|
|
|
|
|
|
|
|
class RopePositionEmbedding(nn.Module): |
|
|
def __init__(self, dim: int, base=10000): |
|
|
super().__init__() |
|
|
inv_freq = 1 / base ** (torch.arange(0, dim, 2).float() / dim) |
|
|
inv_freq = inv_freq.unsqueeze(0) |
|
|
self.register_buffer("inv_freq", inv_freq) |
|
|
|
|
|
def rotate_half(self, x: torch.Tensor): |
|
|
odd = x[..., 1::2] |
|
|
even = x[..., 0::2] |
|
|
return torch.stack((-odd, even), dim=-1).flatten(-2) |
|
|
|
|
|
def apply_rope(self, x: torch.Tensor): |
|
|
x_len = x.shape[2] |
|
|
t = torch.arange(0, x_len, device=x.device, dtype=torch.float32).unsqueeze(1) |
|
|
freq = t * self.inv_freq |
|
|
freq = torch.repeat_interleave(freq, repeats=2, dim=-1)[None, None, :, :] |
|
|
xf = x.float() |
|
|
y = xf * freq.cos() + self.rotate_half(xf) * freq.sin() |
|
|
return y.to(x.dtype) |
|
|
|
|
|
def forward(self, q: torch.Tensor, k: torch.Tensor): |
|
|
return self.apply_rope(q), self.apply_rope(k) |
|
|
|
|
|
|
|
|
class GroupQueryAttention(nn.Module): |
|
|
def __init__( |
|
|
self, |
|
|
num_attention_q_heads, |
|
|
num_attention_kv_heads, |
|
|
dim, |
|
|
qkv_bias, |
|
|
drop_rate, |
|
|
rope_base, |
|
|
): |
|
|
super().__init__() |
|
|
|
|
|
self.head_dim = dim // num_attention_q_heads |
|
|
|
|
|
assert dim % num_attention_q_heads == 0, "dim 必须被 Q 头数整除" |
|
|
assert ( |
|
|
num_attention_q_heads % num_attention_kv_heads == 0 |
|
|
), "Q头数必须是KV头数的整数倍" |
|
|
assert self.head_dim % 2 == 0, "head_dim 必须为偶数以应用 RoPE" |
|
|
|
|
|
self.q_proj = nn.Linear(dim, dim, bias=qkv_bias) |
|
|
self.k_proj = nn.Linear( |
|
|
dim, self.head_dim * num_attention_kv_heads, bias=qkv_bias |
|
|
) |
|
|
self.v_proj = nn.Linear( |
|
|
dim, self.head_dim * num_attention_kv_heads, bias=qkv_bias |
|
|
) |
|
|
self.out_proj = nn.Linear(dim, dim, bias=qkv_bias) |
|
|
|
|
|
self.num_repeat_kv = num_attention_q_heads // num_attention_kv_heads |
|
|
self.drop = nn.Dropout(drop_rate) |
|
|
|
|
|
self.position_embedding = RopePositionEmbedding(self.head_dim, rope_base) |
|
|
|
|
|
self.num_attention_q_heads = num_attention_q_heads |
|
|
self.num_attention_kv_heads = num_attention_kv_heads |
|
|
self.drop_rate = drop_rate |
|
|
|
|
|
def repeat_kv(self, k: torch.Tensor, v: torch.Tensor): |
|
|
k = k.repeat_interleave(self.num_repeat_kv, dim=1) |
|
|
v = v.repeat_interleave(self.num_repeat_kv, dim=1) |
|
|
return k, v |
|
|
|
|
|
def forward(self, x: torch.Tensor): |
|
|
batch_size, seq_len, dim = x.shape |
|
|
Q = ( |
|
|
self.q_proj(x) |
|
|
.reshape(batch_size, seq_len, self.num_attention_q_heads, self.head_dim) |
|
|
.transpose(1, 2) |
|
|
) |
|
|
K = ( |
|
|
self.k_proj(x) |
|
|
.reshape(batch_size, seq_len, self.num_attention_kv_heads, self.head_dim) |
|
|
.transpose(1, 2) |
|
|
) |
|
|
V = ( |
|
|
self.v_proj(x) |
|
|
.reshape(batch_size, seq_len, self.num_attention_kv_heads, self.head_dim) |
|
|
.transpose(1, 2) |
|
|
) |
|
|
|
|
|
Q, K = self.position_embedding(Q, K) |
|
|
|
|
|
K, V = self.repeat_kv(K, V) |
|
|
|
|
|
out = F.scaled_dot_product_attention( |
|
|
Q, K, V, dropout_p=self.drop_rate if self.training else 0.0, is_causal=True |
|
|
) |
|
|
out = out.transpose(1, 2).reshape(batch_size, seq_len, dim) |
|
|
out = self.out_proj(out) |
|
|
out = self.drop(out) |
|
|
return out |
|
|
|
|
|
|
|
|
class FFN(nn.Module): |
|
|
def __init__(self, dim, drop_rate): |
|
|
super().__init__() |
|
|
self.ffn = nn.Sequential( |
|
|
nn.Linear(dim, dim * 4), |
|
|
nn.SiLU(), |
|
|
nn.Linear(dim * 4, dim), |
|
|
nn.Dropout(drop_rate), |
|
|
) |
|
|
|
|
|
def forward(self, x): |
|
|
return self.ffn(x) |
|
|
|
|
|
|
|
|
class DecoderLayer(nn.Module): |
|
|
def __init__( |
|
|
self, |
|
|
num_attention_q_heads, |
|
|
num_attention_kv_heads, |
|
|
dim, |
|
|
qkv_bias, |
|
|
drop_rate, |
|
|
rope_base, |
|
|
): |
|
|
super().__init__() |
|
|
self.norm1 = RMSNorm(dim=dim) |
|
|
self.attn = GroupQueryAttention( |
|
|
num_attention_q_heads=num_attention_q_heads, |
|
|
num_attention_kv_heads=num_attention_kv_heads, |
|
|
dim=dim, |
|
|
qkv_bias=qkv_bias, |
|
|
drop_rate=drop_rate, |
|
|
rope_base=rope_base, |
|
|
) |
|
|
self.norm2 = RMSNorm(dim=dim) |
|
|
self.ffn = FFN(dim=dim, drop_rate=drop_rate) |
|
|
|
|
|
def forward(self, x): |
|
|
x = x + self.attn(self.norm1(x)) |
|
|
x = x + self.ffn(self.norm2(x)) |
|
|
return x |
|
|
|
|
|
|
|
|
class MiniLLM(PreTrainedModel): |
|
|
model_type = "mini-llm" |
|
|
config_class = MiniLLMConfig |
|
|
|
|
|
def __init__(self, config: MiniLLMConfig, pretrain_ckpt=None): |
|
|
super().__init__(config) |
|
|
self.embedding = nn.Embedding(config.vocab_size, config.dim) |
|
|
self.layers = nn.ModuleList([]) |
|
|
for _ in range(config.num_layers): |
|
|
self.layers.append( |
|
|
DecoderLayer( |
|
|
num_attention_q_heads=config.num_attention_q_heads, |
|
|
num_attention_kv_heads=config.num_attention_kv_heads, |
|
|
dim=config.dim, |
|
|
qkv_bias=config.qkv_bias, |
|
|
drop_rate=config.drop_rate, |
|
|
rope_base=config.rope_base, |
|
|
) |
|
|
) |
|
|
self.norm = RMSNorm(dim=config.dim) |
|
|
self.head = nn.Linear(config.dim, config.vocab_size, bias=False) |
|
|
self.apply(self.init_weight) |
|
|
self.head.weight = self.embedding.weight |
|
|
if pretrain_ckpt is not None: |
|
|
self.load_ckpt(pretrain_ckpt) |
|
|
|
|
|
def load_ckpt(self, ckpt_path): |
|
|
ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False) |
|
|
state_dict = ckpt["state_dict"] |
|
|
new_state_dict = {} |
|
|
for k, v in state_dict.items(): |
|
|
new_k = k[len("net._orig_mod.") :] |
|
|
new_state_dict[new_k] = v |
|
|
self.load_state_dict(new_state_dict, strict=True) |
|
|
print(f"load state dict from {ckpt_path}") |
|
|
|
|
|
def init_weight(self, m): |
|
|
if isinstance(m, nn.Linear): |
|
|
nn.init.normal_(m.weight, mean=0, std=0.02) |
|
|
if m.bias is not None: |
|
|
nn.init.constant_(m.bias, 0) |
|
|
elif isinstance(m, RMSNorm): |
|
|
nn.init.constant_(m.weight, 1) |
|
|
elif isinstance(m, nn.Embedding): |
|
|
nn.init.normal_(m.weight, mean=0, std=0.02) |
|
|
|
|
|
def forward(self, input_ids: torch.Tensor): |
|
|
hidden_state = self.embedding(input_ids) |
|
|
for layer in self.layers: |
|
|
hidden_state = layer(hidden_state) |
|
|
|
|
|
hidden_state = self.norm(hidden_state) |
|
|
logits = self.head(hidden_state) |
|
|
return logits |
|
|
|
|
|
def top_k_sample(self, logits, top_k=5): |
|
|
|
|
|
weights, indices = torch.topk(logits, k=top_k, dim=-1) |
|
|
|
|
|
probs = torch.softmax(weights, dim=-1) |
|
|
chosssed_id = torch.multinomial(probs, num_samples=1) |
|
|
new_token = torch.gather(indices, dim=-1, index=chosssed_id) |
|
|
return new_token |
|
|
|
|
|
@torch.no_grad() |
|
|
def chat(self, conversations, tokenizer, max_new_token=256, top_k=5): |
|
|
ids = tokenizer.apply_chat_template( |
|
|
conversations, add_generation_prompt=True, tokenize=True |
|
|
) |
|
|
eos_ids = tokenizer.eos_token_id |
|
|
input_ids = torch.tensor(ids, dtype=torch.long).unsqueeze(0).to(self.device) |
|
|
for _ in range(max_new_token): |
|
|
|
|
|
logits = self(input_ids) |
|
|
last_logits = logits[:, -1] |
|
|
new_token = self.top_k_sample(last_logits, top_k=top_k) |
|
|
input_ids = torch.cat((input_ids, new_token), dim=-1) |
|
|
|
|
|
if new_token.detach()[0].cpu().item() == eos_ids: |
|
|
break |
|
|
|
|
|
output_id = input_ids.detach().cpu()[0].tolist() |
|
|
output_id = output_id[len(ids) :] |
|
|
answer = tokenizer.decode(output_id, skip_special_tokens=True) |
|
|
return answer |
|
|
|
