如何从 Vertex AI 上传并注册模型到 Databricks？

Question

我在 Vertex AI 中微调了 LLM，但我想将模型注册并加载到 Databricks 中，以便我可以从那里运行推理。目前我在 GCS 存储桶中有这些文件：

added_tokens.json
config.json
generation_config.json
merges.txt
pytorch_model.bin
special_tokens_map.json
tokenizer_config.json
training_args.bin
vocab.json

离开此链接，我可以将模型本地保存到 DBFS，但它没有提及我需要上传哪些文件。上传它们后，我假设我可以使用以下代码在模型注册表中注册模型：

mlflow.register_model("runs:/{run_id}/{model-path}", "{registered-model-name}")

我的做法正确吗？我看到了另外两个关于此的问题，但他们没有完全回答我的问题。

Answer 1

我认为你的方法是正确的。

我会将您此处的所有文件上传到 DBFS 中的文件夹中，然后尝试在 Databricks 笔记本中加载模型。

class Net(nn.Module):
// Your Model for which you want to load parameters 

model = Net()
torch.optim.SGD(lr=0.001) #According to your own Configuration.
checkpoint = torch.load(pytorch_model)
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['opt']) 

完成后，使用 mlflow.pytorch.log() 将模型保存到 MLFlow 实验。然后你可以使用

mlflow.register()

或在登录模型时传递参数

registered_model_name

直接注册它。

您可能需要创建一个 MLFlow 自定义类，具体取决于您的模型的复杂程度，这对于 LLM 来说通常是这种情况，因为他们需要额外的组件才能工作。