如何从Vertex AI上传并注册大语言模型至Databricks？

一、操作思路的正确性判断 #

你的核心方向是对的，但缺少关键环节：不能直接上传文件后调用mlflow.register_model完成注册。MLFlow对LLM的注册需要结合适配的模型flavor（比如mlflow.transformers）封装模型，而非仅上传原始文件。

二、需要上传的文件 #

你GCS存储桶里的全部文件都需要同步到DBFS，这些文件共同构成了完整的微调后LLM与分词器：

• 分词器相关：added_tokens.json、special_tokens_map.json、tokenizer_config.json、vocab.json、merges.txt（负责文本编码解码）
• 模型配置相关：config.json、generation_config.json（定义模型结构与推理生成规则）
• 模型权重与训练参数：pytorch_model.bin（微调后的核心权重）、training_args.bin（训练过程参数记录）

三、完整注册与加载步骤 #

1. 同步GCS文件到DBFS
   用Databricks Notebook命令或CLI完成目录同步：

   dbutils.fs.cp("gs://your-bucket/model-dir/*", "dbfs:/databricks/model-store/llm-finetuned/", recurse=True)
   

2. 用MLFlow封装并注册模型（关键步骤）
   通过mlflow.transformers将模型与分词器打包，自动完成注册：

   import mlflow
   from transformers import AutoModelForCausalLM, AutoTokenizer
   
   # 从DBFS加载模型与分词器
   local_model_path = "/dbfs/databricks/model-store/llm-finetuned/"
   model = AutoModelForCausalLM.from_pretrained(local_model_path)
   tokenizer = AutoTokenizer.from_pretrained(local_model_path)
   
   # 启动MLFlow Run并完成模型注册
   with mlflow.start_run() as run:
       mlflow.transformers.log_model(
           transformers_model={"model": model, "tokenizer": tokenizer},
           artifact_path="llm-finetuned-model",
           registered_model_name="vertex-finetuned-llm"
       )
   

   上述代码会直接将模型注册到Databricks模型注册表，无需单独调用mlflow.register_model；若你坚持用该命令，可使用runs:/{run_id}/llm-finetuned-model作为模型路径。

3. 加载注册模型执行推理
   从模型注册表加载模型后即可进行推理：

   import mlflow
   
   # 加载指定版本的注册模型
   loaded_llm = mlflow.transformers.load_model(
       "models:/vertex-finetuned-llm/1"  # 1为模型版本号，可按需修改
   )
   
   # 执行推理示例
   prompt = "请解释大语言模型微调的核心意义"
   inputs = loaded_llm["tokenizer"](prompt, return_tensors="pt")
   outputs = loaded_llm["model"].generate(**inputs, max_new_tokens=150)
   print(loaded_llm["tokenizer"].decode(outputs[0], skip_special_tokens=True))
   

四、注意事项 #

• 确保Databricks集群已安装transformers、mlflow、torch等依赖库
• 大模型建议使用GPU集群，加速模型加载与推理过程
• 检查GCS到DBFS的同步权限，以及MLFlow模型注册表的访问权限

内容的提问来源于stack exchange，提问作者sanminchui