from transformers import AutoModelForCausalLM, AutoTokenizer
# 模型名称
model_name = "Qwen/Qwen-7B-Chat"
# 原生加载：全量FP16加载到GPU，极易OOM
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",  # 自动分配，但无优化
    torch_dtype="auto"
)

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch
from modelscope import snapshot_download
def print_memory_info(stage=""):
    """打印显存使用信息"""
    if torch.cuda.is_available():
        allocated = torch.cuda.memory_allocated() / 1024**3  # 转换为GB
        reserved = torch.cuda.memory_reserved() / 1024**3
        print(f"\n{'='*50}")
        print(f"[{stage}]")
        print(f"  GPU显存占用: {allocated:.2f} GB")
        print(f"  GPU显存预留: {reserved:.2f} GB")
        print(f"{'='*50}")
    else:
        print(f"\n[{stage}] GPU不可用")
# 1. 配置INT4量化参数（核心优化）
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,               # 开启INT4加载
    bnb_4bit_use_double_quant=True,   # 双重量化，进一步压缩显存
    bnb_4bit_quant_type="nf4",        # 专用量化类型，适配大模型
    bnb_4bit_compute_dtype=torch.bfloat16  # 计算精度，4090支持bf16
)
# 2. 加载模型与分词器
model_name = "Qwen/Qwen-7B-Chat"
cache_dir = "/home/model"
print_memory_info("加载前")
local_model_path = snapshot_download(model_name, cache_dir=cache_dir)
tokenizer = AutoTokenizer.from_pretrained(local_model_path, trust_remote_code=True)
print_memory_info("分词器加载后")
model = AutoModelForCausalLM.from_pretrained(
    local_model_path,
    quantization_config=bnb_config,   # 应用量化配置
    device_map="auto",
    trust_remote_code=True
).eval()
print_memory_info("INT4量化模型加载后")
# 3. 显存对比总结
print("\n" + "="*50)
print("显存优化效果对比 (Qwen-7B)")
print("="*50)
print(f"  FP16原始显存: ~14 GB")
print(f"  INT4量化后显存: {torch.cuda.memory_allocated()/1024**3:.2f} GB")
print(f"  显存节省: {14 - torch.cuda.memory_allocated()/1024**3:.2f} GB")
print(f"  压缩比例: {(1 - torch.cuda.memory_allocated()/1024**3/14)*100:.1f}%")
print("="*50)
print("✓ 模型加载完成，显存占用大幅降低！")

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from accelerate import infer_auto_device_map, dispatch_model
import warnings
warnings.filterwarnings("ignore")
def print_memory_info(stage=""):
    """打印显存/内存使用信息"""
    print(f"\n{'='*55}")
    print(f"  📊 [{stage}]")
    print(f"{'='*55}")
    if torch.cuda.is_available():
        gpu_mem = torch.cuda.memory_allocated() / 1024**3
        gpu_mem_reserved = torch.cuda.memory_reserved() / 1024**3
        print(f"  GPU显存占用: {gpu_mem:.2f} GB")
        print(f"  GPU显存预留: {gpu_mem_reserved:.2f} GB")
    import psutil
    ram = psutil.virtual_memory()
    print(f"  CPU内存占用: {ram.used / 1024**3:.2f} GB / {ram.total / 1024**3:.2f} GB")
    print(f"{'='*55}")
