本文将深入探讨如何利用 Google Cloud 的 Vertex AI 和自定义 Handler 部署 PyTorch 大模型。我们将以波兰语语言模型 PLLuM 为例，详细讲解从模型准备、本地测试到最终部署的完整流程。当模型需要特殊预处理、非常规输出格式或标准 serving container 无法满足的独特逻辑时，本文提供的解决方案尤为重要。
1. Vertex AI 与自定义 Handler：掌控推理流程
Vertex AI 作为 Google Cloud 的机器学习平台，提供了强大的模型部署和管理能力。虽然 Model Garden 提供了一键部署许多开源模型（包括 Anthropic、Meta 和 Hugging Face 的模型）的能力，但对于需要自定义逻辑的模型，其灵活性受到限制。这时，自定义 Handler 就派上了用场。
自定义 Handler 是一个 Python 脚本（通常命名为 handler.py），它可以完全控制模型的加载、请求处理和预测结果格式化。使用预构建的 Hugging Face 容器结合自定义 Handler，能让开发者更好地控制推理流程，解决因模型特殊性而带来的部署挑战。
handler.py 需要实现 EndpointHandler 类，该类包含两个关键方法：
__init__: 在模型加载时调用，负责加载模型和 tokenizer 等必要资源。__call__: 在每次收到预测请求时调用，包含预处理、预测和后处理逻辑。
通过自定义 Handler，几乎可以部署任何模型，无论其需求多么特殊。
2. PLLuM 模型：一个需要定制化部署的案例
PLLuM (Polish Language Understanding Model) 是一个强大的波兰语语言模型。它在波兰越来越受欢迎，但无法直接通过 Model Garden 访问。更重要的是，由于其使用自定义 tokenizer，在调用生成函数之前需要实现特定的编码/解码逻辑，因此无法在标准的 Hugging Face 容器中直接工作。
PLLuM 模型很好地诠释了为什么要使用自定义 Handler，通过自定义 Handler，我们可以将自定义 tokenizer 集成到推理流程中，从而成功部署 PLLuM 模型。
3. 环境配置：部署前的准备
在开始部署之前，需要配置相应的环境：
- Google Cloud 项目: 确保拥有启用了 Vertex AI 和 Artifact Registry API 的 Google Cloud 项目。
- Cloud Storage Bucket: 创建一个新的 Cloud Storage bucket 来存储模型文件。
- 权限: 确保您的帐户拥有以下 IAM 角色：Vertex AI User、Artifact Registry Reader 和 Storage Object Admin。
- Docker (可选但推荐): 在部署到云端之前，为了在本地测试模型容器，建议安装并运行 Docker。
- 所需库: 使用 pip 安装必要的 Python 库：
bash
pip install - upgrade - user - quiet 'torch' 'torchvision' 'torchaudio'
pip install - upgrade - user - quiet 'transformers' 'accelerate>=0.26.0'
pip install - upgrade - user - quiet 'google-cloud-aiplatform[prediction]' 'crcmod' 'etils'
完成环境配置后，导入库并初始化 Vertex AI SDK。 示例代码如下:
import json
import torch
import vertexai
from etils import epath
from google.cloud import aiplatform
from google.cloud.aiplatform import Endpoint, Model
from google.cloud.aiplatform.prediction import LocalModel
# 设置你的项目、位置和存储桶详细信息
PROJECT_ID = "your-gcp-project-id"
LOCATION = "your-gcp-location" # 例如: "us-central1"
BUCKET_URI = "gs://your-gcs-bucket-name"
# 初始化 Vertex AI SDK
vertexai.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)
4. 构建 PLLuM 的自定义 Handler
现在我们来构建 PLLuM 模型的自定义 Handler。目标是创建一个简单的文本生成 endpoint。
首先，创建名为 handler.py 的文件，并添加以下代码：
from typing import Any, Dict, List
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import base64
from io import BytesIO
import logging
import sys
# 配置日志输出到 stdout
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
handlers=[logging.StreamHandler(sys.stdout)])
logger = logging.getLogger('huggingface_inference_toolkit')
class EndpointHandler:
def __init__(
self,
model_dir: str = '/opt/huggingface/model',
**kwargs: Any,
) -> None:
self.processor = AutoTokenizer.from_pretrained(model_dir)
self.model = AutoModelForCausalLM.from_pretrained(
model_dir,
torch_dtype=torch.bfloat16,
device_map="auto" # 自动将模型层放置在可用的设备上
).eval()
def __call__(self, data: Dict[str, Any]) -> Dict[str, List[Any]]:
logger.info("Processing new request")
predictions = []
for instance in data['instances']:
logger.info(f"Processing instance: {instance.get('prompt', '')[:100]}...")
if "prompt" not in instance:
error_msg = "Missing prompt in the request body"
logger.info(error_msg)
return {"error": "Missing prompt in the request body"}
inputs = self.processor(
instance["prompt"], return_tensors="pt", return_token_type_ids=False
).to(self.model.device)
input_len = inputs["input_ids"].shape[-1]
logger.info(f"Input processed, length: {input_len}")
with torch.inference_mode():
generation_kwargs = data.get(
"generation_kwargs", {
"max_new_tokens": 100,
"do_sample": False,
"top_k": 50,
"top_p": 0.9,
"temperature": 0.7
}
)
logger.info(f"Generation kwargs: {generation_kwargs}")
generation = self.model.generate(**inputs, **generation_kwargs)
generation = generation[0][input_len:]
response = self.processor.decode(generation, skip_special_tokens=True)
logger.info(f"Generated response: {response[:100]}...")
predictions.append(response)
logger.info(f"Successfully processed {len(predictions)} instances")
return {"predictions": predictions}
__init__ 函数负责加载模型和 tokenizer。__call__ 方法实现了推理逻辑，包括分词、生成和解码。 注意到 __call__ 方法接收一个字典，并且应该以处理多个实例的方式编写。这允许用户在单个请求中发送多个提示。
5. 模型准备：上传到 Cloud Storage
Vertex AI 需要将所有模型工件（模型权重、配置和 handler.py）存储在 Google Cloud Storage 上的一个位置。
- 创建一个包含模型文件和
handler.py的本地目录。 - 将整个目录上传到 GCS bucket。
gsutil -m cp -r /path/to/your/local/model_directory/* gs://your-gcs-bucket-name/model/
为了提高上传速度，建议设置并行复合上传阈值：
gsutil -o GSUtil:parallel_composite_upload_threshold=150M -m cp -r /path/to/your/local/model_directory/* gs://your-gcs-bucket-name/model/
6. 本地测试：使用 LocalModel 模拟云环境
在将模型部署到 GPU 加速的 endpoint 之前，使用 Vertex AI SDK 的 LocalModel 功能在本地机器上模拟云环境，进行快速测试非常重要。
- 获取 CUDA 设备名称:
def get_cuda_device_names():
"""获取 NVIDIA GPU 列表的函数"""
if not torch.cuda.is_available():
return None
return [str(i) for i in range(torch.cuda.device_count())]
- 创建 LocalModel 实例:
local_pllum_model = LocalModel(
serving_container_image_uri="us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-pytorch-inference-cu121.2-3.transformers.4-48.ubuntu2204.py311",
serving_container_ports=[5000],
)
- 创建 LocalEndpoint 实例:
model_uri = epath.Path(BUCKET_URI) / "model"
local_pllum_endpoint = local_pllum_model.deploy_to_local_endpoint(
artifact_uri=str(model_uri), gpu_device_ids=get_cuda_device_names())
local_pllum_endpoint.serve()
- 生成预测:
# EN:"Write a short poem about spring."
prompt = "Napisz krótki wiersz o wiośnie." # @param {type: "string"}
prediction_request = {
"instances": [
{
"prompt": prompt,
"generation_kwargs": {"max_new_tokens": 50, "do_sample": True},
}
]}
vertex_prediction_request = json.dumps(prediction_request)
vertex_prediction_response = local_pllum_endpoint.predict(
request=vertex_prediction_request, headers={"Content-Type": "application/json"})
print(vertex_prediction_response.json()["predictions"])
通过本地测试，可以快速发现代码中的错误，避免长时间等待云端部署结果。
7. 部署到 Vertex AI Endpoint
完成本地测试后，即可将模型部署到 Vertex AI Endpoint。
- 在 Vertex AI Model Registry 中注册模型:
model_uri = epath.Path(BUCKET_URI) / "model"
model = Model.upload(
display_name="cyfragovpl--pllum-12b-it",
artifact_uri=str(model_uri),
serving_container_image_uri="us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-pytorch-inference-cu121.2-3.transformers.4-48.ubuntu2204.py311",
serving_container_ports=[8080],
)
model.wait()
- 将模型部署到 Endpoint:
deployed_model = model.deploy(
endpoint=Endpoint.create(display_name="cyfragovpl--pllum-12b-it-endpoint"),
machine_type="g2-standard-8",
accelerator_type="NVIDIA_L4",
accelerator_count=1,
)
此步骤将花费大约 15-25 分钟。完成后，将拥有一个完全托管的、可扩展的 HTTP endpoint 用于模型。 选择合适的 machine_type 和 accelerator_type 至关重要，需要根据模型的规模和推理需求进行调整。例如，对于更大的模型，可能需要 a2-highgpu-1g 或 a2-ultragpu-1g 这样的机器类型，并配备多个 NVIDIA A100 GPU。
8. 获取实时预测结果
部署完成后，可以使用 Vertex AI SDK、cURL 命令或任何 HTTP 客户端向 endpoint 发送请求。
使用 Vertex AI 的 Python SDK 是最直接的方法：
# EN:"Write a short poem about spring."
prompt = "Napisz krótki wiersz o wiośnie." # @param {type: "string"}
prediction_request = {
"instances": [
{
"prompt": prompt,
"generation_kwargs": {"max_new_tokens": 50, "do_sample": True},
}
]}
prediction = deployed_model.predict(instances=prediction_request["instances"])
print(prediction)
输出将包含从 PLLuM 模型生成的文本，该文本通过自定义 API endpoint 实时提供 🎉。
9. 优化与监控
模型成功部署后，还需要关注性能优化和监控：
- 模型监控: Vertex AI 提供模型监控功能，可以实时监控模型的预测性能，检测潜在的偏差或退化。
- 性能优化: 可以通过模型量化、知识蒸馏等技术来优化模型的推理速度和资源消耗。
- Endpoint 弹性伸缩: 根据流量需求，配置 Endpoint 的自动伸缩，确保服务可用性和响应速度。
结论
通过本文，我们了解了如何利用 Vertex AI 和自定义 Handler 将具有特殊需求的开源 Hugging Face 模型（如 PLLuM）转化为 Google Cloud 上的强大、可扩展的 API。掌握了这项技术，你就能将各种各样的模型投入生产，通过创建简单的自定义 Handler，根据你的确切需求定制推理过程。
自定义 Handler 不仅仅是一种技术，更是一种思维方式。它鼓励我们深入理解模型的内部机制，从而更好地控制模型的部署和推理过程。随着大模型技术的不断发展，自定义 Handler 将在解决复杂模型部署问题中发挥越来越重要的作用。