File size: 3,976 Bytes
33f617f 7812b29 33f617f | 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 | import torch
import torch.nn as nn
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, causal: bool = True):
if not causal:
# Bidirectional (embedding): chỉ chặn key là PAD, mọi token nhìn thấy nhau.
pad = torch.zeros(pad_mask.shape[0], 1, 1, T, device=device)
pad.masked_fill_(~pad_mask[:, None, None, :], float('-inf'))
return pad
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, causal: 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)).
causal=False bật attention hai chiều cho embedding model (LLM2Vec);
causal=True giữ đúng hành vi của LLM gốc (dùng để đo baseline zero-shot)."""
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, causal) if has_padding else None
for block in self.blocks:
x = block(x, cos, sin, attn_mask, causal)
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)
|