FastAPI + Baichuan13B实战：5分钟搭建高性能大模型API服务

在AI应用开发中，将大语言模型快速封装成可调用的API服务是连接模型能力与实际业务的关键一步。今天我们就来手把手教你如何用FastAPI和Baichuan13B搭建一个高性能的本地API服务，整个过程只需5分钟，代码可直接复用。

1. 环境准备与模型加载

首先确保你的开发环境满足以下基础要求：

• Python 3.8+
• CUDA 11.7+（如需GPU加速）
• 至少16GB内存（运行Baichuan13B的最低要求）

安装必要的依赖包：

pip install fastapi uvicorn transformers torch

模型加载是服务启动的关键步骤，这里我们使用Hugging Face的transformers库加载Baichuan13B：

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_path = "baichuan13b_chat"
tokenizer = AutoTokenizer.from_pretrained(
    model_path, 
    use_fast=False,
    trust_remote_code=True
)
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    device_map="auto",
    torch_dtype=torch.float16,
    trust_remote_code=True
)

注意：首次运行时会下载模型权重文件，建议提前准备好模型文件或确保网络通畅

2. FastAPI服务端搭建

FastAPI以其高性能和易用性成为构建API服务的首选框架。下面我们构建一个完整的聊天接口：

from fastapi import FastAPI
from pydantic import BaseModel
import uvicorn

app = FastAPI()

class ChatRequest(BaseModel):
    text: str

@app.post("/chat")
async def chat(query: ChatRequest):
    input_ids = tokenizer([query.text]).input_ids
    outputs = model.generate(
        torch.as_tensor(input_ids).cuda(),
        max_new_tokens=512,
        temperature=0.7,
        top_p=0.9
    )
    response = tokenizer.decode(
        outputs[0][len(input_ids[0]):],
        skip_special_tokens=True
    )
    return {"response": response}

启动服务的命令如下：

uvicorn main:app --host 0.0.0.0 --port 8000 --reload

3. 客户端调用与测试

服务启动后，我们可以通过多种方式测试API：

Python客户端示例：

import requests

response = requests.post(
    "http://localhost:8000/chat",
    json={"text": "请解释量子计算的基本原理"}
)
print(response.json())

cURL测试命令：

curl -X POST "http://localhost:8000/chat" \
-H "Content-Type: application/json" \
-d '{"text":"如何学习深度学习"}'

常见响应格式：

{
  "response": "深度学习的学习路径建议从以下几个方面入手..."
}

4. 性能优化与生产部署

要让API服务达到生产级别，需要考虑以下几个优化点：

1. 批处理支持：

@app.post("/batch_chat")
async def batch_chat(queries: List[ChatRequest]):
    texts = [q.text for q in queries]
    input_ids = tokenizer(texts, padding=True).input_ids
    # 其余处理逻辑...

2. 速率限制（使用FastAPI中间件）：

from fastapi import Request
from fastapi.middleware import Middleware

async def rate_limiter(request: Request):
    # 实现限流逻辑
    pass

3. 日志监控：

import logging
logging.basicConfig(
    filename='api.log',
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

4. GPU内存管理技巧：

• 使用torch.cuda.empty_cache()定期清理缓存
• 设置max_batch_size防止内存溢出
• 考虑使用量化模型减少显存占用

5. 常见问题排查

在实际部署中可能会遇到以下典型问题：

问题1：CUDA内存不足

• 解决方案：减小max_new_tokens或使用fp16量化

model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.float16
)

问题2：响应时间过长

• 优化方法：
  • 启用use_fast版本的tokenizer
  • 使用更高效的采样策略

outputs = model.generate(
    input_ids,
    do_sample=True,
    top_k=50,
    top_p=0.95
)

问题3：并发请求处理不佳

• 改进方案：
  • 增加UVicorn工作线程数

uvicorn main:app --workers 4
  - 使用异步IO优化
```python
@app.post("/chat")
async def chat(query: ChatRequest):
    # 使用异步处理

在实际项目中，我发现最影响性能的往往是tokenizer的处理效率。通过预加载tokenizer并启用fast版本，通常可以获得20-30%的性能提升。另外，对于长时间运行的服务，建议添加健康检查端点：

@app.get("/health")
async def health_check():
    return {"status": "healthy"}