Unity Catalog 联动 SageMaker AI：跨服务治理下的 LLM 微调实战

预计阅读时间：10 分钟

大模型微调从来不只是"跑个训练脚本"的事。数据从哪来、谁有权访问、训练结果归谁管、血缘怎么追踪——这些问题在企业环境里往往比模型本身更棘手。Databricks Unity Catalog 和 Amazon SageMaker AI 各有强项：前者管数据治理与血缘，后者管训练编排与弹性算力。把它们串起来，就能在不动现有服务架构的前提下，把治理链条从数据源一直延伸到模型注册表。

下面拆解这套跨服务微调流程的关键环节，并给出可直接改造的实操示例。

1. 整体流程：四个阶段，一条治理链

这套方案的核心思路是不把治理和训练割裂，而是让 Unity Catalog 作为唯一治理锚点，贯穿全流程：

阶段	执行环境	治理动作
数据准备	EMR Serverless	从 Unity Catalog 读取受控数据集，ETL 后写回
微调训练	SageMaker AI Training Job	拉取预处理数据，运行 Ministral-3-3B-Instruct 微调
模型评估	SageMaker AI（可选本地）	验证微调效果，记录指标
模型注册	Unity Catalog（通过 Databricks）	将训练产物注册回 Catalog，补全血缘

关键点：EMR Serverless 做预处理而非 Databricks Spark，是因为 EMR Serverless 可以在 SageMaker 同一 VPC 内启动，避免数据跨区域搬运；同时它仍然通过 Databricks JDBC/REST 接口读取 Unity Catalog 的表，治理不中断。

2. 数据访问：安全读出，治理不越界

Unity Catalog 的表权限（SELECT / READ_VOLUME）决定了谁能拿数据做微调。在 EMR Serverless 里读 Unity Catalog 数据，有两种常见方式：

JDBC 连接：EMR Spark 通过 Databricks JDBC Driver 直连 Unity Catalog，适合读结构化表。
REST API + 临时凭证：对 Volume 中的文件（如 JSON 训练集），调用 /api/2.1/unity-catalog/volumes/{volume_id}/download 获取临时下载 URL。

下面是 JDBC 方式的 PySpark 代码片段，可在 EMR Serverless 的作业脚本中使用：

# emr_preprocess.py — 在 EMR Serverless 上读取 Unity Catalog 表并做训练集预处理

from pyspark.sql import SparkSession

# JDBC 连接参数（建议从 Secrets Manager 或环境变量注入，不要硬编码）
jdbc_url = "jdbc:spark://<databricks-instance>:443/default;transportMode=http;ssl=1;httpPath=/sql/1.0/warehouses/<warehouse-id>;AuthMech=3;UID=token;PWD=<personal-access-token>"
source_table = "main.fine_tuning.raw_training_data"

spark = SparkSession.builder \
    .appName("UC-Preprocess") \
    .config("spark.jars", "/path/to/databricks-jdbc.jar") \
    .getOrCreate()

# 从 Unity Catalog 读取受控数据
raw_df = spark.read \
    .format("jdbc") \
    .option("url", jdbc_url) \
    .option("dbtable", source_table) \
    .option("fetchsize", "10000") \
    .load()

# 简单预处理：过滤空行、拼接 prompt + response 为微调格式
import json

processed = raw_df.filter("prompt IS NOT NULL AND response IS NOT NULL") \
    .selectExpr("prompt", "response")

# 写成 JSONL 文件（微调输入格式），存到 S3 供 SageMaker 拉取
processed_rdd = processed.rdd.map(
    lambda row: json.dumps({"prompt": row.prompt, "response": row.response})
)
processed_rdd.saveAsTextFile("s3://my-finetune-bucket/preprocessed/ministral-train-jsonl/")

print(f"预处理完成，共 {processed.count()} 条有效样本")

注意：Personal Access Token 应通过 AWS Secrets Manager 注入到 EMR 作业环境，不要写在代码里。生产环境建议用 Service Principal + OAuth，而非个人 Token。

3. 微调训练：SageMaker AI 跑 Ministral-3B

数据落到 S3 后，SageMaker AI Training Job 直接拉取。Ministral-3-3B-Instruct 是 Mistral 系列的小模型，3B 参数量在 SageMaker 的单卡 GPU（如 ml.g5.xlarge）上即可微调，成本可控。

下面是用 SageMaker Python SDK 启动微调的示例：

# sagemaker_finetune.py — 在 SageMaker AI 上微调 Ministral-3B

import sagemaker
from sagemaker.huggingface import HuggingFace

role = sagemaker.get_execution_role()
sagemaker_session = sagemaker.Session()

# 微调超参 — 根据数据量调整
hyperparameters = {
    "model_id": "mistralai/Ministral-3B-Instruct",
    "training_data_path": "/opt/ml/input/data/train",
    "epochs": 3,
    "per_device_train_batch_size": 4,
    "learning_rate": 2e-5,
    "lr_scheduler_type": "cosine",
    "max_seq_length": 2048,
    "lora_r": 16,           # LoRA rank，降低显存占用
    "lora_alpha": 32,
    "lora_dropout": 0.05,
    "bf16": True,            # 使用 bf16 混合精度
}

