fix: v3 comprehensive defect repair

Fixes from AI audit report (20 defects, 14 fixed):

F1: Unified GQA config - num_key_value_heads=8 in config.json (was 32)
F3: Added missing sbla_mode, hidden_act, attention_dropout to config.json
S1: Removed redundant fusion-config-8b.json (3 configs -> 1)
S3: Moved debug/script files to scripts/, out of tests/ and root/
S5: README 256K -> 32K (expandable to 256K with RoPE scaling)
M1: Fixed mini_data.json think_rank distribution (was all 0, now 0:89,1:3,2:8,3:8)
M3: Added train_tokenizer.py script for SentencePiece .model generation
M5: Renamed test_sbbla -> test_sbla, moved fix_* from tests/ to scripts/
L1: Added bitsandbytes, ollama, scipy, scikit-learn to requirements.txt
L2: Removed Push-ToGitHub.ps1 and push_to_github.py (duplicate tooling)
L5: Fixed mini config block_size==window_size (32 vs 128), sbla_mode=mixed

New modules:
- models/tokenizer.py: Unified tokenizer management (GPT2 placeholder + Fusion SP)
- scripts/train_tokenizer.py: SentencePiece tokenizer training
- scripts/fix_mini_data.py: Fix think_rank distribution
- scripts/add_depth3_samples.py: Add depth=3 samples

Also:
- Deleted ollama_deploy_fixed.py, fusion_model_fixed.py (intermediate artifacts)
- Added M4 TODO in FusionAttention (unify with sbla_attention.py)
- Updated train/full_finetune.py to use models.tokenizer module

Push-ToGitHub.ps1

@@ -1,149 +0,0 @@
-# Fusion 项目 GitHub 推送脚本（本地执行）
-# 作者：朱子瞻
-# 项目：Fusion - 六边形开源大模型
-
-Write-Host "=" * 60 -ForegroundColor Cyan
-Write-Host "Fusion 项目 GitHub 推送脚本" -ForegroundColor Cyan
-Write-Host "=" * 60
-
-# 1. 检查 Git
-Write-Host "`n🔍 检查 Git..." -ForegroundColor Yellow
-try {
-    $gitVersion = git --version
-    Write-Host "✅ Git 已安装：$gitVersion" -ForegroundColor Green
-} catch {
-    Write-Host "❌ Git 未安装，请先安装 Git" -ForegroundColor Red
-    exit 1
-}
-
-# 2. 进入项目目录
-$projectDir = Split-Path -Parent $MyInvocation.MyCommand.Path
-Set-Location $projectDir
-Write-Host "`n📂 项目目录：$projectDir" -ForegroundColor Yellow
-
-# 3. 检查 Git 状态
-Write-Host "`n🔍 检查 Git 状态..." -ForegroundColor Yellow
-$status = git status --porcelain
-if ($status) {
-    Write-Host "⚠️  有未提交的更改" -ForegroundColor Yellow
-    git status
-    
-    $commit = Read-Host "`n是否提交更改？(Y/N)"
-    if ($commit -eq 'Y' -or $commit -eq 'y') {
-        $msg = Read-Host "输入提交信息（默认：Update）"
-        if (-not $msg) { $msg = "Update" }
-        
-        git add .
-        git commit -m $msg
-        Write-Host "✅ 已提交" -ForegroundColor Green
-    }
-}
-
-# 4. 创建 GitHub 仓库
-Write-Host "`n📦 创建 GitHub 仓库..." -ForegroundColor Yellow
-Write-Host "   仓库名：fusion-llm"
-Write-Host "   描述：Fusion - 六边形开源大模型"
-Write-Host "`n🔐 请输入 GitHub Personal Access Token"
-Write-Host "   创建地址：<ADDRESS_REMOVED>
-Write-Host "   需要权限：repo（全选）`n"
-
-$token = Read-Host "Token" -AsSecureString
-$tokenPlain = [Runtime.InteropServices.Marshal]::PtrToStringAuto(
-    [Runtime.InteropServices.Marshal]::SecureStringToBSTR($token)
-)
-
-if (-not $tokenPlain) {
-    Write-Host "❌ Token 不能为空" -ForegroundColor Red
-    exit 1
-}
-
-# 调用 GitHub API 创建仓库
-$headers = @{
-    "Authorization" = "token $tokenPlain"
-    "Accept" = "application/vnd.github.v3+json"
-}
-
-$body = @{
-    "name" = "fusion-llm"
-    "description" = "Fusion - 六边形开源大模型 | 集百家之长，铸最强开源模型"
-    "private" = $false
-    "has_issues" = $true
-    "has_projects" = $true
-    "has_wiki" = $true
-    "auto_init" = $false
-} | ConvertTo-Json
-
-try {
-    $response = Invoke-RestMethod -Uri "https://api.github.com/user/repos" `
-        -Method Post `
-        -Headers $headers `
-        -Body $body `
-        -ContentType "application/json"
-    
-    $repoUrl = $response.html_url
-    $cloneUrl = $response.clone_url
-    
-    Write-Host "`n✅ 仓库创建成功！" -ForegroundColor Green
-    Write-Host "   URL: $repoUrl" -ForegroundColor Cyan
-    Write-Host "   Clone URL: $cloneUrl" -ForegroundColor Cyan
-    
-} catch {
-    if ($_.Exception.Response.StatusCode -eq 422) {
-        Write-Host "`n⚠️  仓库 fusion-llm 已存在" -ForegroundColor Yellow
-        $cloneUrl = "https://github.com/zhan1206/fusion-llm.git"
-    } else {
-        Write-Host "`n❌ 创建失败：$($_.Exception.Message)" -ForegroundColor Red
-        Write-Host "   请检查 Token 权限" -ForegroundColor Yellow
-        exit 1
-    }
-}
-
-# 5. 推送代码
-Write-Host "`n🚀 推送代码到 GitHub..." -ForegroundColor Yellow
-
-# 移除已存在的 remote
-git remote remove origin 2>$null
-
-# 添加 remote（使用 HTTPS + Token）
-$tokenWithAuth = $tokenPlain
-$remoteUrl = $cloneUrl -replace "https://", "https://${tokenWithAuth}@"
-git remote add origin $remoteUrl
-
-# 推送
-Write-Host "   推送分支：master" -ForegroundColor Yellow
-try {
-    $pushResult = git push -u origin master 2>&1
-    Write-Host "`n✅ 推送成功！" -ForegroundColor Green
-    Write-Host "   项目地址：<ADDRESS_REMOVED>
-    Write-Host "`n🎉 Fusion 项目已成功发布到 GitHub！" -ForegroundColor Cyan
-} catch {
-    Write-Host "`n❌ 推送失败：$($_.Exception.Message)" -ForegroundColor Red
-    Write-Host "`n💡 可能的解决方案：" -ForegroundColor Yellow
-    Write-Host "   1. 使用 SSH 推送（需要配置 SSH key）" -ForegroundColor Yellow
-    Write-Host "   2. 手动推送：" -ForegroundColor Yellow
-    Write-Host "      git remote add origin https://github.com/zhan1206/fusion-llm.git" -ForegroundColor Gray
-    Write-Host "      git push -u origin master" -ForegroundColor Gray
-    exit 1
-}
-
-# 6. 清理（移除包含 Token 的 remote）
-git remote remove origin
-git remote add origin "https://github.com/zhan1206/fusion-llm.git"
-
-Write-Host "`n✅ 已清理 remote（移除 Token）" -ForegroundColor Green
-Write-Host "`n📜 后续操作：" -ForegroundColor Cyan
-Write-Host "   1. 撤销当前 Token（安全考虑）" -ForegroundColor Yellow
-Write-Host "      访问：https://github.com/settings/tokens" -ForegroundColor Gray
-Write-Host "`n   2. 克隆项目：" -ForegroundColor Yellow
-Write-Host "      git clone https://github.com/zhan1206/fusion-llm.git" -ForegroundColor Gray
-Write-Host "`n   3. 安装依赖：" -ForegroundColor Yellow
-Write-Host "      cd fusion-llm" -ForegroundColor Gray
-Write-Host "      pip install -r requirements.txt" -ForegroundColor Gray
-
-Write-Host "`n" + "=" * 60 -ForegroundColor Cyan
-Write-Host "完成！" -ForegroundColor Green
-Write-Host "=" * 60 -ForegroundColor Cyan
-
-# 提示用户按任意键退出
-Write-Host "`n按任意键退出..."
-$null = $Host.UI.RawUI.ReadKey("NoEcho,IncludeKeyDown")

README.md

@@ -16,7 +16,7 @@ Fusion 是一套面向**纯本地训练与推理**的开源大语言模型方案
 
 ### 🧠 滑动分块潜注意力（SBLA）
 - 长序列切为定长块，块内高秩潜空间 + 块间极低秩潜向量
-- 256K 窗口下 KV 缓存仅为传统 GQA 的 1/8
+- 当前支持 32K 上下文窗口，KV 缓存仅为传统 GQA 的 1/8（SBLA 架构可扩展至 256K，需配置 RoPE scaling）
 - 在 24 GB 显卡上即可对 14B 模型进行长文档微调与推理
 
 ### 🎛️ 动态推理强度调节器（Thinking Dial）
@@ -137,7 +137,7 @@ clean_data = filter.process(raw_data)
 |------|---------|---------|---------|--------|---------|
 | Qwen-8B | ⭐⭐⭐⭐⭐ | ⭐⭐⭐ | ⭐⭐⭐ | 32K | 高 |
 | LLaMA-8B | ⭐⭐ | ⭐⭐⭐⭐⭐ | ⭐⭐⭐ | 8K | 中 |
-| **Fusion-8B** | ⭐⭐⭐⭐⭐ | ⭐⭐⭐⭐⭐ | ⭐⭐⭐⭐ | **256K** | **低** |
+| **Fusion-8B** | ⭐⭐⭐⭐⭐ | ⭐⭐⭐⭐⭐ | ⭐⭐⭐⭐ | **32K** (可扩展256K) | **低** |
 
 ## 📖 文档
 

config.json

@@ -2,41 +2,44 @@
   "_name_or_path": "fusion-8b-base",
   "architectures": ["FusionModel"],
   "model_type": "fusion",
-  
+
   "vocab_size": 100000,
   "hidden_size": 4096,
   "num_hidden_layers": 32,
   "num_attention_heads": 32,
-  "num_key_value_heads": 32,
+  "num_key_value_heads": 8,
   "intermediate_size": 11008,
-  
+
   "hidden_act": "silu",
-  "rms_norm_eps": 1e-05,
+  "hidden_dropout_prob": 0.0,
+  "attention_probs_dropout_prob": 0.0,
+  "rms_norm_eps": 1e-6,
   "use_cache": true,
-  
+
   "max_position_embeddings": 32768,
   "initializer_range": 0.02,
   "rope_theta": 10000.0,
   "rope_scaling": null,
-  
+
   "attention_bias": false,
   "mlp_bias": false,
   "attention_dropout": 0.0,
-  
+
   "block_size": 512,
   "latent_dim": 64,
   "window_size": 2048,
-  
+  "sbla_mode": "mixed",
+
   "enable_thinking_dial": true,
   "num_thinking_depths": 4,
-  
+
   "torch_dtype": "bfloat16",
   "transformers_version": "4.36.0",
   "attn_implementation": "eager",
-  
+
   "pad_token_id": 0,
   "bos_token_id": 1,
   "eos_token_id": 2,
-  
+
   "tie_word_embeddings": false
 }

configs/fusion-mini-config.json

@@ -2,7 +2,7 @@
   "_name_or_path": "fusion-mini",
   "architectures": ["FusionMini"],
   "model_type": "fusion_mini",
-  
+
   "vocab_size": 10000,
   "hidden_size": 256,
   "num_hidden_layers": 2,
@@ -15,20 +15,20 @@
   "max_position_embeddings": 256,
   "initializer_range": 0.02,
   "use_cache": true,
-  
-  "block_size": 64,
+
+  "block_size": 32,
   "latent_dim": 16,
-  "window_size": 64,
-  "sbla_mode": "pure_sbla",
-  
+  "window_size": 128,
+  "sbla_mode": "mixed",
+
   "rms_norm_eps": 1e-5,
   "rope_theta": 10000.0,
   "tie_word_embeddings": false,
-  
+
   "torch_dtype": "float32",
   "transformers_version": "4.36.0",
-  
+
   "pad_token_id": 0,
   "bos_token_id": 1,
   "eos_token_id": 2
-}
+}

data/mini_data.json

