LLM推理优化与部署

1
# 一个 7B 模型的推理开销分析
2
model_params = 7_000_000_000  # 70亿参数
3

4
# FP16 存储（每个参数 2 字节）
5
model_size_gb = model_params * 2 / (1024**3)  # ~13 GB
6

7
# 推理时需要额外显存：
8
# 1. KV-Cache（存储历史 token 的 key/value）
9
# 2. 激活值（中间计算结果）
10
# 3. 临时缓冲区
11

12
# 假设生成 2048 token，batch_size=1
13
num_layers = 32
14
hidden_size = 4096
15
num_heads = 32
16
seq_len = 2048
17

18
kv_cache_size = 2 * num_layers * seq_len * hidden_size * 2 / (1024**3)
19
# 2（key+value）* 32层 * 2048 tokens * 4096维 * 2字节 ≈ 2 GB
20

21
total_memory = model_size_gb + kv_cache_size  # ~15 GB

1
from transformers import AutoModelForCausalLM
2
import torch
3

4
# ===== 方法1: bitsandbytes（最简单）=====
5
from transformers import BitsAndBytesConfig
6

7
bnb_config = BitsAndBytesConfig(
8
    load_in_4bit=True,
9
    bnb_4bit_quant_type="nf4",  # Normal Float 4-bit
10
    bnb_4bit_use_double_quant=True,  # 双重量化
11
    bnb_4bit_compute_dtype=torch.bfloat16
12
)
13

14
model = AutoModelForCausalLM.from_pretrained(
15
    "meta-llama/Llama-2-7b-hf",
16
    quantization_config=bnb_config,
17
    device_map="auto"
18
)
19

20
# 现在模型只占用 ~4GB 显存！

1
from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
2

3
# 量化配置
4
quantize_config = BaseQuantizeConfig(
5
    bits=4,  # 4-bit 量化
6
    group_size=128,  # 分组大小
7
    desc_act=False  # 是否量化激活值
8
)
9

10
# 加载模型
11
model = AutoGPTQForCausalLM.from_pretrained(
12
    "meta-llama/Llama-2-7b-hf",
13
    quantize_config=quantize_config
14
)
15

16
# 使用校准数据量化
17
model.quantize(calibration_dataset)
18

19
# 保存量化模型
20
model.save_quantized("./llama-2-7b-gptq-4bit")
21

22
# 推理
23
quantized_model = AutoGPTQForCausalLM.from_quantized(
24
    "./llama-2-7b-gptq-4bit",
25
    device="cuda:0"
26
)

1
class AttentionWithKVCache(nn.Module):
2
    def __init__(self):
3
        super().__init__()
4
        self.kv_cache = None
5

6
    def forward(self, x, use_cache=False):
7
        batch_size, seq_len, hidden_size = x.shape
8

9
        # 计算 Q, K, V
10
        q = self.q_proj(x)
11
        k = self.k_proj(x)
12
        v = self.v_proj(x)
13

14
        if use_cache:
15
            if self.kv_cache is not None:
16
                # 使用缓存的 K, V
17
                k_cached, v_cached = self.kv_cache
18
                k = torch.cat([k_cached, k], dim=1)
19
                v = torch.cat([v_cached, v], dim=1)
20

21
            # 更新缓存
22
            self.kv_cache = (k, v)
23

24
        # Attention 计算
25
        attn_output = self.attention(q, k, v)
26
        return attn_output

1
# vLLM 的核心优势
2
from vllm import LLM, SamplingParams
3

4
llm = LLM(
5
    model="meta-llama/Llama-2-7b-hf",
6
    tensor_parallel_size=1,
7
    gpu_memory_utilization=0.9  # 使用 90% GPU 显存
8
)
9

10
prompts = [
11
    "Write a story about a robot",
12
    "Explain quantum computing"
13
] * 100  # 200 个请求
14

15
sampling_params = SamplingParams(
16
    temperature=0.8,
17
    top_p=0.95,
18
    max_tokens=512
19
)
20

21
# vLLM 自动做批处理和 KV-Cache 管理
22
outputs = llm.generate(prompts, sampling_params)
23

24
# 吞吐量提升 10-20x！

1
# 标准 Attention（显存占用大）
2
def standard_attention(Q, K, V):
3
    # Q, K, V: [batch, num_heads, seq_len, head_dim]
4
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(head_dim)
5
    # scores: [batch, num_heads, seq_len, seq_len] ← 显存瓶颈！
6
    attn_weights = F.softmax(scores, dim=-1)
