from safetensors.torch import save_file, load_file
import os
import shutil
# ===================== 配置参数 =====================
input_model_dir = r"D:\modelscope\hub\qwen\Qwen1___5-1___8B-Chat"
output_model_dir = r"D:\modelscope\hub\qwen\Qwen1___5-1___8B-Chat-split"
# 手动设置分片数量（1.8B模型设为2即可）
shard_num = 2
# ===================== 核心分片逻辑 =====================
def split_qwen_model_compatible():
    # 1. 创建输出目录
    os.makedirs(output_model_dir, exist_ok=True)
    # 2. 复制必需的配置文件（关键步骤）
    required_files = [
        "config.json", "tokenizer_config.json", 
        "vocab.json", "merges.txt", "generation_config.json"
    ]
    for file_name in required_files:
        src_path = os.path.join(input_model_dir, file_name)
        dst_path = os.path.join(output_model_dir, file_name)
        if os.path.exists(src_path):
            shutil.copy(src_path, dst_path)
            print(f"✅ 复制配置文件：{file_name}")
        else:
            print(f"⚠️  缺失配置文件：{file_name}")
    # 3. 找到原模型的主参数文件
    main_param_file = None
    for file in os.listdir(input_model_dir):
        if file.endswith(".safetensors") and "model" in file:
            main_param_file = os.path.join(input_model_dir, file)
            break
    if not main_param_file:
        raise FileNotFoundError("未找到 .safetensors 格式的参数文件")
    # 4. 加载完整模型参数
    print(f"正在加载模型参数：{main_param_file}")
    model_params = load_file(main_param_file)
    param_keys = list(model_params.keys())
    # 计算每个分片需要存放的参数数量
    keys_per_shard = len(param_keys) // shard_num
    # 5. 手动切分并保存分片文件
    index_data = {"metadata": {}, "weight_map": {}}
    for i in range(shard_num):
        # 切分参数字典
        start_idx = i * keys_per_shard
        end_idx = start_idx + keys_per_shard if i != shard_num -1 else len(param_keys)
        shard_keys = param_keys[start_idx:end_idx]
        shard_params = {k: model_params[k] for k in shard_keys}
        # 定义分片文件名（遵循transformers标准命名）
        shard_name = f"model-{i+1:05d}-of-{shard_num:05d}.safetensors"
        shard_path = os.path.join(output_model_dir, shard_name)
        # 保存当前分片
        save_file(shard_params, shard_path)
        print(f"✅ 生成分片文件：{shard_name}")
        # 填充索引文件数据
        for k in shard_keys:
            index_data["weight_map"][k] = shard_name
    # 6. 手动生成索引文件（关键！框架靠它识别分片）
    index_path = os.path.join(output_model_dir, "model.safetensors.index.json")
    import json
    with open(index_path, "w", encoding="utf-8") as f:
        json.dump(index_data, f, indent=2)
    print(f"✅ 生成索引文件：model.safetensors.index.json")
# ===================== 执行脚本 =====================
if __name__ == "__main__":
    try:
        split_qwen_model_compatible()
    except Exception as e:
        print(f"❌ 分片失败：{str(e)}")

import json
import os
# 配置分片模型路径
model_dir = r"D:\modelscope\hub\qwen\Qwen1___5-1___8B-Chat-split"
def check_shard_integrity():
    # 1. 检查索引文件是否存在
    index_path = os.path.join(model_dir, "model.safetensors.index.json")
    if not os.path.exists(index_path):
        return False, "❌ 缺失分片索引文件 model.safetensors.index.json"
    
    # 2. 读取索引文件，校验分片路径
    with open(index_path, "r", encoding="utf-8") as f:
        index_data = json.load(f)
    
    shard_files = set()
    for _, shard_path in index_data["weight_map"].items():
        shard_files.add(shard_path)
    
    # 3. 检查所有分片文件是否存在
    for shard in shard_files:
        shard_full_path = os.path.join(model_dir, shard)
        if not os.path.exists(shard_full_path):
            return False, f"❌ 缺失分片文件：{shard}"
    
    # 4. 检查配置文件完整性
    required_configs = ["config.json", "tokenizer_config.json", "generation_config.json"]
    for cfg in required_configs:
        if not os.path.exists(os.path.join(model_dir, cfg)):
            return False, f"❌ 缺失配置文件：{cfg}"
    
    return True, "✅ 所有分片文件和配置文件完整！"
