从7B走到173B：AI百川笔记

# 简单的显存需求估算（未考虑优化）
def estimate_vram(params, precision=16):
    # params: 参数量（十亿）
    # precision: 16表示FP16，8表示INT8
    base_size = params * 2  # 基础模型大小（GB）
    activation_overhead = base_size * 0.3  # 激活值开销
    kv_cache = params * 0.2  # KV缓存
    total = base_size + activation_overhead + kv_cache

    if precision == 8:
        total = total * 0.6  # 量化后理论上可减少约40%

    return round(total, 2)

# 各模型的估算值
for model in [7, 13, 53, 173]:
    print(f"Baichuan {model}B FP16: {estimate_vram(model)}GB")
    print(f"Baichuan {model}B INT8: {estimate_vram(model, 8)}GB")

# 环境准备
pip install torch>=2.0.0
pip install transformers>=4.30.0
pip install accelerate

# 模型下载
git clone https://huggingface.co/baichuan-inc/Baichuan-7B

# 推理代码示例
python3 << EOF
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_path = "./Baichuan-7B"
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    device_map="auto",
    torch_dtype=torch.float16,
    trust_remote_code=True
)

# 推理示例
prompt = "简单解释一下什么是机器学习。"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_length=256,
        temperature=0.7,
        top_p=0.9,
        repetition_penalty=1.1
    )

response = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(response)
EOF

# 多卡推理配置
from accelerate import infer_auto_device_map, dispatch_model

model_path = "./Baichuan-53B"
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.float16,
    trust_remote_code=True
)

# 自动设备映射
device_map = infer_auto_device_map(model, max_memory={0: "40GB", 1: "40GB", 2: "40GB", 3: "40GB"})
model = dispatch_model(model, device_map)

# 验证设备分布
for name, param in model.named_parameters():
    print(f"{name}: {param.device}")

# 基础环境准备（多节点）
# 节点配置：8x A100 (40GB) 或类似配置

# 1. 安装依赖
pip install deepspeed
pip install tensor_parallel
pip install vllm  # 高性能推理框架

# 2. 模型分片与加载
python3 << EOF
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from deepspeed import deepspeed

model_path = "./Baichuan-173B"
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

# 使用DeepSpeed进行模型并行
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,
    trust_remote_code=True
)

# 推理配置
deepspeed_config = {
    "tensor_parallel": {"tp_size": 8},
    "fp16": {"enabled": True}
}
model_engine, _, _, _ = deepspeed.initialize(
    model=model,
    config=deepspeed_config
)
EOF

# 3. 启动推理服务（使用vLLM加速）
python -m vllm.entrypoints.api_server \
  --model ./Baichuan-173B \
  --tensor-parallel-size 8 \
  --trust-remote-code \
  --dtype float16 \
  --port 8000

# 错误信息
RuntimeError: CUDA out of memory. Tried to allocate X GiB

# 解决方案
# 1. 使用量化
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    load_in_8bit=True,  # 或者 load_in_4bit=True
    torch_dtype=torch.float16,
    device_map="auto",
    trust_remote_code=True
)

# 2. 调整批处理大小
outputs = model.generate(
    **inputs,
    batch_size=1,  # 降低批处理大小
    max_length=256
)

# 优化推理参数
outputs = model.generate(
    **inputs,
    max_length=256,
    temperature=0.7,  # 控制随机性，0.7-0.9通常较好
    top_p=0.9,        # nucleus sampling
    top_k=50,         # 限制候选词数量
    repetition_penalty=1.1,  # 减少重复
    no_repeat_ngram_size=3   # 避免n-gram重复
)

# 优化方案
# 1. 使用NCCL后端
import os
os.environ["NCCL_IB_DISABLE"] = "0"  # 启用InfiniBand
os.environ["NCCL_SOCKET_IFNAME"] = "eth0"  # 指定网络接口

# 2. 调整梯度累积
# 如果涉及训练，可以累积梯度减少通信频率
gradient_accumulation_steps = 4

# 3. 使用混合精度
from torch.cuda.amp import autocast
with autocast():
    outputs = model(**inputs)

# 解决方案
# 1. 使用KV缓存
outputs = model.generate(
    **inputs,
    use_cache=True,  # 启用KV缓存
    max_length=512
)

# 2. 使用投机采样（ speculative decoding）
# 需要配合小模型
small_model = AutoModelForCausalLM.from_pretrained("./Baichuan-7B")
# 推理时先由小模型生成候选，再由大模型验证

# 解决方案
# 1. 使用延迟加载
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,  # 关键：降低内存峰值
    device_map="auto",
    offload_folder="offload",  # CPU卸载目录
    trust_remote_code=True
)

# 2. 预加载常用权重
# 在服务启动时预加载常用部分的权重

# 使用高性能推理框架
# 1. vLLM（推荐）
python -m vllm.entrypoints.api_server \
  --model ./Baichuan-173B \
  --tensor-parallel-size 8 \
  --gpu-memory-utilization 0.9 \
  --max-model-len 4096 \
  --trust-remote-code

# 2. TensorRT-LLM
# 需要先转换模型
trtllm-build --model_dir ./Baichuan-173B \
  --output_dir ./trt_engine \
  --tp_size 8 \
  --dtype float16

# 3. 分批推理
# 将请求分成小批次并行处理

# 解决方案
# 1. 动态模型选择
def select_model(complexity_score):
    if complexity_score < 0.3:
        return "Baichuan-7B"
    elif complexity_score < 0.7:
        return "Baichuan-13B"
    else:
        return "Baichuan-173B"

# 2. 结果缓存
from functools import lru_cache

@lru_cache(maxsize=1000)
def cached_inference(prompt_hash):
    return model.generate(prompt)

# 3. 提前终止
outputs = model.generate(
    **inputs,
    max_length=256,
    early_stopping=True,
    num_return_sequences=1
)