| |
| """KenT: kiến trúc hồi quy hai luồng (representation X_t / temporal-memory H_t). |
| |
| Mỗi bước nhận ĐÚNG MỘT token (cửa sổ trượt từng token một); lịch sử nằm trong trạng thái H của từng layer. |
| Mỗi KenT layer (hộp "WNN" trong sơ đồ) thực hiện: |
| |
| Z = ResProj_in(Norm_X(X)) # 1 -> M: một vector token thành M resolution, mỗi resolution d feature |
| Z = Z + SwiGLU( DCT_M( RMSNorm(Z) ) ) # WNN Block: DCT-II cố định trên trục M, SwiGLU trên trục d, skip quanh cả cụm |
| H_t = (1 - A) ⊙ H_{t-1} + A ⊙ Z_t # "EMA Learner Matrix": A = sigmoid(ema_logit), MA TRẬN M x d học được / layer |
| X' = X + ResProj_out( DCT_M( Norm_H(H_t) ) ) # DCT-II Proj rồi M -> 1; Add với X vào layer (residual) |
| |
| Mỗi layer xếp K = `num_wnn_blocks` WNN Block nối tiếp (mặc định K = 1, đúng sơ đồ); độ sâu đến từ việc xếp chồng |
| nhiều layer ("xN" trong sơ đồ). Layer cuối cho X' -> norm cuối -> lm_head (tied với embedding). |
| |
| Bộ nhớ H có cơ chế quên: A là ma trận M x d (mỗi vị trí resolution VÀ mỗi feature có hệ số quên riêng), học được và |
| riêng từng layer. Khởi tạo: nội suy theo cấp số nhân từ `ema_alpha_first` (layer đầu, vị trí đầu, feature đầu) tới |
| `ema_alpha_last` (layer cuối, vị trí cuối, feature cuối) theo trung bình của chỉ số layer, vị trí và feature. |
| `ema_per_feature=False` quay về vector M phần tử (chung cho mọi feature), dùng scan theo khối. |
| |
| EMA là một phép hồi quy tuyến tính (đường chéo theo từng kênh) nên vẫn song song hóa được theo thời gian bằng |
| affine-scan log-depth (Hillis-Steele): kết quả bằng đúng hồi quy tuần tự. Trạng thái giải mã chỉ là H cho mỗi |
| layer, shape (M, d): O(1) theo độ dài. |
| |
| Trọng số được lưu ở FP16; EMA và normalization chạy ở FP32. |
| """ |
|
|
| import math |
| from dataclasses import dataclass |
| from functools import lru_cache |
| from typing import Optional, Tuple, Union |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from torch.nn import CrossEntropyLoss |
|
|
| from transformers import GenerationMixin, PretrainedConfig, PreTrainedModel |
| from transformers.utils import ModelOutput |
|
|
|
|
| |
| |
| |
| class KentConfig(PretrainedConfig): |
| r"""Cấu hình KenT. |
| |
| Args: |
| vocab_size: kích thước từ vựng. |
| hidden_size: số chiều của representation stream X_t (n trong sơ đồ). |
| num_hidden_layers: số KenT layer xếp chồng. |
| resolution_size: M, độ dài trục mà DCT-II tác động lên (cũng là độ dài vector hệ số EMA, "d=M" trong sơ đồ). |
| feature_dim: d, số chiều đặc trưng trên mỗi vị trí resolution. |
| num_wnn_blocks: K, số WNN Block (DCT -> SwiGLU -> skip) nối tiếp trong mỗi layer. Sơ đồ: K = 1. |
| intermediate_size: kích thước hidden của SwiGLU (trên trục d). Nên >= 2d để SwiGLU không thành nút cổ chai. |
| rms_norm_eps: epsilon của RMSNorm. |
| use_bias: dùng bias trong Resolution Proj (SwiGLU không bias). |
| initializer_range: độ lệch chuẩn khởi tạo trọng số. |
| ema_alpha_first: hệ số a khởi tạo của layer đầu, vị trí đầu (quên nhanh). |
| ema_alpha_last: hệ số a khởi tạo của layer cuối, vị trí cuối, feature cuối (nhớ lâu); phần còn lại nội suy theo |
| cấp số nhân theo trung bình các chỉ số (layer, vị trí[, feature]). a vẫn là tham số học được. |
| ema_per_feature: True = EMA Learner Matrix (hệ số quên riêng cho mỗi (vị trí M, feature d), ema_logit shape |
| (M, d)); False = vector M phần tử, chung cho mọi feature (ema_logit shape (M,)). |
| ema_block_size: độ dài khối của EMA song song (chỉ dùng khi ema_per_feature=False; ảnh hưởng tốc độ/bộ nhớ). |
| use_cache: trả về KentCache khi suy luận. |
| use_gradient_checkpointing: bật gradient checkpointing theo từng KenT layer khi khởi tạo model |
| (chỉ có tác dụng ở chế độ train; tương đương gọi `model.gradient_checkpointing_enable()`). |
| """ |
|
|
| model_type = "kent" |
|
|
| def __init__( |
| self, |
| vocab_size: int = 50257, |
| hidden_size: int = 256, |
| num_hidden_layers: int = 24, |
| resolution_size: int = 12, |
| feature_dim: int = 96, |
| num_wnn_blocks: int = 1, |
| intermediate_size: int = 384, |
| rms_norm_eps: float = 1e-5, |
| use_bias: bool = True, |
| initializer_range: float = 0.02, |
| ema_alpha_first: float = 0.5, |
| ema_alpha_last: float = 0.01, |
| ema_per_feature: bool = True, |
| ema_block_size: int = 128, |
| use_cache: bool = True, |
| use_gradient_checkpointing: bool = False, |
| pad_token_id: int = 50256, |
| bos_token_id: int = 50256, |
| eos_token_id: int = 50256, |
| tie_word_embeddings: bool = True, |
| **kwargs, |
| ): |
| self.vocab_size = vocab_size |
| self.hidden_size = hidden_size |
| self.num_hidden_layers = num_hidden_layers |
| self.resolution_size = resolution_size |
| self.feature_dim = feature_dim |
| self.num_wnn_blocks = num_wnn_blocks |
| self.intermediate_size = intermediate_size |
| self.rms_norm_eps = rms_norm_eps |
| self.use_bias = use_bias |
| self.initializer_range = initializer_range |
| self.ema_alpha_first = ema_alpha_first |
| self.ema_alpha_last = ema_alpha_last |
| self.ema_per_feature = ema_per_feature |
| self.ema_block_size = ema_block_size |
| self.use_cache = use_cache |
| self.use_gradient_checkpointing = use_gradient_checkpointing |
| super().__init__( |
| pad_token_id=pad_token_id, |
| bos_token_id=bos_token_id, |
| eos_token_id=eos_token_id, |
| tie_word_embeddings=tie_word_embeddings, |
| **kwargs, |
| ) |
|
|
| def alpha_init_vector(self, layer_idx: int): |
| """Vector a (M phần tử) khởi tạo của layer `layer_idx`: giảm theo cấp số nhân theo độ sâu và theo vị trí.""" |
| n_l, n_p = self.num_hidden_layers, self.resolution_size |
| r_l = layer_idx / (n_l - 1) if n_l > 1 else 0.0 |
| ratio = self.ema_alpha_last / self.ema_alpha_first |
| return [float(self.ema_alpha_first * ratio ** (0.5 * (r_l + (p / (n_p - 1) if n_p > 1 else 0.0)))) |
| for p in range(n_p)] |
|
|
| def alpha_init_matrix(self, layer_idx: int): |
| """Ma trận a (M x d) khởi tạo của layer: giảm theo cấp số nhân theo độ sâu, vị trí và feature.""" |
| n_l, n_p, n_d = self.num_hidden_layers, self.resolution_size, self.feature_dim |
| r_l = layer_idx / (n_l - 1) if n_l > 1 else 0.0 |
| ratio = self.ema_alpha_last / self.ema_alpha_first |
| return [[float(self.ema_alpha_first * ratio ** ((r_l + (p / (n_p - 1) if n_p > 1 else 0.0) |
| + (j / (n_d - 1) if n_d > 1 else 0.0)) / 3.0)) |
| for j in range(n_d)] for p in range(n_p)] |
|
|
| def alpha_init(self, layer_idx: int) -> float: |
| """Trung bình của a khởi tạo của layer (để log).""" |
| if self.ema_per_feature: |
| flat = [x for row in self.alpha_init_matrix(layer_idx) for x in row] |
| else: |
| flat = self.alpha_init_vector(layer_idx) |
| return float(sum(flat) / len(flat)) |
|
|
|
|
| |
| |
| |
| @lru_cache(maxsize=32) |
| def _dct2_matrix(n: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor: |
| """Ma trận DCT-II trực chuẩn A_N (cố định, không phải tham số học).""" |
| k = torch.arange(n, dtype=torch.float64).unsqueeze(1) |
| i = torch.arange(n, dtype=torch.float64).unsqueeze(0) |
| a = torch.cos(math.pi / n * (i + 0.5) * k) |
| alpha = torch.full((n, 1), math.sqrt(2.0 / n), dtype=torch.float64) |
| alpha[0, 0] = math.sqrt(1.0 / n) |
| return (alpha * a).to(device=device, dtype=dtype) |
|
|
|
|
| def dct2(z: torch.Tensor) -> torch.Tensor: |
| """D_N(Z) = A_N Z, tác động lên trục -2 (độ dài N) của tensor (..., N, d_h).""" |
| a = _dct2_matrix(z.shape[-2], z.device, z.dtype) |
| return torch.matmul(a, z) |
|
|
|
|
| |
| |
| |
| def _affine_scan(a: torch.Tensor, b: torch.Tensor): |
| """Inclusive scan dọc dim 1 của các ánh xạ affine h -> a*h + b (Hillis-Steele, log K bước). |
| |
| a: (B,K,N) | b: (B,K,N,D). Phần tử k được áp dụng SAU các phần tử < k. |
| Trả về (A, B) sao cho trạng thái sau khối k là A[k]*h_vào + B[k]. |
| """ |
| k = a.shape[1] |
| off = 1 |
| while off < k: |
| b = torch.cat([b[:, :off], a[:, off:, :, None] * b[:, :-off] + b[:, off:]], dim=1) |
| a = torch.cat([a[:, :off], a[:, off:] * a[:, :-off]], dim=1) |
| off *= 2 |
| return a, b |
|
|
|
|
| def ema_scan(u: torch.Tensor, mask: torch.Tensor, h0: torch.Tensor, a: torch.Tensor, block: int = 128): |
| """EMA song song, luôn chạy fp32 (tắt autocast để matmul không bị ép về fp16). |
| |
| a: (M, d) -> hệ số riêng từng kênh (EMA Learner Matrix, scan Hillis-Steele); a: (M,) -> theo khối (chung cho d). |
| """ |
| with torch.autocast(device_type=u.device.type, enabled=False): |
| if a.dim() == 2: |
| return _ema_scan_diag(u.float(), mask, h0.float(), a.float()) |
| return _ema_scan_fp32(u.float(), mask, h0.float(), a.float(), block) |
|
|
|
|
| def _ema_scan_diag(u: torch.Tensor, mask: torch.Tensor, h0: torch.Tensor, a: torch.Tensor): |
| """H_t = (1 - a m_t) ⊙ H_{t-1} + a m_t ⊙ u_t với a: (M, d) khác nhau theo từng kênh. |
| |
| u: (B,T,M,d) fp32 | mask: (B,T) {0,1} (0 = padding: H giữ nguyên) | h0: (B,M,d) fp32. |
| Mỗi bước là ánh xạ affine h -> A_t h + B_t (A_t = 1 - a m_t, B_t = a m_t u_t); inclusive scan log-depth |
| (Hillis-Steele) cho trạng thái sau mọi t. Mọi A_t nằm trong (0, 1] nên ổn định. |
| Trả về H (B,T,M,d) và h_cuối (B,M,d). |
| """ |
| bsz, t, n, d = u.shape |
| am = a.view(1, 1, n, d) * mask.to(u.dtype).view(bsz, t, 1, 1) |
| dec = 1.0 - am |
| acc = am * u |
| off = 1 |
| while off < t: |
| acc = torch.cat([acc[:, :off], dec[:, off:] * acc[:, :-off] + acc[:, off:]], dim=1) |
| dec = torch.cat([dec[:, :off], dec[:, off:] * dec[:, :-off]], dim=1) |
| off *= 2 |
| h = dec * h0.unsqueeze(1) + acc |
| return h, h[:, -1] |
|
|
|
|
| def _ema_scan_fp32(u: torch.Tensor, mask: torch.Tensor, h0: torch.Tensor, a: torch.Tensor, block: int = 128): |
| """H_t = (1 - a m_t) ⊙ H_{t-1} + a m_t ⊙ u_t, tính song song cho mọi t. |
| |
| u: (B,T,N,D) fp32 | mask: (B,T) {0,1} (0 = padding: H giữ nguyên) | h0: (B,N,D) fp32 |
| a: (N,) fp32 trong (0,1): hệ số riêng cho từng vị trí N (= M trong sơ đồ), chung cho mọi feature D. |
| Trả về H (B,T,N,D) và h_cuối (B,N,D). |
| |
| Trong một khối (với từng vị trí n): |
| H_t = sum_{s<=t} a m_s [prod_{s<r<=t}(1 - a m_r)] u_s + [prod_{r<=t}(1 - a m_r)] h_vào, |
| với tích tính qua hiệu của cumsum log (luôn <= 0 nên ổn định). Các khối được nối bằng affine-scan. |
| """ |
| bsz, t, n, d = u.shape |
| blk = max(1, min(block, t)) |
| pad = (-t) % blk |
| m = mask.to(u.dtype) |
| if pad: |
| u = F.pad(u, (0, 0, 0, 0, 0, pad)) |
| m = F.pad(m, (0, pad)) |
| k = (t + pad) // blk |
| uk = u.reshape(bsz, k, blk, n, d).permute(0, 1, 3, 2, 4) |
| am = a.view(1, 1, n, 1) * m.reshape(bsz, k, 1, blk) |
|
|
| cum = torch.cumsum(torch.log1p(-am), dim=-1) |
| tril = torch.ones(blk, blk, dtype=torch.bool, device=u.device).tril() |
| decay = torch.exp((cum.unsqueeze(-1) - cum.unsqueeze(-2)).masked_fill(~tril, float("-inf"))) |
| local = torch.matmul(decay * am.unsqueeze(-2), uk) |
|
|
| a_blk = torch.exp(cum[..., -1]) |
| a_p, b_p = _affine_scan(a_blk, local[:, :, :, -1]) |
| s = a_p.unsqueeze(-1) * h0.unsqueeze(1) + b_p |
| h_in = torch.cat([h0.unsqueeze(1), s[:, :-1]], dim=1) |
| h = local + torch.exp(cum).unsqueeze(-1) * h_in.unsqueeze(3) |
| h = h.permute(0, 1, 3, 2, 4).reshape(bsz, k * blk, n, d)[:, :t] |
| return h, s[:, -1] |
|
|
|
|
| |
| |
| |
| class KentCache: |
| """Trạng thái hồi quy của KenT: H (EMA, fp32) theo layer.""" |
|
|
| def __init__(self, config: KentConfig, batch_size: int = 1, device=None): |
| shape = (batch_size, config.resolution_size, config.feature_dim) |
| self.h = [ |
| torch.zeros(shape, dtype=torch.float32, device=device) |
| for _ in range(config.num_hidden_layers) |
| ] |
| self.steps = 0 |
|
|
| @property |
| def batch_size(self) -> int: |
| return self.h[0].shape[0] |
|
|
| @property |
| def seen_tokens(self) -> int: |
| return self.steps |
|
|
| def get_seq_length(self, layer_idx: int = 0) -> int: |
| return self.steps |
|
|
| def reset(self): |
| for h in self.h: |
| h.zero_() |
| self.steps = 0 |
|
|
| def clone(self) -> "KentCache": |
| new = KentCache.__new__(KentCache) |
| new.h = [h.clone() for h in self.h] |
| new.steps = self.steps |
| return new |
|
|
| def to(self, device): |
| self.h = [h.to(device) for h in self.h] |
| return self |
|
|
| def reorder_cache(self, beam_idx: torch.LongTensor): |
| self.h = [h.index_select(0, beam_idx.to(h.device)) for h in self.h] |
|
|
|
|
| |
| |
| |
| class KentRMSNorm(nn.Module): |
| def __init__(self, dim: int, eps: float = 1e-5): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(dim)) |
| self.eps = eps |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| dtype = x.dtype |
| x = x.float() |
| x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) |
| return (x * self.weight.float()).to(dtype) |
|
|
|
|
| class KentSwiGLU(nn.Module): |
| """SwiGLU trên trục d: down( silu(x W_g) ⊙ (x W_u) ), không bias.""" |
|
|
| def __init__(self, dim: int, hidden: int): |
| super().__init__() |
| self.gate_up = nn.Linear(dim, 2 * hidden, bias=False) |
| self.down = nn.Linear(hidden, dim, bias=False) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| g, u = self.gate_up(x).chunk(2, dim=-1) |
| return self.down(F.silu(g) * u) |
|
|
|
|
| class KentWNNBlock(nn.Module): |
| """WNN Block: Z' = Z + SwiGLU( DCT_M( RMSNorm(Z) ) ). |
| |
| RMSNorm theo trục d; DCT-II cố định tác động lên trục M ("Transpose DCT-II Proj": chuyển trục rồi chiếu); |
| SwiGLU tác động lên trục d (dùng chung cho mọi vị trí M); skip nối quanh cả cụm Norm -> DCT -> SwiGLU. |
| """ |
|
|
| def __init__(self, config: KentConfig): |
| super().__init__() |
| hidden = config.intermediate_size or 4 * config.feature_dim |
| self.norm = KentRMSNorm(config.feature_dim, config.rms_norm_eps) |
| self.mlp = KentSwiGLU(config.feature_dim, hidden) |
|
|
| def forward(self, z: torch.Tensor) -> torch.Tensor: |
| return z + self.mlp(dct2(self.norm(z))) |
|
|
|
|
| class KentLayer(nn.Module): |
| """Một bước KenT: (X_t, H_{t-1}) -> (X_{t+1}, H_t), với K WNN Block và EMA có vector hệ số M phần tử.""" |
|
|
| def __init__(self, config: KentConfig, layer_idx: int = 0): |
| super().__init__() |
| self.n = config.resolution_size |
| self.d = config.feature_dim |
| self.block_size = config.ema_block_size |
| self.norm_x = KentRMSNorm(config.hidden_size, config.rms_norm_eps) |
| self.res_in = nn.Linear(config.hidden_size, self.n * self.d, bias=config.use_bias) |
| self.wnn = nn.ModuleList([KentWNNBlock(config) for _ in range(config.num_wnn_blocks)]) |
| self.norm_h = KentRMSNorm(config.feature_dim, config.rms_norm_eps) |
| self.res_out = nn.Linear(self.n * self.d, config.hidden_size, bias=config.use_bias) |
| a0 = torch.tensor( |
| config.alpha_init_matrix(layer_idx) if config.ema_per_feature else config.alpha_init_vector(layer_idx), |
| dtype=torch.float32) |
| self.ema_logit_init = torch.log(a0 / (1.0 - a0)) |
| self.ema_logit = nn.Parameter(self.ema_logit_init.clone()) |
|
|
| def alpha(self) -> torch.Tensor: |
| """EMA Learner Matrix a (M, d) (hoặc vector (M,)) trong (0,1), tính ở fp32.""" |
| return torch.sigmoid(self.ema_logit.float()).clamp_max(0.999) |
|
|
| def forward(self, x, mask, h0, chunked: bool = False): |
| b, t, _ = x.shape |
| z = self.res_in(self.norm_x(x)).view(b, t, self.n, self.d) |
| for blk in self.wnn: |
| z = blk(z) |
| u = z.float() |
| if chunked: |
| u = u.reshape(1, b * t, self.n, self.d) |
| m = mask.reshape(1, b * t) |
| else: |
| m = mask |
| if h0 is None: |
| h0 = u.new_zeros((u.shape[0], self.n, self.d)) |
| h, h_last = ema_scan(u, m, h0, self.alpha(), self.block_size) |
| h = h.reshape(b, t, self.n, self.d) |
| spec = dct2(self.norm_h(h).to(x.dtype)) |
| dx = self.res_out(spec.reshape(b, t, self.n * self.d)) |
| return x + dx, h_last |
|
|
|
|
| |
| |
| |
| @dataclass |
| class KentOutput(ModelOutput): |
| last_hidden_state: Optional[torch.FloatTensor] = None |
| cache_params: Optional[KentCache] = None |
| hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None |
|
|
|
|
| @dataclass |
| class KentCausalLMOutput(ModelOutput): |
| loss: Optional[torch.FloatTensor] = None |
| logits: Optional[torch.FloatTensor] = None |
| cache_params: Optional[KentCache] = None |
| hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None |
|
|
|
|
| |
| |
| |
| class KentPreTrainedModel(PreTrainedModel): |
| config_class = KentConfig |
| base_model_prefix = "model" |
| supports_gradient_checkpointing = True |
| _no_split_modules = ["KentLayer"] |
| _is_stateful = True |
|
|
| def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None): |
| """Bật checkpoint theo layer; mặc định use_reentrant=False (an toàn với args không phải tensor |
| như h0=None / chunked, và không cần input requires_grad).""" |
| if gradient_checkpointing_kwargs is None: |
| gradient_checkpointing_kwargs = {"use_reentrant": False} |
| super().gradient_checkpointing_enable(gradient_checkpointing_kwargs=gradient_checkpointing_kwargs) |
|
|
| def _init_weights(self, module): |
| std = self.config.initializer_range |
| 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) |
| elif isinstance(module, nn.Embedding): |
| nn.init.normal_(module.weight, mean=0.0, std=std) |
| elif isinstance(module, KentRMSNorm): |
| nn.init.ones_(module.weight) |
| elif isinstance(module, KentLayer): |
| with torch.no_grad(): |
| module.ema_logit.copy_(module.ema_logit_init.to(module.ema_logit.device)) |
| if isinstance(module, nn.Linear) and getattr(module, "_residual_out", False): |
| |
| with torch.no_grad(): |
| module.weight.mul_((2 * self.config.num_hidden_layers) ** -0.5) |
|
|
|
|
| class KentModel(KentPreTrainedModel): |
| def __init__(self, config: KentConfig): |
| super().__init__(config) |
| self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size) |
| self.layers = nn.ModuleList([KentLayer(config, i) for i in range(config.num_hidden_layers)]) |
| for layer in self.layers: |
| layer.res_out._residual_out = True |
| self.norm_f = KentRMSNorm(config.hidden_size, config.rms_norm_eps) |
| self.gradient_checkpointing = False |
| self.post_init() |
| if getattr(config, "use_gradient_checkpointing", False): |
| self.gradient_checkpointing_enable() |
|
|
| def get_input_embeddings(self): |
| return self.embeddings |
|
|
| def set_input_embeddings(self, value): |
| self.embeddings = value |
|
|
| def forward( |
| self, |
| input_ids: Optional[torch.LongTensor] = None, |
| inputs_embeds: Optional[torch.FloatTensor] = None, |
| attention_mask: Optional[torch.Tensor] = None, |
| cache_params: Optional[KentCache] = None, |
| use_cache: Optional[bool] = None, |
| chunked: bool = False, |
| output_hidden_states: Optional[bool] = None, |
| return_dict: Optional[bool] = None, |
| **kwargs, |
| ) -> Union[Tuple, KentOutput]: |
| """ |
| chunked=False: trục batch là các chuỗi độc lập; (B, T). |
| chunked=True : trục batch là K chunk liên tiếp của MỘT tài liệu; (K, C). |
| Cache vào/ra có batch size 1 (trạng thái sau chunk cuối). |
| attention_mask: 1 = token thật, 0 = padding (H giữ nguyên tại vị trí đó). |
| """ |
| use_cache = use_cache if use_cache is not None else (self.config.use_cache and not self.training) |
| output_hidden_states = bool(output_hidden_states) |
| return_dict = return_dict if return_dict is not None else True |
|
|
| if (input_ids is None) == (inputs_embeds is None): |
| raise ValueError("Cần đúng một trong input_ids hoặc inputs_embeds.") |
| x = self.embeddings(input_ids) if inputs_embeds is None else inputs_embeds |
| b, t, _ = x.shape |
|
|
| if attention_mask is None: |
| mask = torch.ones(b, t, dtype=torch.long, device=x.device) |
| else: |
| mask = attention_mask.to(device=x.device, dtype=torch.long) |
| if mask.shape[1] != t: |
| mask = mask[:, -t:] |
|
|
| if cache_params is not None and cache_params.batch_size != (1 if chunked else b): |
| raise ValueError("Batch size của cache không khớp với đầu vào.") |
| if use_cache and cache_params is None: |
| cache_params = KentCache(self.config, 1 if chunked else b, x.device) |
|
|
| all_hidden = () if output_hidden_states else None |
| for i, layer in enumerate(self.layers): |
| if output_hidden_states: |
| all_hidden += (x,) |
| h0 = cache_params.h[i] if cache_params is not None else None |
| if self.gradient_checkpointing and self.training: |
| x, h_last = self._gradient_checkpointing_func(layer.__call__, x, mask, h0, chunked) |
| else: |
| x, h_last = layer(x, mask, h0, chunked) |
| if use_cache: |
| cache_params.h[i] = h_last |
| if use_cache: |
| cache_params.steps += b * t if chunked else t |
|
|
| x = self.norm_f(x) |
| if chunked: |
| x = x.reshape(1, b * t, -1) |
| if all_hidden is not None: |
| all_hidden = tuple(h.reshape(1, b * t, -1) for h in all_hidden) |
| if all_hidden is not None: |
| all_hidden += (x,) |
|
|
| if not return_dict: |
| return tuple(v for v in (x, cache_params, all_hidden) if v is not None) |
| return KentOutput( |
| last_hidden_state=x, |
| cache_params=cache_params if use_cache else None, |
| hidden_states=all_hidden, |
| ) |
|
|
|
|
| class KentForCausalLM(KentPreTrainedModel, GenerationMixin): |
| _tied_weights_keys = {"lm_head.weight": "model.embeddings.weight"} |
|
|
| def __init__(self, config: KentConfig): |
| super().__init__(config) |
| self.model = KentModel(config) |
| self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) |
| self.post_init() |
|
|
| def get_input_embeddings(self): |
| return self.model.get_input_embeddings() |
|
|
| def set_input_embeddings(self, value): |
| self.model.set_input_embeddings(value) |
|
|
| def get_output_embeddings(self): |
| return self.lm_head |
|
|
| def set_output_embeddings(self, new_embeddings): |
| self.lm_head = new_embeddings |
|
|
| def new_cache(self, batch_size: int = 1) -> KentCache: |
| p = next(self.parameters()) |
| return KentCache(self.config, batch_size, p.device) |
|
|
| |
| def prepare_inputs_for_generation( |
| self, |
| input_ids, |
| inputs_embeds=None, |
| attention_mask=None, |
| use_cache=None, |
| cache_params=None, |
| past_key_values=None, |
| **kwargs, |
| ): |
| if cache_params is None and isinstance(past_key_values, KentCache): |
| cache_params = past_key_values |
| if cache_params is not None and cache_params.steps > 0: |
| input_ids = input_ids[:, -1:] |
| inputs_embeds = None |
| inputs = {"inputs_embeds": inputs_embeds} if inputs_embeds is not None else {"input_ids": input_ids} |
| inputs.update( |
| attention_mask=attention_mask, |
| cache_params=cache_params, |
| use_cache=True if use_cache is None else use_cache, |
| ) |
| return inputs |
|
|
| def _update_model_kwargs_for_generation(self, outputs, model_kwargs, is_encoder_decoder=False, **kwargs): |
| model_kwargs["cache_params"] = outputs.get("cache_params", None) |
| model_kwargs.pop("past_key_values", None) |
| if model_kwargs.get("attention_mask", None) is not None: |
| am = model_kwargs["attention_mask"] |
| model_kwargs["attention_mask"] = torch.cat([am, am.new_ones((am.shape[0], 1))], dim=-1) |
| if model_kwargs.get("cache_position", None) is not None: |
| model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + 1 |
| return model_kwargs |
|
|
| def forward( |
| self, |
| input_ids: Optional[torch.LongTensor] = None, |
| inputs_embeds: Optional[torch.FloatTensor] = None, |
| attention_mask: Optional[torch.Tensor] = None, |
| cache_params: Optional[KentCache] = None, |
| past_key_values: Optional[KentCache] = None, |
| use_cache: Optional[bool] = None, |
| chunked: bool = False, |
| labels: Optional[torch.LongTensor] = None, |
| logits_to_keep: int = 0, |
| output_hidden_states: Optional[bool] = None, |
| return_dict: Optional[bool] = None, |
| **kwargs, |
| ) -> Union[Tuple, KentCausalLMOutput]: |
| if cache_params is None and isinstance(past_key_values, KentCache): |
| cache_params = past_key_values |
| return_dict = return_dict if return_dict is not None else True |
| out = self.model( |
| input_ids=input_ids, |
| inputs_embeds=inputs_embeds, |
| attention_mask=attention_mask, |
| cache_params=cache_params, |
| use_cache=use_cache, |
| chunked=chunked, |
| output_hidden_states=output_hidden_states, |
| return_dict=True, |
| ) |
| hidden = out.last_hidden_state |
| if labels is None and logits_to_keep: |
| hidden = hidden[:, -logits_to_keep:] |
| logits = self.lm_head(hidden) |
|
|
| loss = None |
| if labels is not None: |
| shift_logits = logits[..., :-1, :].contiguous().float() |
| shift_labels = labels[..., 1:].contiguous().to(shift_logits.device) |
| loss = CrossEntropyLoss()(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) |
|
|
| if not return_dict: |
| return tuple(v for v in (loss, logits, out.cache_params, out.hidden_states) if v is not None) |
| return KentCausalLMOutput( |
| loss=loss, |
| logits=logits, |
| cache_params=out.cache_params, |
| hidden_states=out.hidden_states, |
| ) |
|
|
|
|
| __all__ = [ |
| "KentConfig", |
| "KentCache", |
| "KentModel", |
| "KentForCausalLM", |
| "KentPreTrainedModel", |
| "dct2", |
| "ema_scan", |
| ] |
|
|