# ==========================================
# 1. 配置路径
# ==========================================
local_model_path = "/home/model/shakechen/Llama-2-13b-chat-hf"
# ==========================================
# 2. INT4 量化配置
# ==========================================
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)
# ==========================================
# 3. 加载 Tokenizer
# ==========================================
print("\n" + "="*55)
print("  🚀 Llama-2-13B 分层分片加载演示")
print("="*55)
print("\n正在加载 Tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(
    local_model_path,
    trust_remote_code=True,
    use_fast=False
)
print_memory_info("分词器加载后")
# ==========================================
# 4. 先加载模型到CPU（不分片）
# ==========================================
print("\n🔄 第一步：加载基础模型...")
model_base = AutoModelForCausalLM.from_pretrained(
    local_model_path,
    quantization_config=bnb_config,
    device_map="cpu",  # 先加载到CPU
    trust_remote_code=True,
    low_cpu_mem_usage=True
)
print_memory_info("模型基础加载（不分片）")
# ==========================================
# 5. 手动分层分片：核心优化
# ==========================================
print("\n🔄 第二步：执行分层分片调度...")
device_map = infer_auto_device_map(
    model_base,
    max_memory={
        0: "20GB",   # GPU 0 分配20GB（预留4GB给计算）
        "cpu": "64GB"# CPU内存分配64GB
    },
    no_split_module_classes=["LlamaDecoderLayer"]
)
print(f"\n📋 分片策略预览:")
layer_count = sum(1 for k in device_map.keys() if 'layer' in k.lower())
print(f"  - 模型层数: {layer_count}")
print(f"  - GPU层数: {sum(1 for v in device_map.values() if v == 0)}")
print(f"  - CPU层数: {sum(1 for v in device_map.values() if v == 'cpu')}")
# 释放基础模型
del model_base
import gc
gc.collect()
# 重新加载并应用分片
model = AutoModelForCausalLM.from_pretrained(
    local_model_path,
    quantization_config=bnb_config,
    device_map=device_map,
    trust_remote_code=True,
    low_cpu_mem_usage=True
)
# 应用分片调度
model = dispatch_model(model, device_map=device_map)
print_memory_info("分层分片加载完成")
# ==========================================
# 6. 显存优化效果对比
# ==========================================
print("\n" + "="*55)
print("  📈 Llama-2-13B 分层分片优化效果")
print("="*55)
print(f"  {'方案':<20} {'GPU显存':<15} {'CPU内存':<15}")
print(f"  {'-'*50}")
print(f"  {'FP16全量加载':<18} {'~26 GB':<15} {'< 1 GB':<15}")
if torch.cuda.is_available():
    gpu_mem = torch.cuda.memory_allocated() / 1024**3
print(f"  {'INT4+分层分片':<18} {f'{gpu_mem:.1f} GB':<15} {'~40 GB':<15}")
print(f"  {'显存节省':<18} {f'{26-gpu_mem:.1f} GB':<15} {'-':<15}")
print(f"  {'压缩比例':<18} {f'{(1-gpu_mem/26)*100:.0f}%':<15} {'-':<15}")
print("="*55)
print("✅ 分层分片加载完成，实现大模型低显存运行！")
# ==========================================
# 7. 测试推理
# ==========================================
prompt = "你是什么模型，简单用中文介绍你自己."
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=50)
    print("\n💬 模型回复:")
    print(tokenizer.decode(outputs[0], skip_special_tokens=True))

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from accelerate import Accelerator
from accelerate import infer_auto_device_map, dispatch_model
import psutil
def print_memory_info(stage=""):
    """打印显存/内存使用信息"""
    print(f"\n{'='*55}")
    print(f"  {stage}")
    print(f"{'='*55}")
    if torch.cuda.is_available():
        gpu_mem = torch.cuda.memory_allocated() / 1024**3
        gpu_total = torch.cuda.get_device_properties(0).total_memory / 1024**3
        print(f"  GPU显存: {gpu_mem:.1f} GB / {gpu_total:.0f} GB (使用率 {gpu_mem/gpu_total*100:.0f}%)")
    ram = psutil.virtual_memory()
    print(f"  CPU内存: {ram.used / 1024**3:.1f} GB / {ram.total / 1024**3:.0f} GB")
    print(f"{'='*55}")
def print_offload_status(device_map, model):
    """打印Offload分配状态"""
    gpu_layers = []
    cpu_layers = []
    total_layers = 0
    
    for name, device in device_map.items():
        # 统计所有参数模块
        if 'weight' in name.lower() or 'bias' in name.lower():
            total_layers += 1
            if device == 0 or (isinstance(device, str) and 'gpu' in device.lower()):
                gpu_layers.append(name)
            elif device == 'cpu':
                cpu_layers.append(name)
    
    # 避免除零
    total = len(gpu_layers) + len(cpu_layers)
    if total == 0:
        total = 1  # 防止除零
    
    gpu_count = len(set(gpu_layers))
    cpu_count = len(set(cpu_layers))
    
    print(f"\n  GPU模块数: {gpu_count}")
    print(f"  CPU模块数: {cpu_count}")
    print(f"  Offload比例: {cpu_count / total * 100:.0f}%")
