共用方式為

實驗追蹤與可觀測性

這很重要

針對單節點任務的 AI 執行時正在 公開預覽中。 多GPU工作負載的分散式訓練API仍處於 測試階段。

本文說明如何使用 MLflow、監控 GPU 健康狀況、查看日誌，以及在 AI 執行時管理模型檢查點。

MLflow 整合

AI 執行環境原生整合 MLflow，用於實驗追蹤、模型記錄及指標視覺化。

設定建議：

  • 將 MLflow 升級到 3.7 或更新版本，並遵循 深度學習的工作流程模式。

  • 啟用 PyTorch Lightning 的自動登入：

    import mlflow
mlflow.pytorch.autolog()

  • 在 API 範圍內封裝模型定型程式碼 mlflow.start_run() ，以自訂您的 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,
)

  • Serverless 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>")

  • 將MLFlowLogger中的參數step設成合理的批次數。 MLflow 的計量步數上限為 1,000 萬步—在大型訓練中記錄每一個批次都可能達到此上限。 請參閱資源限制。

觀看日誌

  • 筆記本輸出 — 訓練代碼的標準輸出與錯誤會顯示在筆記本單元格的輸出中。
  • 驅動程式日誌 — 可透過計算面板存取，用於除錯啟動問題、環境設定問題及執行時錯誤。
  • MLflow 日誌 — 訓練指標、參數與產物可在 MLflow 實驗介面中查看。

模型檢查點

將模型檢查點儲存到 Unity 目錄卷，這些卷提供與其他 Unity 目錄物件相同的治理。 請使用以下路徑格式來參考 Databricks 筆記本中卷中的檔案：

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

把檢查點存到磁碟區的方式與存到本地儲存相同。

以下範例展示了如何將 PyTorch 檢查點寫入 Unity Catalog 卷。

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