Upload 4 files

README.md

@@ -10,10 +10,6 @@ tags:
 - dct
 - swiglu
 - custom_code
-- base
-datasets:
-- HuggingFaceFW/fineweb-edu
-- wikimedia/wikipedia
 ---
 
 # KenT-300M-Base
@@ -29,7 +25,7 @@ H_i[t+1] = (1 - A_i) ⊙ H_i[t] + A_i ⊙ u_i     # Learnable EMA Matrix, A = si
 
 S   = Σ_i H_i[t+1]                            # Add
 y   = S · D                                   # "DCT-II Proj (D)" (DCT ngược)
-out = y + SwiGLU(RMSNorm(y))
+out = x[t] + y + SwiGLU(RMSNorm(y))         # 2 skip: sau phép chiếu D và từ đầu vào gốc x[t]
 logits = LM_head(RMSNorm_f(out))              # tied với embedding
 ```
 
@@ -48,12 +44,10 @@ D là ma trận DCT-II trực chuẩn cố định (hàng = tần số, vector h
 EMA khởi tạo theo cấp số nhân từ 0.5 (tap đầu, feature đầu) đến 0.01 (tap cuối, feature cuối); vẫn học được.
 EMA chạy FP32 bằng affine-scan log-depth, kết quả bằng đúng hồi quy tuần tự.
 
-![KenT diagram](KenT.png)
-
 ```python
 from transformers import AutoModelForCausalLM, AutoTokenizer
 tok = AutoTokenizer.from_pretrained("Vi-LAB/KenT-300M-Base", trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained("Vi-LAB/KenT-300M-Base", trust_remote_code=True)
 ```
 
-Hoặc `python infer.py --model Vi-LAB/KenT-300M-Base --prompt "Today the weather is"`.
+Hoặc `python infer.py --model Vi-LAB/KenT-300M-Base --prompt "Today the weather is"`.

config.json

@@ -19,7 +19,7 @@
   "pad_token_id": 50256,
   "bos_token_id": 50256,
   "eos_token_id": 50256,
-  "dtype": "float16",
+  "dtype": "bfloat16",
   "ema_alpha_first": 0.5,
   "ema_alpha_last": 0.01,
   "use_gradient_checkpointing": false

infer.py

@@ -28,8 +28,8 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
 # --------------------------------------------------------------------------- #
 def load(model_id: str, device: Optional[str] = None, dtype: Optional[torch.dtype] = None):
     device = device or ("cuda" if torch.cuda.is_available() else "cpu")
-    if dtype is None:  # FP16 trên GPU; CPU dùng FP32 vì nhiều phép FP16 chậm/không hỗ trợ
-        dtype = torch.float16 if device.startswith("cuda") else torch.float32
+    if dtype is None:  # BF16 trên GPU; CPU dùng FP32 vì nhiều phép FP16 chậm/không hỗ trợ
+        dtype = torch.bfloat16 if device.startswith("cuda") else torch.float32
     tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
     model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True)
     model = model.to(device=device, dtype=dtype).eval()

model.py

@@ -11,7 +11,7 @@ Mỗi bước nhận ĐÚNG MỘT token x[t] (vector 1 x D); lịch sử nằm h
 
     S   = sum_i H_i[t+1]                         # Add
     y   = S · D                                  # "DCT-II Proj (D)": DCT ngược (D trực chuẩn, D^T D = I)
-    out = y + SwiGLU( RMSNorm(y) )               # RMSNorm -> SwiGLU -> Add (skip lấy sau phép chiếu D)
+    out = x[t] + y + SwiGLU( RMSNorm(y) )        # RMSNorm -> SwiGLU -> Add: skip lấy sau phép chiếu D VÀ skip từ đầu vào gốc x[t]
     logits = LM_head( RMSNorm_f(out) )           # head buộc chung trọng số với embedding
 
 Quy ước hàng: x là vector hàng nên DCT-II là x·D^T (D[k, i] = α_k cos(π(i+½)k/D), hàng = tần số k, cột = vị trí i)
@@ -22,7 +22,7 @@ bằng DCT ngược.
 toàn mô hình chỉ là H shape (n+1, D): O(1) theo độ dài chuỗi.
 
 EMA là hồi quy tuyến tính theo từng kênh nê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ự. EMA và RMSNorm chạy ở FP32; trọng số lưu ở FP16.
+(Hillis-Steele): kết quả bằng đúng hồi quy tuần tự. EMA và RMSNorm chạy ở FP32; trọng số lưu ở BF16.
 
 RMSNorm cuối (`norm_f`) trước LM head không có trong sơ đồ; thêm vào để logits ổn định.
 """
@@ -417,7 +417,8 @@ class KentModel(KentPreTrainedModel):
         # --- Add -> DCT-II Proj (D) -> RMSNorm -> SwiGLU -> Add
         s = h.sum(dim=2).to(feats[0].dtype)
         y = idct2(s)
-        y = y + self.out_mlp(self.out_norm(y))
+        x0 = feats[0].reshape(1, b * t, d) if chunked else feats[0]                     # đầu vào gốc x[t] (embedding)
+        y = x0 + y + self.out_mlp(self.out_norm(y))                                     # Add: skip sau D + skip đầu vào gốc
         y = self.norm_f(y)                                                              # (1, K*C, D) nếu chunked
         if chunked:
             y = y.reshape(1, b * t, d)