试验跟踪和可观测性

AI 运行时原生集成了 MLflow，用于实验跟踪，并内置了 GPU 资源面板，用于监控利用率、内存占用和温度。 使用 MLflow 记录指标和运行，在笔记本和 MLflow UI 中查看训练输出，将模型检查点保存到 Unity 目录卷，并在代码运行时跟踪 GPU 运行状况。

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 时，每次调用 .distributed() 都会自动创建 MLflow 试验运行。 如果在当前活动的 MLflow 运行中调用，则会在该活动父运行下创建一个嵌套的子运行。

  import mlflow
  
  with mlflow.start_run() as outer_run:
      ...
      run_train.distributed()  # creates a nested child run under outer_run
  

• 若要自定义所使用的.distributed()试验，请调用mlflow.set_experiment()之前调用.distributed()或设置MLFLOW_EXPERIMENT_NAME环境变量。 默认试验名称为 /Users/{WORKSPACE_USER}/{notebook-name}. 始终使用绝对路径。

  import mlflow
  mlflow.set_experiment("/Users/<username>/my-experiment")
  run_train.distributed()
  

  Alternatively:

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

• 若要恢复以前的 MLflow 运行，请使用 mlflow.start_run(run_id="<previous-run-id>")。

• 若要使用 .distributed() 恢复之前的 MLflow 运行，请在调用它之前设置 MLFLOW_RUN_ID：

  os.environ["MLFLOW_RUN_ID"] = "<previous-run-id>"
  run_train.distributed()
  

• 将 step 中的 MLFlowLogger 参数设置为合理的批次数。 MLflow 的指标步数上限为 1000 万，因此，在大规模训练运行中记录每一个批次时，可能会达到这一上限。 请参阅资源限制。

查看日志

• 笔记本输出：训练代码的标准输出和错误显示在笔记本单元输出中。
• MLflow 日志：MLflow 试验 UI 显示训练指标、参数和项目。

模型检查点

对于分布式训练，请将模型检查点异步保存到 Unity 目录卷，该卷提供与其他 Unity 目录对象相同的治理。 将来自 UCVolumeWriter 包的 UCVolumeReader 和 serverless_gpu.data 与 Torch Distributed Checkpoint (DCP) API 结合使用。 这些存储后端通过快速本地目录（/tmp，在无服务器 GPU 节点上由 NVMe 支持）暂存所有 I/O，并上传到或从 Unity Catalog 卷下载，这比直接将检查点分片写入 FUSE 装载要快。 保留元数据原子性：编写器仅在其数据分片完成上传后发布 .metadata 文件。

检查点的频率要足够高，以限制中断后丢失的工作，但又不能太高，以免 I/O 开销减慢训练速度。 以每隔 30 分钟到 1 小时设置一个检查点为目标，并根据步骤时间和检查点大小调整间隔。

要在训练继续进行时在后台上传检查点，请将 UCVolumeWriter 作为 storage_writer 传递给 dcp.async_save。 异步保存需要进程组中有一个 CPU 后端，因此请使用 torch.distributed.init_process_group(backend="cpu:gloo,cuda:nccl", ...) 对其进行初始化：

import torch.distributed.checkpoint as dcp
from serverless_gpu.data import UCVolumeWriter

checkpoint_path = "/Volumes/my_catalog/my_schema/model/checkpoints"
writer = UCVolumeWriter(checkpoint_path)

future = dcp.async_save(state_dict, storage_writer=writer)
# ...continue training...
future.result()  # blocks until the upload lands on the UC volume

使用 UCVolumeReader 加载检查点：

from serverless_gpu.data import UCVolumeReader

reader = UCVolumeReader(checkpoint_path)
dcp.load(state_dict, storage_reader=reader)

数据管道检查点机制

模型检查点会保存模型和优化器的状态，但不包含数据管道在数据集中的位置，因此恢复后的运行无法快进到上次停止时的确切样本。 在恢复方式中要考虑到这一点：从时期边界重新开始，或者在自己的训练状态中跟踪已处理的样本或分片，以便在恢复时跳过它们。

监视 GPU 资源

在 AI 运行时上运行代码时，使用 GPU 资源 窗格监视 GPU 运行状况和利用率。 该窗格支持单节点工作负荷和多节点工作负荷。

若要打开窗格，请将笔记本连接到 AI 运行时，然后单击右侧窗格中的 GPU 资源。

该窗格显示每个 GPU 的以下指标：

• GPU 利用率百分比
• GPU 内存使用情况
• 温度

该窗格每 10 秒轮询一次指标，最多保留 2 小时的历史记录。 单击“ 刷新 以立即提取最新值。 处于非活动状态 5 分钟后，窗格将暂停;重新打开它以恢复监视。

多用户协作

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

Azure Databricks 中的全局限制

请参阅资源限制。