from transformers import AutoModelForCausalLM
import torch
 
model = AutoModelForCausalLM.from_pretrained(
    "gpt2",
    torch_dtype=torch.float16,   # 或 torch.bfloat16
)

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    torch_dtype=torch.float16,
    device_map="auto",           # 自动分配到可用 GPU，不够则溢出到 CPU
)

浮点权重值 w ∈ [w_min, w_max]
映射到整数 q ∈ [-128, 127]

量化：q = round(w / scale + zero_point)
反量化：w ≈ scale * (q - zero_point)

from transformers import AutoModelForCausalLM, BitsAndBytesConfig
import torch
 
# load_in_4bit=True 背后做了什么：
#   1. 将模型权重从 fp16/fp32 量化为 NF4 格式（4 bit）
#   2. 计算时反量化为 fp16 再做矩阵乘法（compute_dtype）
#   3. 激活值仍然以 fp16 存储
#   整体效果：显存减少约 75%，推理速度接近 fp16
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,  # 计算时临时转回 fp16
    bnb_4bit_use_double_quant=True,        # 对量化参数再做一次量化，进一步省显存
    bnb_4bit_quant_type="nf4",             # NF4 比 fp4 精度更高
)
 
model = AutoModelForCausalLM.from_pretrained(
    "facebook/opt-125m",
    quantization_config=bnb_config,
    device_map="auto",
)

无缓存（第 n+1 步）：计算 n+1 个 token 的 K、V，做全量 attention
有缓存（第 n+1 步）：只计算第 n+1 个 token 的 K、V，与缓存的 n 组 K、V 拼接后做 attention

# attention_mask: 1 表示真实 token，0 表示 padding token
# tokenizer 会自动生成
inputs = tokenizer(
    texts,
    padding=True,          # 自动 padding 到 batch 内最长序列
    truncation=True,
    max_length=512,
    return_tensors="pt",
)
 
# inputs["attention_mask"] 的形状：(batch_size, seq_len)
# 值为 1 的位置参与 attention，值为 0 的位置被 mask 掉
outputs = model(**inputs)

from vllm import LLM, SamplingParams
 
llm = LLM(model="meta-llama/Llama-2-7b-hf")
sampling_params = SamplingParams(temperature=0.7, max_tokens=256)
 
outputs = llm.generate(["Tell me about transformers."], sampling_params)
print(outputs[0].outputs[0].text)

python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Llama-2-7b-hf \
    --port 8000

# 安装 Ollama（Linux）
curl -fsSL https://ollama.com/install.sh | sh
 
# 下载并运行模型（自动处理量化和显存分配）
ollama run llama3
 
# Python 调用
import ollama
response = ollama.chat(model='llama3', messages=[
    {'role': 'user', 'content': 'What is a transformer model?'}
])
print(response['message']['content'])

import time
import torch
 
def benchmark(model, inputs, num_runs=100, warmup=10):
    """
    测量模型推理的延迟和吞吐量。
    
    warmup：前几次推理用于 JIT 编译和 GPU 预热，不计入统计。
    """
    # GPU 操作是异步的，需要 synchronize() 确保计时准确
    if torch.cuda.is_available():
        torch.cuda.synchronize()
 
    # 预热
    for _ in range(warmup):
        with torch.no_grad():
            model(**inputs)
 
    if torch.cuda.is_available():
        torch.cuda.synchronize()
 
    # 正式测量
    latencies = []
    for _ in range(num_runs):
        start = time.perf_counter()
        with torch.no_grad():
            outputs = model(**inputs)
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        end = time.perf_counter()
        latencies.append((end - start) * 1000)  # 转换为毫秒
 
    import numpy as np
    return {
        "mean_ms": np.mean(latencies),
        "p50_ms":  np.percentile(latencies, 50),
        "p99_ms":  np.percentile(latencies, 99),
    }

# 查看当前显存占用
print(f"已分配显存：{torch.cuda.memory_allocated() / 1024**2:.1f} MB")
print(f"缓存显存：  {torch.cuda.memory_reserved() / 1024**2:.1f} MB")
 
# 峰值显存（从进程启动以来的最大值）
print(f"峰值显存：  {torch.cuda.max_memory_allocated() / 1024**2:.1f} MB")
 
# 重置峰值统计（用于对比不同阶段的显存峰值）
torch.cuda.reset_peak_memory_stats()

from torch.profiler import profile, ProfilerActivity
 
with profile(activities=[ProfilerActivity.CUDA], profile_memory=True) as prof:
    with torch.no_grad():
        model(**inputs)
 
print(prof.key_averages().table(sort_by="cuda_memory_usage", row_limit=10))