# 执行校验
is_valid, msg = check_shard_integrity()
print(msg)

from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
model_dir = r"D:\modelscope\hub\qwen\Qwen1___5-1___8B-Chat-split"
def test_shard_loading():
    try:
        # 1. 加载配置文件（先不加载参数，快速验证配置）
        config = AutoConfig.from_pretrained(model_dir)
        print(f"✅ 配置文件加载成功，模型维度：{config.hidden_size}")
        # 2. 加载分词器
        tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
        print(f"✅ 分词器加载成功，词表大小：{len(tokenizer)}")
        # 3. 加载分片模型（关键：device_map='auto' 启用按需加载）
        model = AutoModelForCausalLM.from_pretrained(
            model_dir,
            device_map="cpu",  # 自动分配CPU/GPU显存，实现分片按需加载
            low_cpu_mem_usage=True,
            trust_remote_code=True
        )
        print(f"✅ 分片模型加载成功，模型设备：{model.device}")
        return True
    
    except Exception as e:
        print(f"❌ 模型加载失败：{str(e)}")
        return False
# 执行加载测试
test_shard_loading()

from transformers import AutoModelForCausalLM, AutoTokenizer
import time
model_dir = r"D:\modelscope\hub\qwen\Qwen1___5-1___8B-Chat-split"
# 加载模型和分词器
tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_dir,
    device_map="auto",
    low_cpu_mem_usage=True,
    trust_remote_code=True
)
# 测试案例
test_prompts = [
    "你好，请介绍一下自己",
    "解释什么是大模型分片存储",
    "写一段关于人工智能的短文（50字以内）"
]
# 执行推理
for i, prompt in enumerate(test_prompts):
    print(f"\n===== 测试案例 {i+1} =====")
    # 构造Qwen的对话格式
    messages = [{"role": "user", "content": prompt}]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    
    # 计时推理过程
    start_time = time.time()
    inputs = tokenizer([text], return_tensors="pt").to(model.device)
    outputs = model.generate(
        **inputs,
        max_new_tokens=150,
        temperature=0.7,
        do_sample=True
    )
    end_time = time.time()
    
    # 解码输出
    response = tokenizer.decode(outputs[0], skip_special_tokens=True).split("assistant\n")[-1]
    print(f"用户输入：{prompt}")
    print(f"模型回复：{response}")
    print(f"响应耗时：{end_time - start_time:.2f} 秒")

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model_dir = r"D:\modelscope\hub\qwen\Qwen1___5-1___8B-Chat-split"
# 加载模型
model = AutoModelForCausalLM.from_pretrained(
    model_dir,
    device_map="auto",
    low_cpu_mem_usage=True,
    trust_remote_code=True
)
tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
# 测试显存占用
def test_resource_usage():
    # 空输入时的基础显存占用
    if torch.cuda.is_available():
        base_memory = torch.cuda.memory_allocated() / 1024 / 1024
        print(f"✅ GPU基础显存占用：{base_memory:.2f} MB")
    
    # 推理时的峰值显存占用
    prompt = "写一个大模型本地部署的教程大纲"
    messages = [{"role": "user", "content": prompt}]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer([text], return_tensors="pt").to(model.device)
    
    if torch.cuda.is_available():
        torch.cuda.reset_peak_memory_stats()
    outputs = model.generate(**inputs, max_new_tokens=200)
    
    if torch.cuda.is_available():
        peak_memory = torch.cuda.max_memory_allocated() / 1024 / 1024
        print(f"✅ 推理峰值显存占用：{peak_memory:.2f} MB")
    else:
        print("✅ 使用CPU推理，无显存占用")
test_resource_usage()

from transformers import AutoModelForCausalLM, AutoTokenizer
# 指向配置文件和索引文件的目录
model_config_dir = r"D:\qwen-split-config"
tokenizer = AutoTokenizer.from_pretrained(model_config_dir, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_config_dir,
    device_map="auto",
    low_cpu_mem_usage=True,
    trust_remote_code=True
)
# 测试推理
prompt = "你好，详细介绍你自己"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=50)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))