通过

试验跟踪和可观测性

重要

用于单节点任务的 AI 运行时为 公共预览版。 多 GPU 工作负载的分布式训练 API 仍为 Beta 版。

本文介绍如何使用 MLflow、监视 GPU 运行状况、查看日志和管理 AI 运行时上的模型检查点。

MLflow 集成

AI 运行时以本机方式与 MLflow 集成，用于试验跟踪、模型日志记录和指标可视化。

设置建议：

  • 将 MLflow 升级到版本 3.7 或更高版本，并遵循 深度学习工作流模式。

  • 为 PyTorch Lightning 启用自动记录：

    import mlflow
mlflow.pytorch.autolog()

  • 通过将模型训练代码封装在 mlflow.start_run() API 范围内，来自定义 MLflow 运行名称。 这样你就可以控制运行名称，并能够从以前的运行中重启。可以使用run_name参数在mlflow.start_run(run_name="your-custom-name")中或在支持 MLflow 的第三方库（例如 Hugging Face Transformers）中自定义运行名称。 否则，默认运行名称为 jobTaskRun-xxxxx。

    from transformers import TrainingArguments
args = TrainingArguments(
    report_to="mlflow",
    run_name="llama7b-sft-lr3e5",  # <-- MLflow run name
    logging_steps=50,
)

  • 无服务器 GPU API 会自动启动具有默认名称 /Users/{WORKSPACE_USER}/{get_notebook_name()}的 MLflow 试验。 用户可以使用环境变量 MLFLOW_EXPERIMENT_NAME覆盖它。始终对 MLFLOW_EXPERIMENT_NAME 环境变量使用绝对路径：

    import os
os.environ["MLFLOW_EXPERIMENT_NAME"] = "/Users/<username>/my-experiment"

  • 继续先前的训练，方法是使用之前运行的MLFLOW_RUN_ID设置：

    mlflow.start_run(run_id="<previous-run-id>")

  • 将 step 参数 MLFlowLogger 设置为合理的批号。 MLflow 的度量步骤限制为 1000 万，在大型训练运行中记录每一个批次可能达到此限制。 请参阅资源限制。

查看日志

  • 笔记本输出 - 训练代码的标准输出和错误显示在笔记本单元格输出中。
  • 驱动程序日志 - 可通过计算面板进行访问，以便调试启动问题、环境设置问题和运行时错误。
  • MLflow 日志 - 可在 MLflow 试验 UI 中查看训练指标、参数和项目。

模型检查点

将模型检查点保存到 Unity 目录卷，该卷提供与其他 Unity 目录对象相同的治理。 使用以下路径格式从 Databricks 笔记本中引用卷内的文件：

/Volumes/<catalog>/<schema>/<volume>/<path>/<file-name>

将检查点保存到卷中，与保存到本地存储的方式相同。

以下示例演示如何将 PyTorch 检查点写入 Unity 目录卷：

import torch

checkpoint = {
    "epoch": epoch,  # last finished epoch
    "model_state_dict": model.state_dict(),  # weights & buffers
    "optimizer_state_dict": optimizer.state_dict(),  # optimizer state
    "loss": loss,  # optional current loss
    "metrics": {"val_acc": val_acc},  # optional metrics
    # Add scheduler state, RNG state, and other metadata as needed.
}
checkpoint_path = "/Volumes/my_catalog/my_schema/model/checkpoints/ckpt-0001.pt"
torch.save(checkpoint, checkpoint_path)

此方法也适用于分布式检查点。 以下示例显示了使用 Torch 分布式检查点 API 的分布式模型检查点：

import torch.distributed.checkpoint as dcp

def save_checkpoint(self, checkpoint_path):
    state_dict = self.get_state_dict(model, optimizer)
    dcp.save(state_dict, checkpoint_id=checkpoint_path)

trainer.save_checkpoint("/Volumes/my_catalog/my_schema/model/checkpoints")

多用户协作

  • 为了确保所有用户都可以访问共享代码（例如帮助程序模块或环境 YAML 文件），请将它们存储在其中 /Workspace/Shared ，而不是用户特定的文件夹，例如 /Workspace/Users/<your_email>/。
  • 对于处于活动开发中的代码，请使用特定于用户的文件夹中 /Workspace/Users/<your_email>/ 的 Git 文件夹并推送到远程 Git 存储库。 这样，多个用户就可以拥有特定于用户的克隆和分支，同时仍使用远程 Git 存储库进行版本控制。 请参阅有关在 Databricks 上使用 Git 的 最佳做法 。
  • 协作者可以 共享和评论 笔记本。

Databricks 中的全局限制

请参阅资源限制。

  • Last updated on