Keep a single inference source tree with original filenames

src/model/DecoderBlock.py → DecoderBlock.py

@@ -1,6 +1,6 @@
 import torch.nn as nn
-from src.model.SwiGLU import SwiGLU
-from src.model.GroupedQueryAttention import GroupedQueryAttention
+from .SwiGLU import SwiGLU
+from .GroupedQueryAttention import GroupedQueryAttention
 
 class DecoderBlock(nn.Module):
     def __init__(self, d_model: int, num_heads: int, num_kv_heads: int, ff_dim: int, dropout: float):

src/model/GroupedQueryAttention.py → GroupedQueryAttention.py

@@ -2,7 +2,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torch.nn.attention import sdpa_kernel, SDPBackend
-from src.model.RotaryPositionalEmbedding import RotaryPositionalEmbedding
+from .RotaryPositionalEmbedding import RotaryPositionalEmbedding
 
 _SDPA_BACKENDS = [
     SDPBackend.CUDNN_ATTENTION,

README.md

@@ -68,7 +68,7 @@ Mặc định, mô hình sinh tối đa 200 token bằng beam search với 5 bea
 
 ### Source và adapter Hugging Face
 
-Thư mục `src/` chứa nguyên vẹn các file Python từ project huấn luyện. `generation_sai.py` là bản sao byte-for-byte của `src/inference/generate.py`. Các module `runtime_*.py` lấy trực tiếp từ source model; thay đổi duy nhất ở các file cần thiết là chuyển import `src.*` thành relative import để cơ chế `trust_remote_code` có thể tải chúng.
+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.
 
@@ -115,7 +115,7 @@ answer = model.generate(
 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 `generation_sai.py`.
+`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.
 

src/model/TransformerModel.py → TransformerModel.py

@@ -1,8 +1,8 @@
 import torch
 import torch.nn as nn
-from src.inference.generate import generate
-from src.model.DecoderBlock import DecoderBlock
-from src.model.RotaryPositionalEmbedding import RotaryPositionalEmbedding
+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):

generation_sai.py

@@ -1,196 +0,0 @@
-import torch
-from pathlib import Path
-
-def build_input(user_input, tokenizer):
-    BOS      = tokenizer.piece_to_id("[BOS]")
-    EOS      = tokenizer.piece_to_id("[EOS]")
-    IM_START = tokenizer.piece_to_id("<|im_start|>")
-    IM_END   = tokenizer.piece_to_id("<|im_end|>")
-
-    user_ids = (
-        [IM_START]
-        + tokenizer.encode("user\n" + user_input.strip().lower(), out_type=int)
-        + [IM_END]
-    )
-    assistant_prefix = [IM_START] + tokenizer.encode("model\n", out_type=int)
-    input_ids = [BOS] + user_ids + assistant_prefix
-
-    return input_ids, len(input_ids), EOS, IM_END
-
-def decode_output(best, start, EOS, IM_END, tokenizer):
-    output_tokens = best["seq"][start:]
-    while output_tokens and output_tokens[-1] in (EOS, IM_END):
-        output_tokens.pop()
-    return tokenizer.decode(output_tokens)
-
-def score(seq, log_prob, start):
-    out_len = max(len(seq) - start, 1)
-    return log_prob / (out_len ** 1.0)
-
-def apply_penalty(logits, seq, penalty):
-    if penalty == 1.0:
-        return logits
-    logits = logits.clone()
-    for tid in set(seq):
-        if logits[tid] < 0:
-            logits[tid] *= penalty
-        else:
-            logits[tid] /= penalty
-    return logits
-
-def get_banned_tokens(seq, n):
-    banned = set()
-    if n > 0 and len(seq) >= n:
-        prefix = tuple(seq[-(n - 1):])
-        for i in range(len(seq) - n + 1):
-            if tuple(seq[i:i + n - 1]) == prefix:
-                banned.add(seq[i + n - 1])
-    return banned
-
-def forward_init(model, input_ids, max_beam_size, max_new_tokens, device):
-    prompt_len = len(input_ids)
-    max_total  = prompt_len + max_new_tokens
-
-    kv_buffers = model.init_cache(max_beam_size, max_total, device)
-
-    prompt_tensor = torch.tensor([input_ids], dtype=torch.long, device=device)
-
-    with torch.inference_mode():
-        first_logits, present_cache = model.prefill(prompt_tensor, kv_cache=None)
-        for i, (k, v) in enumerate(present_cache):
-            kv_buffers[i][0][:, :, :prompt_len, :].copy_(k.expand(max_beam_size, -1, -1, -1))
-            kv_buffers[i][1][:, :, :prompt_len, :].copy_(v.expand(max_beam_size, -1, -1, -1))
-
-    return first_logits[0], kv_buffers, prompt_len
-
-def forward_step(model, last_tokens, kv_buffers, cache_len):
-    with torch.inference_mode():
-        return model.decode_step(last_tokens, kv_buffers, cache_len)
-
-def beam_core(model, input_ids, start, score_fn, penalty_fn, max_new_tokens, beam_size, no_repeat_ngram, penalty, device, EOS, IM_END, max_seq_len, early_stop=True, patience=10):
-
-    first_logits, kv_buffers, cache_len = forward_init(
-        model, input_ids, beam_size, max_new_tokens, device
-    )
-
-    first_logits = penalty_fn(first_logits, input_ids, penalty)
-    first_lp = torch.clamp(torch.log_softmax(first_logits, -1), -1e9, 0.0)
-    topk_lp, topk_tok = torch.topk(first_lp, beam_size)
-
-    seqs = [input_ids + [int(t)] for t in topk_tok.tolist()]
-    log_probs = topk_lp.tolist()
-    dones = [int(t) in (EOS, IM_END) for t in topk_tok.tolist()]
-    unique_sets = [set(input_ids) | {int(t)} for t in topk_tok.tolist()]
-    completed = []
-    K = beam_size * 3
-    patience_counter = 0
-
-    for _ in range(max_new_tokens - 1):
-        if all(dones):
-            break
-
-        last = torch.tensor([[seqs[i][-1]] for i in range(len(seqs))], dtype=torch.long, device=device)
-        logits_batch = forward_step(model, last, kv_buffers, cache_len)
-        cache_len += 1
-
-        n_beams = len(seqs)
-        if penalty != 1.0:
-            pen_mask = torch.ones(n_beams, logits_batch.shape[-1], device=device)
-            for i, uid in enumerate(unique_sets):
-                if dones[i]:
-                    continue
-                idx = torch.tensor(list(uid), dtype=torch.long, device=device)
-                pen_mask[i, idx] = penalty
-            neg = logits_batch < 0
-            logits_batch = torch.where(neg, logits_batch * pen_mask, logits_batch / pen_mask)
-
-        lp_batch = torch.clamp(torch.log_softmax(logits_batch, -1), -1e9, 0.0)
-        topk_lp_b, topk_tok_b = torch.topk(lp_batch, K, dim=-1)
-        topk_lp_b = topk_lp_b.tolist()
-        topk_tok_b = topk_tok_b.tolist()
-
-        candidates = []
-        for beam_i in range(n_beams):
-            if dones[beam_i]:
-                continue
-            banned = get_banned_tokens(seqs[beam_i], no_repeat_ngram)
-            count = 0
-            for l, t in zip(topk_lp_b[beam_i], topk_tok_b[beam_i]):
-                if count >= beam_size:
-                    break
-                if t in banned:
-                    continue
-                new_seq = seqs[beam_i] + [t]
-                done = t in (EOS, IM_END) or len(new_seq) >= max_seq_len
-                candidates.append({
-                    "beam_src": beam_i,
-                    "seq": new_seq,
-                    "log_prob": log_probs[beam_i] + l,
-                    "done": done
-                })
-                count += 1
-
-        if not candidates:
-            break
-
-        candidates.sort(key=lambda x: score_fn(x["seq"], x["log_prob"], start), reverse=True)
-
-        kept, src = [], []
-        for c in candidates:
-            if c["done"]:
-                completed.append(c)
-            elif len(kept) < beam_size:
-                kept.append(c)
-                src.append(c["beam_src"])
-
-        if not kept:
-            break
-
-        if len(completed) >= beam_size:
-            best_done = max(completed, key=lambda x: score_fn(x["seq"], x["log_prob"], start))
-            best_alive = max(kept, key=lambda x: score_fn(x["seq"], x["log_prob"], start))
-            if score_fn(best_done["seq"], best_done["log_prob"], start) >= \
-               score_fn(best_alive["seq"], best_alive["log_prob"], start):
-                if early_stop:
-                    break
-                else:
-                    patience_counter += 1
-                    if patience_counter >= patience:
-                        break
-            else:
-                patience_counter = 0
-
-        if src != list(range(len(src))):
-            src_t = torch.tensor(src, dtype=torch.long, device=device)
-            for k_buf, v_buf in kv_buffers:
-                k_tmp = k_buf[src_t].clone()
-                v_tmp = v_buf[src_t].clone()
-                k_buf[:len(src)].copy_(k_tmp)
-                v_buf[:len(src)].copy_(v_tmp)
-
-        unique_sets = [unique_sets[src[i]] | {c["seq"][-1]} for i, c in enumerate(kept)]
-        seqs = [c["seq"] for c in kept]
-        log_probs = [c["log_prob"] for c in kept]
-        dones = [False] * len(kept)
-
-    pool = completed if completed else [{"seq": seqs[i], "log_prob": log_probs[i]} for i in range(len(seqs))]
-    return max(pool, key=lambda x: score_fn(x["seq"], x["log_prob"], start))
-
-# ================= PUBLIC API =================
-def generate(model, user_input, tokenizer,
-                   max_new_tokens=200, beam_size=5,
-                   no_repeat_ngram=3, penalty=1.2,
-                   early_stop=False, patience=30):
-
-    device = model.lm_head.weight.device
-    ids, start, EOS, IM_END = build_input(user_input, tokenizer)
-
-    best = beam_core(
-        model, ids, start,
-        score, apply_penalty,
-        max_new_tokens, beam_size, no_repeat_ngram, penalty,
-        device, EOS, IM_END, model.max_seq_len,
-        early_stop=early_stop, patience=patience
-    )
-
-    return decode_output(best, start, EOS, IM_END, tokenizer)

modeling_sai.py

@@ -3,12 +3,12 @@ import torch.nn as nn
 from transformers import PreTrainedModel
 
 from .configuration_sai import SAIConfig
-from .generation_sai import generate as _source_generate
-from .runtime_decoder import DecoderBlock as _SourceDecoderBlock
-from .runtime_gqa import GroupedQueryAttention as _SourceGroupedQueryAttention
-from .runtime_rope import RotaryPositionalEmbedding as _SourceRoPE
-from .runtime_swiglu import SwiGLU as _SourceSwiGLU
-from .runtime_transformer import TransformerModel
+from .generate import generate as _source_generate
+from .DecoderBlock import DecoderBlock as _SourceDecoderBlock
+from .GroupedQueryAttention import GroupedQueryAttention as _SourceGroupedQueryAttention
+from .RotaryPositionalEmbedding import RotaryPositionalEmbedding as _SourceRoPE
+from .SwiGLU import SwiGLU as _SourceSwiGLU
+from .TransformerModel import TransformerModel
 
 
 class SAIForCausalLM(PreTrainedModel):

runtime_decoder.py

@@ -1,28 +0,0 @@
-import torch.nn as nn
-from .runtime_swiglu import SwiGLU
-from .runtime_gqa import GroupedQueryAttention
-
-class DecoderBlock(nn.Module):
-    def __init__(self, d_model: int, num_heads: int, num_kv_heads: int, ff_dim: int, dropout: float):
-        super().__init__()
-        self.gqa = GroupedQueryAttention(d_model, num_heads, num_kv_heads, dropout)
-        self.ffn = SwiGLU(d_model, ff_dim)
-        self.norm1 = nn.RMSNorm(d_model, eps=1e-6)
-        self.norm2 = nn.RMSNorm(d_model, eps=1e-6)
-        self.drop = nn.Dropout(dropout)
-
-    def forward(self, x, cos, sin, attn_mask=None):
-        x = x + self.drop(self.gqa(self.norm1(x), cos, sin, attn_mask))
-        x = x + self.drop(self.ffn(self.norm2(x)))
-        return x
-
-    def prefill(self, x, cos, sin):
-        attn, kv = self.gqa.prefill(self.norm1(x), cos, sin)
-        x = x + self.drop(attn)
-        x = x + self.drop(self.ffn(self.norm2(x)))
-        return x, list(kv)
-
-    def forward_with_cache(self, x, kv, cache_len: int, cos, sin):
-        x = x + self.drop(self.gqa.forward_with_cache(self.norm1(x), kv, cache_len, cos, sin))
-        x = x + self.drop(self.ffn(self.norm2(x)))
-        return x

runtime_gqa.py

@@ -1,96 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torch.nn.attention import sdpa_kernel, SDPBackend
-from .runtime_rope import RotaryPositionalEmbedding
-
-_SDPA_BACKENDS = [
-    SDPBackend.CUDNN_ATTENTION,
-    SDPBackend.FLASH_ATTENTION,
-    SDPBackend.EFFICIENT_ATTENTION,
-    SDPBackend.MATH,
-]
-
-class GroupedQueryAttention(nn.Module):
-    def __init__(self, d_model: int, num_heads: int, num_kv_heads: int, dropout: float):
-        super().__init__()
-        assert d_model % num_heads == 0
-        assert num_heads % num_kv_heads == 0, \
-            "num_heads phải chia hết cho num_kv_heads"
-
-        self.num_heads = num_heads
-        self.num_kv_heads = num_kv_heads
-        self.num_groups = num_heads // num_kv_heads
-        self.d_k = d_model // num_heads
-        self.dropout_rate = dropout
-        self.q_dim = num_heads * self.d_k
-        self.kv_dim = num_kv_heads * self.d_k
-        self.wqkv = nn.Linear(d_model, self.q_dim + 2 * self.kv_dim, bias=False)
-        self.wo = nn.Linear(self.q_dim, d_model, bias=False)
-
-    def _project_qkv(self, x: torch.Tensor):
-        B, T, _ = x.shape
-        qkv = self.wqkv(x)
-        q, k, v = qkv.split([self.q_dim, self.kv_dim, self.kv_dim], dim=-1)
-        q = q.view(B, T, self.num_heads, self.d_k)
-        k = k.view(B, T, self.num_kv_heads, self.d_k)
-        v = v.view(B, T, self.num_kv_heads, self.d_k)
-        return q, k, v
-
-    def _merge(self, out: torch.Tensor) -> torch.Tensor:
-        B, _, T, _ = out.shape
-        return self.wo(out.transpose(1, 2).contiguous().view(B, T, self.num_heads * self.d_k))
-
-    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, attn_mask=None) -> torch.Tensor:
-        q, k, v = self._project_qkv(x)
-
-        q = RotaryPositionalEmbedding.apply_rope(q, cos, sin)
-        k = RotaryPositionalEmbedding.apply_rope(k, cos, sin)
-
-        q = q.transpose(1, 2)
-        k = k.transpose(1, 2)
-        v = v.transpose(1, 2)
-
-        dropout_p = self.dropout_rate if self.training else 0.0
-
-        with sdpa_kernel(_SDPA_BACKENDS):
-            out = F.scaled_dot_product_attention(
-                q, k, v, attn_mask=attn_mask, is_causal=(attn_mask is None), dropout_p=dropout_p,
-                enable_gqa=True,
-            )
-        return self._merge(out)
-
-    def prefill(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
-        q, k, v = self._project_qkv(x)
-
-        q = RotaryPositionalEmbedding.apply_rope(q, cos, sin)
-        k = RotaryPositionalEmbedding.apply_rope(k, cos, sin)
-
-        q = q.transpose(1, 2)
-        k = k.transpose(1, 2)
-        v = v.transpose(1, 2)
-
-        with sdpa_kernel(_SDPA_BACKENDS):
-            out = F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=True)
-        return self._merge(out), (k, v)
-
-    def forward_with_cache(self, x: torch.Tensor, past_kv, cache_len: int, cos: torch.Tensor, sin: torch.Tensor):
-        B, T, _ = x.shape
-        q, k, v = self._project_qkv(x)
-
-        q = RotaryPositionalEmbedding.apply_rope(q, cos, sin)
-        k = RotaryPositionalEmbedding.apply_rope(k, cos, sin)
-
-        q = q.transpose(1, 2)
-        k = k.transpose(1, 2)
-        v = v.transpose(1, 2)
-
-        past_kv[0][:B, :, cache_len:cache_len + T, :] = k
-        past_kv[1][:B, :, cache_len:cache_len + T, :] = v
-
-        k_full = past_kv[0][:B, :, :cache_len + T, :]
-        v_full = past_kv[1][:B, :, :cache_len + T, :]
-
-        with sdpa_kernel(_SDPA_BACKENDS):
-            out = F.scaled_dot_product_attention(q, k_full, v_full, is_causal=False, enable_gqa=True)
-        return self._merge(out)

runtime_rope.py

@@ -1,33 +0,0 @@
-import torch
-import torch.nn as nn
-
-class RotaryPositionalEmbedding(nn.Module):
-    def __init__(self, head_dim: int, max_seq_len: int, base: int = 10_000):
-        super().__init__()
-        assert head_dim % 2 == 0, "head_dim phải chẵn"
-        inv_freq = 1.0 / (
-            base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)
-        )
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self._build_cache(max_seq_len)
-
-    def _build_cache(self, seq_len: int):
-        pos = torch.arange(seq_len, dtype=torch.float32)
-        freqs = torch.outer(pos, self.inv_freq)  # (seq_len, head_dim/2)
-        emb = torch.cat([freqs, freqs], dim=-1)  # (seq_len, head_dim)
-        self.register_buffer("cos_cached", emb.cos(), persistent=False)
-        self.register_buffer("sin_cached", emb.sin(), persistent=False)
-
-    def get_cos_sin(self, positions: torch.Tensor):
-        cos = self.cos_cached[positions][None, :, None, :]
-        sin = self.sin_cached[positions][None, :, None, :]
-        return cos, sin
-
-    @staticmethod
-    def _rotate_half(x: torch.Tensor) -> torch.Tensor:
-        x1, x2 = x.chunk(2, dim=-1)
-        return torch.cat((-x2, x1), dim=-1)
-
-    @staticmethod
-    def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
-        return x * cos + RotaryPositionalEmbedding._rotate_half(x) * sin

runtime_swiglu.py

@@ -1,14 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-class SwiGLU(nn.Module):
-    def __init__(self, d_model: int, ff_dim: int):
-        super().__init__()
-        self.ff_dim = ff_dim
-        self.gate_up = nn.Linear(d_model, 2 * ff_dim, bias=False)
-        self.down = nn.Linear(ff_dim, d_model, bias=False)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        gate, up = self.gate_up(x).chunk(2, dim=-1)
-        return self.down(up * F.silu(gate))

runtime_transformer.py

@@ -1,114 +0,0 @@
-import torch
-import torch.nn as nn
-from .generation_sai import generate
-from .runtime_decoder import DecoderBlock
-from .runtime_rope 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)

src/data/Dataset.py

@@ -1,260 +0,0 @@
-import torch
-import numpy as np
-import gc
-import json
-from pathlib import Path
-from src.utils.utils import log_progress
-
-class Dataset(torch.utils.data.Dataset):
-    def __init__(self, X, Y, lengths, indices, loss_masks=None):
-        self.X = X
-        self.Y = Y
-        self.lengths = lengths
-        self.indices = indices
-        self.loss_masks = loss_masks
-
-    def __len__(self):
-        return len(self.indices)
-
-    def __getitem__(self, idx):
-        real_idx = self.indices[idx]
-        return (
-            self.X[real_idx],
-            self.Y[real_idx],
-            self.lengths[real_idx],
-            self.loss_masks[real_idx] if self.loss_masks is not None else None
-        )
-    
-    @classmethod
-    def create_dataloader(cls, X, Y, lengths, batch_size, max_seq_len, num_workers, shuffle, loss_masks=None):
-        log_progress(f"Đang tạo dataset từ {len(X)} samples...")
-        indices = np.arange(len(X))
-        dataset = cls(X, Y, lengths, indices, loss_masks)
-        bucket_size = 128
-        PAD_ID = 0
-
-        def collate_fn(batch):
-            X_batch = [item[0] for item in batch]
-            Y_batch = [item[1] for item in batch]
-            lm_batch = [item[3] for item in batch]
-
-            max_len = max(len(x) for x in X_batch)
-            if max_len > (max_seq_len - (max_seq_len%bucket_size)):
-                bucket = max_seq_len
-            else:
-                bucket = int(np.ceil(max_len / bucket_size) * bucket_size)
-
-            bsz = len(X_batch)
-            X_padded = torch.zeros((bsz, bucket), dtype=torch.long)
-            Y_padded = torch.zeros((bsz, bucket), dtype=torch.long)
-            has_padding = any(len(x) < bucket for x in X_batch)
-
-            if loss_masks is not None:
-                loss_mask_padded = torch.zeros((bsz, bucket), dtype=torch.float)
-
-            for i, (x, y) in enumerate(zip(X_batch, Y_batch)):
-                x_len = len(x)
-                y_len = len(y)
-
-                X_padded[i, :x_len] = torch.as_tensor(x, dtype=torch.long)
-                Y_padded[i, :y_len] = torch.as_tensor(y, dtype=torch.long)
-
-                if loss_masks is not None:
-                    lm = lm_batch[i]
-                    lm_len = len(lm)
-                    loss_mask_padded[i, :lm_len] = torch.as_tensor(lm, dtype=torch.float)
-
-            attention_mask = (X_padded != PAD_ID).float()
-
-            if loss_masks is None:
-                sample_weight = (Y_padded != PAD_ID).float()
-            else:
-                sample_weight = loss_mask_padded
-
-            return X_padded, Y_padded, sample_weight, attention_mask, has_padding
-
-        dataloader = torch.utils.data.DataLoader(
-            dataset,
-            batch_size=batch_size,
-            shuffle=shuffle,
-            collate_fn=collate_fn,
-            num_workers=num_workers,
-            pin_memory=torch.cuda.is_available(),
-            persistent_workers = True,
-            drop_last=False
-        )
-
-        log_progress(f"Dataset được tạo với batch_size={batch_size}")
-        return dataloader
-
-def split_train_val_test(X, Y, loss_masks, lengths, train_ratio, val_ratio, seed=54):
-    total_sample = len(X)
-    rng = np.random.default_rng(seed)
-    indices = rng.permutation(total_sample)
-
-    train_end = int(total_sample * train_ratio)
-    val_end = int(total_sample * (train_ratio + val_ratio))
-
-    train_idx = indices[:train_end]
-    val_idx = indices[train_end:val_end]
-    test_idx = indices[val_end:]
-
-    X_train, Y_train, lengths_train = X[train_idx], Y[train_idx], lengths[train_idx]
-    X_val, Y_val, lengths_val = X[val_idx], Y[val_idx], lengths[val_idx]
-    X_test, Y_test, lengths_test = X[test_idx], Y[test_idx], lengths[test_idx]
-    
-    if loss_masks is not None:
-        mask_train = loss_masks[train_idx]
-        mask_val = loss_masks[val_idx]
-        mask_test = loss_masks[test_idx]
-        
-        return (X_train, Y_train, mask_train, lengths_train, 
-                X_val, Y_val, mask_val, lengths_val, 
-                X_test, Y_test, mask_test, lengths_test)
-    else:
-        return (X_train, Y_train, None, lengths_train,
-                X_val, Y_val, None, lengths_val,
-                X_test, Y_test, None, lengths_test)
-
-def _load_single_npz(path, load_mask=False):
-    """Bản gốc: đọc 1 file .npz duy nhất."""
-    f = np.load(path)
-    def reconstruct(name):
-        flat, offsets = f[f"{name}_flat"], f[f"{name}_offsets"]
-        return np.array([flat[offsets[i]:offsets[i+1]] for i in range(len(offsets)-1)], dtype=object)
-    X = reconstruct("X")
-    Y = reconstruct("Y")
-    lengths = f["lengths"]
-    if load_mask:
-        M = reconstruct("M")
-        f.close()
-        return X, Y, M, lengths
-    f.close()
-    return X, Y, lengths
-
-def _load_manifest_shards(manifest_path, load_mask=False):
-    """Đọc manifest.json rồi ghép nhiều shard .npz lại thành 1 dataset logic."""
-    with open(manifest_path, "r", encoding="utf-8") as f:
-        manifest = json.load(f)
-
-    manifest_dir = Path(manifest_path).parent
-    shards = manifest["shards"]
-
-    X_list, Y_list, M_list, L_list = [], [], [], []
-    for shard in shards:
-        shard_path = manifest_dir / shard["file"]
-        if load_mask:
-            X_s, Y_s, M_s, L_s = _load_single_npz(shard_path, load_mask=True)
-            M_list.append(M_s)
-        else:
-            X_s, Y_s, L_s = _load_single_npz(shard_path, load_mask=False)
-        X_list.append(X_s)
-        Y_list.append(Y_s)
-        L_list.append(L_s)
-
-    X = np.concatenate(X_list)
-    Y = np.concatenate(Y_list)
-    lengths = np.concatenate(L_list)
-
-    if load_mask:
-        M = np.concatenate(M_list)
-        return X, Y, M, lengths
-    return X, Y, lengths
-
-def load_npz(path, load_mask=False):
-    """
-    - path là file .json -> manifest nhiều shard, tự đọc & ghép tất cả.
-    - path là file .npz -> đọc 1 file như bản gốc.
-    """
-    path = Path(path)
-    if path.suffix == ".json":
-        return _load_manifest_shards(path, load_mask=load_mask)
-    return _load_single_npz(path, load_mask=load_mask)
-
-def load_data(data_type, main_data, sub_data=None, seed=54):
-    if data_type == "pretrain":
-        X, Y, lengths = load_npz(main_data)
-        return X, Y, lengths
-
-    elif data_type == "continued_pretrain":
-        X_main, Y_main, L_main = load_npz(main_data)
-
-        if sub_data is not None:
-            X_sub, Y_sub, L_sub = load_npz(sub_data)
-
-            n_continued = len(X_main)
-            n_pretrain_needed = 3 * n_continued
-            total_pretrain = len(X_sub)
-            
-            rng = np.random.default_rng(seed)
-            shuffled_indices = rng.permutation(total_pretrain)
-            
-            n_samples = min(n_pretrain_needed, total_pretrain)
-            selected_indices = shuffled_indices[:n_samples]
-            
-            X_sub_sampled = X_sub[selected_indices]
-            Y_sub_sampled = Y_sub[selected_indices]
-            L_sub_sampled = L_sub[selected_indices]
-            
-            del X_sub, Y_sub, L_sub, shuffled_indices, selected_indices
-            gc.collect()
-            
-            X_combined = np.concatenate([X_main, X_sub_sampled])
-            Y_combined = np.concatenate([Y_main, Y_sub_sampled])
-            L_combined = np.concatenate([L_main, L_sub_sampled])
-            
-            combined_indices = rng.permutation(len(X_combined))
-            X = X_combined[combined_indices]
-            Y = Y_combined[combined_indices]
-            lengths = L_combined[combined_indices]
-            
-        else:
-            X, Y, lengths = X_main, Y_main, L_main
-
-        return X, Y, lengths
-
-    elif data_type == "sft1":
-        X, Y, loss_mask, lengths = load_npz(main_data, load_mask=True)
-        return X, Y, loss_mask, lengths
-    
-    elif data_type == "sft2":
-        X_main, Y_main, M_main, L_main = load_npz(main_data, load_mask=True)
-
-        if sub_data is not None:
-            X_sub, Y_sub, M_sub, L_sub = load_npz(sub_data, load_mask=True)
-
-            n_sub = len(X_main) * 3
-            total_sub = len(X_sub)
-
-            rng = np.random.default_rng(seed)
-            shuffled_indices = rng.permutation(total_sub)
-
-            n_samples = min(n_sub, total_sub)
-            selected_indices = shuffled_indices[:n_samples]
-
-            X_sub_sampled = X_sub[selected_indices]
-            Y_sub_sampled = Y_sub[selected_indices]
-            M_sub_sampled = M_sub[selected_indices]
-            L_sub_sampled = L_sub[selected_indices]
-
-            del X_sub, Y_sub, M_sub, L_sub, shuffled_indices, selected_indices
-            gc.collect()
-
-            X_combined = np.concatenate([X_main, X_sub_sampled])
-            Y_combined = np.concatenate([Y_main, Y_sub_sampled])
-            M_combined = np.concatenate([M_main, M_sub_sampled])
-            L_combined = np.concatenate([L_main, L_sub_sampled])
-
-            combined_indices = rng.permutation(len(X_combined))
-            X = X_combined[combined_indices]
-            Y = Y_combined[combined_indices]
-            loss_mask = M_combined[combined_indices]
-            lengths = L_combined[combined_indices]
-
-        else:
-            X, Y, loss_mask, lengths = X_main, Y_main, M_main, L_main
-
-        return X, Y, loss_mask, lengths
-
-    else:
-        raise ValueError(f"Unknown data_type: {data_type}")

src/tokenizer/encoders.py

@@ -1,224 +0,0 @@
-import json
-from pathlib import Path
-import numpy as np
-import sentencepiece as spm
-
-current_file = Path(__file__).resolve()
-project_root = current_file.parent.parent.parent
-config_file = project_root / "config" / "base.json"
-tokenizer_dir = project_root / "src" / "tokenizer"
-data_dir = project_root / "data"
-raw_dir = data_dir / "raw"
-processed_dir = data_dir / "processed"
-processed_dir.mkdir(parents=True, exist_ok=True)
-
-with open(config_file, 'r') as f:
-    config = json.load(f)
-max_seq_len = config['max_seq_len']
-
-tokenizer_file = tokenizer_dir / "tokenizer.model"
-sp = spm.SentencePieceProcessor()
-sp.load(str(tokenizer_file))
-
-BOS = sp.piece_to_id("[BOS]")
-EOS = sp.piece_to_id("[EOS]")
-IM_START = sp.piece_to_id("<|im_start|>")
-IM_END = sp.piece_to_id("<|im_end|>")
-
-def iter_text_jsonl(path, key="text"):
-    with open(path, "r", encoding="utf-8") as f:
-        for line in f:
-            line = line.strip()
-            if not line:
-                continue
-            try:
-                obj = json.loads(line)
-            except json.JSONDecodeError:
-                continue
-            if isinstance(obj, dict) and key in obj:
-                text = obj[key].strip()
-                if text:
-                    yield text
-
-def iter_sft_conversations(path):
-    """Đọc file .json chứa 1 mảng các object {instruction, history: [{user, model}, ...]}."""
-    with open(raw_dir / path, "r", encoding="utf-8") as f:
-        data = json.load(f)
-    if not isinstance(data, list):
-        data = [data]
-    for obj in data:
-        if not isinstance(obj, dict):
-            continue
-        history = obj.get("history", [])
-        if not history:
-            continue
-        yield obj
-
-def _encode_pretrain(text):
-    tokens = sp.encode(text.lower(), out_type=int)
-    if len(tokens) < 2 or len(tokens) + 2 > max_seq_len:
-        return None
-    x = [BOS] + tokens
-    y = tokens + [EOS]
-    return x, y
-
-def _encode_sft(sample):
-    instruction = sample.get("instruction", "").strip()
-    history = sample.get("history", [])
-    if not history:
-        return None
-
-    all_tokens = []
-    mask = []
-
-    if instruction:
-        instr_ids = (
-            [IM_START]
-            + sp.encode("system\n" + instruction.lower(), out_type=int)
-            + [IM_END]
-        )
-        all_tokens += instr_ids
-        mask += [0] * len(instr_ids)
-
-    for turn in history:
-        user_text = turn.get("user", "").strip()
-        model_text = turn.get("model", "").strip()
-        if not model_text:
-            continue
-
-        user_ids = (
-            [IM_START]
-            + sp.encode("user\n" + user_text.lower(), out_type=int)
-            + [IM_END]
-        )
-        model_ids = (
-            [IM_START]
-            + sp.encode("model\n" + model_text.lower(), out_type=int)
-            + [IM_END]
-        )
-
-        all_tokens += user_ids + model_ids
-        mask += [0] * len(user_ids) + [1] * len(model_ids)
-
-    if not all_tokens:
-        return None
-
-    input_ids = [BOS] + all_tokens + [EOS]
-    if len(input_ids) > max_seq_len:
-        return None
-
-    target_ids = input_ids[1:]
-    mask = mask + [1]
-
-    assert len(mask) == len(target_ids), (
-        f"Độ dài mask ({len(mask)}) khác target_ids ({len(target_ids)})"
-    )
-    return input_ids, target_ids, mask
-
-def _two_pass_write(path, encoded_iter_fn, has_mask=False):
-    """Logic 2-pass gốc, ghi ra .npz (giống bản đầu tiên)."""
-    x_offsets = [0]
-    y_offsets = [0]
-    m_offsets = [0]
-    lengths = []
-
-    for idx, item in enumerate(encoded_iter_fn()):
-        if idx % 100_000 == 0:
-            print(f"  [pass 1] {idx:,} dòng...")
-        if has_mask:
-            x, y, m = item
-        else:
-            x, y = item
-        x_offsets.append(x_offsets[-1] + len(x))
-        y_offsets.append(y_offsets[-1] + len(y))
-        if has_mask:
-            m_offsets.append(m_offsets[-1] + len(m))
-        lengths.append(len(x))
-
-    n = len(lengths)
-    if n == 0:
-        return 0, 0
-
-    total_x = x_offsets[-1]
-    total_y = y_offsets[-1]
-
-    x_flat = np.empty(total_x, dtype=np.int32)
-    y_flat = np.empty(total_y, dtype=np.int32)
-    if has_mask:
-        m_flat = np.empty(m_offsets[-1], dtype=np.int8)
-
-    x_off_arr = np.array(x_offsets, dtype=np.int64)
-    y_off_arr = np.array(y_offsets, dtype=np.int64)
-    len_arr = np.array(lengths, dtype=np.int32)
-    if has_mask:
-        m_off_arr = np.array(m_offsets, dtype=np.int64)
-
-    for idx, item in enumerate(encoded_iter_fn()):
-        if idx % 100_000 == 0:
-            print(f"  [pass 2] {idx:,}/{n:,} dòng...")
-        if has_mask:
-            x, y, m = item
-            m_flat[m_offsets[idx]: m_offsets[idx + 1]] = m
-        else:
-            x, y = item
-        x_flat[x_offsets[idx]: x_offsets[idx + 1]] = x
-        y_flat[y_offsets[idx]: y_offsets[idx + 1]] = y
-
-    out = {
-        "X_flat": x_flat, "X_offsets": x_off_arr,
-        "Y_flat": y_flat, "Y_offsets": y_off_arr,
-        "lengths": len_arr,
-    }
-    if has_mask:
-        out["M_flat"] = m_flat
-        out["M_offsets"] = m_off_arr
-
-    np.savez_compressed(path, **out)
-    return n, int(len_arr.sum())
-
-def _write_manifest(manifest_path, shards):
-    """Sổ mục lục các shard .npz để lúc load ghép lại thành 1 dataset logic."""
-    manifest = {
-        "shards": shards,
-        "total_samples": sum(s["n_samples"] for s in shards),
-        "total_tokens": sum(s["n_tokens"] for s in shards),
-    }
-    with open(manifest_path, "w", encoding="utf-8") as f:
-        json.dump(manifest, f, ensure_ascii=False, indent=2)
-    return manifest
-
-# ---- Pretrain: mỗi file pretrain_data_00X.jsonl -> 1 shard .npz riêng ----
-pretrain_paths = sorted(raw_dir.glob("pretrain_data_*.jsonl"))
-print(f"Pretrain - tìm thấy {len(pretrain_paths)} file shard...")
-
-pretrain_shards = []
-for p in pretrain_paths:
-    shard_name = p.stem  # ví dụ: pretrain_data_001
-    print(f"Pretrain shard {shard_name} - pass 1 & 2...")
-    n, total = _two_pass_write(
-        processed_dir / f"{shard_name}_ids.npz",
-        lambda p=p: filter(None, (_encode_pretrain(t) for t in iter_text_jsonl(p))),
-    )
-    print(f"✅ {shard_name}: {n} samples | Tổng số token: {total:,}")
-    pretrain_shards.append({
-        "name": shard_name,
-        "file": f"{shard_name}_ids.npz",
-        "n_samples": n,
-        "n_tokens": total,
-    })
-
-manifest = _write_manifest(processed_dir / "pretrain_manifest.json", pretrain_shards)
-print(f"✅ Pretrain tổng: {manifest['total_samples']} samples | {manifest['total_tokens']:,} token | manifest: pretrain_manifest.json")
-
-# cont_path = raw_dir / "continued_pretrain_data.jsonl"
-# print("Continued Pretrain - pass 1 & 2...")
-# n, total = _two_pass_write(processed_dir / "continued_pretrain_data_ids.npz", lambda: filter(None, (_encode_pretrain(t) for t in iter_text_jsonl(cont_path))))
-# print(f"✅ Continued Pretrain: {n} samples | Tổng số token: {total:,}")
-
-print("SFT1 - pass 1 & 2...")
-n, total = _two_pass_write(processed_dir / "SFT1_data_ids.npz", lambda: filter(None, (_encode_sft(s) for s in iter_sft_conversations("SFT_1.json"))), has_mask=True)
-print(f"✅ Đã lưu SFT1: {n} samples | Tổng số token: {total:,}")
-
-# print("SFT2 - pass 1 & 2...")
-# n, total = _two_pass_write(processed_dir / "SFT2_data_ids.npz", lambda: filter(None, (_encode_sft(s) for s in iter_sft_conversations("SFT_2.json"))), has_mask=True)
-# print(f"✅ Đã lưu SFT2: {n} samples | Tổng số token: {total:,}")

src/tokenizer/tokenizer.model

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:3eb2e8f55e0472fbe6c5f597d601746934e69542b5238a18852132761c03a9ac
-size 156890

src/tokenizer/train_tokenizer.py

@@ -1,103 +0,0 @@
-import json
-import time
-from pathlib import Path
-import sentencepiece as spm
-
-# Khởi động đồng hồ đo thời gian toàn bộ tiến trình
-start_total_time = time.time()
-
-# 1. Giữ nguyên logic tìm thư mục gốc tự động
-current_file = Path(__file__).resolve()
-project_root = current_file.parent
-
-RAW_DATA_ABSOLUTE_PATH = Path("/Users/thongbui.nd/Documents/Thong Bui/dev_2/data/raw/pretrain_data.jsonl")
-
-vocab_size = 10000
-input_text_path = project_root / "spm_input.txt"
-
-# =========================================================
-# BƯỚC 1+2: ĐỌC, SÀN LỌC VÀ GHI THẲNG RA FILE TẠM
-# =========================================================
-print(f"🚀 Bước 1+2: Đang xử lý dữ liệu và ghi ra file tạm: {input_text_path.name}...")
-start_step1 = time.time()
-
-total_lines_read = 0
-malformed_lines = 0
-empty_lines = 0
-valid_count = 0
-
-# Đọc từ file hardcode tuyệt đối, ghi vào file tạm ở project_root
-with open(RAW_DATA_ABSOLUTE_PATH, "r", encoding="utf-8") as fin, \
-     open(input_text_path, "w", encoding="utf-8") as fout:
-
-    for line in fin:
-        total_lines_read += 1
-
-        if total_lines_read % 100000 == 0:
-            print(f"   [Tiến trình] Đã quét qua {total_lines_read:,} dòng dữ liệu thô...")
-
-        line = line.strip()
-        if not line:
-            empty_lines += 1
-            continue
-
-        try:
-            obj = json.loads(line)
-        except json.JSONDecodeError:
-            malformed_lines += 1
-            continue
-
-        if isinstance(obj, dict) and "text" in obj:
-            text = obj["text"].strip()
-            if text:
-                fout.write(text.lower() + "\n")
-                valid_count += 1
-            else:
-                empty_lines += 1
-
-print(f"📊 Thống kê dữ liệu đã đọc:")
-print(f"   - Tổng số dòng quét qua: {total_lines_read:,}")
-print(f"   - Số dòng trống/không có text: {empty_lines:,}")
-print(f"   - Số dòng lỗi định dạng JSON: {malformed_lines:,}")
-print(f"   - Số dòng hợp lệ đưa vào huấn luyện: {valid_count:,}")
-print(f"⏱️  Thời gian xử lý Bước 1+2: {time.time() - start_step1:.2f} giây\n")
-
-# =========================================================
-# BƯỚC 3: HUẤN LUYỆN SENTENCEPIECE
-# =========================================================
-print("🧠 Bước 3: Đang huấn luyện SentencePiece (Unigram)...")
-print("   * Lưu ý: Thư viện SentencePiece có bộ log C++ nội bộ riêng, chuẩn bị hiển thị bên dưới... *")
-start_step3 = time.time()
-
-spm.SentencePieceTrainer.train(
-    input=str(input_text_path),
-    model_prefix=str(project_root / "tokenizer"),
-    vocab_size=vocab_size,
-    model_type="unigram",
-    character_coverage=0.9999,
-    pad_piece="[PAD]",
-    unk_piece="[UNK]",
-    bos_piece="[BOS]",
-    eos_piece="[EOS]",
-    user_defined_symbols=["<|im_start|>", "<|im_end|>"],
-    normalization_rule_name='identity',
-    # split_by_whitespace=False,
-    treat_whitespace_as_suffix=False,
-    input_sentence_size=1000000,
-    shuffle_input_sentence=True,
-    remove_extra_whitespaces=False,
-    split_digits=True,
-    byte_fallback=True,
-    max_sentencepiece_length=8
-)
-
-print(f"\nThời gian huấn luyện SentencePiece: {time.time() - start_step3:.2f} giây\n")
-
-# =========================================================
-# HOÀN TẤT TIẾN TRÌNH
-# =========================================================
-print("🎉 Đã huấn luyện xong!")
-print(f"📁 Xuất file thành công tại:")
-print(f"   - Model: {project_root / 'tokenizer.model'}")
-print(f"   - Vocab: {project_root / 'tokenizer.vocab'}")
-print(f"🚀 Tổng thời gian thực hiện toàn bộ script: {time.time() - start_total_time:.2f} giây.")

src/training/PenaltyEngine.py

@@ -1,188 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from abc import ABC, abstractmethod
-from typing import Optional, NamedTuple
-
-"""
-WrongTokenMarginPenalty: Khi mô hình dự đoán sai
-    - Ép mô hình tăng logit của token đúng (nếu detach_max = False thì trực tiếp giảm thêm logit của token sai nhưng không nên dùng)
-
-WrongTokenEntropyPenalty: Khi mô hình đã dự đoán sai, mà lại còn tự tin vào cái token sai đấy
-    - Dấu hiệu: entropy thấp
-    - Công thức: penalty = max(0, min_entropy - H) * WrongMask
-    - Không quan tâm đến token đúng, chỉ cần gradient kéo toàn bộ phân phối xác suất phẳng hơn là được.
-
-FocalOverconfidencePenalty: Khi mô hình dự đoán sai
-    - Khuếch đại cái sai theo hệ số gamma, phạt cho xác suất của token sai giảm xuống
-    - Xác suất mô hình dành cho token sai càng lơn, phạt càng mạnh
-
-Các trường hợp:
-    - tất cả z đều thấp, và model dự đoán đúng: chưa xử lý
-    - tất cả z đều thấp, và model dự đoán sai: dùng WrongTokenMarginPenalty để kéo logit của token đúng lên
-    - tất cả z đều cao, và model dự đoán sai: dùng WrongTokenMarginPenalty để kéo logit của token đúng lên và FocalOverconfidencePenalty để giảm xác suất của token sai
-    - tất cả z đều cao, và model dự đoán đúng: chưa xử lý
-    - entropy thấp, model dự đoán đúng: chưa xử lý
-    - entropy thấp, model dự đoán sai: dùng WrongTokenMarginPenalty để kéo logit của token đúng lên và WrongTokenEntropyPenalty để kéo phẳng phân phối và FocalOverconfidencePenalty để giảm xác suất của token sai
-"""
-
-class PenaltyContext(NamedTuple):
-    logits: torch.Tensor
-    inputs: torch.Tensor
-    targets: torch.Tensor
-    loss_mask: Optional[torch.Tensor]
-    max_logits: torch.Tensor
-    predicted: torch.Tensor
-    target_logits: torch.Tensor
-    wrong_mask: torch.Tensor
-    log_probs: Optional[torch.Tensor]
-
-def _safe_masked_mean(tensor: torch.Tensor, mask: Optional[torch.Tensor]) -> torch.Tensor:
-    if mask is not None:
-        tensor = tensor * mask
-        denom = mask.sum().clamp(min=1e-6)
-    else:
-        denom = torch.tensor(tensor.numel(), dtype=tensor.dtype, device=tensor.device).clamp(min=1e-6)
-    return tensor.sum() / denom
-
-class PenaltyRule(nn.Module, ABC):
-    needs_log_probs: bool = False
-
-    @abstractmethod
-    def forward(self, ctx: PenaltyContext) -> torch.Tensor:
-        """Trả về scalar penalty."""
-
-class WrongTokenMarginPenalty(PenaltyRule):
-    needs_log_probs = False
-    def __init__(self, weight: float = 1.0, detach_max: bool = True):
-        super().__init__()
-        self.weight = weight
-        self.detach_max = detach_max
-
-    def forward(self, ctx: PenaltyContext) -> torch.Tensor:
-        max_l = ctx.max_logits.detach() if self.detach_max else ctx.max_logits
-        margin = F.relu(max_l - ctx.target_logits)
-        penalty = margin * ctx.wrong_mask
-
-        return self.weight * _safe_masked_mean(penalty, ctx.loss_mask)
-
-class WrongTokenEntropyPenalty(PenaltyRule):
-    needs_log_probs = True
-    def __init__(self, weight: float = 0.01, min_entropy: float = 0.5):
-        super().__init__()
-        self.weight = weight
-        self.min_entropy = min_entropy
-
-    def forward(self, ctx: PenaltyContext) -> torch.Tensor:
-        assert ctx.log_probs is not None, \
-            "WrongTokenEntropyPenalty cần log_probs; thêm needs_log_probs=True"
-
-        probs = ctx.log_probs.exp()
-        entropy = -(probs * ctx.log_probs).sum(dim=-1)
-
-        low_entropy = F.relu(self.min_entropy - entropy)
-
-        combined_mask = ctx.wrong_mask * (ctx.loss_mask if ctx.loss_mask is not None else 1.0)
-        penalty = low_entropy * combined_mask
-
-        return self.weight * _safe_masked_mean(penalty, ctx.loss_mask)
-
-class FocalOverconfidencePenalty(PenaltyRule):
-    """
-    Lấy cảm hứng từ Focal Loss nhưng đảo chiều:
-    phạt nặng khi model quá chắc vào token SAI.
-
-        penalty = (p_wrong)^gamma * log(p_wrong + ε)  khi sai
-
-    Với gamma > 1: gradient TĂNG khi overconfident
-    (ngược lại với p(1-p) trong softmax thông thường).
-
-    Cần log_probs để lấy p_max một cách hiệu quả.
-    """
-    needs_log_probs = True
-
-    def __init__(self, weight: float = 1.0, gamma: float = 2.0):
-        super().__init__()
-        self.weight = weight
-        self.gamma = gamma
-
-    def forward(self, ctx: PenaltyContext) -> torch.Tensor:
-        assert ctx.log_probs is not None
-
-        # xác suất của token được predict (token sai)
-        log_p_wrong = ctx.log_probs.gather(
-            dim=-1,
-            index=ctx.predicted.unsqueeze(-1)
-        ).squeeze(-1)
-        p_wrong = log_p_wrong.exp()
-
-        # focal penalty: nặng khi p cao, nhẹ khi p thấp
-        penalty = (p_wrong ** self.gamma) * (-log_p_wrong)
-        penalty = penalty * ctx.wrong_mask
-
-        return self.weight * _safe_masked_mean(penalty, ctx.loss_mask)
-
-class PenaltyEngine(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.rules = nn.ModuleList()
-
-    @classmethod
-    def build_default(cls) -> "PenaltyEngine":
-        return (
-            cls()
-            .add_rule(WrongTokenMarginPenalty(weight=0.1, detach_max=True))
-            .add_rule(WrongTokenEntropyPenalty(weight=0.001, min_entropy=0.5))
-            .add_rule(FocalOverconfidencePenalty(weight=0.05, gamma=2.0))
-        )
-
-    def add_rule(self, rule: PenaltyRule) -> "PenaltyEngine":
-        self.rules.append(rule)
-        return self
-
-    def _build_context(
-        self,
-        logits: torch.Tensor,
-        inputs: torch.Tensor,
-        targets: torch.Tensor,
-        loss_mask: Optional[torch.Tensor],
-    ) -> PenaltyContext:
-
-        max_logits, predicted = logits.max(dim=-1)
-        target_logits = logits.gather(
-            dim=-1, index=targets.unsqueeze(-1)
-        ).squeeze(-1)
-
-        wrong_mask = (predicted != targets).float()
-
-        needs_log = any(getattr(r, "needs_log_probs", False) for r in self.rules)
-        log_probs = F.log_softmax(logits, dim=-1) if needs_log else None
-
-        return PenaltyContext(
-            logits=logits,
-            inputs=inputs,
-            targets=targets,
-            loss_mask=loss_mask,
-            max_logits=max_logits,
-            predicted=predicted,
-            target_logits=target_logits,
-            wrong_mask=wrong_mask,
-            log_probs=log_probs,
-        )
-
-    def forward(
-        self,
-        logits: torch.Tensor,
-        inputs: torch.Tensor,
-        targets: torch.Tensor,
-        loss_mask: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        if not self.rules:
-            return torch.tensor(0.0, device=logits.device)
-
-        ctx = self._build_context(logits, inputs, targets, loss_mask)
-
-        total = torch.tensor(0.0, device=logits.device)
-        for rule in self.rules:
-            total = total + rule(ctx)
-        return total

src/training/finetune.py

@@ -1,581 +0,0 @@
-from pathlib import Path
-import torch
-import torch.optim as optim
-from torch.amp import GradScaler, autocast
-import torch._dynamo
-import json
-import gc
-import argparse
-import logging
-import math
-from src.utils.utils import (get_step_lr_lambda, freeze_layers, unfreeze_all_layers, log_progress, load_checkpoint, save_checkpoint, estimate_training_flops, print_training_flops_summary, TflopsBenchmarker, KernelLogger)
-from src.utils.chunked_loss import chunked_lm_loss, chunked_sft_loss
-from src.data.Dataset import Dataset, split_train_val_test, load_data
-from src.training.PenaltyEngine import PenaltyEngine, WrongTokenMarginPenalty, WrongTokenEntropyPenalty, FocalOverconfidencePenalty
-from src.model.TransformerModel import TransformerModel
-
-parser = argparse.ArgumentParser()
-parser.add_argument("--model", type=str, required=True, help="Model size: 35M, 100M or 500M")
-parser.add_argument(
-    "--phase", type=str, required=True,
-    choices=["sft1", "sft2", "sft1_resume", "sft2_resume", "full"],
-    help="Training phase: sft1 | sft2 | sft1_resume | sft2_resume | full"
-)
-parser.add_argument(
-    "--profile-kernels", action="store_true",
-    help="Profile kernel tại step 500 và lưu compact JSON"
-)
-parser.add_argument(
-    "--limited-max-autotune", action="store_true",
-    help="Bật max-autotune giới hạn cho forward_features (chỉ ATEN/TRITON)"
-)
-args = parser.parse_args()
-torch._dynamo.config.cache_size_limit = 24
-logging.getLogger("torch._inductor.select_algorithm").setLevel(logging.ERROR)
-model_size = args.model
-
-current_file = Path(__file__).resolve()
-project_root = current_file.parent.parent.parent
-config_dir = project_root / "config"
-data_dir = project_root / "data"
-model_dir = project_root / "model"
-src_dir = project_root / "src"
-base_config_file = config_dir / "base.json"
-model_config_file = config_dir / f"{args.model}.json"
-model_dir.mkdir(parents=True, exist_ok=True)
-data_processed_dir = project_root / "data" / "processed"
-SFT1_data_ids_file = data_processed_dir / "SFT1_data_ids.npz"
-SFT2_data_ids_file = data_processed_dir / "SFT2_data_ids.npz"
-
-pretrained_save_path = model_dir / f"pretrained_{model_size}.pt"
-sft1_save_path = model_dir / f"sft1_{model_size}.pt"
-sft1_ckpt_path = model_dir / f"sft1_{model_size}.ckpt.pt"
-sft2_save_path = model_dir / f"sft2_{model_size}.pt"
-sft2_ckpt_path = model_dir / f"sft2_{model_size}.ckpt.pt"
-
-# Peak logits mỗi chunk ~ LM_LOSS_CHUNK_SIZE * vocab_size * 4 bytes.
-LM_LOSS_CHUNK_SIZE = 4096
-
-def _build_val_test_loaders(phase_name, main_data, sub_data, train_ratio, val_ratio, batch_size, num_workers):
-    X, Y, loss_mask, lengths = load_data(phase_name, main_data, sub_data)
-
-    X_train, Y_train, mask_train, len_train, \
-    X_val, Y_val, mask_val, len_val, \
-    X_test, Y_test, mask_test, len_test = split_train_val_test(
-        X, Y, loss_mask, lengths, train_ratio, val_ratio
-    )
-
-    val_ds = Dataset.create_dataloader(X_val, Y_val, len_val, batch_size, max_seq_len, num_workers, shuffle=False, loss_masks=mask_val)
-    test_ds = Dataset.create_dataloader(X_test, Y_test, len_test, batch_size, max_seq_len, num_workers, shuffle=False, loss_masks=mask_test)
-
-    train_size = len(X_train)
-    val_size = len(X_val)
-    test_size = len(X_test)
-
-    del X_train, Y_train, mask_train, len_train
-    del X_val, Y_val, mask_val, len_val
-    del X_test, Y_test, mask_test, len_test
-    del X, Y, loss_mask, lengths
-    gc.collect()
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-
-    return val_ds, test_ds, train_size, val_size, test_size
-
-def _build_train_loader_epoch(phase_name, main_data, sub_data, train_ratio, val_ratio, batch_size, num_workers, epoch=0):
-    X, Y, loss_mask, lengths = load_data(phase_name, main_data, sub_data, seed=epoch)
-
-    X_train, Y_train, mask_train, len_train, \
-    X_val, Y_val, mask_val, len_val, \
-    X_test, Y_test, mask_test, len_test = split_train_val_test(
-        X, Y, loss_mask, lengths, train_ratio, val_ratio
-    )
-
-    train_ds = Dataset.create_dataloader(X_train, Y_train, len_train, batch_size, max_seq_len, num_workers, shuffle=True, loss_masks=mask_train)
-
-    del X_val, Y_val, mask_val, len_val
-    del X_test, Y_test, mask_test, len_test
-    del X_train, Y_train, mask_train, len_train
-    del X, Y, loss_mask, lengths
-    gc.collect()
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-
-    return train_ds
-
-def finetune(model_raw, optimizer, device, main_data, sub_data, num_epochs, model_save_path, train_ratio, val_ratio, batch_size, num_workers, phase_name, penalty_engine, resample_per_epoch=False, resume_checkpoint_path: Path = None, profile_kernels=False):
-    print(f"╠════════════════════════════════════════════════════════════════════════════════════╣")
-    print(f"║                               BẮT ĐẦU LOAD {phase_name.upper():<4} DATA                               ║")
-    print(f"╠════════════════════════════════════════════════════════════════════════════════════╣")
-
-    if not resample_per_epoch:
-        X, Y, loss_mask, lengths = load_data(phase_name, main_data, sub_data)
-
-        X_train, Y_train, mask_train, len_train, \
-        X_val, Y_val, mask_val, len_val, \
-        X_test, Y_test, mask_test, len_test = split_train_val_test(
-            X, Y, loss_mask, lengths, train_ratio, val_ratio
-        )
-
-        log_progress(f"[{phase_name}] Train: {len(X_train)}, Val: {len(X_val)}, Test: {len(X_test)}")
-
-        train_ds = Dataset.create_dataloader(X_train, Y_train, len_train, batch_size, max_seq_len, num_workers, shuffle=True, loss_masks=mask_train)
-        val_ds = Dataset.create_dataloader(X_val, Y_val, len_val, batch_size, max_seq_len, num_workers, shuffle=False, loss_masks=mask_val)
-        test_ds = Dataset.create_dataloader(X_test, Y_test, len_test, batch_size, max_seq_len, num_workers, shuffle=False, loss_masks=mask_test)
-
-        del X_train, Y_train, mask_train, len_train
-        del X_val, Y_val, mask_val, len_val
-        del X_test, Y_test, mask_test, len_test
-        gc.collect()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-
-        total_steps = math.ceil(len(train_ds) / accumulation_steps) * num_epochs
-
-    else:
-        val_ds, test_ds, train_size, val_size, test_size = _build_val_test_loaders(
-            phase_name, main_data, sub_data, train_ratio, val_ratio, batch_size, num_workers
-        )
-        log_progress(f"[{phase_name}] Train: ~{train_size}, Val: {val_size}, Test: {test_size}")
-
-        train_ds = _build_train_loader_epoch(phase_name, main_data, sub_data, train_ratio, val_ratio, batch_size, num_workers, epoch=0)
-        steps_per_epoch = len(train_ds)
-        total_steps = math.ceil(steps_per_epoch / accumulation_steps) * num_epochs
-
-    warmup_steps = int(total_steps * 0.15)
-    lr_lambda = get_step_lr_lambda(warmup_steps, total_steps)
-    scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
-    log_progress(f"Step-based LR: warmup={warmup_steps} steps, total={total_steps} steps")
-
-    best_val_loss = float("inf")
-    global_step = 0
-    start_epoch = 0
-    use_penalty = penalty_engine is not None and len(penalty_engine.rules) > 0
-    checkpoint_path = model_save_path.with_suffix(".ckpt.pt")
-    if resume_checkpoint_path is not None:
-        if resume_checkpoint_path.exists():
-            start_epoch, global_step, best_val_loss = load_checkpoint(
-                resume_checkpoint_path, model_raw, optimizer, scheduler, device
-            )
-        else:
-            log_progress(f"[WARNING] Checkpoint not found at {resume_checkpoint_path}. Starting from scratch.")
-
-    flops_info = estimate_training_flops(
-        model=model_raw,
-        num_layers=num_layers,
-        max_seq_len=max_seq_len,
-        d_model=d_model,
-        epochs=num_epochs,
-        batches_per_epoch=len(train_ds),
-        batch_size=batch_size,
-    )
-    print_training_flops_summary(flops_info, num_epochs)
-    bench = TflopsBenchmarker(
-        flops_per_token=flops_info["flops_per_token"],
-        total_flops_needed=flops_info["total_flops"],
-        total_batches_per_epoch=flops_info["total_batches_per_epoch"],
-        epochs=num_epochs,
-        device=device,
-        warmup_steps=100,
-        target_steps=500,
-    )
-    kernel_logger = KernelLogger(enabled=profile_kernels, log_step=500)
-
-    for epoch in range(start_epoch, num_epochs):
-
-        if resample_per_epoch and epoch > 0:
-            log_progress(f"[{phase_name}] Epoch {epoch+1}: Re-sampling train data...")
-            train_ds = _build_train_loader_epoch(
-                phase_name, main_data, sub_data, train_ratio, val_ratio, batch_size, num_workers, epoch
-            )
-
-        model_raw.train()
-        train_loss = torch.zeros((), device=device)
-        batch_count = 0
-        total_batches = len(train_ds)
-        penalty = torch.zeros((), device=device)
-
-        optimizer.zero_grad(set_to_none=True)
-        for batch_idx, (X_batch, Y_batch, sample_weight, attention_mask, has_padding) in enumerate(train_ds):
-            X_batch = X_batch.to(device, non_blocking=True)
-            Y_batch = Y_batch.to(device, non_blocking=True)
-            sample_weight = sample_weight.to(device, non_blocking=True)
-            attention_mask = attention_mask.to(device, non_blocking=True)
-
-            _is_bench_step = bench.is_bench_step(epoch, start_epoch, batch_idx)
-            if _is_bench_step:
-                bench.on_step_begin(batch_idx, X_batch.shape[0] * X_batch.shape[1])
-
-            _is_kernel_log = kernel_logger.should_log(epoch, start_epoch, batch_idx)
-            if _is_kernel_log:
-                kernel_logger.start()
-
-            with autocast(device_type=device.type, dtype=amp_dtype, enabled=amp_enabled):
-                hidden = model_features(
-                    X_batch,
-                    attention_mask=attention_mask,
-                    has_padding=has_padding,
-                )
-                hidden_flat = hidden.reshape(-1, hidden.size(-1))
-                inputs_flat = X_batch.reshape(-1)
-                targets_flat = Y_batch.reshape(-1)
-                weight_flat = sample_weight.reshape(-1)
-                num_valid_tokens = weight_flat.sum()
-
-                ce_sum, penalty_sum = chunked_sft_loss(
-                    hidden_flat,
-                    model_raw.lm_head.weight,
-                    inputs_flat,
-                    targets_flat,
-                    weight_flat,
-                    penalty_engine=penalty_engine if use_penalty else None,
-                    chunk_size=LM_LOSS_CHUNK_SIZE,
-                )
-                ce_loss = ce_sum / (num_valid_tokens + 1e-8)
-                penalty = penalty_sum / (num_valid_tokens + 1e-8)
-                loss = ce_loss + penalty
-
-            scaled_loss = loss / accumulation_steps
-            scaler.scale(scaled_loss).backward()
-            penalty_for_log = penalty.detach()
-
-            if _is_kernel_log:
-                kernel_logger.stop_and_report(batch_idx, epoch, device)
-
-            if (batch_idx + 1) % accumulation_steps == 0 or (batch_idx + 1) == total_batches:
-                scaler.unscale_(optimizer)
-                torch.nn.utils.clip_grad_norm_(model_raw.parameters(), max_norm=1.0)
-                scaler.step(optimizer)
-                scaler.update()
-                optimizer.zero_grad(set_to_none=True)
-                scheduler.step()
-                global_step += 1
-
-            train_loss += loss.detach()
-            batch_count += 1
-            current_lr = optimizer.param_groups[0]["lr"]
-
-            if (batch_idx + 1) % 10 == 0 or (batch_idx + 1) == total_batches:
-                avg_loss = (train_loss / batch_count).item()
-                penalty_str = f" | penalty: {penalty_for_log.item():.4f}" if use_penalty else ""
-                print(
-                    f"\r{phase_name} | Epoch {epoch+1}/{num_epochs} "
-                    f"- Step {global_step}/{total_steps} "
-                    f"- loss: {avg_loss:.4f}{penalty_str} "
-                    f"- lr: {current_lr:.2e}",
-                    end=""
-                )
-
-            if _is_bench_step:
-                bench.on_step_end(batch_idx)
-
-            del X_batch, Y_batch, hidden, hidden_flat, inputs_flat, targets_flat
-            del weight_flat, num_valid_tokens, loss, scaled_loss, ce_loss, penalty
-            del penalty_for_log, ce_sum, penalty_sum, sample_weight, attention_mask
-
-        print()
-        train_loss = (train_loss / len(train_ds)).item()
-
-        model_raw.eval()
-        val_loss = torch.zeros((), device=device)
-        with torch.no_grad():
-            for X_batch, Y_batch, sample_weight, attention_mask, has_padding in val_ds:
-                X_batch = X_batch.to(device, non_blocking=True)
-                Y_batch = Y_batch.to(device, non_blocking=True)
-                sample_weight = sample_weight.to(device, non_blocking=True)
-                attention_mask = attention_mask.to(device, non_blocking=True)
-
-                with autocast(device_type=device.type, dtype=amp_dtype, enabled=amp_enabled):
-                    hidden = model_features(
-                        X_batch,
-                        attention_mask=attention_mask,
-                        has_padding=has_padding,
-                    )
-                    hidden_flat = hidden.reshape(-1, hidden.size(-1))
-                    targets_flat = Y_batch.reshape(-1)
-                    weight_flat = sample_weight.reshape(-1)
-                    num_valid_tokens = weight_flat.sum()
-                    loss_sum = chunked_lm_loss(
-                        hidden_flat,
-                        model_raw.lm_head.weight,
-                        targets_flat,
-                        weight_flat,
-                        chunk_size=LM_LOSS_CHUNK_SIZE,
-                    )
-                    loss = loss_sum / (num_valid_tokens + 1e-8)
-                val_loss += loss.detach()
-
-        val_loss = (val_loss / len(val_ds)).item()
-        log_progress(f"{phase_name} Epoch {epoch+1}/{num_epochs} | Train Loss: {train_loss:.4f} | Val Loss (CE only): {val_loss:.4f}")
-
-        if val_loss < best_val_loss:
-            best_val_loss = val_loss
-            torch.save(model_raw.state_dict(), model_save_path)
-            log_progress(f"Epoch {epoch+1}: Saving model and checkpoint")
-
-        save_checkpoint(
-            checkpoint_path, epoch, global_step,
-            model_raw, optimizer, scheduler, best_val_loss
-        )
-
-    print(f"╠════════════════════════════════════════════════════════════════════════════════════╣")
-    print(f"║                            ĐÁNH GIÁ {phase_name.upper():<4} TRÊN TEST SET                             ║")
-    print(f"╠════════════════════════════════════════════════════════════════════════════════════╣")
-
-    model_raw.load_state_dict(torch.load(model_save_path, map_location=device))
-    model_raw.eval()
-    test_loss = torch.zeros((), device=device)
-    with torch.no_grad():
-        for X_batch, Y_batch, sample_weight, attention_mask, has_padding in test_ds:
-            X_batch = X_batch.to(device, non_blocking=True)
-            Y_batch = Y_batch.to(device, non_blocking=True)
-            sample_weight = sample_weight.to(device, non_blocking=True)
-            attention_mask = attention_mask.to(device, non_blocking=True)
-
-            with autocast(device_type=device.type, dtype=amp_dtype, enabled=amp_enabled):
-                hidden = model_features(
-                    X_batch,
-                    attention_mask=attention_mask,
-                    has_padding=has_padding,
-                )
-                hidden_flat = hidden.reshape(-1, hidden.size(-1))
-                targets_flat = Y_batch.reshape(-1)
-                weight_flat = sample_weight.reshape(-1)
-                num_valid_tokens = weight_flat.sum()
-                loss_sum = chunked_lm_loss(
-                    hidden_flat,
-                    model_raw.lm_head.weight,
-                    targets_flat,
-                    weight_flat,
-                    chunk_size=LM_LOSS_CHUNK_SIZE,
-                )
-                loss = loss_sum / (num_valid_tokens + 1e-8)
-            test_loss += loss.detach()
-
-    test_loss = (test_loss / len(test_ds)).item()
-    log_progress(f"{phase_name} Test Loss (CE only): {test_loss:.4f}")
-    print(f"╠════════════════════════════════════════════════════════════════════════════════════╣")
-
-    return test_loss
-
-with open(base_config_file, 'r') as f:
-    config = json.load(f)
-with open(model_config_file, 'r') as f:
-    config.update(json.load(f))
-
-vocab_size = config["vocab_size"]
-max_seq_len = config["max_seq_len"]
-d_model = config["d_model"]
-num_heads = config["num_heads"]
-num_kv_heads = config['num_kv_heads']
-num_layers = config["num_layers"]
-ff_dim = config["ff_dim"]
-dropout = config["dropout"]
-sft1_epochs = config["sft1_epochs"]
-sft2_epochs = config["sft2_epochs"]
-batch_size = config["batch_size"]
-train_ratio = config["train_ratio"]
-val_ratio = config["val_ratio"]
-sft1_learning_rate = config["sft1_learning_rate"]
-sft1_learning_weight_decay = config["sft1_learning_weight_decay"]
-sft2_learning_rate = config["sft2_learning_rate"]
-accumulation_steps = config["accumulation_steps"]
-sft2_learning_weight_decay = config["sft2_learning_weight_decay"]
-freeze = config["freeze"]
-num_workers = config["num_workers"]
-penalty_engine = (PenaltyEngine()
-    .add_rule(WrongTokenMarginPenalty(weight=config["penalty_margin_weight"], detach_max=config["penalty_margin_detach_max"]))
-    .add_rule(WrongTokenEntropyPenalty(weight=config["penalty_entropy_weight"], min_entropy=config["penalty_entropy_min_entropy"]))
-    .add_rule(FocalOverconfidencePenalty(weight=config["penalty_focal_weight"], gamma=config["penalty_focal_gamma"]))
-)
-
-if torch.cuda.is_available():
-    device = torch.device("cuda")
-    torch.backends.cudnn.benchmark = True
-elif torch.backends.mps.is_available():
-    device = torch.device("mps")
-else:
-    device = torch.device("cpu")
-
-amp_enabled = device.type == "cuda"
-amp_dtype = torch.bfloat16 if amp_enabled and torch.cuda.is_bf16_supported() else torch.float16
-scaler = GradScaler("cuda", enabled=amp_enabled and amp_dtype == torch.float16)
-log_progress(f"Sử dụng device: {device}")
-if amp_enabled:
-    precision_name = "BF16" if amp_dtype == torch.bfloat16 else "FP16 + GradScaler"
-    log_progress(f"CUDA mixed precision: {precision_name}")
-
-model_raw = TransformerModel(vocab_size, d_model, num_heads, num_kv_heads, num_layers, ff_dim, max_seq_len, dropout).to(device)
-# Giữ alias `model` để không thay khung dispatch phase hiện tại.
-model = model_raw
-if args.limited_max_autotune:
-    compile_options = {
-        "max_autotune": True,
-        "max_autotune_gemm_backends": "ATEN,TRITON",
-        "max_autotune_gemm_search_space": "DEFAULT",
-        "epilogue_fusion": True,
-        "triton.cudagraphs": False,
-    }
-    model_features = torch.compile(
-        model_raw.forward_features,
-        dynamic=True,
-        options=compile_options,
-    )
-    log_progress("torch.compile: limited max-autotune (ATEN,TRITON; DEFAULT; dynamic=True; CUDA Graph off)")
-else:
-    model_features = torch.compile(model_raw.forward_features, dynamic=True)
-
-phase = args.phase
-if phase == "sft1":
-    log_progress("Load model từ pretrain...")
-    model.load_state_dict(torch.load(pretrained_save_path, map_location=device))
-    freeze_layers(model, freeze)
-
-    optimizer_sft1 = optim.AdamW(
-        filter(lambda p: p.requires_grad, model.parameters()),
-        lr=sft1_learning_rate, weight_decay=sft1_learning_weight_decay,
-        fused=amp_enabled,
-    )
-
-    test_loss = finetune(
-        model, optimizer_sft1, device, SFT1_data_ids_file, sub_data=None,
-        num_epochs=sft1_epochs,
-        model_save_path=sft1_save_path,
-        train_ratio=train_ratio, val_ratio=val_ratio,
-        batch_size=batch_size, phase_name="sft1",
-        num_workers=num_workers,
-        penalty_engine=penalty_engine,
-        resample_per_epoch=False,
-        resume_checkpoint_path=None,
-        profile_kernels=args.profile_kernels,
-    )
-    log_progress(f"SFT1 Test Loss: {test_loss:.4f}")
-
-elif phase == "sft1_resume":
-    log_progress("SFT1 resume: khởi tạo model skeleton trước khi load checkpoint...")
-    freeze_layers(model, freeze)
-
-    optimizer_sft1 = optim.AdamW(
-        filter(lambda p: p.requires_grad, model.parameters()),
-        lr=sft1_learning_rate, weight_decay=sft1_learning_weight_decay,
-        fused=amp_enabled,
-    )
-
-    test_loss = finetune(
-        model, optimizer_sft1, device, SFT1_data_ids_file, sub_data=None,
-        num_epochs=sft1_epochs,
-        model_save_path=sft1_save_path,
-        train_ratio=train_ratio, val_ratio=val_ratio,
-        batch_size=batch_size, phase_name="sft1",
-        num_workers=num_workers,
-        penalty_engine=penalty_engine,
-        resample_per_epoch=False,
-        resume_checkpoint_path=sft1_ckpt_path,
-        profile_kernels=args.profile_kernels,
-    )
-    log_progress(f"SFT1 Test Loss: {test_loss:.4f}")
-
-elif phase == "sft2":
-    log_progress("Load model từ sft1...")
-    model.load_state_dict(torch.load(sft1_save_path, map_location=device))
-    unfreeze_all_layers(model)
-    freeze_layers(model, freeze)
-
-    optimizer_sft2 = optim.AdamW(
-        filter(lambda p: p.requires_grad, model.parameters()),
-        lr=sft2_learning_rate, weight_decay=sft2_learning_weight_decay,
-        fused=amp_enabled,
-    )
-
-    test_loss = finetune(
-        model, optimizer_sft2, device, SFT2_data_ids_file, SFT1_data_ids_file,
-        num_epochs=sft2_epochs,
-        model_save_path=sft2_save_path,
-        train_ratio=train_ratio, val_ratio=val_ratio,
-        batch_size=batch_size, phase_name="sft2",
-        num_workers=num_workers,
-        penalty_engine=penalty_engine,
-        resample_per_epoch=True,
-        resume_checkpoint_path=None,
-        profile_kernels=args.profile_kernels,
-    )
-    log_progress(f"SFT2 Test Loss: {test_loss:.4f}")
-
-elif phase == "sft2_resume":
-    log_progress("SFT2 resume: khởi tạo model skeleton trước khi load checkpoint...")
-    unfreeze_all_layers(model)
-    freeze_layers(model, freeze)
-
-    optimizer_sft2 = optim.AdamW(
-        filter(lambda p: p.requires_grad, model.parameters()),
-        lr=sft2_learning_rate, weight_decay=sft2_learning_weight_decay,
-        fused=amp_enabled,
-    )
-
-    test_loss = finetune(
-        model, optimizer_sft2, device, SFT2_data_ids_file, SFT1_data_ids_file,
-        num_epochs=sft2_epochs,
-        model_save_path=sft2_save_path,
-        train_ratio=train_ratio, val_ratio=val_ratio,
-        batch_size=batch_size, phase_name="sft2",
-        num_workers=num_workers,
-        penalty_engine=penalty_engine,
-        resample_per_epoch=True,
-        resume_checkpoint_path=sft2_ckpt_path,
-        profile_kernels=args.profile_kernels,
-    )
-    log_progress(f"SFT2 Test Loss: {test_loss:.4f}")
-
-elif phase == "full":
-    # ── SFT1 ──
-    log_progress("Load model từ pretrain...")
-    model.load_state_dict(torch.load(pretrained_save_path, map_location=device))
-    freeze_layers(model, freeze)
-
-    optimizer_sft1 = optim.AdamW(
-        filter(lambda p: p.requires_grad, model.parameters()),
-        lr=sft1_learning_rate, weight_decay=sft1_learning_weight_decay,
-        fused=amp_enabled,
-    )
-
-    test_loss = finetune(
-        model, optimizer_sft1, device, SFT1_data_ids_file, sub_data=None,
-        num_epochs=sft1_epochs,
-        model_save_path=sft1_save_path,
-        train_ratio=train_ratio, val_ratio=val_ratio,
-        batch_size=batch_size, phase_name="sft1",
-        num_workers=num_workers,
-        penalty_engine=penalty_engine,
-        resample_per_epoch=False,
-        resume_checkpoint_path=None,
-        profile_kernels=args.profile_kernels,
-    )
-    log_progress(f"SFT1 Test Loss: {test_loss:.4f}")
-
-    # ── SFT2 ──
-    log_progress("Load model từ sft1...")
-    model.load_state_dict(torch.load(sft1_save_path, map_location=device))
-    unfreeze_all_layers(model)
-    freeze_layers(model, freeze)
-
-    optimizer_sft2 = optim.AdamW(
-        filter(lambda p: p.requires_grad, model.parameters()),
-        lr=sft2_learning_rate, weight_decay=sft2_learning_weight_decay,
-        fused=amp_enabled,
-    )
-
-    test_loss = finetune(
-        model, optimizer_sft2, device, SFT2_data_ids_file, SFT1_data_ids_file,
-        num_epochs=sft2_epochs,
-        model_save_path=sft2_save_path,
-        train_ratio=train_ratio, val_ratio=val_ratio,
-        batch_size=batch_size, phase_name="sft2",
-        num_workers=num_workers,
-        penalty_engine=penalty_engine,
-        resample_per_epoch=True,
-        resume_checkpoint_path=None,
-        profile_kernels=args.profile_kernels,
-    )
-    log_progress(f"SFT2 Test Loss: {test_loss:.4f}")
-
-log_progress(f"Model cuối cùng lưu tại: {model_dir / (phase.split('_')[0] + f'_{model_size}.pt')}")

src/training/pretrain.py

@@ -1,421 +0,0 @@
-from pathlib import Path
-import torch
-import torch.nn as nn
-import torch.optim as optim
-from torch.amp import GradScaler, autocast
-import torch._dynamo
-import logging
-import json
-import gc
-import argparse
-import math
-from src.utils.utils import (get_step_lr_lambda, log_progress, load_checkpoint, save_checkpoint, estimate_training_flops, print_training_flops_summary, TflopsBenchmarker, KernelLogger)
-from src.utils.chunked_loss import chunked_lm_loss
-from src.data.Dataset import Dataset, split_train_val_test, load_data
-from src.model.TransformerModel import TransformerModel
-
-parser = argparse.ArgumentParser()
-parser.add_argument("--model", type=str, required=True, help="Model size: 100M or 500M")
-parser.add_argument(
-    "--phase", type=str, required=True,
-    choices=["pretrain", "pretrain_resume", "continued_pretrain", "continued_pretrain_resume", "full"],
-    help="Training phase: pretrain | pretrain_resume | continued_pretrain | continued_pretrain_resume | full"
-)
-parser.add_argument(
-    "--profile-kernels", action="store_true",
-    help="Chỉ chạy 1 lần duy nhất, log tên kernel CPU/CUDA tại 1 step rồi thôi (dùng để debug performance)"
-)
-parser.add_argument(
-    "--limited-max-autotune", action="store_true",
-    help="Bật max-autotune giới hạn cho forward_features (chỉ ATEN/TRITON, không exhaustive/CUDA Graph)"
-)
-args = parser.parse_args()
-torch._dynamo.config.cache_size_limit = 24
-logging.getLogger("torch._inductor.select_algorithm").setLevel(logging.ERROR)
-model_size = args.model
-
-current_file = Path(__file__).resolve()
-project_root = current_file.parent.parent.parent
-config_dir = project_root / "config"
-data_dir = project_root / "data"
-model_dir = project_root / "model"
-src_dir = project_root / "src"
-base_config_file = config_dir / "base.json"
-model_config_file = config_dir / f"{args.model}.json"
-model_dir.mkdir(parents=True, exist_ok=True)
-data_processed_dir = project_root / "data" / "processed"
-pretrain_tokenized_file = data_processed_dir / "pretrain_manifest.json"
-continued_pretrain_tokenized_file = data_processed_dir / "continued_pretrain_data_ids.npz"
-
-pretrained_save_path = model_dir / f"pretrained_{model_size}.pt"
-pretrained_ckpt_path = model_dir / f"pretrained_{model_size}.ckpt.pt"
-continued_pretrained_save_path = model_dir / f"continued_pretrained_{model_size}.pt"
-continued_pretrained_ckpt_path = model_dir / f"continued_pretrained_{model_size}.ckpt.pt"
-
-LM_LOSS_CHUNK_SIZE = 4096
-
-def train_loop(data_type, tokenized_file, epochs, learning_rate, weight_decay, num_workers, extra_file=None, model_save_path=None, resume_checkpoint_path=None, profile_kernels=False):
-    print("╠════════════════════════════════════════════════════════════════════════════════════╣")
-    print("║                                BAT ĐAU LOAD DATA                                   ║")
-    print("╠════════════════════════════════════════════════════════════════════════════════════╣")
-
-    if extra_file is None:
-        X, Y, lengths = load_data(data_type, tokenized_file)
-    else:
-        X, Y, lengths = load_data(data_type, tokenized_file, extra_file)
-
-    X_train, Y_train, _, lengths_train, X_val, Y_val, _, lengths_val, X_test, Y_test, _, lengths_test = split_train_val_test(X, Y, None, lengths, train_ratio, val_ratio)
-    log_progress(f"Train: {len(X_train)}, Val: {len(X_val)}, Test: {len(X_test)}")
-
-    train_ds = Dataset.create_dataloader(X_train, Y_train, lengths_train, batch_size, max_seq_len, num_workers, shuffle=True)
-    val_ds = Dataset.create_dataloader(X_val, Y_val, lengths_val, batch_size, max_seq_len, num_workers, shuffle=False)
-    test_ds = Dataset.create_dataloader(X_test, Y_test, lengths_test, batch_size, max_seq_len, num_workers, shuffle=False)
-
-    del X_train, Y_train, lengths_train, X_val, Y_val, lengths_val, X_test, Y_test, lengths_test
-    gc.collect()
-
-    global optimizer, scaler
-    optimizer = optim.AdamW(model_raw.parameters(), lr=learning_rate, weight_decay=weight_decay, fused=True)
-
-    total_steps = math.ceil(len(train_ds) / accumulation_steps) * epochs
-    warmup_steps = total_steps // 10
-
-    lr_lambda = get_step_lr_lambda(warmup_steps, total_steps)
-    scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
-
-    best_val_loss = float('inf')
-    global_step = 0
-    start_epoch = 0
-
-    checkpoint_path = model_save_path.with_suffix(".ckpt.pt") if model_save_path is not None else None
-    if resume_checkpoint_path is not None:
-        if resume_checkpoint_path.exists():
-            start_epoch, global_step, best_val_loss = load_checkpoint(
-                resume_checkpoint_path, model_raw, optimizer, scheduler, device
-            )
-        else:
-            log_progress(f"[WARNING] Checkpoint not found at {resume_checkpoint_path}. Starting from scratch.")
-
-    flops_info = estimate_training_flops(model=model_raw, num_layers=num_layers, max_seq_len=max_seq_len, d_model=d_model, epochs=epochs, batches_per_epoch=len(train_ds), batch_size=batch_size)
-    print_training_flops_summary(flops_info, epochs)
-    bench = TflopsBenchmarker(flops_per_token=flops_info["flops_per_token"], total_flops_needed=flops_info["total_flops"], total_batches_per_epoch=flops_info["total_batches_per_epoch"], epochs=epochs, device=device, warmup_steps=100, target_steps=500)
-    kernel_logger = KernelLogger(enabled=profile_kernels, log_step=500)
-
-    for epoch in range(start_epoch, epochs):
-        model_raw.train()
-        # Gom loss trên GPU (tensor), chỉ .item() khi thực sự cần in/log ->
-        # tránh sync CPU-GPU mỗi micro-batch, kết quả loss cuối cùng không đổi.
-        train_loss = torch.zeros((), device=device)
-        batch_count = 0
-        total_batches = len(train_ds)
-        optimizer.zero_grad()
-
-        for batch_idx, (X_batch, Y_batch, sample_weight, attention_mask, has_padding) in enumerate(train_ds):
-            X_batch = X_batch.to(device, non_blocking=True)
-            Y_batch = Y_batch.to(device, non_blocking=True)
-            sample_weight = sample_weight.to(device, non_blocking=True)
-            attention_mask = attention_mask.to(device, non_blocking=True)
-
-            _is_bench_step = bench.is_bench_step(epoch, start_epoch, batch_idx)
-            if _is_bench_step:
-                bench.on_step_begin(batch_idx, X_batch.shape[0] * X_batch.shape[1])
-
-            _is_kernel_log = kernel_logger.should_log(epoch, start_epoch, batch_idx)
-            if _is_kernel_log:
-                kernel_logger.start()
-
-            with autocast(device_type=device.type, dtype=amp_dtype, enabled=amp_enabled):
-                hidden = model_features(X_batch, attention_mask=attention_mask, has_padding=has_padding)
-                hidden_flat = hidden.reshape(-1, hidden.size(-1))
-                targets_flat = Y_batch.reshape(-1)
-                weight_flat = sample_weight.reshape(-1)
-                num_valid_tokens = weight_flat.sum()
-                loss_sum = chunked_lm_loss(
-                    hidden_flat, model_raw.lm_head.weight, targets_flat, weight_flat,
-                    chunk_size=LM_LOSS_CHUNK_SIZE,
-                )
-                loss = loss_sum / (num_valid_tokens + 1e-8)
-
-            scaled_loss = loss / accumulation_steps
-            scaler.scale(scaled_loss).backward()
-
-            if _is_kernel_log:
-                kernel_logger.stop_and_report(batch_idx, epoch, device)
-
-            if (batch_idx + 1) % accumulation_steps == 0 or (batch_idx + 1) == total_batches:
-                scaler.unscale_(optimizer)
-                torch.nn.utils.clip_grad_norm_(model_raw.parameters(), max_norm=1.0)
-                scaler.step(optimizer)
-                scaler.update()
-
-                optimizer.zero_grad()
-                scheduler.step()
-                global_step += 1
-
-            train_loss += loss.detach()
-            batch_count += 1
-            current_lr = optimizer.param_groups[0]['lr']
-
-            if (batch_idx + 1) % 10 == 0 or (batch_idx + 1) == total_batches:
-                avg_loss = (train_loss / batch_count).item()
-                print(f"\rEpoch {epoch+1}/{epochs} - Step {global_step}/{total_steps} - loss: {avg_loss:.4f} - lr: {current_lr:.2e}", end='')
-
-            if _is_bench_step:
-                bench.on_step_end(batch_idx)
-
-            del X_batch, Y_batch, hidden, hidden_flat, loss, loss_sum, sample_weight, attention_mask
-
-        print()
-
-        train_loss = (train_loss / len(train_ds)).item()
-        model_raw.eval()
-        val_loss = torch.zeros((), device=device)
-
-        with torch.no_grad():
-            for X_batch, Y_batch, sample_weight, attention_mask, has_padding in val_ds:
-                X_batch = X_batch.to(device, non_blocking=True)
-                Y_batch = Y_batch.to(device, non_blocking=True)
-                sample_weight = sample_weight.to(device, non_blocking=True)
-                attention_mask = attention_mask.to(device, non_blocking=True)
-
-                with autocast(device_type=device.type, dtype=amp_dtype, enabled=amp_enabled):
-                    hidden = model_features(X_batch, attention_mask=attention_mask, has_padding=has_padding)
-                    hidden_flat = hidden.reshape(-1, hidden.size(-1))
-                    targets_flat = Y_batch.reshape(-1)
-                    weight_flat = sample_weight.reshape(-1)
-                    num_valid_tokens = weight_flat.sum()
-                    loss_sum = chunked_lm_loss(
-                        hidden_flat, model_raw.lm_head.weight, targets_flat, weight_flat,
-                        chunk_size=LM_LOSS_CHUNK_SIZE,
-                    )
-
-                loss = loss_sum / (num_valid_tokens + 1e-8)
-                val_loss += loss.detach()
-
-                del X_batch, Y_batch, hidden, hidden_flat, loss, loss_sum, sample_weight, attention_mask
-
-        val_loss = (val_loss / len(val_ds)).item()
-        log_progress(f"Epoch {epoch+1}/{epochs} Train Loss: {train_loss:.4f} Val Loss: {val_loss:.4f}")
-
-        if val_loss < best_val_loss:
-            best_val_loss = val_loss
-            if model_save_path is not None:
-                torch.save(model_raw.state_dict(), model_save_path)
-                print(f"Epoch {epoch+1}: val_loss improved to {val_loss:.5f}, saving model to {model_save_path}")
-            else:
-                torch.save(model_raw.state_dict(), pretrained_save_path)
-                print(f"Epoch {epoch+1}: val_loss improved to {val_loss:.5f}, saving model to default path")
-
-        if checkpoint_path is not None:
-            save_checkpoint(
-                checkpoint_path, epoch, global_step,
-                model_raw, optimizer, scheduler, best_val_loss
-            )
-            log_progress(f"Checkpoint saved → {checkpoint_path}")
-
-    print("╠════════════════════════════════════════════════════════════════════════════════════╣")
-    print("║                               DANH GIA TREN TEST SET                               ║")
-    print("╠════════════════════════════════════════════════════════════════════════════════════╣")
-
-    model_raw.eval()
-    test_loss = torch.zeros((), device=device)
-
-    with torch.no_grad():
-        for X_batch, Y_batch, sample_weight, attention_mask, has_padding in test_ds:
-            X_batch = X_batch.to(device, non_blocking=True)
-            Y_batch = Y_batch.to(device, non_blocking=True)
-            sample_weight = sample_weight.to(device, non_blocking=True)
-            attention_mask = attention_mask.to(device, non_blocking=True)
-
-            with autocast(device_type=device.type, dtype=amp_dtype, enabled=amp_enabled):
-                hidden = model_features(X_batch, attention_mask=attention_mask, has_padding=has_padding)
-                hidden_flat = hidden.reshape(-1, hidden.size(-1))
-                targets_flat = Y_batch.reshape(-1)
-                weight_flat = sample_weight.reshape(-1)
-                num_valid_tokens = weight_flat.sum()
-                loss_sum = chunked_lm_loss(
-                    hidden_flat, model_raw.lm_head.weight, targets_flat, weight_flat,
-                    chunk_size=LM_LOSS_CHUNK_SIZE,
-                )
-
-            loss = loss_sum / (num_valid_tokens + 1e-8)
-            test_loss += loss.detach()
-
-            del X_batch, Y_batch, hidden, hidden_flat, loss, loss_sum, sample_weight, attention_mask
-
-    test_loss = (test_loss / len(test_ds)).item()
-    log_progress(f"Test Loss: {test_loss:.4f}")
-    print("╠════════════════════════════════════════════════════════════════════════════════════╣")
-
-    return test_loss
-
-with open(base_config_file, 'r') as f:
-    config = json.load(f)
-with open(model_config_file, 'r') as f:
-    config.update(json.load(f))
-
-vocab_size = config['vocab_size']
-max_seq_len = config['max_seq_len']
-d_model = config['d_model']
-num_heads = config['num_heads']
-num_kv_heads = config['num_kv_heads']
-num_layers = config['num_layers']
-ff_dim = config['ff_dim']
-dropout = config['dropout']
-pretrain_epochs = config['pretrain_epochs']
-continued_pretrain_epochs = config['continued_pretrain_epochs']
-batch_size = config['batch_size']
-train_ratio = config['train_ratio']
-val_ratio = config['val_ratio']
-pretrain_learning_rate = config['pretrain_learning_rate']
-continued_pretrain_learning_rate = config['continued_pretrain_learning_rate']
-accumulation_steps = config['accumulation_steps']
-pretrain_weight_decay = config['pretrain_weight_decay']
-continued_pretrain_weight_decay = config['continued_pretrain_weight_decay']
-num_workers = config['num_workers']
-
-if torch.cuda.is_available():
-    device = torch.device("cuda")
-    torch.backends.cudnn.benchmark = True
-elif torch.backends.mps.is_available():
-    device = torch.device("mps")
-else:
-    device = torch.device("cpu")
-
-amp_enabled = device.type == "cuda"
-amp_dtype = torch.bfloat16 if amp_enabled and torch.cuda.is_bf16_supported() else torch.float16
-scaler = GradScaler("cuda", enabled=amp_enabled and amp_dtype == torch.float16)
-
-if amp_enabled:
-    precision_name = "BF16" if amp_dtype == torch.bfloat16 else "FP16 + GradScaler"
-    log_progress(f"CUDA mixed precision: {precision_name}")
-
-model_raw = TransformerModel(vocab_size, d_model, num_heads, num_kv_heads, num_layers, ff_dim, max_seq_len, dropout).to(device)
-model = torch.compile(model_raw, dynamic=True)
-if args.limited_max_autotune:
-    compile_options = {
-        "max_autotune": True,
-        "max_autotune_gemm_backends": "ATEN,TRITON",
-        "max_autotune_gemm_search_space": "DEFAULT",
-        "epilogue_fusion": True,
-        "triton.cudagraphs": False,
-    }
-    model_features = torch.compile(
-        model_raw.forward_features,
-        dynamic=True,
-        options=compile_options,
-    )
-    log_progress("torch.compile: limited max-autotune (ATEN,TRITON; DEFAULT; dynamic=True; CUDA Graph off)")
-else:
-    model_features = torch.compile(model_raw.forward_features, dynamic=True)
-optimizer = optim.AdamW(model_raw.parameters(), lr=pretrain_learning_rate, weight_decay=pretrain_weight_decay, fused=True)
-
-print("╠════════════════════════════════════════════════════════════════════════════════════╣")
-print("║                                 BAT ĐAU TRAINING                                   ║")
-print("╠════════════════════════════════════════════════════════════════════════════════════╣")
-
-phase = args.phase
-
-if phase == "pretrain":
-    pretrain_test_loss = train_loop(
-        data_type="pretrain",
-        tokenized_file=pretrain_tokenized_file,
-        epochs=pretrain_epochs,
-        learning_rate=pretrain_learning_rate,
-        weight_decay=pretrain_weight_decay,
-        num_workers=num_workers,
-        model_save_path=pretrained_save_path,
-        resume_checkpoint_path=None,
-        profile_kernels=args.profile_kernels,
-    )
-    log_progress(f"Pretrain Test Loss: {pretrain_test_loss:.4f}")
-
-elif phase == "pretrain_resume":
-    log_progress("Pretrain resume: khởi tạo model skeleton trước khi load checkpoint...")
-    pretrain_test_loss = train_loop(
-        data_type="pretrain",
-        tokenized_file=pretrain_tokenized_file,
-        epochs=pretrain_epochs,
-        learning_rate=pretrain_learning_rate,
-        weight_decay=pretrain_weight_decay,
-        num_workers=num_workers,
-        model_save_path=pretrained_save_path,
-        resume_checkpoint_path=pretrained_ckpt_path,
-        profile_kernels=args.profile_kernels,
-    )
-    log_progress(f"Pretrain Test Loss: {pretrain_test_loss:.4f}")
-
-elif phase == "continued_pretrain":
-    log_progress("Load best model từ pretrain để tiếp tục training...")
-    model_raw.load_state_dict(torch.load(pretrained_save_path, map_location=device))
-    model_raw.to(device)
-    optimizer = optim.AdamW(model_raw.parameters(), lr=continued_pretrain_learning_rate, weight_decay=continued_pretrain_weight_decay, fused=True)
-    log_progress(f"Reset optimizer với learning rate: {continued_pretrain_learning_rate}")
-
-    continued_pretrain_test_loss = train_loop(
-        data_type="continued_pretrain",
-        tokenized_file=continued_pretrain_tokenized_file,
-        epochs=continued_pretrain_epochs,
-        learning_rate=continued_pretrain_learning_rate,
-        weight_decay=continued_pretrain_weight_decay,
-        num_workers=num_workers,
-        extra_file=pretrain_tokenized_file,
-        model_save_path=continued_pretrained_save_path,
-        resume_checkpoint_path=None,
-        profile_kernels=args.profile_kernels,
-    )
-    log_progress(f"Continued Pretrain Test Loss: {continued_pretrain_test_loss:.4f}")
-
-elif phase == "continued_pretrain_resume":
-    log_progress("Continued pretrain resume: khởi tạo model skeleton trước khi load checkpoint...")
-    optimizer = optim.AdamW(model_raw.parameters(), lr=continued_pretrain_learning_rate, weight_decay=continued_pretrain_weight_decay, fused=True)
-
-    continued_pretrain_test_loss = train_loop(
-        data_type="continued_pretrain",
-        tokenized_file=continued_pretrain_tokenized_file,
-        epochs=continued_pretrain_epochs,
-        learning_rate=continued_pretrain_learning_rate,
-        weight_decay=continued_pretrain_weight_decay,
-        num_workers=num_workers,
-        extra_file=pretrain_tokenized_file,
-        model_save_path=continued_pretrained_save_path,
-        resume_checkpoint_path=continued_pretrained_ckpt_path,
-        profile_kernels=args.profile_kernels,
-    )
-    log_progress(f"Continued Pretrain Test Loss: {continued_pretrain_test_loss:.4f}")
-
-elif phase == "full":
-    pretrain_test_loss = train_loop(
-        data_type="pretrain",
-        tokenized_file=pretrain_tokenized_file,
-        epochs=pretrain_epochs,
-        learning_rate=pretrain_learning_rate,
-        weight_decay=pretrain_weight_decay,
-        num_workers=num_workers,
-        model_save_path=pretrained_save_path,
-        resume_checkpoint_path=None,
-        profile_kernels=args.profile_kernels,
-    )
-
-    log_progress("Đang load best model từ pretrain để tiếp tục training...")
-    model_raw.load_state_dict(torch.load(pretrained_save_path, map_location=device))
-    model_raw.to(device)
-    optimizer = optim.AdamW(model_raw.parameters(), lr=continued_pretrain_learning_rate, weight_decay=continued_pretrain_weight_decay, fused=True)
-    log_progress(f"Reset optimizer với learning rate: {continued_pretrain_learning_rate}")
-
-    continued_pretrain_test_loss = train_loop(
-        data_type="continued_pretrain",
-        tokenized_file=continued_pretrain_tokenized_file,
-        epochs=continued_pretrain_epochs,
-        learning_rate=continued_pretrain_learning_rate,
-        weight_decay=continued_pretrain_weight_decay,
-        num_workers=num_workers,
-        extra_file=pretrain_tokenized_file,
-        model_save_path=continued_pretrained_save_path,
-        resume_checkpoint_path=None,
-        profile_kernels=args.profile_kernels,
-    )
-
-    log_progress(f"Hoàn thành training!")
-    log_progress(f"Pretrain Test Loss: {pretrain_test_loss:.4f}")
-    log_progress(f"Continued Pretrain Test Loss: {continued_pretrain_test_loss:.4f}")

src/utils/chunked_loss.py

@@ -1,90 +0,0 @@
-import torch
-import torch.nn.functional as F
-from torch.utils.checkpoint import checkpoint
-
-def chunked_lm_loss(hidden: torch.Tensor, weight: torch.Tensor, targets: torch.Tensor,
-                     token_weight: torch.Tensor, chunk_size: int = 4096) -> torch.Tensor:
-    N = hidden.shape[0]
-    total_loss = hidden.new_zeros((), dtype=torch.float32)
-
-    def _chunk_fn(h_chunk, t_chunk, w_chunk):
-        logits_chunk = F.linear(h_chunk, weight).float()
-        loss_per_token = F.cross_entropy(logits_chunk, t_chunk, reduction='none')
-        return (loss_per_token * w_chunk).sum()
-
-    for start in range(0, N, chunk_size):
-        end = min(start + chunk_size, N)
-        chunk_args = (
-            hidden[start:end],
-            targets[start:end],
-            token_weight[start:end],
-        )
-        if torch.is_grad_enabled():
-            chunk_loss = checkpoint(
-                _chunk_fn, *chunk_args, use_reentrant=False,
-            )
-        else:
-            chunk_loss = _chunk_fn(*chunk_args)
-        total_loss = total_loss + chunk_loss
-
-    return total_loss
-
-
-def chunked_sft_loss(
-    hidden: torch.Tensor,
-    weight: torch.Tensor,
-    inputs: torch.Tensor,
-    targets: torch.Tensor,
-    token_weight: torch.Tensor,
-    penalty_engine=None,
-    chunk_size: int = 4096,
-):
-    """Tính CE + penalty theo chunk mà không vật lý hóa full logits.
-
-    PenaltyEngine trả về mean theo loss mask. Nhân lại với số token
-    hợp lệ trong từng chunk rồi chia một lần ở caller giữ nguyên
-    công thức mean trên toàn batch.
-    """
-    N = hidden.shape[0]
-    total_ce = hidden.new_zeros((), dtype=torch.float32)
-    total_penalty = hidden.new_zeros((), dtype=torch.float32)
-
-    def _chunk_fn(h_chunk, i_chunk, t_chunk, w_chunk):
-        logits_chunk = F.linear(h_chunk, weight).float()
-        loss_per_token = F.cross_entropy(
-            logits_chunk, t_chunk, reduction="none"
-        )
-        ce_sum = (loss_per_token * w_chunk).sum()
-
-        if penalty_engine is None:
-            penalty_sum = logits_chunk.new_zeros(())
-        else:
-            penalty_mean = penalty_engine(
-                logits=logits_chunk.unsqueeze(0),
-                inputs=i_chunk.unsqueeze(0),
-                targets=t_chunk.unsqueeze(0),
-                loss_mask=w_chunk.unsqueeze(0),
-            )
-            penalty_sum = penalty_mean * w_chunk.sum()
-
-        return ce_sum, penalty_sum
-
-    for start in range(0, N, chunk_size):
-        end = min(start + chunk_size, N)
-        chunk_args = (
-            hidden[start:end],
-            inputs[start:end],
-            targets[start:end],
-            token_weight[start:end],
-        )
-        if torch.is_grad_enabled():
-            ce_sum, penalty_sum = checkpoint(
-                _chunk_fn, *chunk_args, use_reentrant=False,
-            )
-        else:
-            ce_sum, penalty_sum = _chunk_fn(*chunk_args)
-
-        total_ce = total_ce + ce_sum
-        total_penalty = total_penalty + penalty_sum
-
-    return total_ce, total_penalty

src/utils/utils.py

@@ -1,331 +0,0 @@
-import torch
-import time
-import json
-from pathlib import Path
-
-def get_step_lr_lambda(warmup_steps, total_steps):
-    def lr_lambda(current_step):
-        if current_step < warmup_steps:
-            return float(current_step) / float(max(1, warmup_steps))
-        elif current_step < total_steps * 0.4:
-            return 1.0
-        else:
-            progress = (current_step - total_steps * 0.4) / (total_steps * 0.3)
-            return max(0.1, 1.0 - 0.9 * progress)
-    return lr_lambda
-
-def freeze_layers(model, freeze):
-    for idx in freeze:
-        for param in model.blocks[idx].parameters():
-            param.requires_grad = False
-
-def unfreeze_all_layers(model):
-    for param in model.parameters():
-        param.requires_grad = True
-
-def log_progress(text):
-    fixed_width = 82
-    formatted_text = f"║ {text:<{fixed_width}} ║"
-    print(formatted_text)
-
-def save_checkpoint(path: Path, epoch: int, global_step: int, model, optimizer, scheduler, best_val_loss: float):
-    """Save full training state so training can be resumed exactly."""
-    torch.save({
-        "epoch": epoch,
-        "global_step": global_step,
-        "model_state_dict": model.state_dict(),
-        "optimizer_state_dict": optimizer.state_dict(),
-        "scheduler_state_dict": scheduler.state_dict(),
-        "best_val_loss": best_val_loss,
-    }, path)
-
-
-def load_checkpoint(path: Path, model, optimizer, scheduler, device):
-    """Load full training state. Returns (start_epoch, global_step, best_val_loss)."""
-    ckpt = torch.load(path, map_location=device)
-    model.load_state_dict(ckpt["model_state_dict"])
-    optimizer.load_state_dict(ckpt["optimizer_state_dict"])
-    scheduler.load_state_dict(ckpt["scheduler_state_dict"])
-    start_epoch = ckpt["epoch"] + 1
-    global_step = ckpt["global_step"]
-    best_val_loss = ckpt["best_val_loss"]
-    log_progress(f"Resumed from checkpoint: epoch {ckpt['epoch']+1}, step {global_step}, best_val_loss {best_val_loss:.5f}")
-    return start_epoch, global_step, best_val_loss
-
-def format_time(seconds: float) -> str:
-    h = int(seconds // 3600)
-    m = int((seconds % 3600) // 60)
-    s = seconds % 60
-    return f"{h}h {m}m {s:.1f}s"
-
-def estimate_training_flops(
-    model: torch.nn.Module,
-    num_layers: int,
-    max_seq_len: int,
-    d_model: int,
-    epochs: int,
-    batches_per_epoch: int,
-    batch_size: int,
-) -> dict:
-    n_params = sum(p.numel() for p in model.parameters())
-    flops_per_token = 6 * n_params + 6 * num_layers * max_seq_len * d_model
-    total_token_slots = epochs * batches_per_epoch * batch_size * max_seq_len
-    total_flops = flops_per_token * total_token_slots
-    total_tflops = total_flops / 1e12
-
-    return {
-        "n_params": n_params,
-        "flops_per_token": flops_per_token,
-        "total_batches_per_epoch": batches_per_epoch,
-        "total_token_slots": total_token_slots,
-        "total_flops": total_flops,
-        "total_tflops": total_tflops,
-    }
-
-def print_training_flops_summary(flops_info: dict, epochs: int) -> None:
-    """In bảng tổng FLOPs cần cho training."""
-    print("╠════════════════════════════════════════════════════════════════════════════════════╣")
-    print("║                      TINH TONG FLOPs CAN CHO TOAN BO TRAINING                      ║")
-    print("╠════════════════════════════════════════════════════════════════════════════════════╣")
-    print(f"So tham so model: {flops_info['n_params']:,}")
-    print(
-        f"Tong so TFLOP can cho toan bo qua trinh training "
-        f"({epochs} epoch, {flops_info['total_batches_per_epoch']} batch/epoch): "
-        f"{flops_info['total_tflops']:,.2f} TFLOP"
-    )
-    print("╠════════════════════════════════════════════════════════════════════════════════════╣")
-
-class TflopsBenchmarker:
-    def __init__(
-        self,
-        flops_per_token: float,
-        total_flops_needed: float,
-        total_batches_per_epoch: int,
-        epochs: int,
-        device: torch.device,
-        warmup_steps: int = 100,
-        target_steps: int = 500,
-    ):
-        self.flops_per_token = flops_per_token
-        self.total_flops_needed = total_flops_needed
-        self.total_batches_per_epoch = total_batches_per_epoch
-        self.epochs = epochs
-        self.device = device
-
-        self.warmup_steps = min(warmup_steps, total_batches_per_epoch)
-        self.target_steps = min(target_steps, total_batches_per_epoch)
-        self.reported = self.target_steps <= self.warmup_steps
-
-        self.tokens = 0
-        self.start_time = None
-
-    def is_bench_step(self, epoch: int, start_epoch: int, batch_idx: int) -> bool:
-        return (
-            (not self.reported)
-            and (epoch == start_epoch)
-            and (self.warmup_steps <= batch_idx < self.target_steps)
-        )
-
-    def on_step_begin(self, batch_idx: int, batch_tokens: int) -> None:
-        """Gọi khi bắt đầu step nằm trong vùng đo."""
-        if batch_idx == self.warmup_steps:
-            if self.device.type == "cuda":
-                torch.cuda.synchronize()
-            self.start_time = time.perf_counter()
-        self.tokens += batch_tokens
-
-    def on_step_end(self, batch_idx: int) -> None:
-        """Gọi ở step cuối vùng đo → in kết quả."""
-        if batch_idx != self.target_steps - 1:
-            return
-
-        if self.device.type == "cuda":
-            torch.cuda.synchronize()
-
-        elapsed = time.perf_counter() - self.start_time
-        measured_steps = self.target_steps - self.warmup_steps
-        flops = self.flops_per_token * self.tokens
-        tflops_per_sec = (flops / elapsed) / 1e12 if elapsed > 0 else 0.0
-        avg_time_per_step = elapsed / measured_steps if measured_steps > 0 else 0.0
-        token_slots_per_sec = self.tokens / elapsed if elapsed > 0 else 0.0
-
-        print()
-        print("╠════════════════════════════════════════════════════════════════════════════════════╣")
-        print(
-            f"[DO TFLOPS] Do {measured_steps} step training that "
-            f"(step {self.warmup_steps + 1} → {self.target_steps}), "
-            f"mat {elapsed:.2f}s (da bo qua {self.warmup_steps} step warm-up dau)"
-        )
-        print(f"[DO TOC DO] Trung binh moi step: {avg_time_per_step * 1000:,.2f} ms/step")
-        print(f"[DO TOC DO] Thong luong tensor thuc te: {token_slots_per_sec:,.0f} token-slot/s")
-
-        if tflops_per_sec > 0:
-            print(
-                f"[DO TFLOPS] TFLOPS uoc tinh theo cong thuc model: "
-                f"{tflops_per_sec:,.2f} TFLOPS"
-            )
-
-            total_steps_all_epochs = self.epochs * self.total_batches_per_epoch
-
-            est_time_actual = avg_time_per_step * total_steps_all_epochs
-            print(f"[Cach 1 - Thuc te] Uoc tinh thoi gian train (ngoai suy toc do do duoc): {format_time(est_time_actual)}")
-
-            # Cách 2 – lý thuyết
-            est_time_theory = self.total_flops_needed / (tflops_per_sec * 1e12)
-            print(f"[Cach 2 - Ly thuyet] Uoc tinh thoi gian train (Tong FLOPs / TFLOPS do duoc): {format_time(est_time_theory)}")
-        else:
-            print("[CẢNH BÁO] Thời gian đo quá nhỏ để tính TFLOPS đáng tin cậy.")
-
-        print("╠════════════════════════════════════════════════════════════════════════════════════╣")
-        self.reported = True
-
-class KernelLogger:
-    def __init__(self, enabled: bool = False, log_step: int = 500):
-        self.enabled = enabled
-        self.log_step = log_step
-        self.reported = not enabled
-        self._prof_ctx = None
-
-    def should_log(self, epoch: int, start_epoch: int, batch_idx: int) -> bool:
-        return (
-            (not self.reported)
-            and (epoch == start_epoch)
-            and (batch_idx == self.log_step)
-        )
-
-    def start(self):
-        self._prof_ctx = torch.profiler.profile(
-            activities=[
-                torch.profiler.ProfilerActivity.CPU,
-                torch.profiler.ProfilerActivity.CUDA,
-            ],
-            record_shapes=True,
-            profile_memory=True,
-            with_stack=True,
-            with_flops=True,
-        )
-        self._prof_ctx.__enter__()
-
-    def stop_and_report(self, batch_idx: int, epoch: int, device: torch.device):
-        if device.type == "cuda":
-            torch.cuda.synchronize()
-
-        self._prof_ctx.__exit__(None, None, None)
-
-        # Gom theo operator + input shape để chẩn đoán GEMM/Linear shape
-        # mà không phải lưu timeline gồm hàng triệu event.
-        key_avgs = self._prof_ctx.key_averages(group_by_input_shape=True)
-
-        def _self_cuda_us(e):
-            if hasattr(e, "self_cuda_time_total"):
-                return e.self_cuda_time_total
-            return getattr(e, "self_device_time_total", 0)
-
-        def _cuda_total_us(e):
-            if hasattr(e, "cuda_time_total"):
-                return e.cuda_time_total
-            return getattr(e, "device_time_total", 0)
-
-        total_self_cuda = sum(_self_cuda_us(e) for e in key_avgs)
-        total_self_cpu = sum(e.self_cpu_time_total for e in key_avgs)
-
-        rows = []
-        for e in key_avgs:
-            self_cuda = _self_cuda_us(e)
-            cuda_total = _cuda_total_us(e)
-
-            cuda_pct = (
-                self_cuda / total_self_cuda * 100
-                if total_self_cuda > 0 else 0.0
-            )
-
-            cpu_pct = (
-                e.self_cpu_time_total / total_self_cpu * 100
-                if total_self_cpu > 0 else 0.0
-            )
-
-            if cuda_pct >= 0.1 or cpu_pct >= 1.0:
-                rows.append({
-                    "name": e.key,
-                    "input_shapes": e.input_shapes,
-                    "self_cuda_ms": self_cuda / 1000,
-                    "cuda_total_ms": cuda_total / 1000,
-                    "cuda_pct": cuda_pct,
-                    "self_cpu_ms": e.self_cpu_time_total / 1000,
-                    "cpu_total_ms": e.cpu_time_total / 1000,
-                    "cpu_pct": cpu_pct,
-                    "calls": e.count,
-                    "flops": getattr(e, "flops", 0),
-                    "self_cpu_memory_bytes": getattr(e, "self_cpu_memory_usage", 0),
-                    "self_device_memory_bytes": getattr(e, "self_device_memory_usage", 0),
-                    "source_stack": list(getattr(e, "stack", []))[:5],
-                })
-
-        rows.sort(
-            key=lambda r: (r["self_cuda_ms"], r["self_cpu_ms"]),
-            reverse=True,
-        )
-        rows = rows[:20]
-
-        print()
-        print("╠════════════════════════════════════════════════════════════════════════════════════╣")
-        print(f"[LOG KERNEL] Top kernel tại step {batch_idx} (epoch {epoch+1})")
-        print(
-            f"{'Name':<55}"
-            f"{'Self CUDA':>12}"
-            f"{'CUDA Total':>12}"
-            f"{'CUDA %':>9}"
-            f"{'CPU %':>9}"
-            f"{'Calls':>9}"
-        )
-
-        for row in rows:
-            print(
-                f"{row['name'][:55]:<55}"
-                f"{row['self_cuda_ms']:>9.2f}ms"
-                f"{row['cuda_total_ms']:>9.2f}ms"
-                f"{row['cuda_pct']:>8.2f}%"
-                f"{row['cpu_pct']:>8.2f}%"
-                f"{row['calls']:>9}"
-            )
-
-        device_info = {"type": device.type}
-        if device.type == "cuda":
-            device_info.update({
-                "name": torch.cuda.get_device_name(device),
-                "capability": list(torch.cuda.get_device_capability(device)),
-            })
-
-        report = {
-            "format": "compact-kernel-profile-v1",
-            "torch_version": torch.__version__,
-            "device": device_info,
-            "epoch": epoch + 1,
-            "batch_idx": batch_idx,
-            "display_step": batch_idx + 1,
-            "totals": {
-                "self_cuda_ms": total_self_cuda / 1000,
-                "self_cpu_ms": total_self_cpu / 1000,
-            },
-            "filters": {
-                "min_cuda_pct": 0.1,
-                "min_cpu_pct": 1.0,
-                "max_rows": 20,
-            },
-            "operators": rows,
-        }
-
-        trace_file = Path(
-            f"torch_profile_epoch{epoch+1}_step{batch_idx}_summary.json"
-        )
-        trace_file.write_text(
-            json.dumps(report, indent=2, ensure_ascii=False),
-            encoding="utf-8",
-        )
-
-        print("────────────────────────────────────────────────────────────────────────────────────")
-        print(f"Compact profiler JSON: {trace_file}")
-        print("╠════════════════════════════════════════════════════════════════════════════════════╣")
-
-        self.reported = True
-        self._prof_ctx = None