用 MLflow 管理几十个实验版本的实践经验

某次项目里，我在两周内跑了超过 50 个训练实验，超参数记在一个共享文档里，结果对比写在另一个 Excel 里。想找"上周五那个 lr=2e-4、rank=16 的实验"，翻了将近半小时才找到。更糟糕的是，对应的 checkpoint 在服务器上用时间戳命名，已经不知道对应哪次。从那之后我开始认真用 MLflow，这篇文章记录我实际用下来觉得有用的部分。

基础搭建

MLflow 有两种用法：本地存储和远程 tracking server。对于个人或小团队，本地模式足够：

# 启动本地 UI（默认监听 5000 端口）
mlflow ui --host 0.0.0.0 --port 5000

# 在代码里指定 tracking URI
import mlflow
mlflow.set_tracking_uri("http://localhost:5000")

如果是多人协作，推荐搭一个共享的 tracking server，把 artifact 存到对象存储（如 MinIO 或 S3）：

mlflow server \
  --host 0.0.0.0 \
  --port 5000 \
  --backend-store-uri postgresql://user:pass@localhost/mlflow \
  --default-artifact-root s3://your-bucket/mlflow-artifacts

命名规范

我用的实验命名格式是 {项目代号}-{任务类型}，比如 defect-v2-detection、qwen2-lora-sft。run name 则是 {日期}-{关键变量摘要}，例如 0321-lr2e4-r8-bs16。这样在 UI 里扫一眼就能知道大概是什么配置。

超参和指标管理

每次训练开始时，我会把完整的配置字典一次性 log 进去，包括模型结构、数据路径、训练参数：

import mlflow

with mlflow.start_run(run_name="0321-lr2e4-r8-bs16"):
    # 一次性记录所有超参
    mlflow.log_params({
        "model": "Qwen2-7B",
        "lora_r": 8,
        "lora_alpha": 16,
        "learning_rate": 2e-4,
        "batch_size": 16,
        "num_epochs": 3,
        "warmup_ratio": 0.05,
        "data_size": 120000,
    })

    # 训练循环里记录指标
    for step, (train_loss, eval_loss) in enumerate(train_loop()):
        mlflow.log_metrics({
            "train_loss": train_loss,
            "eval_loss": eval_loss,
        }, step=step)

    # 结束时记录最终指标
    mlflow.log_metric("best_eval_loss", best_loss)
    mlflow.log_metric("final_accuracy", final_acc)

小技巧： 把训练用的配置文件（YAML/JSON）作为 artifact 上传，这样即使代码改动，也能还原当时的完整配置：mlflow.log_artifact("config.yaml")

实验对比

MLflow UI 的 parallel coordinates（平行坐标图）是对比多个实验最直观的工具。选中多个 run，选择想对比的参数和指标，能一眼看出哪个参数对 eval loss 影响最大。

也可以用 Python API 做程序化对比，比如找出某个实验组里 eval_loss 最低的 run：

import mlflow

client = mlflow.MlflowClient()

# 找出实验下所有 run，按 eval_loss 排序
runs = client.search_runs(
    experiment_ids=["1"],
    filter_string="metrics.eval_loss < 1.5",
    order_by=["metrics.eval_loss ASC"],
    max_results=10,
)

for run in runs:
    print(
        f"Run: {run.info.run_name} | "
        f"lr={run.data.params.get('learning_rate')} | "
        f"eval_loss={run.data.metrics.get('eval_loss'):.4f}"
    )

模型 Artifact 管理

每次实验结束，我会把最优 checkpoint 和 tokenizer 一起记录进 MLflow Model Registry：

from transformers import AutoModelForCausalLM, AutoTokenizer
import mlflow.transformers

model = AutoModelForCausalLM.from_pretrained(best_checkpoint_path)
tokenizer = AutoTokenizer.from_pretrained(best_checkpoint_path)

with mlflow.start_run(run_id=current_run_id):
    mlflow.transformers.log_model(
        transformers_model={"model": model, "tokenizer": tokenizer},
        artifact_path="model",
        registered_model_name="qwen2-sft-v2",
    )

注册之后，可以在 Model Registry 里把表现最好的版本标记为 Production，部署时直接从 registry 加载，不用再记住某个服务器路径。

最佳实践小结

失败的实验也要记录：用 tag 标注失败原因（如 mlflow.set_tag("status", "diverged")），方便日后排除相同的坏配置。
run name 要人类可读：时间戳 + 关键变量，日后扫一眼能大概知道是什么实验。
log artifact 不嫌多：除了 checkpoint，把训练曲线图、验证集预测样例一并存进去，review 实验时省时间。
定期清理：明确失败或被覆盖的实验定期归档，保持 UI 整洁，不然实验一多照样找不到。

MLflow 上手不难，但要用好需要建立一套自己的规范并坚持执行。如果你在用 W&B 或者自研的实验管理工具，思路是一样的——关键是让每次实验的参数、结果、产物都可以被找到和复现。有问题欢迎交流。