實驗追蹤與可觀測性

這很重要

這項功能目前處於 公開預覽版。

實驗追蹤與可觀察性內建於 AI 執行環境中。 MLflow 是執行參數、指標、GPU 系統指標、日誌和產物的單一存放點。 每次執行都會記錄在可與團隊分享的 MLflow 實驗中,而內建的 GPU 資源 窗格會在程式碼執行期間顯示即時的 GPU 使用率、記憶體用量和溫度。

小提示

  • MLflow 是 AI Runtime 實驗的統一介面:指標、參數、系統指標、日誌和成品檔案。
  • 透過 Databricks CLI 提交的工作負載會自動建立一個 MLflow 執行。 在 notebook 和指令碼中,呼叫 mlflow.start_run() 或 mlflow.autolog()。
  • 內建的 GPU 資源面板顯示使用率、記憶體和溫度。

MLflow 為深度學習提供的功能

  • 指標與參數:記錄訓練損失、評估指標、學習率與超參數,並在 MLflow UI 中比較各次執行。
  • 系統指標:GPU、CPU 和記憶體使用率會與訓練指標一併記錄在該次執行的 系統指標 分頁中。
  • 日誌:作業執行期間的驅動程式輸出,顯示於該次執行的 日誌 分頁。
  • 成品與模型:將模型檔案、設定檔及其他輸出與該次執行一併儲存。 工件可以儲存在 Unity Catalog 磁碟區中。
  • 分享與協作:實驗是工作空間的物件。 授予隊友實驗的存取權限,以便分享執行記錄並比較結果。 請參閱 使用 MLflow 實驗組織訓練回合。
  • 框架整合功能:Hugging Face Transformers、PyTorch Lightning 和其他函式庫可直接將記錄寫入 MLflow。

關於 MLflow 3 中的深度學習模式,請參見 MLflow 3 深度學習工作流程。

我需要加入 MLflow 的程式碼嗎?

這取決於你提交工作負載的方式:

執行方式 MLflow run 會自動建立嗎? 你新增的內容
Databricks CLI(databricks air run) Yes. 工作負載 YAML 中的 experiment_name 會設定實驗,系統指標與日誌則被擷取,無需程式碼。 Optional. 在 MLFLOW_RUN_ID 中記錄該次執行的自訂指標。 請參考 Track Runs with MLflow 和 Jobs run 頁面。
無伺服器 GPU API (@distributed) Yes. 每次 .distributed() 呼叫都會產生一次跑分。 Optional. 在函式內記錄自訂計量。
在單一節點上的筆記本程式或指令碼 No. 無伺服器運算不會自動啟用自動記錄。 呼叫 mlflow.start_run() 並記錄指標,或呼叫 mlflow.autolog()。

開始使用

使用 MLflow 3.7 及以上版本。 以下範例可直接複製到 Notebook 儲存格或 Python 指令碼中。

從訓練循環中記錄指標

import mlflow

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

with mlflow.start_run(run_name="baseline-lr3e-4"):
    mlflow.log_params({"learning_rate": 3e-4, "batch_size": 32, "epochs": 3})
    for epoch in range(3):
        train_loss = train_one_epoch(model, train_loader, optimizer)  # your training code
        val_loss = evaluate(model, val_loader)
        mlflow.log_metrics({"train_loss": train_loss, "val_loss": val_loss}, step=epoch)

使用自動記錄

對於 PyTorch Lightning,請在訓練前呼叫 mlflow.pytorch.autolog()。 如需其他受支援的函式庫,請致電 mlflow.autolog()。

import mlflow

mlflow.pytorch.autolog()

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

來自 Hugging Face Transformers 的日誌

設定 report_to="mlflow"。 run_name 參數會設定 MLflow 執行的名稱。

from transformers import TrainingArguments

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

來自多張 GPU 的日誌

在分散式訓練中,每個程序都會執行你的訓練程式碼。 僅記錄 rank 0 的日誌,這樣每個指標只會被記錄一次:

import os

import mlflow

if int(os.environ.get("RANK", "0")) == 0:
    mlflow.log_metric("train_loss", loss, step=step)

最佳做法

  • 將 step 設為有意義的值,例如全域批次或 epoch,並按照固定間隔記錄(例如每 50 步一次),而不是每個批次都記錄。 MLflow 會限制每次執行中的指標步數上限。 請參閱資源限制。
  • 使用絕對實驗路徑,例如 /Users/<username>/my-experiment 或 /Workspace/Shared/<team>/my-experiment。 把你想分享的實驗放進共用資料夾。
  • 若要繼續先前的執行,請傳入其 ID:mlflow.start_run(run_id="<previous-run-id>")。

無伺服器 GPU API

當你使用 Serverless GPU API 時,每次呼叫 .distributed() 都會自動建立一筆 MLflow 執行記錄。 預設的實驗是 /Users/{WORKSPACE_USER}/{notebook-name}。

  • 如果你在一個作用中的 MLflow 執行內呼叫 .distributed(),它會在該執行下建立一個巢狀的子執行:

    import mlflow
    
    with mlflow.start_run() as outer_run:
        run_train.distributed()  # creates a nested child run under outer_run
    
  • 若要使用不同的實驗,請呼叫 mlflow.set_experiment() 之前 .distributed(),或設定 MLFLOW_EXPERIMENT_NAME 環境變數。 一律使用絕對路徑。

    import os
    
    import mlflow
    
    mlflow.set_experiment("/Users/<username>/my-experiment")
    # or: os.environ["MLFLOW_EXPERIMENT_NAME"] = "/Users/<username>/my-experiment"
    run_train.distributed()
    
  • 要繼續之前的運算,請在呼叫.distributed()前設定MLFLOW_RUN_ID:

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

觀看日誌

  • Notebook 輸出:您的訓練程式碼所產生的標準輸出和錯誤輸出會顯示在 notebook 儲存格的輸出內容中。
  • MLflow 日誌:MLflow 實驗介面會顯示訓練指標、參數與產物。

如果你無法查看日誌

MLflow 執行頁面上的 日誌 標籤會串流與 MLflow 執行相關的 Databricks 工作執行日誌,因此存取受該工作權限所控制。 如果分頁顯示 「你無法存取這些日誌」,代表你沒有足夠的權限。

存取 MLflow 中的該執行,並不表示也能存取該作業。 你可以保留 MLflow 實驗的權限,但仍然會被拒絕接收這些日誌。 要取得存取權限,請有「 可管理 」權限的使用者或工作區管理員至少授權你「 可檢視 」。 請參閱 「控制存取」 以了解工作權限的授予方式。

監控 GPU 資源

GPU 資源面板是筆記型電腦會話時的便利功能。 它能顯示即時 GPU 健康狀況和使用率,無需設定 MLflow 即可,因此當你的筆記本工作階段沒有建立 MLflow 實驗時,它特別有用。 若想永久記錄與執行相關的 GPU、CPU 和記憶體指標,請使用 MLflow System 指標 標籤。 面板支援單節點與多節點工作負載。

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

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

面板顯示每個 GPU 的以下指標:

  • GPU 使用率百分比
  • GPU 記憶體使用率
  • 溫度

面板每 10 秒輪詢一次指標,並保留最多 2 小時的歷史紀錄。 點擊 重新整理圖示。重新整理 以立即取得最新值。 五分鐘不動後,窗格會暫停;重新開啟它以恢復監控。

Azure Databricks 中的全域限制

請參閱資源限制。