print("\n" + "="*55)
print("  CPU Offload 实践应用演示")
print("  模型: Qwen3.6-27B (INT4量化)")
print("="*55)
# 1. 初始化加速器
accelerator = Accelerator()
# 2. INT4量化配置
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)
# 3. 加载模型（本地路径）
model_name = "Qwen/Qwen3.6-27B"  # 模型名称标识
local_model_path = "/home/model/Qwen/Qwen3___6-27B"  # 本地模型路径
tokenizer = AutoTokenizer.from_pretrained(local_model_path, trust_remote_code=True, use_fast=False)
print_memory_info("初始状态")
# ========================================
# 核心：CPU Offload 实现
# ========================================
print("\n" + "-"*55)
print("  Step 1: 将模型完整加载到CPU内存")
print("-"*55)
print("  原理: 模型参数先存放在CPU RAM中，不占用GPU显存")
# 加载模型
model = AutoModelForCausalLM.from_pretrained(
    local_model_path,
    quantization_config=bnb_config,
    trust_remote_code=True,
    device_map="cpu",  # 关键：加载到CPU而非GPU
    low_cpu_mem_usage=True
)
print_memory_info("CPU加载完成")
print("  效果: 模型完整在CPU内存中，GPU显存几乎未占用")
# ========================================
print("\n" + "-"*55)
print("  Step 2: 创建分层分片策略")
print("-"*55)
print("  原理: 根据层级重要性分配存储位置")
print("        核心Attention层 -> GPU (高频访问)")
print("        FFN层/Embed层 -> CPU (低频访问)")
device_map = infer_auto_device_map(
    model,
    max_memory={
        0: "22GB",   # GPU: 放核心层
        "cpu": "128GB"  # CPU: 放非核心层
    },
    no_split_module_classes=["LlamaDecoderLayer", "Qwen2DecoderLayer"]
)
print(f"\n  分片详情:")
print_offload_status(device_map, model)
# ========================================
print("\n" + "-"*55)
print("  Step 3: 应用分层调度 (dispatch_model)")
print("-"*55)
print("  原理: 按需加载/卸载，推理时只激活需要的层")
model = dispatch_model(model, device_map=device_map)
print_memory_info("分层调度完成")
print("  效果: GPU只承载22GB，剩余计算需求通过CPU交换")
# ========================================
print("\n" + "-"*55)
print("  Step 4: 加速器优化")
print("-"*55)
model = accelerator.prepare(model)
print_memory_info("加速器准备完成")
# ========================================
# 最终效果总结
# ========================================
print("\n" + "="*55)
print("  📈 CPU Offload 优化效果总结")
print("="*55)
if torch.cuda.is_available():
    gpu_mem = torch.cuda.memory_allocated() / 1024**3
    print(f"\n  模型规模: Qwen3.6-27B (约270亿参数)")
    print(f"  原始显存需求: ~54 GB (FP16)")
    print(f"  优化后GPU显存: {gpu_mem:.1f} GB")
    print(f"  显存节省: {54 - gpu_mem:.1f} GB ({((54-gpu_mem)/54*100):.0f}%)")
    print(f"\n  关键技术:")
    print(f"    1. INT4量化: 参数体积缩减75%")
    print(f"    2. CPU Offload: 非核心层卸载到CPU内存")
    print(f"    3. 分层调度: 按需加载，减少峰值显存")
    print(f"\n  适用场景: 消费级GPU (如RTX 4090 24GB) 运行超大模型")
print("\n" + "="*55)
print("  CPU Offload实践完成 - 超大模型稳定运行!")
print("="*55)

import torch
import time
def monitor_gpu_memory():
    while True:
        # 获取4090显存使用情况
        gpu_mem_used = torch.cuda.memory_allocated(0) / 1024**3  # 已用显存（GB）
        gpu_mem_total = torch.cuda.get_device_properties(0).total_memory / 1024**3  # 总显存
        gpu_mem_free = gpu_mem_total - gpu_mem_used  # 剩余显存
        
        print(f"RTX 4090显存：已用 {gpu_mem_used:.2f}GB | 剩余 {gpu_mem_free:.2f}GB")
        time.sleep(2)
# 后台启动监控
import threading
threading.Thread(target=monitor_gpu_memory, daemon=True).start()