7
    output = torch.matmul(attn_weights, V)
8
    return output
9

10
# Flash Attention（显存友好）
11
from flash_attn import flash_attn_func
12

13
def flash_attention(Q, K, V):
14
    # 分块计算，只存储最终结果
15
    output = flash_attn_func(Q, K, V, causal=True)
16
    return output
17

18
# 性能提升：
19
# - 速度: 2-4x
20
# - 显存: 减少 5-20x

1
class MultiQueryAttention(nn.Module):
2
    def __init__(self, hidden_size, num_heads):
3
        super().__init__()
4
        self.num_heads = num_heads
5
        self.head_dim = hidden_size // num_heads
6

7
        # Q 有多个头
8
        self.q_proj = nn.Linear(hidden_size, hidden_size)
9

10
        # K, V 只有一个头
11
        self.k_proj = nn.Linear(hidden_size, self.head_dim)
12
        self.v_proj = nn.Linear(hidden_size, self.head_dim)
13

14
    def forward(self, x):
15
        batch_size, seq_len, _ = x.shape
16

17
        # Q: [batch, num_heads, seq_len, head_dim]
18
        q = self.q_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim)
19
        q = q.transpose(1, 2)
20

21
        # K, V: [batch, 1, seq_len, head_dim]
22
        k = self.k_proj(x).unsqueeze(1)
23
        v = self.v_proj(x).unsqueeze(1)
24

25
        # K, V 广播到所有头
26
        k = k.expand(-1, self.num_heads, -1, -1)
27
        v = v.expand(-1, self.num_heads, -1, -1)
28

29
        # 标准 attention
30
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
31
        attn = F.softmax(scores, dim=-1)
32
        output = torch.matmul(attn, v)
33

34
        return output

1
from vllm import LLM, SamplingParams
2

3
# 初始化
4
llm = LLM(
5
    model="meta-llama/Llama-2-7b-chat-hf",
6
    tensor_parallel_size=2,  # 2 卡张量并行
7
    dtype="float16",
8
    max_model_len=4096
9
)
10

11
# 批量推理
12
prompts = ["Hello"] * 1000
13
outputs = llm.generate(prompts, SamplingParams(temperature=0.8, max_tokens=100))
14

15
# 吞吐量：~2000 tokens/s（单卡 A100）

1
import tensorrt_llm
2
from tensorrt_llm.runtime import ModelRunner
3

4
# 构建引擎（需要预先编译）
5
# 1. 导出模型
6
# 2. 量化（可选）
7
# 3. 构建 TensorRT 引擎
8

9
# 加载引擎
10
runner = ModelRunner.from_dir("./llama-7b-engine")
11

12
# 推理
13
input_ids = tokenizer.encode("Hello, world!")
14
outputs = runner.generate(
15
    input_ids,
16
    max_new_tokens=100,
17
    temperature=0.8
18
)
19

20
# 吞吐量：~3000 tokens/s（单卡 A100，INT8）

1
from transformers import pipeline
2

3
# DeepSpeed Pipeline
4
pipe = pipeline(
5
    "text-generation",
6
    model="meta-llama/Llama-2-70b-hf",
7
    device_map="auto",  # 自动分配层到 GPU
8
    max_length=200
9
)
10

11
# 层分配示例（8 卡）：
12
# GPU 0-1: Layers 0-9
13
# GPU 2-3: Layers 10-19
14
# GPU 4-5: Layers 20-29
15
# GPU 6-7: Layers 30-39

1
from ray import serve
2
from transformers import AutoModelForCausalLM, AutoTokenizer
3

4
@serve.deployment(num_replicas=4, ray_actor_options={"num_gpus": 1})
5
class LLMDeployment:
6
    def __init__(self):
7
        self.model = AutoModelForCausalLM.from_pretrained("gpt2")
8
        self.tokenizer = AutoTokenizer.from_pretrained("gpt2")
9

10
    def __call__(self, request):
11
        prompt = request.query_params["prompt"]
12
        inputs = self.tokenizer(prompt, return_tensors="pt")
13
        outputs = self.model.generate(**inputs, max_length=100)
14
        return self.tokenizer.decode(outputs[0])
15

16
# 部署
17
serve.run(LLMDeployment.bind())
18

19
# 现在有 4 个副本自动负载均衡

1
用户请求
2
   ↓
3
[Nginx/负载均衡器]
4
   ↓
5
[API Gateway]
6
   ↓
7
[请求队列 (Redis/RabbitMQ)]
8
   ↓
