實驗追蹤與可觀測性

AI Runtime 與 MLflow 整合,提供實驗追蹤、可觀察性及除錯工具。 MLflow 提供一個單一空間來存放執行的參數、指標、GPU 系統指標、日誌和成品,無論你是在筆記本中開發、用 databricks air CLI 提交工作負載,或是排程工作。

MLflow 提供的功能

MLflow 實驗會將相關的訓練執行分組,每次執行都會提供訓練配置與結果的持久記錄。

  • 指標與參數:記錄訓練損失、評估指標、學習率與超參數,然後跨次跑比較。
  • 系統指標:在 系統指標 標籤中查看管理運行的 GPU、CPU 與記憶體使用率。
  • 日誌檢視:在該執行的 日誌 索引標籤中查看由作業支援的工作負載的驅動程式輸出。請參閱 在 MLflow 日誌標籤中檢視日誌。
  • 產物與模型:與此次執行一起儲存模型檔案、配置及其他輸出。 對於大型或持久性成品,請使用 Unity Catalog volume。
  • 分享與協作:授與實驗的存取權限,讓其他工作區使用者能查看並比較該實驗的各次執行。 請參閱 使用 MLflow 實驗組織訓練回合。
  • 框架整合:使用 MLflow 與 Hugging Face Transformers、PyTorch Lightning 及其他支援函式庫整合。

對於由作業支援的工作負載,例如以 databricks air run 提交的工作負載,Azure Databricks 工作執行會補充 MLflow,記錄執行狀態、運算、嘗試次數及驅動程式輸出。

如何建立 MLflow 執行紀錄

你如何開始或加入 MLflow 執行,取決於你如何開始工作負載:

你如何執行工作負載 MLflow run 的啟動方式 如何選擇實驗項目
Databricks CLI(databricks air run) 每次嘗試都會自動產生一次執行。 在工作負載 YAML 中設定 experiment_name。
搭載 Databricks AI Runtime SDK 的筆記本(.distributed()) 每次呼叫都會自動建立一次 run。 請在MLFLOW_EXPERIMENT_NAME之前致電mlflow.set_experiment(),或設定.distributed()。
程式碼直接在附帶的筆記本中執行 不會自動建立任何執行作業。 使用 mlflow.start_run() 或 autologging。 在筆記本中呼叫 mlflow.set_experiment()。
已排程 ai_runtime_task 該任務會自動產生執行。 在任務組態中設定 experiment。 請參見 AI 執行時任務。

當你的程式碼開始或加入正確的執行後,使用 MLflow 追蹤 API 來記錄自訂的訓練資料。 對於支援的框架,你可以使用自動記錄,而非新增個別的日誌呼叫。 關於深度學習模式,請參見 MLflow 3 深度學習工作流程。

在配置 MLflow 追蹤時,請採用以下做法:

  • 設定 step 為有意義的值,例如全域批次或紀元。 請按間隔記錄,而不是每個批次都記錄。 MLflow 限制每次執行的度量步數。 請參閱資源限制。
  • 對於分散式訓練,只由 rank 0 記錄一個指標,讓每個值只會被記錄一次。
  • 對於 Notebook 工作流程,請使用絕對實驗路徑,例如 /Users/<username>/my-experiment 或 /Workspace/Shared/<team>/my-experiment。 把你想分享的實驗放到共用工作區資料夾裡。
  • 將大型且耐用的檢查點儲存在 Unity Catalog 磁碟區中。 關於檢查點模式,請參見 「提升 AI 執行時的訓練效能與韌性」。

Note

本頁範例請使用 MLflow 3.7 及以上版本。

追蹤 Databricks CLI 工作負載

當你提交含有 databricks air run 的工作負載時,AI Runtime 會為執行細節建立 Azure Databricks 工作執行,以及為實驗資料建立 MLflow 執行。 原始提交與重試共用一個工作執行,而每次嘗試會在同一實驗中建立獨立的 MLflow 執行,方便比較各次嘗試。

GPU、CPU 與記憶體系統的指標會自動為每次受管理的執行擷取,並顯示在 MLflow 系統的指標 標籤中。

MLflow 執行時的系統指標分頁,顯示 GPU、CPU 和記憶體的指標。

設定並記錄執行

使用工作負載 YAML 選擇實驗,並可選擇設定執行名稱和產物位置。 以下範例展示了所需的 MLflow 欄位:

experiment_name: my-training
mlflow_run_name: baseline-lr3e5
mlflow_artifact_location: /Volumes/main/default/mlflow-artifacts/my-training

關於欄位需求、預設值及支援的值,請參閱 Workload YAML 參考。

AI 執行階段透過 MLFLOW_RUN_ID 公開受管理的執行 ID。 要將你自己的參數、指標和產物記錄到該次執行,請用 MLflow 追蹤 API 重新開啟:

import os

import mlflow

# Log from rank 0 only because all nodes share the same MLFLOW_RUN_ID.
if os.environ.get("RANK", "0") == "0":
    with mlflow.start_run(run_id=os.environ["MLFLOW_RUN_ID"]):
        mlflow.log_param("learning_rate", 3e-4)
        for step, loss in enumerate(training_losses):
            mlflow.log_metric("train_loss", loss, step=step)

