File size: 5,148 Bytes
fb6b7f8 be8ee5f fb6b7f8 | 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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | import torch
import torch.nn as nn
from .generate import generate
from .DecoderBlock import DecoderBlock
from .RotaryPositionalEmbedding import RotaryPositionalEmbedding
class TransformerModel(nn.Module):
def __init__(self, vocab_size: int, d_model: int, num_heads: int, num_kv_heads: int, num_layers: int, ff_dim: int, max_seq_len: int, dropout: float, pad_token_id: int = 0):
super().__init__()
self.d_model = d_model
self.num_heads = num_heads
self.num_kv_heads = num_kv_heads
self.num_layers = num_layers
self.pad_token_id = pad_token_id
self.max_seq_len = max_seq_len
self.embed = nn.Embedding(vocab_size, d_model)
d_k = d_model // num_heads
self.rope = RotaryPositionalEmbedding(d_k, max_seq_len)
self.blocks = nn.ModuleList([
DecoderBlock(d_model, num_heads, num_kv_heads, ff_dim, dropout)
for _ in range(num_layers)
])
self.norm = nn.RMSNorm(d_model, eps=1e-6)
self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
self.lm_head.weight = self.embed.weight
causal = torch.triu(
torch.full((max_seq_len, max_seq_len), float('-inf')), diagonal=1
)
self.register_buffer("causal_mask", causal, persistent=False)
self._init_weights()
def _init_weights(self):
std = self.d_model ** -0.5
nn.init.normal_(self.embed.weight, mean=0.0, std=std)
for module in self.modules():
if isinstance(module, nn.Linear):
nn.init.normal_(module.weight, mean=0.0, std=std)
if module.bias is not None:
nn.init.zeros_(module.bias)
def _build_attn_mask(self, T: int, pad_mask, device):
attn_mask = self.causal_mask[:T, :T][None, None, :, :]
if pad_mask is not None:
pad = torch.zeros(pad_mask.shape[0], 1, 1, T, device=device)
pad.masked_fill_(~pad_mask[:, None, None, :], float('-inf'))
attn_mask = attn_mask + pad
return attn_mask
def forward_features(self, input_ids: torch.Tensor, attention_mask=None, has_padding: bool = True) -> torch.Tensor:
"""Giống forward() nhưng DỪNG TRƯỚC lm_head — dùng cho training loss chunked
(tránh vật lý hóa logits full (B*T, vocab_size))."""
pad_mask = attention_mask.bool() if attention_mask is not None \
else (input_ids != self.pad_token_id)
B, T = input_ids.shape
x = self.embed(input_ids)
pos = torch.arange(T, device=input_ids.device)
cos, sin = self.rope.get_cos_sin(pos)
# Chỉ build mask (và cộng pad-bias) khi batch này thực sự có token PAD.
# Nếu không có PAD, causal mask thuần == causal+pad mask về mặt toán học
# (pad-bias toàn số 0) -> bỏ qua an toàn, không đổi kết quả.
# has_padding được tính sẵn trên CPU trong collate_fn nên không tốn sync GPU ở đây.
attn_mask = self._build_attn_mask(T, pad_mask, x.device) if has_padding else None
for block in self.blocks:
x = block(x, cos, sin, attn_mask)
return self.norm(x) # (B, T, d_model) — CHƯA qua lm_head
def forward(self, input_ids: torch.Tensor, attention_mask=None, has_padding: bool = True) -> torch.Tensor:
x = self.forward_features(input_ids, attention_mask, has_padding)
return self.lm_head(x)
def init_cache(self, batch_size: int, max_gen_len: int, device: torch.device):
d_k = self.d_model // self.num_heads
return [
[
torch.empty(batch_size, self.num_kv_heads, max_gen_len, d_k, device=device),
torch.empty(batch_size, self.num_kv_heads, max_gen_len, d_k, device=device),
]
for _ in self.blocks
]
def prefill(self, input_ids: torch.Tensor, kv_cache=None):
B, T = input_ids.shape
x = self.embed(input_ids)
pos = torch.arange(T, device=input_ids.device)
cos, sin = self.rope.get_cos_sin(pos)
new_cache = []
for i, block in enumerate(self.blocks):
x, kv = block.prefill(x, cos, sin)
if kv_cache is not None:
kv_cache[i][0][:B, :, :T, :] = kv[0]
kv_cache[i][1][:B, :, :T, :] = kv[1]
new_cache.append(kv_cache[i])
else:
new_cache.append(kv)
logits = self.lm_head(self.norm(x))[:, -1, :]
return logits, new_cache
def decode_step(self, token_ids: torch.Tensor, kv_cache, cache_len: int):
token_ids = token_ids.view(-1, 1)
x = self.embed(token_ids)
pos = torch.arange(cache_len, cache_len + 1, device=token_ids.device)
cos, sin = self.rope.get_cos_sin(pos)
for block, kv in zip(self.blocks, kv_cache):
x = block.forward_with_cache(x, kv, cache_len, cos, sin)
return self.lm_head(self.norm(x))[:, 0, :]
def generate_response(self, user_input, tokenizer, **kwargs):
return generate(self, user_input, tokenizer, **kwargs)
|