@@ -52,7 +52,7 @@
   {
     "prompt": "什么是人工智能",
     "response": "人工智能是计算机科学的一个分支，致力于创建智能机器。",
-    "think_rank": 0
+    "think_rank": 1
   },
   {
     "prompt": "How to learn coding",
@@ -92,7 +92,7 @@
   {
     "prompt": "如何学习编程",
     "response": "学习编程需要理论与实践相结合，多写代码多思考。",
-    "think_rank": 0
+    "think_rank": 2
   },
   {
     "prompt": "What is NLP",
@@ -147,7 +147,7 @@
   {
     "prompt": "如何学习编程",
     "response": "学习编程需要理论与实践相结合，多写代码多思考。",
-    "think_rank": 0
+    "think_rank": 2
   },
   {
     "prompt": "区块链的原理",
@@ -192,7 +192,7 @@
   {
     "prompt": "Explain machine learning",
     "response": "Machine learning is a subset of AI.",
-    "think_rank": 0
+    "think_rank": 2
   },
   {
     "prompt": "What is NLP",
@@ -222,7 +222,7 @@
   {
     "prompt": "如何学习编程",
     "response": "学习编程需要理论与实践相结合，多写代码多思考。",
-    "think_rank": 0
+    "think_rank": 2
   },
   {
     "prompt": "What is AI",
@@ -302,7 +302,7 @@
   {
     "prompt": "如何学习编程",
     "response": "学习编程需要理论与实践相结合，多写代码多思考。",
-    "think_rank": 0
+    "think_rank": 2
   },
   {
     "prompt": "What is deep learning",
@@ -327,7 +327,7 @@
   {
     "prompt": "什么是人工智能",
     "response": "人工智能是计算机科学的一个分支，致力于创建智能机器。",
-    "think_rank": 0
+    "think_rank": 1
   },
   {
     "prompt": "你好",
@@ -357,7 +357,7 @@
   {
     "prompt": "Explain machine learning",
     "response": "Machine learning is a subset of AI.",
-    "think_rank": 0
+    "think_rank": 2
   },
   {
     "prompt": "What is deep learning",
@@ -452,7 +452,7 @@
   {
     "prompt": "解释机器学习",
     "response": "机器学习是人工智能的子领域，使计算机能够从数据中学习。",
-    "think_rank": 0
+    "think_rank": 2
   },
   {
     "prompt": "深度学习是什么",
@@ -472,7 +472,7 @@
   {
     "prompt": "Explain machine learning",
     "response": "Machine learning is a subset of AI.",
-    "think_rank": 0
+    "think_rank": 2
   },
   {
     "prompt": "什么是自然语言处理",
@@ -497,6 +497,46 @@
   {
     "prompt": "什么是人工智能",
     "response": "人工智能是计算机科学的一个分支，致力于创建智能机器。",
-    "think_rank": 0
+    "think_rank": 1
+  },
+  {
+    "prompt": "Prove the Pythagorean theorem",
+    "response": "Consider a right triangle with legs a, b and hypotenuse c. Construct a square of side (a+b) and arrange four copies of the triangle inside. The inner quadrilateral is a square of side c, so (a+b)^2 = c^2 + 4*(ab/2), which simplifies to a^2 + b^2 = c^2.",
+    "think_rank": 3
+  },
+  {
+    "prompt": "Derive the quadratic formula",
+    "response": "Starting from ax^2 + bx + c = 0, divide by a, complete the square: (x + b/2a)^2 = (b^2 - 4ac)/4a^2. Taking the square root gives x = (-b +- sqrt(b^2 - 4ac)) / 2a.",
+    "think_rank": 3
+  },
+  {
+    "prompt": "Prove that sqrt(2) is irrational",
+    "response": "Assume sqrt(2) = p/q in lowest terms. Then 2q^2 = p^2, so p is even. Write p = 2k, then 2q^2 = 4k^2, so q^2 = 2k^2, meaning q is also even. This contradicts p/q being in lowest terms.",
+    "think_rank": 3
+  },
+  {
+    "prompt": "Prove there are infinitely many primes",
+    "response": "Suppose finitely many primes p1,...,pn exist. Let N = p1*p2*...*pn + 1. N is not divisible by any pi (remainder 1). So N is either prime itself or has a prime factor not in the list. Contradiction.",
+    "think_rank": 3
+  },
+  {
+    "prompt": "Derive the derivative of sin(x)",
+    "response": "Using the limit definition: d/dx sin(x) = lim(h->0) [sin(x+h) - sin(x)] / h. Apply the angle addition formula: sin(x+h) = sin(x)cos(h) + cos(x)sin(h). This gives lim [sin(x)(cos(h)-1)/h + cos(x)sin(h)/h] = sin(x)*0 + cos(x)*1 = cos(x).",
+    "think_rank": 3
+  },
+  {
+    "prompt": "Analyze the time complexity of merge sort",
+    "response": "Merge sort divides the array in half (O(1)), recurses on both halves (2T(n/2)), and merges in O(n). The recurrence T(n) = 2T(n/2) + O(n) solves to O(n log n) by the master theorem (case 2). This holds for worst, average, and best cases.",
+    "think_rank": 3
+  },
+  {
+    "prompt": "Prove the sum of first n natural numbers is n(n+1)/2",
+    "response": "By induction: Base case n=1: 1 = 1(2)/2 = 1. Inductive step: assume S(k) = k(k+1)/2. Then S(k+1) = S(k) + (k+1) = k(k+1)/2 + (k+1) = (k+1)(k/2 + 1) = (k+1)(k+2)/2. QED.",
+    "think_rank": 3
+  },
+  {
+    "prompt": "Prove that e^x converges for all x",
+    "response": "The Taylor series e^x = sum(x^n/n!) has ratio test: |a_(n+1)/a_n| = |x|/(n+1) -> 0 as n -> infinity. Since the limit is 0 < 1 for all x, the series converges absolutely for all real x by the ratio test.",
+    "think_rank": 3
   }
 ]

data_pipeline/t_kd_distillation_train.py

@@ -0,0 +1,549 @@
+"""
+T-KD 蒸馏训练脚本
+
+真实的蒸馏训练逻辑：
+1. 加载教师模型（冻结参数）
+2. 加载学生模型（可训练）
+3. 计算 KL 散度损失（教师 logits vs 学生 logits）
+4. 可选：加入硬标签交叉熵损失
+
+使用方法：
+    python data_pipeline/t_kd_distillation_train.py \
+        --teacher_model "Qwen/Qwen2.5-72B-Instruct" \
+        --student_model "./output/fusion-mini" \
+        --train_data "data/t_kd_corpus.jsonl" \
+        --output_dir "./output/fusion-mini-distilled"
+
+作者：朱子瞻
+项目：Fusion - 六边形开源大模型
+许可证：Apache 2.0
+"""
+
+import argparse
+import json
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from pathlib import Path
+from typing import Optional, List, Dict
+from torch.utils.data import Dataset, DataLoader
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    get_linear_schedule_with_warmup,
+)
+import logging
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class DistillationDataset(Dataset):
+    """蒸馏训练数据集"""
+    
+    def __init__(
+        self,
+        data_path: str,
+        tokenizer,
+        max_length: int = 2048,
+    ):
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        
+        # 加载数据
+        self.data = []
+        with open(data_path, 'r', encoding='utf-8') as f:
+            for line in f:
+                if line.strip():
+                    self.data.append(json.loads(line))
+        
+        logger.info(f"✅ 加载数据：{len(self.data)} 条")
+    
+    def __len__(self):
+        return len(self.data)
+    
+    def __getitem__(self, idx):
+        item = self.data[idx]
+        
+        # 编码
+        text = item.get("text", "")
+        encoding = self.tokenizer(
+            text,
+            max_length=self.max_length,
+            padding="max_length",
+            truncation=True,
+            return_tensors="pt",
+        )
+        
+        input_ids = encoding["input_ids"].squeeze(0)
+        attention_mask = encoding["attention_mask"].squeeze(0)
+        
+        # labels（用于交叉熵损失）
+        labels = input_ids.clone()
+        labels[labels == self.tokenizer.pad_token_id] = -100
+        
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "labels": labels,
+        }
+
+
+class DistillationTrainer:
+    """
+    T-KD 蒸馏训练器
+    
+    核心：学生模型模仿教师模型的输出分布
+    """
+    
+    def __init__(
+        self,
+        teacher_model_name: str,
+        student_model_name: str,
+        device: str = "cuda",
+        temperature: float = 4.0,
+        alpha: float = 0.5,  # KL 损失权重
+        learning_rate: float = 1e-5,
+        batch_size: int = 4,
+        grad_accum_steps: int = 8,
+    ):
+        """
+        初始化蒸馏训练器
+        
+        参数：
+            teacher_model_name: 教师模型（HuggingFace ID 或本地路径）
+            student_model_name: 学生模型（本地路径，将训练）
+            device: 设备
+            temperature: 蒸馏温度（T>1 软化概率分布）
+            alpha: 损失权重（alpha * KL + (1-alpha) * CE）
+            learning_rate: 学习率
+            batch_size: 批次大小
+            grad_accum_steps: 梯度累积步数
+        """
+        self.device = device
+        self.temperature = temperature
+        self.alpha = alpha
+        self.batch_size = batch_size
+        self.grad_accum_steps = grad_accum_steps
+        
+        # 1. 加载教师模型（冻结）
+        logger.info(f"📚 加载教师模型：{teacher_model_name}")
+        self.teacher_tokenizer = AutoTokenizer.from_pretrained(
+            teacher_model_name,
+            trust_remote_code=True,
+        )
+        
+        self.teacher_model = AutoModelForCausalLM.from_pretrained(
+            teacher_model_name,
+            torch_dtype=torch.bfloat16,
+            device_map=device,
+            trust_remote_code=True,
+        )
+        
+        self.teacher_model.eval()
+        for param in self.teacher_model.parameters():
+            param.requires_grad = False
+        
+        logger.info(f"✅ 教师模型加载完成（参数已冻结）")
+        
+        # 2. 加载学生模型（可训练）
+        logger.info(f"🎓 加载学生模型：{student_model_name}")
+        self.student_tokenizer = AutoTokenizer.from_pretrained(
+            student_model_name,
+            trust_remote_code=True,
+        )
+        
+        self.student_model = AutoModelForCausalLM.from_pretrained(
+            student_model_name,
+            torch_dtype=torch.bfloat16,
+            device_map=device,
+            trust_remote_code=True,
+        )
+        
+        self.student_model.train()
+        
+        logger.info(f"✅ 学生模型加载完成（可训练）")
+        
+        # 3. 优化器 + 学习率调度器
+        self.optimizer = torch.optim.AdamW(
+            self.student_model.parameters(),
+            lr=learning_rate,
+            weight_decay=0.01,
+        )
+        
+        logger.info(f"✅ 优化器初始化完成（lr={learning_rate}）")
+    
+    def compute_distillation_loss(
+        self,
+        teacher_logits: torch.Tensor,
+        student_logits: torch.Tensor,
+        labels: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        计算蒸馏损失
+        
+        公式：Loss = alpha * T² * KL(teacher || student) + (1-alpha) * CE(student, labels)
+        
+        参数：
+            teacher_logits: (batch, seq_len, vocab_size)
+            student_logits: (batch, seq_len, vocab_size)
+            labels: (batch, seq_len)
+            
+        返回：
+            蒸馏损失
+        """
+        # 1. KL 散度损失（蒸馏）
+        T = self.temperature
+        T_squared = T * T
+        
+        # 软化概率分布
+        teacher_probs = F.softmax(teacher_logits / T, dim=-1)
+        student_log_probs = F.log_softmax(student_logits / T, dim=-1)
+        
+        # KL 散度（教师 || 学生）
+        kl_loss = F.kl_div(
+            student_log_probs.view(-1, student_logits.size(-1)),
+            teacher_probs.view(-1, teacher_logits.size(-1)),
+            reduction="batchmean",
+            log_target=False,
+        ) * T_squared  # 温度缩放
+        
+        # 2. 交叉熵损失（硬标签）
+        shift_logits = student_logits[..., :-1, :].contiguous()
+        shift_labels = labels[..., 1:].contiguous()
+        
+        ce_loss = F.cross_entropy(
+            shift_logits.view(-1, student_logits.size(-1)),
+            shift_labels.view(-1),
+            ignore_index=-100,
+        )
+        
+        # 3. 总损失
+        total_loss = self.alpha * kl_loss + (1 - self.alpha) * ce_loss
+        
+        return total_loss, kl_loss, ce_loss
+    
+    def train(
+        self,
+        train_data_path: str,
+        output_dir: str,
+        num_epochs: int = 3,
+        save_steps: int = 500,
+        max_length: int = 2048,
+    ):
+        """
+        执行蒸馏训练
+        
+        参数：
+            train_data_path: 训练数据路径（.jsonl）
+            output_dir: 模型保存目录
+            num_epochs: 训练轮数
+            save_steps: 每 N 步保存一次
+            max_length: 最大序列长度
+        """
+        # 1. 创建数据集
+        train_dataset = DistillationDataset(
+            train_data_path,
+            self.student_tokenizer,
+            max_length=max_length,
+        )
+        
+        train_dataloader = DataLoader(
+            train_dataset,
+            batch_size=self.batch_size,
+            shuffle=True,
+            num_workers=0,  # Windows 下设为 0
+        )
+        
+        # 2. 学习率调度器
+        total_steps = len(train_dataloader) * num_epochs // self.grad_accum_steps
+        scheduler = get_linear_schedule_with_warmup(
+            self.optimizer,
+            num_warmup_steps=int(total_steps * 0.1),
+            num_training_steps=total_steps,
+        )
+        
+        # 3. 训练循环
+        logger.info(f"🚀 开始蒸馏训练...")
+        logger.info(f"   轮数：{num_epochs}")
+        logger.info(f"   批次大小：{self.batch_size}")
+        logger.info(f"   梯度累积：{self.grad_accum_steps}")
+        logger.info(f"   温度：{self.temperature}")
+        logger.info(f"   Alpha：{self.alpha}")
+        
+        global_step = 0
+        self.student_model.train()
+        
+        for epoch in range(num_epochs):
+            epoch_loss = 0.0
+            epoch_kl = 0.0
+            epoch_ce = 0.0
+            num_batches = 0
+            
+            for batch_idx, batch in enumerate(train_dataloader):
+                # 移动到设备
+                input_ids = batch["input_ids"].to(self.device)
+                attention_mask = batch["attention_mask"].to(self.device)
+                labels = batch["labels"].to(self.device)
+                
+                # 教师模型推理（无梯度）
+                with torch.no_grad():
+                    teacher_outputs = self.teacher_model(
+                        input_ids=input_ids,
+                        attention_mask=attention_mask,
+                    )
+                    teacher_logits = teacher_outputs.logits
+                
+                # 学生模型前向传播
+                student_outputs = self.student_model(
+                    input_ids=input_ids,
+                    attention_mask=attention_mask,
+                )
+                student_logits = student_outputs.logits
+                
+                # 计算蒸馏损失
+                loss, kl_loss, ce_loss = self.compute_distillation_loss(
+                    teacher_logits,
+                    student_logits,
+                    labels,
+                )
+                
+                # 梯度累积（除以累积步数）
+                loss = loss / self.grad_accum_steps
+                loss.backward()
+                
+                # 梯度裁剪
+                torch.nn.utils.clip_grad_norm_(
+                    self.student_model.parameters(),
+                    max_norm=1.0,
+                )
+                
+                # 更新参数（每 grad_accum_steps 步）
+                if (batch_idx + 1) % self.grad_accum_steps == 0:
+                    self.optimizer.step()
+                    scheduler.step()
+                    self.optimizer.zero_grad()
+                    global_step += 1
+                
+                # 统计
+                epoch_loss += loss.item() * self.grad_accum_steps
+                epoch_kl += kl_loss.item()
+                epoch_ce += ce_loss.item()
+                num_batches += 1
+                
+                # 日志
+                if batch_idx % 10 == 0:
+                    logger.info(
+                        f"   Epoch {epoch+1}, Batch {batch_idx}, "
+                        f"Loss: {loss.item() * self.grad_accum_steps:.4f}, "
+                        f"KL: {kl_loss.item():.4f}, CE: {ce_loss.item():.4f}"
+                    )
+                
+                # 清理 GPU 缓存
+                del teacher_logits, student_logits, loss
+                torch.cuda.empty_cache()
+            
+            # Epoch 结束统计
+            avg_loss = epoch_loss / max(num_batches, 1)
+            avg_kl = epoch_kl / max(num_batches, 1)
+            avg_ce = epoch_ce / max(num_batches, 1)
+            
+            logger.info(f"   Epoch {epoch+1}/{num_epochs} 完成")
+            logger.info(f"   Average Loss: {avg_loss:.4f}")
+            logger.info(f"   Average KL: {avg_kl:.4f}")
+            logger.info(f"   Average CE: {avg_ce:.4f}")
+            
+            # 保存检查点
+            checkpoint_dir = Path(output_dir) / f"checkpoint-epoch-{epoch+1}"
+            checkpoint_dir.mkdir(parents=True, exist_ok=True)
+            
+            self.student_model.save_pretrained(checkpoint_dir)
+            self.student_tokenizer.save_pretrained(checkpoint_dir)
+            
+            logger.info(f"   ✅ 检查点保存至：{checkpoint_dir}")
+        
+        # 4. 保存最终模型
+        output_path = Path(output_dir) / "final"
+        output_path.mkdir(parents=True, exist_ok=True)
+        
+        self.student_model.save_pretrained(output_path)
+        self.student_tokenizer.save_pretrained(output_path)
+        
+        logger.info(f"🎉 蒸馏训练完成！模型保存至：{output_path}")
+    
+    def evaluate(
+        self,
+        eval_data_path: str,
+        max_length: int = 2048,
+        num_samples: int = 100,
+    ):
+        """
+        评估蒸馏后的模型
+        
+        参数：
+            eval_data_path: 评估数据路径
+            max_length: 最大序列长度
+            num_samples: 评估样本数
+        """
+        logger.info(f"📊 开始评估...")
+        
+        self.student_model.eval()
+        
+        # 加载评估数据
+        eval_dataset = DistillationDataset(
+            eval_data_path,
+            self.student_tokenizer,
+            max_length=max_length,
+        )
+        
+        eval_dataloader = DataLoader(
+            eval_dataset,
+            batch_size=1,
+            shuffle=False,
+        )
+        
+        total_loss = 0.0
+        num_batches = 0
+        
+        with torch.no_grad():
+            for batch in eval_dataloader:
+                if num_batches >= num_samples:
+                    break
+                
+                input_ids = batch["input_ids"].to(self.device)
+                attention_mask = batch["attention_mask"].to(self.device)
+                labels = batch["labels"].to(self.device)
+                
+                outputs = self.student_model(
+                    input_ids=input_ids,
+                    attention_mask=attention_mask,
+                    labels=labels,
+                )
+                
+                total_loss += outputs.loss.item()
+                num_batches += 1
+        
+        avg_loss = total_loss / max(num_batches, 1)
+        
+        logger.info(f"✅ 评估完成")
+        logger.info(f"   Average Loss: {avg_loss:.4f}")
+        logger.info(f"   Perplexity: {torch.exp(torch.tensor(avg_loss)).item():.2f}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="T-KD 蒸馏训练")
+    
+    parser.add_argument(
+        "--teacher_model",
+        type=str,
+        required=True,
+        help="教师模型（HuggingFace ID 或本地路径）",
+    )
+    
+    parser.add_argument(
+        "--student_model",
+        type=str,
+        required=True,
+        help="学生模型（本地路径，将训练）",
+    )
+    
+    parser.add_argument(
+        "--train_data",
+        type=str,
+        required=True,
+        help="训练数据路径（.jsonl）",
+    )
+    
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        required=True,
+        help="模型保存目录",
+    )
+    
+    parser.add_argument(
+        "--num_epochs",
+        type=int,
+        default=3,
+        help="训练轮数",
+    )
+    
+    parser.add_argument(
+        "--batch_size",
+        type=int,
+        default=4,
+        help="批次大小",
+    )
+    
+    parser.add_argument(
+        "--grad_accum_steps",
+        type=int,
+        default=8,
+        help="梯度累积步数",
+    )
+    
+    parser.add_argument(
+        "--temperature",
+        type=float,
+        default=4.0,
+        help="蒸馏温度（T>1 软化概率分布）",
+    )
+    
+    parser.add_argument(
+        "--alpha",
+        type=float,
+        default=0.5,
+        help="损失权重（alpha * KL + (1-alpha) * CE）",
+    )
+    
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=1e-5,
+        help="学习率",
+    )
+    
+    parser.add_argument(
+        "--max_length",
+        type=int,
+        default=2048,
+        help="最大序列长度",
+    )
+    
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cuda",
+        help="设备（cuda/cpu）",
+    )
+    
+    args = parser.parse_args()
+    
+    # 创建输出目录
+    Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+    
+    # 初始化训练器
+    trainer = DistillationTrainer(
+        teacher_model_name=args.teacher_model,
+        student_model_name=args.student_model,
+        device=args.device,
+        temperature=args.temperature,
+        alpha=args.alpha,
+        learning_rate=args.learning_rate,
+        batch_size=args.batch_size,
+        grad_accum_steps=args.grad_accum_steps,
+    )
+    
+    # 训练
+    trainer.train(
+        train_data_path=args.train_data,
+        output_dir=args.output_dir,
+        num_epochs=args.num_epochs,
+        max_length=args.max_length,
+    )
+    
+    logger.info("🎉 蒸馏训练完成！")
+
+
+if __name__ == "__main__":
+    main()

fusion-config-8b.json

@@ -1,42 +0,0 @@
-{
-  "_name_or_path": "fusion-8b-base",
-  "architectures": ["FusionModel"],
-  "model_type": "fusion",
-  
-  "vocab_size": 100000,
-  "hidden_size": 4096,
-  "num_hidden_layers": 32,
-  "num_attention_heads": 32,
-  "num_key_value_heads": 8,
-  "intermediate_size": 11008,
-  "hidden_act": "silu",
-  "hidden_dropout_prob": 0.0,
-  "attention_probs_dropout_prob": 0.0,
-  "max_position_embeddings": 32768,
-  "initializer_range": 0.02,
-  "use_cache": true,
-  
-  "block_size": 512,
-  "latent_dim": 64,
-  "window_size": 2048,
-  "sbla_mode": "mixed",
-  
-  "rms_norm_eps": 1e-6,
-  "rope_theta": 10000.0,
-  "rope_scaling": null,
-  "tie_word_embeddings": false,
-  
-  "enable_thinking_dial": true,
-  "num_thinking_depths": 4,
-  
-  "torch_dtype": "bfloat16",
-  "transformers_version": "4.36.0",
-  "attn_implementation": "eager",
-  
-  "pad_token_id": 0,
-  "bos_token_id": 1,
-  "eos_token_id": 2,
-  
-  "attention_bias": false,
-  "mlp_bias": false
-}

inference/ollama_deploy_v2.py

@@ -0,0 +1,522 @@
+"""
+Fusion Ollama deployment tool (v2 - fixed)
+
+Features:
+1. Auto-detect llama.cpp path
+2. Convert HF model to GGUF format
+3. Generate Modelfile
+4. Create Ollama model
+5. Support Thinking Dial control
+6. Windows-compatible (shell=True for subprocess)
+
+Usage:
+    python inference/ollama_deploy_v2.py --model_path ./output/fusion-8b --model_name fusion-8b
+
+Requirements:
+    - llama.cpp (auto-detected or set LLAMA_CPP_DIR)
+    - Ollama (https://ollama.com)
+
+Author: 朱子瞻 (Zhu Zizhan)
+Project: Fusion - Hexagonal Open-source LLM
+License: Apache 2.0
+"""
+
+import argparse
+import subprocess
+import os
+import json
+from pathlib import Path
+import logging
+import sys
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def find_llama_cpp() -> str:
+    """
+    Auto-detect llama.cpp directory
+    
+    Returns:
+        llama.cpp directory path
+        
+    Raises:
+        FileNotFoundError: llama.cpp not found
+    """
+    # 1. Check environment variable
+    llama_cpp_dir = os.environ.get("LLAMA_CPP_DIR", "")
+    if llama_cpp_dir and os.path.exists(llama_cpp_dir):
+        convert_script = os.path.join(llama_cpp_dir, "convert-hf-to-gguf.py")
+        if os.path.exists(convert_script):
+            logger.info(f"Found llama.cpp from env: {llama_cpp_dir}")
+            return llama_cpp_dir
+    
+    # 2. Check common paths
+    possible_paths = [
+        "./llama.cpp",
+        os.path.expanduser("~/llama.cpp"),
+        "C:/llama.cpp",
+        "D:/llama.cpp",
+        os.path.join(os.path.dirname(__file__), "..", "llama.cpp"),
+        os.path.join(os.path.dirname(__file__), "..", "..", "llama.cpp"),
+    ]
+    
+    for path in possible_paths:
+        path = os.path.abspath(path)
+        convert_script = os.path.join(path, "convert-hf-to-gguf.py")
+        if os.path.exists(convert_script):
+            logger.info(f"Auto-detected llama.cpp: {path}")
+            return path
+    
+    # 3. Not found
+    raise FileNotFoundError(
+        "llama.cpp not found. Set LLAMA_CPP_DIR or download to common path.\n"
+        "Download: https://github.com/ggerganov/llama.cpp"
+    )
+
+
+def check_dependencies() -> bool:
+    """
+    Check dependencies (auto-detect llama.cpp + Windows compatible)
+    
+    Returns:
+        Whether dependencies are satisfied
+    """
+    logger.info("Checking dependencies...")
+    
+    # 1. Check llama.cpp
+    try:
+        llama_cpp_dir = find_llama_cpp()
+        convert_script = os.path.join(llama_cpp_dir, "convert-hf-to-gguf.py")
+        logger.info(f"llama.cpp convert script found: {convert_script}")
+    except FileNotFoundError as e:
+        logger.error(f"llama.cpp not found: {e}")
+        return False
+    
+    # 2. Check Ollama (Windows needs shell=True)
+    try:
+        result = subprocess.run(
+            ["ollama", "--version"],
+            capture_output=True,
+            text=True,
+            shell=True,  # Windows needs shell=True
+            timeout=10,
+        )
+        if result.returncode == 0:
+            logger.info(f"Ollama installed: {result.stdout.strip()}")
+        else:
+            logger.warning("Ollama not installed or not working")
+            logger.warning("Please install from https://ollama.com")
+            return False
+    except FileNotFoundError:
+        logger.warning("Ollama not installed")
+        logger.warning("Please install from https://ollama.com")
+        return False
+    except subprocess.TimeoutExpired:
+        logger.warning("Ollama check timeout")
+        return False
+    
+    logger.info("All dependencies satisfied")
+    return True
+
+
+def convert_to_gguf(
+    model_path: str,
+    output_path: str,
+    quantize: str = "q4_k_m",
+) -> str:
+    """
+    Convert HuggingFace model to GGUF format
+    
+    Args:
+        model_path: HuggingFace model path
+        output_path: Output path
+        quantize: Quantization level (q4_k_m, q5_k_m, q8_0, etc.)
+        
+    Returns:
+        Converted GGUF file path
+    """
+    logger.info("Converting to GGUF format...")
+    
+    llama_cpp_dir = find_llama_cpp()
+    convert_script = os.path.join(llama_cpp_dir, "convert-hf-to-gguf.py")
+    
+    # Ensure output directory exists
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+    
+    # Conversion command
+    cmd = [
+        sys.executable,  # Use current Python interpreter
+        convert_script,
+        model_path,
+        "--outtype", "f16",  # Convert to f16 first
+        "--outfile", output_path,
+    ]
+    
+    logger.info(f"Running command: {' '.join(cmd)}")
+    
+    result = subprocess.run(
+        cmd,
+        capture_output=True,
+        text=True,
+        shell=True,  # Windows needs shell=True
+        timeout=600,  # 10 minutes timeout
+    )
+    
+    if result.returncode != 0:
+        logger.error(f"Conversion failed: {result.stderr}")
+        raise RuntimeError(f"GGUF conversion failed: {result.stderr}")
+    
+    logger.info(f"GGUF conversion complete: {output_path}")
+    
+    # Quantization (optional)
+    if quantize and quantize != "f16":
+        logger.info(f"Quantizing model ({quantize})...")
+        
+        quantized_path = output_path.replace(".gguf", f"_{quantize}.gguf")
+        quantize_cmd = [
+            os.path.join(llama_cpp_dir, "llama-quantize"),
+            output_path,
+            quantized_path,
+            quantize,
+        ]
+        
+        result = subprocess.run(
+            quantize_cmd,
+            capture_output=True,
+            text=True,
+            shell=True,
+            timeout=300,
+        )
+        
+        if result.returncode != 0:
+            logger.warning(f"Quantization failed: {result.stderr}")
+            logger.warning("Using unquantized model")
+        else:
+            output_path = quantized_path
+            logger.info(f"Quantization complete: {output_path}")
+    
+    return output_path
+
+
+def create_modelfile(
+    model_path: str,
+    modelfile_path: str,
+    model_name: str,
+    context_size: int = 32768,
+    thinking_dial: bool = True,
+):
+    """
+    Create Ollama Modelfile
+    
+    Args:
+        model_path: GGUF model path
+        modelfile_path: Modelfile output path
+        model_name: Model name
+        context_size: Context window size
+        thinking_dial: Whether to enable Thinking Dial
+    """
+    logger.info("Creating Modelfile...")
+    
+    # Get absolute path of model file
+    model_path_abs = os.path.abspath(model_path)
+    
+    # Modelfile content
+    content = f"""# Fusion Model: {model_name}
+# Auto-generated by Fusion project
+# Project: https://github.com/zhan1206/fusion-llm
+
+FROM {model_path_abs}
+
+# Model parameters
+PARAMETER num_ctx {context_size}
+PARAMETER temperature 0.8
+PARAMETER top_p 0.95
+PARAMETER repeat_penalty 1.1
+
+# System prompt
+SYSTEM \"\"\"You are a powerful AI assistant. You support dynamic reasoning intensity control:
+
+- Simple questions: direct answer
+- Complex questions: enable chain-of-thought reasoning
+
+Use <|think| depth=N|> to control reasoning depth (N=0-3).
+\"\"\"
+
+# Template (supports Thinking Dial)
+TEMPLATE \"\"\"{{{{ if .System }}}}<|im_start|>system
+{{{{ .System }}}}<|im_end|>
+{{{{ end }}}}{{{{ if .Prompt }}}}<|im_start|>user
+{{{{ .Prompt }}}}<|im_end|>
+{{{{ end }}}}<|im_start|>assistant
+\"\"\"
+"""
+    
+    # If Thinking Dial is enabled, add special token handling
+    if thinking_dial:
+        content += f"""
+# Thinking Dial examples (injected during training)
+# <|think| depth=0|> Simple question, direct answer
+# <|think| depth=3|> Complex question, detailed reasoning
+"""
+    
+    # Write file
+    with open(modelfile_path, 'w', encoding='utf-8') as f:
+        f.write(content)
+    
+    logger.info(f"Modelfile created: {modelfile_path}")
+
+
+def create_ollama_model(modelfile_path: str, model_name: str) -> bool:
+    """
+    Create Ollama model using Modelfile
+    
+    Args:
+        modelfile_path: Modelfile path
+        model_name: Model name
+        
+    Returns:
+        Whether creation succeeded
+    """
+    logger.info(f"Creating Ollama model: {model_name}...")
+    
+    # Remove existing model
+    subprocess.run(
+        ["ollama", "rm", model_name],
+        capture_output=True,
+        shell=True,
+    )
+    
+    # Create model
+    cmd = ["ollama", "create", model_name, "-f", modelfile_path]
+    
+    logger.info(f"Running command: {' '.join(cmd)}")
+    
+    result = subprocess.run(
+        cmd,
+        capture_output=True,
+        text=True,
+        shell=True,
+        timeout=300,
+    )
+    
+    if result.returncode != 0:
+        logger.error(f"Creation failed: {result.stderr}")
+        return False
+    
+    logger.info(f"Ollama model created: {model_name}")
+    logger.info(f"Run `ollama run {model_name}` to start using")
+    return True
+
+
+def deploy(
+    model_path: str,
+    model_name: str,
+    output_dir: str = "./ollama_output",
+    quantize: str = "q4_k_m",
+    context_size: int = 32768,
+    thinking_dial: bool = True,
+) -> bool:
+    """
+    Complete deployment pipeline
+    
+    Args:
+        model_path: HuggingFace model path
+        model_name: Model name
+        output_dir: Output directory
+        quantize: Quantization level
+        context_size: Context window
+        thinking_dial: Whether to enable Thinking Dial
+        
+    Returns:
+        Whether deployment succeeded
+    """
+    logger.info("Starting Ollama deployment...")
+    logger.info(f"Model path: {model_path}")
+    logger.info(f"Model name: {model_name}")
+    
+    # 1. Check dependencies
+    if not check_dependencies():
+        logger.error("Dependency check failed")
+        return False
+    
+    # 2. Create output directory
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # 3. Convert to GGUF
+    gguf_path = os.path.join(output_dir, f"{model_name}.gguf")
+    try:
+        gguf_path = convert_to_gguf(
+            model_path=model_path,
+            output_path=gguf_path,
+            quantize=quantize,
+        )
+    except RuntimeError as e:
+        logger.error(f"GGUF conversion failed: {e}")
+        return False
+    
+    # 4. Create Modelfile
+    modelfile_path = os.path.join(output_dir, "Modelfile")
+    create_modelfile(
+        model_path=gguf_path,
+        modelfile_path=modelfile_path,
+        model_name=model_name,
+        context_size=context_size,
+        thinking_dial=thinking_dial,
+    )
+    
+    # 5. Create Ollama model
+    if not create_ollama_model(
+        modelfile_path=modelfile_path,
+        model_name=model_name,
+    ):
+        logger.error("Ollama model creation failed")
+        return False
+    
+    # 6. Generate usage example
+    example_path = os.path.join(output_dir, "USAGE.md")
+    generate_usage_example(model_name, example_path)
+    
+    logger.info("Deployment complete!")
+    logger.info(f"Run: `ollama run {model_name}`")
+    logger.info(f"Examples: see {example_path}")
+    
+    return True
+
+
+def generate_usage_example(model_name: str, output_path: str):
+    """
+    Generate usage example document
+    """
+    content = f"""# Fusion Model Usage Examples
+
+## 1. Basic Usage
+
+```bash
+# Start model
+ollama run {model_name}
+
+# Input question in interactive interface
+> Explain quantum entanglement
+```
+
+## 2. Thinking Dial Control
+
+Fusion supports dynamic reasoning intensity control. Add control token before question:
+
+```bash
+# depth=0: direct answer (casual chat, translation)
+> <|think| depth=0|> How's the weather today?
+
+# depth=1: simple reasoning
+> <|think| depth=1|> Calculate 123 * 456
+
+# depth=2: medium reasoning
+> <|think| depth=2|> Prove Pythagorean theorem
+
+# depth=3: deep reasoning (chain-of-thought)
+> <|think| depth=3|> Solve this algorithm problem: ...
+```
+
+## 3. REST API
+
+Ollama provides OpenAI-compatible API:
+
+```bash
+# Start Ollama service
+ollama serve
+
+# Call API
+curl <a href="http://localhost:11434/api/generate">http://localhost:11434/api/generate</a> -d {{{{
+  "model": "{model_name}",
+  "prompt": "Explain machine learning",
+  "stream": false
+}}}}
+```
+
+## 4. Python Call
+
+```python
+import ollama
+
+# Basic call
+response = ollama.generate(
+    model="{model_name}",
+    prompt="Explain quantum entanglement",
+)
+
+print(response["response"])
+
+# With Thinking Dial
+response = ollama.generate(
+    model="{model_name}",
+    prompt="<|think| depth=2|> Prove Pythagorean theorem",
+)
+
+print(response["response"])
+```
+
+## 5. Parameter Tuning
+
+Adjust generation parameters in Ollama:
+
+```bash
+# Temperature (creativity)
+ollama run {model_name} --temperature 0.9
+
+# Context window
+ollama run {model_name} --num_ctx 16384
+
+# Top-p sampling
+ollama run {model_name} --top_p 0.95
+```
+
+---
+
+**Tip**: See Ollama docs for more: https://ollama.com/docs
+"""
+    
+    with open(output_path, 'w', encoding='utf-8') as f:
+        f.write(content)
+    
+    logger.info(f"Usage example generated: {output_path}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Fusion Ollama One-Click Deployment (v2)")
+    
+    parser.add_argument("--model_path", type=str, required=True,
+                        help="HuggingFace model path")
+    parser.add_argument("--model_name", type=str, required=True,
+                        help="Ollama model name (e.g., fusion-8b)")
+    parser.add_argument("--output_dir", type=str, default="./ollama_output",
+                        help="Output directory")
+    parser.add_argument("--quantize", type=str, default="q4_k_m",
+                        choices=["q4_k_m", "q5_k_m", "q8_0", "f16"],
+                        help="Quantization level")
+    parser.add_argument("--context_size", type=int, default=32768,
+                        help="Context window size")
+    parser.add_argument("--no_thinking_dial", action="store_false",
+                        dest="thinking_dial",
+                        help="Disable Thinking Dial")
+    
+    args = parser.parse_args()
+    
+    # Execute deployment
+    success = deploy(
+        model_path=args.model_path,
+        model_name=args.model_name,
+        output_dir=args.output_dir,
+        quantize=args.quantize,
+        context_size=args.context_size,
+        thinking_dial=args.thinking_dial,
+    )
+    
+    if success:
+        logger.info("Deployment successful!")
+    else:
+        logger.error("Deployment failed")
+
+
+if __name__ == "__main__":
+    main()

models/fusion_model.py

@@ -118,9 +118,14 @@ class RMSNorm(nn.Module):
 
 class FusionAttention(nn.Module):
     """
-    Fusion 注意力层 - 集成 SBLA
+    Fusion Attention Layer with integrated SBLA.
     
-    在 FusionMiniLayer 中使用，实际集成到 fusion_model.py
+    NOTE (M4): This is a standalone reimplementation. The canonical SBLA logic
+    lives in sbla_attention.py (SBLAttention class). Future work should unify
+    by having FusionAttention delegate to SBLAttention instead of duplicating
+    mask building and block latent computation logic.
+    
+    See: models/sbla_attention.py::SBLAttention
     """
     
     def __init__(self, config: FusionConfig):
@@ -220,7 +225,9 @@ class FusionAttention(nn.Module):
         self,
         hidden_states: torch.Tensor,
         attention_mask: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
+        past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        use_cache: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
         batch_size, seq_len, _ = hidden_states.shape
         device = hidden_states.device
         
@@ -229,6 +236,14 @@ class FusionAttention(nn.Module):
         K = self.k_proj(hidden_states).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
         V = self.v_proj(hidden_states).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
         
+        # KV Cache 逻辑
+        if past_key_value is not None:
+            past_k, past_v = past_key_value
+            K = torch.cat([past_k, K], dim=2)
+            V = torch.cat([past_v, V], dim=2)
+        
+        present_key_value = (K, V) if use_cache else None
+        
         # 构建注意力掩码
         causal_mask = self._build_causal_mask(seq_len, device)
         window_mask = self._build_window_mask(seq_len, self.block_size, device)
@@ -274,7 +289,7 @@ class FusionAttention(nn.Module):
         output = output_std + gate_value * latent_expanded
         
         output = self.LayerNorm(output)
-        return self.dropout(output)
+        return self.dropout(output), present_key_value
 
 
 class FusionLayer(nn.Module):
@@ -298,11 +313,18 @@ class FusionLayer(nn.Module):
         self,
         hidden_states: torch.Tensor,
         attention_mask: Optional[torch.Tensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        use_cache: bool = False,
         **kwargs,
-    ) -> Tuple[torch.Tensor, None]:
+    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
         residual = hidden_states
         hidden_states = self.input_layernorm(hidden_states)
-        attn_output = self.attention(hidden_states, attention_mask)
+        attn_output, present_key_value = self.attention(
+            hidden_states, 
+            attention_mask,
+            past_key_value=past_key_value if past_key_value is not None else None,
+            use_cache=use_cache,
+        )
         hidden_states = residual + self.dropout(attn_output)
         
         residual = hidden_states
@@ -312,7 +334,7 @@ class FusionLayer(nn.Module):
         ffn_output = self.down_proj(gate * up)
         hidden_states = residual + self.dropout(ffn_output)
         
-        return hidden_states, None
+        return hidden_states, present_key_value
 
 
 class FusionModel(PreTrainedModel, GenerationMixin):
@@ -381,9 +403,23 @@ class FusionModel(PreTrainedModel, GenerationMixin):
             float_mask = attention_mask.to(dtype=hidden_states.dtype)
             attention_mask = (1.0 - float_mask) * torch.finfo(hidden_states.dtype).min
         
-        # Transformer 层
-        for layer in self.layers:
-            hidden_states, _ = layer(hidden_states, attention_mask=attention_mask)
+        # Transformer 层（支持 KV Cache）
+        past_key_values = kwargs.get("past_key_values", None)
+        use_cache = kwargs.get("use_cache", False) or (past_key_values is not None)
+        
+        present_key_values = () if use_cache else None
+        
+        for i, layer in enumerate(self.layers):
+            layer_past = past_key_values[i] if past_key_values is not None else None
+            layer_outputs, cache = layer(
+                hidden_states,
+                attention_mask=attention_mask,
+                past_key_value=layer_past,
+                use_cache=use_cache,
+            )
+            hidden_states = layer_outputs
+            if use_cache:
+                present_key_values = present_key_values + (cache,)
         
         # Final norm
         hidden_states = self.norm(hidden_states)
@@ -399,6 +435,9 @@ class FusionModel(PreTrainedModel, GenerationMixin):
             loss_fct = nn.CrossEntropyLoss()
             loss = loss_fct(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1))
         
+        if use_cache:
+            return {"loss": loss, "logits": logits, "past_key_values": present_key_values}
+        
         if not return_dict:
             return (loss, logits) if loss is not None else (logits,)
         
@@ -438,8 +477,7 @@ class FusionModel(PreTrainedModel, GenerationMixin):
             )
             
             logits = outputs["logits"]
-            if "past_key_values" in outputs:
-                past_key_values = outputs["past_key_values"]
+            past_key_values = outputs.get("past_key_values", None)
             
             next_token_logits = logits[:, -1, :] / max(temperature, 1e-8)
             

models/tokenizer.py

@@ -0,0 +1,146 @@
+"""
+Fusion Tokenizer - Unified tokenizer management
+
+Handles the gap between GPT2 (vocab_size=50257) and Fusion's target vocab (100K).
+Current status: Uses GPT2 as placeholder until SentencePiece model is trained.
+
+Usage:
+    from models.tokenizer import get_tokenizer
+    tokenizer = get_tokenizer("gpt2")  # placeholder
+    tokenizer = get_tokenizer("fusion", vocab_size=100000)  # future: SentencePiece
+
+Author: Zhu Zizhan
+Project: Fusion-LLM
+License: Apache 2.0
+"""
+
+import json
+import os
+from pathlib import Path
+from typing import Optional
+
+try:
+    from transformers import AutoTokenizer, PreTrainedTokenizer
+except ImportError:
+    AutoTokenizer = None
+    PreTrainedTokenizer = None
+
+
+# Fusion special tokens
+FUSION_SPECIAL_TOKENS = {
+    "pad_token": "<|pad|>",
+    "bos_token": "<|start|>",
+    "eos_token": "<|end|>",
+    "think_tokens": ["<|think_depth_0|>", "<|think_depth_1|>", "<|think_depth_2|>", "<|think_depth_3|>"],
+}
+
+
+def get_tokenizer(
+    tokenizer_type: str = "gpt2",
+    vocab_size: int = 50257,
+    tokenizer_dir: Optional[str] = None,
+) -> "PreTrainedTokenizer":
+    """
+    Get a tokenizer for Fusion models.
+
+    Args:
+        tokenizer_type: "gpt2" (placeholder) or "fusion" (SentencePiece, if available)
+        vocab_size: Target vocabulary size
+        tokenizer_dir: Directory containing tokenizer files (for "fusion" type)
+
+    Returns:
+        A HuggingFace PreTrainedTokenizer instance
+
+    Notes:
+        - "gpt2" mode: Uses GPT2 BPE tokenizer (vocab_size=50257). This is a
+          PLACEHOLDER. The model config should set vocab_size=50257 when using this.
+        - "fusion" mode: Loads a SentencePiece tokenizer from tokenizer_dir.
+          Requires tokenizer.model file to exist. Falls back to GPT2 if not found.
+    """
+    if AutoTokenizer is None:
+        raise ImportError("transformers is required: pip install transformers")
+
+    if tokenizer_type == "fusion":
+        sp_model_path = None
+        if tokenizer_dir:
+            sp_model_path = Path(tokenizer_dir) / "tokenizer.model"
+        else:
+            # Try project root
+            for candidate in ["tokenizers", ".", "data"]:
+                p = Path(candidate) / "tokenizer.model"
+                if p.exists():
+                    sp_model_path = p
+                    break
+
+        if sp_model_path and sp_model_path.exists():
+            tokenizer = AutoTokenizer.from_pretrained(
+                str(sp_model_path.parent),
+                tokenizer_type="SentencePiece",
+            )
+            tokenizer = _add_fusion_special_tokens(tokenizer)
+            return tokenizer
+        else:
+            import warnings
+            warnings.warn(
+                "Fusion SentencePiece tokenizer not found. "
+                "Falling back to GPT2 tokenizer. "
+                "Set vocab_size=50257 in model config to match.",
+                UserWarning,
+            )
+            tokenizer_type = "gpt2"
+            vocab_size = 50257
+
+    if tokenizer_type == "gpt2":
+        tokenizer = AutoTokenizer.from_pretrained("gpt2")
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer = _add_fusion_special_tokens(tokenizer)
+
+        # Verify vocab size consistency
+        actual_vocab = len(tokenizer)
+        if actual_vocab != vocab_size:
+            import warnings
+            warnings.warn(
+                f"GPT2 tokenizer vocab_size={actual_vocab}, but config specifies {vocab_size}. "
+                f"Using actual tokenizer size ({actual_vocab}). "
+                f"Update model config vocab_size to match.",
+                UserWarning,
+            )
+        return tokenizer
+
+    raise ValueError(f"Unknown tokenizer_type: {tokenizer_type}")
+
+
+def _add_fusion_special_tokens(tokenizer: "PreTrainedTokenizer") -> "PreTrainedTokenizer":
+    """Add Fusion-specific special tokens to any tokenizer."""
+    special_tokens_dict = {
+        "pad_token": FUSION_SPECIAL_TOKENS["pad_token"],
+        "additional_special_tokens": FUSION_SPECIAL_TOKENS["think_tokens"],
+    }
+    tokenizer.add_special_tokens(special_tokens_dict)
+    return tokenizer
+
+
+def get_effective_vocab_size(tokenizer_type: str = "gpt2", requested_vocab: int = 100000) -> int:
+    """
+    Return the effective vocab size that should be used in model config.
+    This ensures model embedding size matches the actual tokenizer.
+    """
+    if tokenizer_type == "gpt2":
+        return 50257 + len(FUSION_SPECIAL_TOKENS["think_tokens"]) + 1  # ~50262
+    if tokenizer_type == "fusion":
+        return requested_vocab
+    return requested_vocab
+
+
+if __name__ == "__main__":
+    print("[Fusion Tokenizer] Testing tokenizer creation...")
+    tok = get_tokenizer("gpt2")
+    print(f"  Type: GPT2 (placeholder)")
+    print(f"  Vocab size: {len(tok)}")
+    print(f"  Pad token: {tok.pad_token}")
+    print(f"  Think tokens: {FUSION_SPECIAL_TOKENS['think_tokens']}")
+    print(f"  Effective vocab: {get_effective_vocab_size('gpt2')}")
+    test_text = "Hello, Fusion! <|think_depth_2|>"
+    encoded = tok.encode(test_text)
+    decoded = tok.decode(encoded)
+    print(f"  Encode/decode test: '{test_text}' -> {encoded} -> '{decoded}'")

push_to_github.py

@@ -1,228 +0,0 @@
-"""
-Fusion 项目 GitHub 推送脚本
-
-使用方法（在你本地电脑运行）：
-1. 安装依赖：pip install requests
-2. 运行脚本：python push_to_github.py
-3. 按提示输入 GitHub token（不会显示在聊天中）
-
-注意：Token 需要 repo 权限（全选）
-创建 Token：https://github.com/settings/tokens
-
-作者：朱子瞻
-项目：Fusion - 六边形开源大模型
-"""
-
-import os
-import json
-import subprocess
-import getpass
-import requests
-from pathlib import Path
-
-
-def create_github_repo(token: str, repo_name: str = "fusion-llm", private: bool = False):
-    """
-    使用 GitHub API 创建仓库
-    
-    参数：
-        token: GitHub Personal Access Token
-        repo_name: 仓库名称
-        private: 是否私有（默认 False，公开）
-    
-    返回：
-        仓库 URL
-    """
-    print(f"\n📦 创建 GitHub 仓库：{repo_name}...")
-    
-    url = "https://api.github.com/user/repos"
-    
-    headers = {
-        "Authorization": f"token {token}",
-        "Accept": "application/vnd.github.v3+json",
-    }
-    
-    data = {
-        "name": repo_name,
-        "description": "Fusion - 六边形开源大模型 | 集百家之长，铸最强开源模型",
-        "private": private,
-        "has_issues": True,
-        "has_projects": True,
-        "has_wiki": True,
-        "auto_init": False,  # 不要自动创建 README
-    }
-    
-    response = requests.post(url, headers=headers, json=data)
-    
-    if response.status_code == 201:
-        repo_data = response.json()
-        repo_url = repo_data["html_url"]
-        clone_url = repo_data["clone_url"]
-        
-        print(f"✅ 仓库创建成功！")
-        print(f"   URL: {repo_url}")
-        print(f"   克隆 URL (HTTPS): {clone_url}")
-        
-        return clone_url
-        
-    elif response.status_code == 422:
-        # 仓库已存在
-        print(f"⚠️  仓库 {repo_name} 已存在")
-        return f"https://github.com/zhan1206/{repo_name}.git"
-        
-    else:
-        print(f"❌ 创建失败：{response.status_code}")
-        print(f"   错误信息：{response.text}")
-        return None
-
-
-def push_to_github(repo_url: str, project_dir: str, use_ssh: bool = False):
-    """
-    推送代码到 GitHub
-    
-    参数：
-        repo_url: 仓库 URL
-        project_dir: 项目目录
-        use_ssh: 是否使用 SSH（默认 False，使用 HTTPS）
-    """
-    print(f"\n🚀 推送代码到 GitHub...")
-    
-    # 切换到项目目录
-    os.chdir(project_dir)
-    
-    # 如果已经设置了 remote，先删除
-    subprocess.run(
-        ["git", "remote", "remove", "origin"],
-        capture_output=True,
-    )
-    
-    # 添加 remote
-    if use_ssh:
-        # SSH 格式
-        ssh_url = repo_url.replace(
-            "https://github.com/",
-            "git@github.com:",
-        )
-        remote_url = ssh_url
-    else:
-        # HTTPS 格式
-        remote_url = repo_url
-    
-    print(f"   Remote URL: {remote_url}")
-    
-    result = subprocess.run(
-        ["git", "remote", "add", "origin", remote_url],
-        capture_output=True,
-        text=True,
-    )
-    
-    if result.returncode != 0:
-        print(f"❌ 添加 remote 失败：{result.stderr}")
-        return False
-    
-    # 推送代码
-    print(f"   推送分支：master")
-    
-    result = subprocess.run(
-        ["git", "push", "-u", "origin", "master"],
-        capture_output=True,
-        text=True,
-    )
-    
-    if result.returncode == 0:
-        print(f"✅ 推送成功！")
-        print(f"\n🎉 项目已发布：{repo_url.replace('.git', '')}")
-        return True
-    else:
-        print(f"❌ 推送失败：{result.stderr}")
-        
-        # 如果是 HTTPS 且失败，提示使用 SSH
-        if not use_ssh:
-            print(f"\n💡 提示：如果 HTTPS 推送失败，可以尝试使用 SSH：")
-            print(f"   1. 生成 SSH key：ssh-keygen -t ed25519 -C \"your_email@example.com\"")
-            print(f"   2. 添加 SSH key 到 GitHub：https://github.com/settings/keys")
-            print(f"   3. 重新运行脚本，输入 'y' 使用 SSH")
-        
-        return False
-
-
-def main():
-    print("=" * 60)
-    print("Fusion 项目 GitHub 推送脚本")
-    print("=" * 60)
-    
-    # 1. 获取项目目录
-    project_dir = os.path.dirname(os.path.abspath(__file__))
-    print(f"\n📂 项目目录：{project_dir}")
-    
-    # 2. 检查 Git 状态
-    os.chdir(project_dir)
-    
-    result = subprocess.run(
-        ["git", "status", "--porcelain"],
-        capture_output=True,
-        text=True,
-    )
-    
-    if result.stdout:
-        print(f"\n⚠️  有未提交的更改：")
-        print(result.stdout)
-        
-        commit = input("\n是否提交这些更改？(y/N): ").strip().lower()
-        
-        if commit == 'y':
-            subprocess.run(["git", "add", "."])
-            commit_msg = input("输入提交信息（默认：Update）：") or "Update"
-            subprocess.run(["git", "commit", "-m", commit_msg])
-            print(f"✅ 已提交")
-        else:
-            print(f"⚠️  取消推送")
-            return
-    
-    # 3. 获取 GitHub Token（安全输入，不显示）
-    print(f"\n🔐 输入 GitHub Personal Access Token")
-    print(f"   创建 Token：https://github.com/settings/tokens")
-    print(f"   需要权限：repo（全选）")
-    print(f"   （输入时不会显示，这是正常的）")
-    
-    token = getpass.getpass("Token: ")
-    
-    if not token:
-        print(f"❌ Token 不能为空")
-        return
-    
-    # 4. 创建 GitHub 仓库
-    repo_name = "fusion-llm"
-    repo_url = create_github_repo(token, repo_name, private=False)
-    
-    if not repo_url:
-        print(f"❌ 创建仓库失败")
-        return
-    
-    # 5. 推送代码
-    use_ssh = input("\n使用 SSH 推送？(y/N): ").strip().lower() == 'y'
-    
-    success = push_to_github(repo_url, project_dir, use_ssh=use_ssh)
-    
-    if not success and not use_ssh:
-        # 如果 HTTPS 失败，询问是否尝试 SSH
-        retry_ssh = input("\n是否尝试使用 SSH 推送？(y/N): ").strip().lower()
-        
-        if retry_ssh == 'y':
-            # 修改 URL 为 SSH 格式
-            ssh_url = repo_url.replace(
-                "https://github.com/",
-                "git@github.com:",
-            )
-            success = push_to_github(ssh_url, project_dir, use_ssh=True)
-    
-    if success:
-        print(f"\n🎉 完成！项目已成功推送到 GitHub")
-        print(f"   仓库地址：<ADDRESS_REMOVED>
-        print(f"   克隆命令：git clone {repo_url.replace('.git', '')}.git")
-    else:
-        print(f"\n❌ 推送失败，请检查错误信息")
-
-
-if __name__ == "__main__":
-    main()

requirements.txt

@@ -8,6 +8,7 @@ datasets>=2.16.0
 
 # PEFT (LoRA/QLoRA)
 peft>=0.7.0
+bitsandbytes>=0.41.0
 
 # 分布式训练
 deepspeed>=0.12.0
@@ -25,8 +26,12 @@ tokenizers>=0.15.0
 langid>=1.1.6
 
 # 推理部署
-# 注意：Ollama 需要从 https://ollama.com/ 下载二进制客户端
-# 或使用 ollama-python SDK: pip install ollama
+ollama>=0.1.0
+# 注意：Ollama 也需要从 https://ollama.com/ 下载二进制客户端
+
+# 量化工具依赖
+scipy>=1.11.0
+scikit-learn>=1.3.0
 
 # 评估（可选）
 evaluate>=0.4.0

scripts/add_depth3_samples.py

@@ -0,0 +1,42 @@
+#!/usr/bin/env python3
+"""Add depth=3 samples to mini_data.json for Thinking Dial coverage."""
+
+import json
+
+DEPTH_3_SAMPLES = [
+    {"prompt": "Prove the Pythagorean theorem", "response": "Consider a right triangle with legs a, b and hypotenuse c. Construct a square of side (a+b) and arrange four copies of the triangle inside. The inner quadrilateral is a square of side c, so (a+b)^2 = c^2 + 4*(ab/2), which simplifies to a^2 + b^2 = c^2.", "think_rank": 3},
+    {"prompt": "Derive the quadratic formula", "response": "Starting from ax^2 + bx + c = 0, divide by a, complete the square: (x + b/2a)^2 = (b^2 - 4ac)/4a^2. Taking the square root gives x = (-b +- sqrt(b^2 - 4ac)) / 2a.", "think_rank": 3},
+    {"prompt": "Prove that sqrt(2) is irrational", "response": "Assume sqrt(2) = p/q in lowest terms. Then 2q^2 = p^2, so p is even. Write p = 2k, then 2q^2 = 4k^2, so q^2 = 2k^2, meaning q is also even. This contradicts p/q being in lowest terms.", "think_rank": 3},
+    {"prompt": "Prove there are infinitely many primes", "response": "Suppose finitely many primes p1,...,pn exist. Let N = p1*p2*...*pn + 1. N is not divisible by any pi (remainder 1). So N is either prime itself or has a prime factor not in the list. Contradiction.", "think_rank": 3},
+    {"prompt": "Derive the derivative of sin(x)", "response": "Using the limit definition: d/dx sin(x) = lim(h->0) [sin(x+h) - sin(x)] / h. Apply the angle addition formula: sin(x+h) = sin(x)cos(h) + cos(x)sin(h). This gives lim [sin(x)(cos(h)-1)/h + cos(x)sin(h)/h] = sin(x)*0 + cos(x)*1 = cos(x).", "think_rank": 3},
+    {"prompt": "Analyze the time complexity of merge sort", "response": "Merge sort divides the array in half (O(1)), recurses on both halves (2T(n/2)), and merges in O(n). The recurrence T(n) = 2T(n/2) + O(n) solves to O(n log n) by the master theorem (case 2). This holds for worst, average, and best cases.", "think_rank": 3},
+    {"prompt": "Prove the sum of first n natural numbers is n(n+1)/2", "response": "By induction: Base case n=1: 1 = 1(2)/2 = 1. Inductive step: assume S(k) = k(k+1)/2. Then S(k+1) = S(k) + (k+1) = k(k+1)/2 + (k+1) = (k+1)(k/2 + 1) = (k+1)(k+2)/2. QED.", "think_rank": 3},
+    {"prompt": "Prove that e^x converges for all x", "response": "The Taylor series e^x = sum(x^n/n!) has ratio test: |a_(n+1)/a_n| = |x|/(n+1) -> 0 as n -> infinity. Since the limit is 0 < 1 for all x, the series converges absolutely for all real x by the ratio test.", "think_rank": 3},
+]
+
+def main():
+    with open('data/mini_data.json', 'r', encoding='utf-8') as f:
+        data = json.load(f)
+
+    old_dist = {}
+    for item in data:
+        r = item.get('think_rank', 0)
+        old_dist[r] = old_dist.get(r, 0) + 1
+    print(f"Before: {old_dist}")
+
+    data.extend(DEPTH_3_SAMPLES)
+
+    new_dist = {}
+    for item in data:
+        r = item.get('think_rank', 0)
+        new_dist[r] = new_dist.get(r, 0) + 1
+    print(f"After: {new_dist}")
+
+    with open('data/mini_data.json', 'w', encoding='utf-8') as f:
+        json.dump(data, f, ensure_ascii=False, indent=2)
+
+    print(f"Total: {len(data)} items")
+
+
+if __name__ == '__main__':
+    main()

{tests → scripts}/create_mini_data.py

@@ -67,10 +67,20 @@ def create_mini_dataset(output_path: str, num_samples: int = 100):
         else:
             prompt, response = random.choice(english_samples)
         
+        # Assign think_rank based on content depth
+        if any(kw in prompt for kw in ["Prove", "Derive", "Analyze", "\u8bc1\u660e", "\u63a8\u5bfc", "\u5206\u6790"]):
+            think_rank = 3
+        elif any(kw in prompt for kw in ["Explain", "How", "Why", "\u89e3\u91ca", "\u5982\u4f55", "\u4e3a\u4ec0\u4e48"]):
+            think_rank = 2
+        elif any(kw in prompt for kw in ["Write", "Implement", "\u5199", "\u5b9e\u73b0"]):
+            think_rank = 1
+        else:
+            think_rank = 0
+
         data.append({
             "prompt": prompt,
             "response": response,
-            "think_rank": 0,  # mini 模型不使用 thinking dial
+            "think_rank": think_rank,
         })
     
     # 保存为 JSON

scripts/fix_mini_data.py

@@ -0,0 +1,61 @@
+#!/usr/bin/env python3
+"""Fix mini_data.json: distribute think_rank 0-3 based on prompt content."""
+
+import json
+import re
+
+# Keywords suggesting different thinking depths
+DEPTH_3_KEYWORDS = ['prove', 'theorem', 'proof', 'derive', 'mathematical', 'complex',
+                    'prove', 'derive', 'calculate', 'analyze deeply',
+                    '\u8bc1\u660e', '\u63a8\u5bfc', '\u5b9a\u7406', '\u590d\u6742', '\u6df1\u5165\u5206\u6790']
+DEPTH_2_KEYWORDS = ['explain', 'why', 'how does', 'compare', 'difference',
+                    'algorithm', 'design', 'optimize',
+                    '\u89e3\u91ca', '\u4e3a\u4ec0\u4e48', '\u5982\u4f55', '\u6bd4\u8f83', '\u7b97\u6cd5', '\u8bbe\u8ba1', '\u4f18\u5316']
+DEPTH_1_KEYWORDS = ['write', 'implement', 'code', 'function', 'create',
+                    '\u5199', '\u5b9e\u73b0', '\u7f16\u5199', '\u4ee3\u7801', '\u521b\u5efa']
+
+
+def assign_depth(item):
+    text = (item.get('prompt', '') + ' ' + item.get('response', '')).lower()
+    for kw in DEPTH_3_KEYWORDS:
+        if kw.lower() in text:
+            return 3
+    for kw in DEPTH_2_KEYWORDS:
+        if kw.lower() in text:
+            return 2
+    for kw in DEPTH_1_KEYWORDS:
+        if kw.lower() in text:
+            return 1
+    return 0
+
+
+def main():
+    with open('data/mini_data.json', 'r', encoding='utf-8') as f:
+        data = json.load(f)
+
+    # Count current distribution
+    old_dist = {}
+    for item in data:
+        r = item.get('think_rank', 0)
+        old_dist[r] = old_dist.get(r, 0) + 1
+    print(f"Before fix: {old_dist}")
+
+    # Fix
+    for item in data:
+        item['think_rank'] = assign_depth(item)
+
+    # Count new distribution
+    new_dist = {}
+    for item in data:
+        r = item.get('think_rank', 0)
+        new_dist[r] = new_dist.get(r, 0) + 1
+    print(f"After fix: {new_dist}")
+
+    with open('data/mini_data.json', 'w', encoding='utf-8') as f:
+        json.dump(data, f, ensure_ascii=False, indent=2)
+
+    print(f"Fixed {len(data)} items")
+
+
+if __name__ == '__main__':
+    main()

scripts/train_tokenizer.py

@@ -0,0 +1,155 @@
+#!/usr/bin/env python3
+"""
+Train a SentencePiece tokenizer for Fusion models.
+
+This creates the tokenizer.model file needed by Fusion's 100K vocabulary.
+Training data should be a plain text file with one sentence per line.
+
+Usage:
+    python scripts/train_tokenizer.py --input data/tokenizer_train.txt --vocab_size 100000 --output tokenizers/
+
+Requirements:
+    pip install sentencepiece
+
+Author: Zhu Zizhan
+Project: Fusion-LLM
+License: Apache 2.0
+"""
+
+import argparse
+import os
+import sys
+
+
+def train_tokenizer(input_path: str, vocab_size: int, output_dir: str, model_type: str = "unigram"):
+    """Train a SentencePiece tokenizer."""
+    try:
+        import sentencepiece as spm
+    except ImportError:
+        print("[ERROR] sentencepiece not installed. Run: pip install sentencepiece")
+        sys.exit(1)
+
+    if not os.path.exists(input_path):
+        print(f"[ERROR] Training data not found: {input_path}")
+        print("Create a plain text file with one sentence per line.")
+        print("For bilingual (zh+en) tokenizer, mix Chinese and English text.")
+        sys.exit(1)
+
+    os.makedirs(output_dir, exist_ok=True)
+    model_prefix = os.path.join(output_dir, "tokenizer")
+
+    print(f"[Tokenizer] Training SentencePiece model...")
+    print(f"  Input: {input_path}")
+    print(f"  Vocab size: {vocab_size}")
+    print(f"  Model type: {model_type}")
+    print(f"  Output: {output_dir}/")
+
+    # Special tokens for Fusion
+    control_symbols = [
+        "<|pad|>", "<|start|>", "<|end|>",
+        "<|think_depth_0|>", "<|think_depth_1|>",
+        "<|think_depth_2|>", "<|think_depth_3|>",
+    ]
+
+    spm.SentencePieceTrainer.train(
+        input=input_path,
+        model_prefix=model_prefix,
+        vocab_size=vocab_size,
+        model_type=model_type,
+        character_coverage=0.9995,  # High coverage for CJK
+        input_sentence_size=10000000,
+        shuffle_input_sentence=True,
+        control_symbols=control_symbols,
+        unk_id=0,
+        bos_id=1,
+        eos_id=2,
+        pad_id=3,
+        byte_fallback=True,  # Important for multilingual
+        split_by_unicode_script=True,
+        allow_whitespace_only_pieces=True,
+    )
+
+    model_path = os.path.join(output_dir, "tokenizer.model")
+    vocab_path = os.path.join(output_dir, "tokenizer.vocab")
+
+    print(f"\n[Done] Tokenizer trained successfully!")
+    print(f"  Model: {model_path}")
+    print(f"  Vocab: {vocab_path}")
+
+    # Verify
+    sp = spm.SentencePieceProcessor()
+    sp.load(model_path)
+    print(f"  Actual vocab size: {sp.get_piece_size()}")
+
+    # Test
+    test_zh = "Fusion是一个开源大语言模型"
+    test_en = "Fusion is an open-source language model"
+    print(f"\n  Test encode (zh): {test_zh}")
+    print(f"    -> {sp.encode(test_zh)}")
+    print(f"  Test encode (en): {test_en}")
+    print(f"    -> {sp.encode(test_en)}")
+
+
+def create_sample_training_data(output_path: str, num_lines: int = 100000):
+    """Create sample training data for tokenizer (for testing only)."""
+    import random
+
+    print(f"[Sample Data] Creating sample training data: {output_path}")
+    samples_zh = [
+        "人工智能是计算机科学的一个重要分支",
+        "深度学习使用多层神经网络来模拟人脑",
+        "自然语言处理帮助计算机理解人类语言",
+        "机器学习使计算机能够从数据中学习",
+        "大语言模型在文本生成任务中表现出色",
+        "Transformer架构彻底改变了自然语言处理领域",
+        "注意力机制是现代深度学习的核心组件",
+        "预训练语言模型通过大规模语料学习语言知识",
+        "微调技术使模型适应特定下游任务",
+        "开源模型推动了人工智能技术的普及",
+    ]
+    samples_en = [
+        "Artificial intelligence is a branch of computer science",
+        "Deep learning uses multi-layer neural networks",
+        "Natural language processing helps computers understand language",
+        "Machine learning enables computers to learn from data",
+        "Large language models excel at text generation tasks",
+        "The Transformer architecture revolutionized NLP",
+        "Attention mechanisms are core components of modern deep learning",
+        "Pre-trained models learn language knowledge from large corpora",
+        "Fine-tuning adapts models to specific downstream tasks",
+        "Open-source models promote the democratization of AI",
+    ]
+
+    with open(output_path, 'w', encoding='utf-8') as f:
+        for _ in range(num_lines):
+            if random.random() > 0.5:
+                f.write(random.choice(samples_zh) + "\n")
+            else:
+                f.write(random.choice(samples_en) + "\n")
+
+    print(f"  Generated {num_lines} lines")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Train Fusion SentencePiece tokenizer")
+    parser.add_argument("--input", type=str, default="data/tokenizer_train.txt",
+                        help="Path to training data (one sentence per line)")
+    parser.add_argument("--vocab_size", type=int, default=100000,
+                        help="Vocabulary size (default: 100000)")
+    parser.add_argument("--output", type=str, default="tokenizers/",
+                        help="Output directory")
+    parser.add_argument("--model_type", type=str, default="unigram",
+                        choices=["unigram", "bpe"], help="Model type")
+    parser.add_argument("--create_sample_data", action="store_true",
+                        help="Create sample training data for testing")
+    args = parser.parse_args()
+
+    if args.create_sample_data:
+        create_sample_training_data(args.input)
+        return
+
+    train_tokenizer(args.input, args.vocab_size, args.output, args.model_type)
+
+
+if __name__ == "__main__":
+    main()

test_train_loop.py

@@ -1,38 +0,0 @@
-"""Small training loop to verify loss decreases"""
-import sys
-sys.path.insert(0, ".")
-import torch
-from models.fusion_model import FusionModel, FusionConfig
-
-config = FusionConfig(
-    vocab_size=10000,
-    hidden_size=256,
-    num_hidden_layers=2,
-    num_attention_heads=4,
-    intermediate_size=512,
-    block_size=64,
-    latent_dim=16,
-    sbla_mode="pure_sbla",
-    max_position_embeddings=256,
-)
-
-model = FusionModel(config)
-model.train()
-
-optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
-batch_size, seq_len = 4, 32
-
-print("[DEBUG] Small training loop (5 steps)...")
-for step in range(5):
-    input_ids = torch.randint(0, 10000, (batch_size, seq_len))
-    attention_mask = torch.ones(batch_size, seq_len)
-    
-    optimizer.zero_grad()
-    outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)
-    loss = outputs["loss"]
-    loss.backward()
-    optimizer.step()
-    
-    print(f"Step {step}: loss = {loss.item():.4f}")
-
-print("\nTraining loop successful - loss is decreasing!")

test_training.py

@@ -1,38 +0,0 @@
-"""Test training"""
-import sys
-sys.path.insert(0, ".")
-import torch
-from models.fusion_model import FusionModel, FusionConfig
-
-config = FusionConfig(
-    vocab_size=10000,
-    hidden_size=256,
-    num_hidden_layers=2,
-    num_attention_heads=4,
-    intermediate_size=512,
-    block_size=64,
-    latent_dim=16,
-    sbla_mode="pure_sbla",
-    max_position_embeddings=256,
-)
-
-model = FusionModel(config)
-model.train()
-
-batch_size, seq_len = 2, 32
-input_ids = torch.randint(0, 10000, (batch_size, seq_len))
-attention_mask = torch.ones(batch_size, seq_len)
-
-outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)
-print(f"Loss: {outputs['loss'].item():.4f}")
-
-loss = outputs["loss"]
-loss.backward()
-print(f"Gradients exist: {model.embeddings.weight.grad is not None}")
-
-optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
-optimizer.zero_grad()
-outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)
-outputs["loss"].backward()
-optimizer.step()
-print("Optimizer step successful!")

tests/debug_attn.py

@@ -1,74 +0,0 @@
-"""Debug attention step by step"""
-import sys
-sys.path.insert(0, ".")
-import torch
-import torch.nn.functional as F
-import math
-
-print("[DEBUG] Step-by-step attention debugging...")
-
-from models.fusion_model import FusionConfig
-
-config = FusionConfig(
-    vocab_size=10000,
-    hidden_size=256,
-    num_hidden_layers=2,
-    num_attention_heads=4,
-    intermediate_size=512,
-    block_size=64,
-    latent_dim=16,
-    sbla_mode="pure_sbla",
-    max_position_embeddings=256,
-)
-
-# Manual attention computation
-batch_size, seq_len = 2, 32
-hidden_states = torch.randn(batch_size, seq_len, 256)
-device = hidden_states.device
-
-# Q/K/V
-q_proj = torch.nn.Linear(256, 256, bias=False)
-k_proj = torch.nn.Linear(256, 256, bias=False)
-v_proj = torch.nn.Linear(256, 256, bias=False)
-
-Q = q_proj(hidden_states).view(batch_size, seq_len, 4, 64).transpose(1, 2)
-K = k_proj(hidden_states).view(batch_size, seq_len, 4, 64).transpose(1, 2)
-V = v_proj(hidden_states).view(batch_size, seq_len, 4, 64).transpose(1, 2)
-
-print(f"Q shape: {Q.shape}, has_nan: {torch.isnan(Q).any()}")
-print(f"K shape: {K.shape}, has_nan: {torch.isnan(K).any()}")
-print(f"V shape: {V.shape}, has_nan: {torch.isnan(V).any()}")
-
-# Compute attention scores
-attn_scores = torch.matmul(Q, K.transpose(-1, -2)) / math.sqrt(64)
-print(f"Attn scores: min={attn_scores.min():.4f}, max={attn_scores.max():.4f}, has_nan: {torch.isnan(attn_scores).any()}")
-
-# Check for -inf in scores
-print(f"Scores has -inf: {torch.isinf(attn_scores).any()}")
-print(f"Scores has inf: {torch.isinf(attn_scores).any()}")
-
-# Check causal mask
-causal_mask = torch.triu(torch.ones(seq_len, seq_len, device=device, dtype=torch.bool), diagonal=1)
-causal_mask = causal_mask.float().masked_fill(causal_mask, float('-inf'))
-print(f"Causal mask: min={causal_mask.min():.4f}, max={causal_mask.max():.4f}")
-
-# Add causal mask
-attn_scores_masked = attn_scores + causal_mask
-print(f"Scores after mask: min={attn_scores_masked.min():.4f}, has_nan: {torch.isnan(attn_scores_masked).any()}")
-
-# Softmax
-attn_probs = F.softmax(attn_scores_masked, dim=-1)
-print(f"Attn probs: min={attn_probs.min():.4f}, max={attn_probs.max():.4f}, has_nan: {torch.isnan(attn_probs).any()}")
-
-# Context
-context = torch.matmul(attn_probs, V)
-print(f"Context: min={context.min():.4f}, max={context.max():.4f}, has_nan: {torch.isnan(context).any()}")
-
-# Reshape
-context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, 256)
-print(f"Context reshaped: min={context.min():.4f}, max={context.max():.4f}, has_nan: {torch.isnan(context).any()}")
-
-# Output projection
-out_proj = torch.nn.Linear(256, 256, bias=False)
-output = out_proj(context)
-print(f"Output: min={output.min():.4f}, max={output.max():.4f}, has_nan: {torch.isnan(output).any()}")

tests/debug_layer.py

@@ -1,63 +0,0 @@
-"""Debug layer by layer"""
-import sys
-sys.path.insert(0, ".")
-import torch
-
-print("[DEBUG] Testing layer by layer...")
-
-from models.fusion_model import FusionModel, FusionConfig, FusionAttention
-
-config = FusionConfig(
-    vocab_size=10000,
-    hidden_size=256,
-    num_hidden_layers=2,
-    num_attention_heads=4,
-    intermediate_size=512,
-    block_size=64,
-    latent_dim=16,
-    sbla_mode="pure_sbla",
-    max_position_embeddings=256,
-)
-
-model = FusionModel(config)
-model.eval()
-
-# Get the attention layer
-attn = model.layers[0].attention
-
-batch_size, seq_len = 2, 32
-hidden_states = torch.randn(batch_size, seq_len, 256)
-
-# Test attention
-attn_out = attn.forward(hidden_states)
-print(f"Attention output: min={attn_out.min():.4f}, max={attn_out.max():.4f}, has_nan: {torch.isnan(attn_out).any()}")
-
-# Test RMSNorm
-norm = model.layers[0].input_layernorm
-norm_out = norm.forward(hidden_states)
-print(f"RMSNorm output: min={norm_out.min():.4f}, max={norm_out.max():.4f}, has_nan: {torch.isnan(norm_out).any()}")
-
-# Test FFN
-ffn = model.layers[0]
-residual = hidden_states
-norm1_out = ffn.input_layernorm(hidden_states)
-attn_out = ffn.attention(norm1_out)
-after_attn = residual + attn_out
-print(f"After attention residual: min={after_attn.min():.4f}, max={after_attn.max():.4f}, has_nan: {torch.isnan(after_attn).any()}")
-
-norm2_out = ffn.post_attention_layernorm(after_attn)
-print(f"Post-attention norm: min={norm2_out.min():.4f}, max={norm2_out.max():.4f}, has_nan: {torch.isnan(norm2_out).any()}")
-
-gate = torch.nn.functional.silu(ffn.gate_proj(norm2_out))
-up = ffn.up_proj(norm2_out)
-print(f"Gate: min={gate.min():.4f}, max={gate.max():.4f}, has_nan: {torch.isnan(gate).any()}")
-print(f"Up: min={up.min():.4f}, max={up.max():.4f}, has_nan: {torch.isnan(up).any()}")
-
-gate_up = gate * up
-print(f"Gate*Up: min={gate_up.min():.4f}, max={gate_up.max():.4f}, has_nan: {torch.isnan(gate_up).any()}")
-
-ffn_out = ffn.down_proj(gate_up)
-print(f"FFN output: min={ffn_out.min():.4f}, max={ffn_out.max():.4f}, has_nan: {torch.isnan(ffn_out).any()}")
-
-final = after_attn + ffn_out
-print(f"Final layer output: min={final.min():.4f}, max={final.max():.4f}, has_nan: {torch.isnan(final).any()}")

tests/debug_lm.py

@@ -1,54 +0,0 @@
-"""Debug norm and lm_head"""
-import sys
-sys.path.insert(0, ".")
-import torch
-
-print("[DEBUG] Testing norm and lm_head...")
-
-from models.fusion_model import FusionModel, FusionConfig
-
-config = FusionConfig(
-    vocab_size=10000,
-    hidden_size=256,
-    num_hidden_layers=2,
-    num_attention_heads=4,
-    intermediate_size=512,
-    block_size=64,
-    latent_dim=16,
-    sbla_mode="pure_sbla",
-    max_position_embeddings=256,
-)
-
-model = FusionModel(config)
-model.eval()
-
-# Simulate hidden states after all layers
-batch_size, seq_len = 2, 32
-hidden_states = torch.randn(batch_size, seq_len, 256)
-
-# Final norm
-norm_out = model.norm(hidden_states)
-print(f"Final norm output: min={norm_out.min():.4f}, max={norm_out.max():.4f}, has_nan: {torch.isnan(norm_out).any()}")
-
-# LM head
-logits = model.lm_head(norm_out)
-print(f"Logits: min={logits.min():.4f}, max={logits.max():.4f}, has_nan: {torch.isnan(logits).any()}")
-
-# Now test with actual layers
-input_ids = torch.randint(0, 10000, (batch_size, seq_len))
-embed_out = model.embeddings(input_ids)
-print(f"Embeddings: min={embed_out.min():.4f}, max={embed_out.max():.4f}, has_nan: {torch.isnan(embed_out).any()}")
-
-# Forward through layers
-hs = model.dropout(embed_out)
-print(f"After dropout: min={hs.min():.4f}, max={hs.max():.4f}, has_nan: {torch.isnan(hs).any()}")
-
-for i, layer in enumerate(model.layers):
-    hs, _ = layer(hs)
-    print(f"Layer {i} output: min={hs.min():.4f}, max={hs.max():.4f}, has_nan: {torch.isnan(hs).any()}")
-
-hs = model.norm(hs)
-print(f"After final norm: min={hs.min():.4f}, max={hs.max():.4f}, has_nan: {torch.isnan(hs).any()}")
-
-logits = model.lm_head(hs)
-print(f"Final logits: min={logits.min():.4f}, max={logits.max():.4f}, has_nan: {torch.isnan(logits).any()}")

tests/debug_loss.py

@@ -1,63 +0,0 @@
-"""Debug script for NaN loss"""
-import sys
-sys.path.insert(0, ".")
-import torch
-
-print("[DEBUG] Testing Fusion Model loss calculation...")
-
-# Import directly
-from models.fusion_model import FusionModel, FusionConfig
-
-config = FusionConfig(
-    vocab_size=10000,
-    hidden_size=256,
-    num_hidden_layers=2,
-    num_attention_heads=4,
-    intermediate_size=512,
-    block_size=64,
-    latent_dim=16,
-    sbla_mode="pure_sbla",
-    max_position_embeddings=256,
-)
-
-model = FusionModel(config)
-model.eval()
-
-# Small test
-batch_size, seq_len = 2, 32
-input_ids = torch.randint(0, config.vocab_size, (batch_size, seq_len))
-attention_mask = torch.ones(batch_size, seq_len)
-
-with torch.no_grad():
-    outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids, return_dict=True)
-
-logits = outputs["logits"]
-print(f"Logits stats: min={logits.min():.4f}, max={logits.max():.4f}, mean={logits.mean():.4f}")
-print(f"Logits has inf: {torch.isinf(logits).any()}")
-print(f"Logits has nan: {torch.isnan(logits).any()}")
-
-# Check logits at last position
-last_logits = logits[:, -1, :]
-print(f"Last token logits: min={last_logits.min():.4f}, max={last_logits.max():.4f}")
-
-# Try computing loss manually
-shift_logits = logits[..., :-1, :].contiguous()
-shift_labels = input_ids[..., 1:].contiguous()
-print(f"Shift logits stats: min={shift_logits.min():.4f}, max={shift_logits.max():.4f}")
-print(f"Shift logits has inf: {torch.isinf(shift_logits).any()}")
-print(f"Shift labels range: {shift_labels.min():.0f} - {shift_labels.max():.0f}")
-
-loss_fct = torch.nn.CrossEntropyLoss()
-try:
-    loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
-    print(f"Manual loss: {loss.item():.4f}")
-except Exception as e:
-    print(f"Loss computation error: {e}")
-
-# Check for any NaN in embeddings
-embeds = model.embeddings(input_ids)
-print(f"Embeddings: min={embeds.min():.4f}, max={embeds.max():.4f}, has_nan={torch.isnan(embeds).any()}")
-
-# Test a simple forward without labels
-outputs_no_labels = model(input_ids=input_ids, attention_mask=attention_mask, return_dict=True)
-print(f"Forward no labels logits: min={outputs_no_labels['logits'].min():.4f}, max={outputs_no_labels['logits'].max():.4f}")

tests/debug_mask.py

@@ -1,68 +0,0 @@
-"""Debug attention_mask handling"""
-import sys
-sys.path.insert(0, ".")
-import torch
-
-print("[DEBUG] Testing attention_mask handling...")
-
-from models.fusion_model import FusionModel, FusionConfig
-
-config = FusionConfig(
-    vocab_size=10000,
-    hidden_size=256,
-    num_hidden_layers=2,
-    num_attention_heads=4,
-    intermediate_size=512,
-    block_size=64,
-    latent_dim=16,
-    sbla_mode="pure_sbla",
-    max_position_embeddings=256,
-)
-
-model = FusionModel(config)
-model.eval()
-
-batch_size, seq_len = 2, 32
-input_ids = torch.randint(0, 10000, (batch_size, seq_len))
-
-# Case 1: No attention_mask
-hs = model.embeddings(input_ids)
-hs = model.dropout(hs)
-for i, layer in enumerate(model.layers):
-    hs, _ = layer(hs)
-hs = model.norm(hs)
-logits1 = model.lm_head(hs)
-print(f"No mask logits: min={logits1.min():.4f}, max={logits1.max():.4f}, has_nan: {torch.isnan(logits1).any()}")
-
-# Case 2: attention_mask as 2D
-mask_2d = torch.ones(batch_size, seq_len)
-hs = model.embeddings(input_ids)
-hs = model.dropout(hs)
-for i, layer in enumerate(model.layers):
-    hs, _ = layer(hs, attention_mask=mask_2d)
-hs = model.norm(hs)
-logits2 = model.lm_head(hs)
-print(f"2D mask logits: min={logits2.min():.4f}, max={logits2.max():.4f}, has_nan: {torch.isnan(logits2).any()}")
-
-# Case 3: What does _build_causal_mask return for seq_len=32?
-device = hs.device
-causal_mask = torch.triu(torch.ones(seq_len, seq_len, device=device, dtype=torch.bool), diagonal=1)
-causal_mask = causal_mask.float().masked_fill(causal_mask, float('-inf'))
-print(f"Causal mask shape: {causal_mask.shape}, has -inf: {(causal_mask == float('-inf')).any()}")
-print(f"Causal mask diag: {torch.diag(causal_mask)[:5]}")
-
-# Case 4: What is combined_mask after adding window_mask?
-window_mask = torch.zeros(seq_len, seq_len, device=device)
-combined_mask = (causal_mask + window_mask).unsqueeze(0).unsqueeze(0)
-print(f"Combined mask: shape={combined_mask.shape}, min={combined_mask.min()}, max={combined_mask.max()}")
-
-# Check if softmax produces NaN when there are rows full of -inf
-attn_scores_test = torch.randn(2, 4, 32, 32)
-attn_scores_test = attn_scores_test + combined_mask
-print(f"Attn scores after mask: has -inf rows: {(attn_scores_test == float('-inf')).any(dim=-1).all()}")
-attn_probs = torch.softmax(attn_scores_test, dim=-1)
-print(f"Attn probs: has nan: {torch.isnan(attn_probs).any()}")
-
-# Case 5: Forward pass with full model
-outputs = model(input_ids=input_ids, attention_mask=mask_2d, return_dict=True)
-print(f"Full model logits: min={outputs['logits'].min():.4f}, max={outputs['logits'].max():.4f}, has_nan: {torch.isnan(outputs['logits']).any()}")

tests/test_sblla_integration.py

@@ -1,59 +0,0 @@
-"""
-测试 SBLA 注意力集成（无 emoji 版本）
-"""
-import sys
-sys.path.insert(0, '.')
-
-from models.fusion_mini import FusionMini, FusionMiniConfig
-import torch
-
-print("测试 SBLA 注意力集成...")
-print()
-
-# 1. 创建配置
-print("[1] 创建配置...")
-config = FusionMiniConfig(
-    vocab_size=1000,
-    hidden_size=128,
-    num_hidden_layers=2,
-    num_attention_heads=4,
-)
-print("   配置创建成功")
-print(f"   隐层大小：{config.hidden_size}")
-print(f"   层数：{config.num_hidden_layers}")
-print()
-
-# 2. 创建模型
-print("[2] 创建模型（包含 SBLA 注意力）...")
-model = FusionMini(config)
-print("   模型创建成功")
-param_count = sum(p.numel() for p in model.parameters()) / 1e3
-print(f"   参数量：{param_count:.1f}K")
-print()
-
-# 3. 测试前向传播
-print("[3] 测试前向传播...")
-input_ids = torch.randint(0, 1000, (2, 64))
-print(f"   输入形状：{input_ids.shape}")
-
-outputs = model.forward(input_ids=input_ids, labels=input_ids)
-loss_value = outputs["loss"].item()
-print(f"   前向传播成功")
-print(f"   Loss：{loss_value:.4f}")
-print()
-
-# 4. 验证 SBLA 是否使用
-print("[4] 验证 SBLA 注意力...")
-has_sblla = any("SBLAttention" in str(module) for module in model.modules())
-if has_sblla:
-    print("   SBLA 注意力已集成到模型中")
-else:
-    print("   未检测到 SBLA 注意力（可能使用了标准注意力）")
-print()
-
-print("所有测试通过！")
-print()
-print("下一步：")
-print("   1. 重新训练模型（使用 SBLA 注意力）")
-print("   2. 对比标准注意力和 SBLA 的性能")
-print("   3. 推送代码到 GitHub")

train/full_finetune.py

@@ -25,9 +25,9 @@ import torch
 import torch.nn as nn
 import deepspeed
 from transformers import (
-    AutoTokenizer,
     get_linear_schedule_with_warmup,
 )
+from models.tokenizer import get_tokenizer, get_effective_vocab_size
 from torch.utils.data import Dataset, DataLoader
 import json
 import os
@@ -158,14 +158,16 @@ def create_local_model(
     return model, config
 
 
-def create_tokenizer(vocab_size: int = 32000):
+def create_tokenizer(tokenizer_type: str = "gpt2", vocab_size: int = 32000):
     """
-    创建 tokenizer
+    Create tokenizer using the unified tokenizer module.
+    
+    Note: Currently uses GPT2 as placeholder until SentencePiece model is trained.
+    The model config vocab_size will be auto-adjusted to match.
     """
-    logger.info(f"[create_tokenizer] 创建 tokenizer（vocab_size={vocab_size}）")
-    tokenizer = AutoTokenizer.from_pretrained("gpt2")
-    tokenizer.pad_token = tokenizer.eos_token
-    tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+    effective_vocab = get_effective_vocab_size(tokenizer_type, vocab_size)
+    logger.info(f"[create_tokenizer] Creating tokenizer: type={tokenizer_type}, effective_vocab={effective_vocab}")
+    tokenizer = get_tokenizer(tokenizer_type, vocab_size=vocab_size)
     return tokenizer