尋找執行與記錄

日誌可透過 Databricks CLI 與 MLflow 執行頁面取得。 請使用以下指令來尋找工作負載、檢查工作負載,以及串流或下載日誌:

databricks air list --all-status
databricks air get <job-run-id>
databricks air logs <job-run-id>
databricks air logs <job-run-id> --node 2
databricks air logs <job-run-id> --download-to ./logs/

追蹤 Notebook 工作負載

你可以從筆記本直接在已附加的工作階段中執行程式碼,或透過 Databricks AI 執行時 SDK 的 @distributed 裝飾器執行分散式訓練。

直接在筆記本中執行程式碼

直接在附加 notebook 中執行的程式碼不會自動建立 MLflow 執行。 選擇一個實驗,然後開始執行或啟用自動記錄:

import mlflow

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

with mlflow.start_run(run_name="baseline"):
    train()

要啟用 PyTorch Lightning 自動記錄,請使用以下模式:

mlflow.pytorch.autolog()

with mlflow.start_run(run_name="lightning-baseline"):
    trainer.fit(model, datamodule=datamodule)

對於 Hugging Face Transformers,請如下配置 TrainingArguments:

from transformers import TrainingArguments

args = TrainingArguments(
    output_dir="/Volumes/<catalog>/<schema>/<volume>/checkpoints",
    report_to="mlflow",
    run_name="llama7b-sft-lr3e5",
    logging_steps=50,
)

使用 Databricks AI Runtime SDK

若要從筆記本進行分散式訓練,請使用來自 Databricks AI Runtime SDK 的 @distributed 裝飾器(databricks-sdk-air)。 SDK 在 GPU 環境第 6 版及以上版本中可於 databricks.air 命名空間下使用,或在 GPU 環境第 5 版及以下版本中以 serverless_gpu 提供。

from databricks.air import distributed

@distributed()
def train():
    # Training code
    ...

train.distributed()

每次 .distributed() 呼叫都會自動建立一個 MLflow run。 預設的實驗是 /Users/{WORKSPACE_USER}/{notebook-name}。 要選擇另一個實驗,請在呼叫 .distributed()前設定:

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

當你在一個主動執行中呼叫 .distributed() 時,AI 執行時會建立一個巢狀子執行:

with mlflow.start_run() as parent_run:
    train.distributed()

若要繼續先前的執行,請在 .distributed() 之前設定 MLFLOW_RUN_ID:

import os

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

使用儲存格輸出中的執行連結或筆記本 實驗 面板開啟 MLflow 執行項目。

查看日誌

  • 筆記本輸出:標準輸出和錯誤輸出會顯示在筆記本儲存格的輸出中。
  • MLflow 日誌:對於由作業支援的筆記本工作負載(如 .distributed(remote=True) 呼叫),MLflow Logs 標籤會顯示相關作業執行的輸出。

監控即時 GPU 資源

NotebookGPU 資源面板顯示單節點與多節點工作負載的即時 GPU 使用率、記憶體與溫度。 當你需要與一次執行關聯的持久記錄時,可以使用 MLflow 系統指標 分頁。

要打開窗格,請將筆記本連接到 AI 執行環境,然後在右側窗格中點選晶片圖示。GPU 資源。

GPU 資源面板顯示每張 GPU 的使用率、記憶體和溫度指標。

面板每 10 秒輪詢一次指標,並保留最多 2 小時的歷史紀錄。 點擊 「重新整理」圖示。重新整理 以取得最新的數值。 閒置五分鐘後,窗格會暫停。 重新開啟它以繼續監控。

在 MLflow Logs 標籤中查看日誌

MLflow 執行頁面上的 記錄 標籤會顯示由 job 支援的 AI Runtime 工作負載輸出。 要開啟它,請前往執行頁面並點選 日誌。

對於新的執行,日誌 索引標籤會取代先前儲存在執行作業中已棄用的 成品 索引標籤內的日誌檔案。 先前建立的執行作業會把日誌檔案保留在 Artifacts 分頁。

在 記錄 索引標籤中,你可以:

  • 搜尋:只顯示包含搜尋詞的行。
  • 依節點、等級及時間範圍篩選:選擇多節點工作負載中的單一節點,僅顯示資訊、警告或錯誤線,或縮小視圖至時間範圍。
  • 追蹤運行中的工作負載:當執行仍在進行中,新的輸出會在到達時顯示。 往上捲動以暫停跟隨。
  • 複製並下載:複製已載入的行,或下載所有符合當前篩選器的行。

日誌保留

日誌可在 日誌 分頁中查看 60 天。

排解日誌存取問題

MLflow 執行頁面上的 Logs 標籤會串流與 Azure Databricks 工作執行的輸出。 存取 MLflow 實驗並不代表能取得該作業。

如果分頁顯示 你無法存取這些日誌,請有 可管理 權限的使用者或工作區管理員,至少授權你在該工作上具備 可檢視 權限。 請參閱控制對作業的存取。

Azure Databricks 中的全域限制

請參閱資源限制。