# 训练数据 S3 位置（与 EMR 预处理输出对齐）
train_data_s3_uri = "s3://my-finetune-bucket/preprocessed/ministral-train-jsonl/"

estimator = HuggingFace(
    entry_point="train.py",          # 你的训练脚本，需实现 LoRA 微调逻辑
    source_dir="src/",               # 本地训练代码目录
    instance_type="ml.g5.2xlarge",   # 单卡 A10G，足够跑 3B LoRA
    instance_count=1,
    role=role,
    transformers_version="4.45",
    pytorch_version="2.4",
    py_version="py311",
    hyperparameters=hyperparameters,
    sagemaker_session=sagemaker_session,
)

estimator.fit({"train": train_data_s3_uri})

print(f"训练完成，模型 S3 位置: {estimator.model_data}")

src/train.py 需自行实现，核心逻辑是加载 Ministral-3B → 应用 LoRA → 读 JSONL → 训练 → 合并权重 → 保存。可参考 HuggingFace PEFT 库的 LoraConfig + SFTTrainer 模式，此处不展开。

4. 血缘回写：把训练产物注册回 Unity Catalog

微调完成后，模型权重在 S3。下一步是把这条"从哪个表的数据 → 经什么参数 → 产出哪个模型"的血缘关系写回 Unity Catalog，让后续审计和模型使用都有据可查。

Databricks 提供了 REST API 注册模型到 Unity Catalog 的 Model Registry：

# 在 SageMaker 训练完成后，通过 Databricks REST API 注册模型版本

# 1. 先确保 Unity Catalog 中已创建模型注册表（可在 Databricks UI 或 API 中操作）
#    CREATE MODEL main.fine_tuning.ministral_3b_finetuned;

# 2. 注册新版本，关联 S3 模型路径
curl -X POST \
  "https://<databricks-instance>/api/2.1/unity-catalog/models/main.fine_tuning.ministral_3b_finetuned/versions" \
  -H "Authorization: Bearer <databricks-token>" \
  -H "Content-Type: application/json" \
  -d '{
    "source": "s3://my-finetune-bucket/output/ministral-3b-lora-merged/",
    "comment": "LoRA fine-tune on main.fine_tuning.raw_training_data, epochs=3, lr=2e-5, lora_r=16",
    "run_id": "<mlflow-run-id-if-available>"
  }'

血缘的关键信息藏在 comment 字段里——记录了源数据表名、训练参数、S3 路径。更严谨的做法是用 MLflow Tracking 在训练过程中自动记录参数和指标，然后通过 mlflow.register_model 直接注册到 Unity Catalog，血缘会自动关联到 MLflow Run。

5. 落地时的取舍与风险

这套方案不是万能的，有几个现实边界需要正视：

网络与权限是最容易卡住的地方。 EMR Serverless 和 SageMaker Training Job 需要在同一 VPC 内运行，且要有到 Databricks JDBC 端口和 S3 的网络通路。跨账户场景下，S3 Bucket Policy 和 VPC Peering 都要提前配好，调试成本不低。

EMR Serverless 的冷启动。 首次提交作业时，EMR Serverless 的 Application 启动需要 1-3 分钟。如果微调是频繁触发（比如每天跑），考虑保持 Application 活跃或改用 Databricks 自有 Spark 做预处理，牺牲一点网络隔离换取启动速度。

模型注册的粒度。 Unity Catalog Model Registry 目前对"哪个数据表的哪次快照产出了这个模型"的自动血缘追踪还不够细。手动在 comment 里写源表名是务实的过渡方案，但长期应依赖 MLflow Run 的自动血缘能力。

LoRA 合并后的部署。 微调产出的是合并后的完整权重（或 LoRA adapter），部署时需要选择 SageMaker Endpoint 或 Databricks Model Serving。两者都能从 Unity Catalog 拉取模型，但部署配置不同，需要提前规划。

快速检查清单，上线前逐项确认：

✅ Unity Catalog 表权限已授予微调使用的 Service Principal
✅ EMR Serverless Application 与 SageMaker 在同一 VPC
✅ S3 Bucket Policy 允许 EMR 写、SageMaker 读
✅ Databricks JDBC/REST 凭证通过 Secrets Manager 注入，未硬编码
✅ 训练脚本已实现 LoRA + bf16，单卡显存可承载
✅ 模型注册回 Unity Catalog 时，comment 包含源表名与关键超参
✅ 部署路径（SageMaker Endpoint 或 Databricks Serving）已选定

跨服务治理的核心不是"哪个平台更好"，而是让治理锚点唯一、让数据不搬家、让血缘不断链。Unity Catalog + SageMaker AI 的组合，在"各干各的擅长事"的前提下，把这条链串了起来。