README.md

---
language:
- vi
license: apache-2.0
library_name: transformers
model_name: SAI_35M
pipeline_tag: text-generation
tags:
- custom-code
- vietnamese
- causal-lm
---

# SAI_35M — Mô hình ngôn ngữ **decoder-only Transformer** thuần Việt

## Kiến trúc mô hình

| Cấu hình | `d_model` | Attention heads | KV heads | Layers | FFN dimension |
|---|---:|---:|---:|---:|---:|
| 35M | 512 | 8 | 8 | 8 | 2.048 |
| 100M | 768 | 12 | 6 | 12 | 3.072 |
| 500M | 1.024 | 16 | 8 | 32 | 4.096 |

### Tokenizer

Sử dụng thư viện **SentencePiece Unigram** với vocabulary 10.000 token. Tokenizer được huấn luyện trên dữ liệu tiếng Việt đã chuyển về chữ thường, giữ nguyên khoảng trắng và bật byte fallback để hạn chế token không xác định.

### Embedding

Mỗi token ID được ánh xạ thành một vector có kích thước `d_model` bằng `nn.Embedding`. Khác với Transformer trong bài báo *Attention Is All You Need*, vốn cộng sinusoidal positional encoding trực tiếp vào token embedding, SAI sử dụng **Rotary Positional Embedding (RoPE)**. RoPE mã hoá vị trí tương đối bằng cách xoay các vector query và key trong attention. Context tối đa của mô hình là 2.048 token.

Trọng số của output language-model head được chia sẻ với ma trận embedding (weight tying). Cách này giảm số tham số và buộc biểu diễn đầu vào, đầu ra cùng nằm trong một không gian token.

### Grouped Query Attention

Query, key và value được chiếu qua một lớp tuyến tính gộp không dùng bias. Nhiều query head có thể dùng chung một cặp key/value head, nhờ đó giảm kích thước KV cache so với Multi-Head Attention thông thường.

- Cấu hình 35M có 8 query head và 8 KV head (tương đương Multi-Head Attention).
- Cấu hình 100M có 12 query head và 6 KV head.
- Cấu hình 500M có 16 query head và 8 KV head.

Attention được tính bằng `scaled_dot_product_attention` của PyTorch, cho phép tự chọn cuDNN Attention, Flash Attention, memory-efficient attention hoặc math backend phù hợp. Causal mask ngăn mô hình nhìn thấy token tương lai; padding mask loại bỏ token đệm khi huấn luyện.

### Decoder Block

Mỗi decoder block dùng kiến trúc **pre-norm** với hai nhánh residual:

```text
x = x + Dropout(GQA(RMSNorm(x)))
x = x + Dropout(SwiGLU(RMSNorm(x)))
```

Feed-forward network dùng SwiGLU: một phép chiếu sinh đồng thời nhánh `gate` và `up`, sau đó tính `up × SiLU(gate)` và chiếu trở lại `d_model`. Tất cả lớp tuyến tính trong attention và SwiGLU đều không dùng bias. Sau chồng decoder block, mô hình áp dụng một RMSNorm cuối trước language-model head.

So với Transformer gốc, SAI chỉ giữ phần decoder và thay LayerNorm bằng RMSNorm, ReLU FFN bằng SwiGLU. Transformer gốc áp dụng normalization sau residual connection (post-norm), trong khi SAI chuẩn hoá đầu vào trước attention và feed-forward network (pre-norm).

### Training

Quá trình huấn luyện gồm bốn giai đoạn: pretraining, continued pretraining, SFT1 và SFT2. Pretraining tối ưu next-token cross-entropy trên toàn bộ chuỗi; SFT dùng loss mask để chỉ học trên token thuộc câu trả lời.

Loss được chia thành từng phần nhỏ để giảm lượng bộ nhớ cần dùng khi huấn luyện. Pipeline cũng hỗ trợ mixed precision, gradient accumulation và `torch.compile` để tăng tốc trên GPU.

### Inference

Prompt của người dùng được đóng gói bằng chính hàm `build_input` của dự án, bao gồm chuyển nội dung về chữ thường và thêm các token vai trò. Bước **prefill** xử lý toàn bộ prompt một lần và lưu key/value của từng decoder block vào KV buffer cấp phát sẵn. Các bước sau gọi `decode_step` cho token mới và cập nhật trực tiếp buffer, thay vì dùng `GenerationMixin` hoặc `DynamicCache` của Transformers.

Mặc định, mô hình sinh tối đa 200 token bằng beam search với 5 beam, repetition penalty `1.2` và no-repeat 3-gram. Quá trình sinh dừng khi gặp `[EOS]`, `<|im_end|>`, đạt giới hạn context hoặc thoả điều kiện dừng sớm.

### Source và adapter Hugging Face

Repo chỉ giữ một bộ code inference với tên file nguyên bản: `TransformerModel.py`, `DecoderBlock.py`, `GroupedQueryAttention.py`, `RotaryPositionalEmbedding.py`, `SwiGLU.py` và `generate.py`. Nội dung thuật toán lấy trực tiếp từ project gốc; thay đổi duy nhất là chuyển các import `src.*` thành relative import để cơ chế `trust_remote_code` tải được chúng.

`modeling_sai.py` chỉ là adapter mỏng: đọc `SAIConfig`, giữ nguyên tên tensor trong state dict, chuyển lời gọi sang `TransformerModel` gốc và materialize RoPE/causal-mask buffer sau cơ chế meta-device của `from_pretrained()`. Adapter không cài đặt lại attention, decoder, KV cache hay beam search.

## Tải mô hình và chạy thử

Checkpoint hiện được phát hành là phiên bản **35M**. Mô hình có thể tải trực tiếp từ Hugging Face bằng thư viện Transformers.
 
```bash
pip install torch transformers sentencepiece
```

Ví dụ chạy thử:

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


model_id = "thongbuind/SAI_35M"

tokenizer = AutoTokenizer.from_pretrained(
    model_id,
    trust_remote_code=True,
)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    trust_remote_code=True,
)
device = "cuda" if torch.cuda.is_available() else (
    "mps" if torch.backends.mps.is_available() else "cpu"
)
model = model.to(device).eval()

answer = model.generate(
    "hướng dẫn tôi cách nấu cháo gà",
    tokenizer,
    max_new_tokens=200,
    beam_size=5,
    penalty=1.2,
    no_repeat_ngram=3,
    early_stop=False,
    patience=30,
)
print(answer)
```

`AutoModelForCausalLM` và `AutoTokenizer` chỉ được dùng làm lớp tương thích để tải config, trọng số và SentencePiece tokenizer. Thuật toán inference không gọi `GenerationMixin.generate()`, `DynamicCache`, `GenerationConfig`, `apply_chat_template()` hay beam search mặc định của Transformers. Phương thức `SAIForCausalLM.generate()` gọi trực tiếp beam search trong `generate.py`.

Model cần nằm trọn trên một device vì KV buffer của thuật toán gốc được cấp phát trên device chứa `lm_head`. Không dùng `device_map="auto"` để chia model qua nhiều device.

Trong lần chạy đầu tiên, Transformers sẽ tự tải model, tokenizer và config; các lần sau sẽ sử dụng bộ nhớ đệm cục bộ.