9
[推理服务集群]
10
  ├─ vLLM Instance 1 (GPU 0-1)
11
  ├─ vLLM Instance 2 (GPU 2-3)
12
  └─ vLLM Instance 3 (GPU 4-5)
13
   ↓
14
[响应缓存 (Redis)]
15
   ↓
16
[监控 (Prometheus + Grafana)]

1
apiVersion: apps/v1
2
kind: Deployment
3
metadata:
4
  name: llm-inference
5
spec:
6
  replicas: 3
7
  selector:
8
    matchLabels:
9
      app: llm-inference
10
  template:
11
    metadata:
12
      labels:
13
        app: llm-inference
14
    spec:
15
      containers:
16
      - name: vllm
17
        image: vllm/vllm-openai:latest
18
        args:
19
          - --model
20
          - meta-llama/Llama-2-7b-chat-hf
21
          - --tensor-parallel-size
22
          - "2"
23
        resources:
24
          limits:
25
            nvidia.com/gpu: 2
26
        ports:
27
        - containerPort: 8000
28
---
29
apiVersion: v1
30
kind: Service
31
metadata:
32
  name: llm-service
33
spec:
34
  selector:
35
    app: llm-inference
36
  ports:
37
    - protocol: TCP
38
      port: 80
39
      targetPort: 8000
40
  type: LoadBalancer

1
from prometheus_client import Counter, Histogram, start_http_server
2
import time
3

4
# 定义指标
5
request_count = Counter('llm_requests_total', 'Total requests')
6
request_duration = Histogram('llm_request_duration_seconds', 'Request duration')
7
tokens_generated = Counter('llm_tokens_generated_total', 'Total tokens generated')
8

9
def monitored_inference(prompt):
10
    start = time.time()
11
    request_count.inc()
12

13
    try:
14
        # 推理
15
        output = model.generate(prompt)
16
        tokens_generated.inc(len(output))
17

18
        # 记录延迟
19
        duration = time.time() - start
20
        request_duration.observe(duration)
21

22
        return output
23
    except Exception as e:
24
        error_count.inc()
25
        raise
26

27
# 启动 Prometheus 服务器
28
start_http_server(8001)

1
# AWS Spot Instance 部署
2
# 成本降低 70%，但可能被中断
3

4
import boto3
5

6
ec2 = boto3.client('ec2')
7

8
# 请求 Spot Instance
9
response = ec2.request_spot_instances(
10
    InstanceCount=4,
11
    Type='persistent',
12
    LaunchSpecification={
13
        'ImageId': 'ami-xxx',  # 包含 vLLM 的 AMI
14
        'InstanceType': 'g5.12xlarge',  # 4x A10G
15
        'KeyName': 'my-key',
16
        'SecurityGroups': ['llm-sg']
17
    },
18
    SpotPrice='2.50'  # 最高出价 $2.50/小时（按需价格 $5/小时）
19
)

1
class DynamicBatcher:
2
    def __init__(self, max_batch_size=32, max_wait_ms=100):
3
        self.max_batch_size = max_batch_size
4
        self.max_wait_ms = max_wait_ms
5
        self.queue = []
6

7
    async def add_request(self, request):
8
        self.queue.append(request)
9

10
        # 条件满足时触发批处理
11
        if len(self.queue) >= self.max_batch_size:
12
            return await self.process_batch()
13

14
        # 或等待超时
15
        await asyncio.sleep(self.max_wait_ms / 1000)
16
        if self.queue:
17
            return await self.process_batch()
18

19
    async def process_batch(self):
20
        batch = self.queue[:self.max_batch_size]
21
        self.queue = self.queue[self.max_batch_size:]
22

23
        # 批量推理（吞吐量提升 10x）
24
        results = model.generate(batch)
25
        return results

1
import redis
2
import hashlib
3

4
cache = redis.Redis(host='localhost', port=6379, db=0)
5

6
def cached_inference(prompt, **kwargs):
7
    # 生成缓存键
8
    cache_key = hashlib.md5(f"{prompt}{kwargs}".encode()).hexdigest()
9

10
    # 检查缓存
11
    cached_result = cache.get(cache_key)
12
    if cached_result:
13
        return cached_result.decode()
14

15
    # 推理
16
    result = model.generate(prompt, **kwargs)
17

18
    # 存入缓存（TTL 1小时）
19
    cache.setex(cache_key, 3600, result)
20

21
    return result
22

23
# 缓存命中率 30-50% 时，成本降低显著