feat: v1.3-lite-flash 训练代码 + 新词表(vocab_size=47725)

- 用 CharTokenizer.build_vocab 从 977K 数据生成词表 (vocab_size=47725)
- train_cnb.py: batch=224 + accum=2 (等效448) + AMP + torch.compile
- build_vocab.py: 词表生成脚本
- vocab_full.json: 含特殊token完整词表
- 模型 41.85M参数, L40 GPU训练

llm_chat_improved.py

@@ -7,6 +7,7 @@ LLM对话AI训练代码 - 改进版
 
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 import torch.optim as optim
 from torch.utils.data import Dataset, DataLoader
 import math
@@ -235,21 +236,13 @@ class Attention(nn.Module):
         xk = xk.transpose(1, 2)
         xv = xv.transpose(1, 2)
         
-        # 计算注意力分数
-        scores = torch.matmul(xq, xk.transpose(-2, -1)) / math.sqrt(self.d_head)
-        
-        # 应用掩码
-        if mask is not None:
-            scores = scores.masked_fill(mask == 0, float('-inf'))
-        
-        # Softmax
-        attn_weights = torch.softmax(scores, dim=-1)
-        attn_weights = self.dropout(attn_weights)
-        
-        # 注意力输出
-        output = torch.matmul(attn_weights, xv)
+        # 融合注意力（SDPA / FlashAttention-2）- 纯因果掩码，等价于原下三角 mask
+        dropout_p = self.dropout.p if self.training else 0.0
+        output = F.scaled_dot_product_attention(
+            xq, xk, xv, attn_mask=None, dropout_p=dropout_p, is_causal=True)
+
         output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
-        
+
         return self.wo(output)
 
 
@@ -583,6 +576,9 @@ class LLMTrainer:
     def __init__(self, model: nn.Module, device: str = 'cpu', lr: float = 0.0001, use_amp: bool = None):
         self.model = model.to(device)
         self.device = device
+        # 火力全开：TF32 加速 fp32 矩阵乘法（L40/A100/H100 原生支持）
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
         self.criterion = nn.CrossEntropyLoss(ignore_index=0)
         self.optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
         self.scheduler = optim.lr_scheduler.CosineAnnealingLR(self.optimizer, T_max=100)
@@ -618,7 +614,9 @@ class LLMTrainer:
             
             # 生成因果掩码
             mask = self.model.generate_mask(input_ids.size(1)).to(self.device)
-            
+
+            # CUDA graphs step 边界标记（torch.compile reduce-overhead 必需，避免输出 buffer 被覆盖）
+            torch.compiler.cudagraph_mark_step_begin()
             # 前向传播
             if self.use_amp:
                 with torch.cuda.amp.autocast():
@@ -677,30 +675,35 @@ class LLMTrainer:
                 target_ids = batch['target_ids'].to(self.device, non_blocking=True)
                 
                 mask = self.model.generate_mask(input_ids.size(1)).to(self.device)
+                torch.compiler.cudagraph_mark_step_begin()
                 outputs = self.model(input_ids, mask)
                 outputs = outputs.reshape(-1, outputs.size(-1))
                 tgt = target_ids.reshape(-1)
                 
                 loss = self.criterion(outputs, tgt)
                 total_loss += loss.item()
-        
+
+        if len(dataloader) == 0:
+            return float('inf')
         return total_loss / len(dataloader)
     
     def save_model(self, filepath: str, vocab: dict, epoch: int = None):
+        # 兼容 torch.compile：取原始（未编译）模块，避免 state_dict 带 _orig_mod. 前缀
+        core = getattr(self.model, '_orig_mod', self.model)
         save_dict = {
-            'model_state_dict': self.model.state_dict(),
+            'model_state_dict': core.state_dict(),
             'optimizer_state_dict': self.optimizer.state_dict(),
             'vocab': vocab,
             'model_config': {
-                'vocab_size': self.model.vocab_size,
-                'd_model': self.model.d_model,
-                'n_heads': self.model.n_heads,
-                'n_layers': self.model.n_layers,
-                'd_ff': self.model.d_ff,
-                'max_len': self.model.max_len,
-                'use_memory': self.model.use_memory,
-                'use_thinking': self.model.use_thinking,
-                'use_moe': self.model.use_moe
+                'vocab_size': core.vocab_size,
+                'd_model': core.d_model,
+                'n_heads': core.n_heads,
+                'n_layers': core.n_layers,
+                'd_ff': core.d_ff,
+                'max_len': core.max_len,
+                'use_memory': core.use_memory,
+                'use_thinking': core.use_thinking,
+                'use_moe': core.use_moe
             }
         }
         if epoch is not None:
@@ -711,7 +714,8 @@ class LLMTrainer:
     
     def load_model(self, filepath: str):
         checkpoint = torch.load(filepath, map_location=self.device)
-        self.model.load_state_dict(checkpoint['model_state_dict'])
+        core = getattr(self.model, '_orig_mod', self.model)
+        core.load_state_dict(checkpoint['model_state_dict'])
         self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
         vocab = checkpoint['vocab']
         print(f"模型已从 {filepath} 加载")