在筆記本中進行分散式訓練

Important

這項功能位於測試版 (Beta) 中。

Databricks AI 執行時 SDK（databricks-sdk-air）中的 @distributed 裝飾器是從 Databricks 筆記本執行分散式訓練最方便的方式。為你的訓練函式加上裝飾器並呼叫它，AI Runtime 就會在你的筆記本所連線節點上的所有 GPU 上執行該函式。同樣的程式碼可以從單一 GPU 擴展到多 GPU，無需叢集配置，也無需分散式啟動器配置。

本頁重點：

@distributed 裝飾器可在筆記本中，於你的節點上的每個 GPU 上執行訓練函式。
它支援 PyTorch DDP、FSDP 和 DeepSpeed，並將單 GPU 程式碼以最小改動轉為多 GPU。
將筆電連接到 8xH100 加速器，進行完整的多 GPU 訓練。預設情況下，@distributed 會使用節點上的所有 GPU。

Note

本頁介紹使用 Databricks AI Runtime SDK 從 Databricks 筆記本進行分散式訓練。若要從本地機器提交分散式訓練工作負載，請使用用於 AI Runtime 的 Databricks CLI 指令。請參閱搭配 AI 執行階段使用 Databricks CLI。

快速入門

本頁範例使用 Databricks AI Runtime SDK，該軟體已預先安裝於 GPU 環境版本 6 及以上。在版本 5 及以下，改用套件 serverless_gpu 。請參閱從 serverless_gpu 遷移。要將筆記本連接到加速器，請參考步驟 1：將筆記本連接到 GPU。

使用 @distributed 裝飾你的訓練函式，然後以 .distributed() 呼叫它：

from databricks.air import distributed

# By default, @distributed uses every GPU on the node your notebook is connected
# to. To use fewer, pass num_accelerators (for example, num_accelerators=2).
@distributed()
def train():
    import os
    import torch
    import torch.distributed as dist

    # Bind this process to its own GPU before training.
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    device = torch.device(f"cuda:{local_rank}")
    dist.init_process_group("nccl")
    # ... build the model and data on `device`, then run your training loop ...
    dist.destroy_process_group()

train.distributed()

分散式訓練需要一個 8xH100 加速器，該加速器在單一節點配備 8 顆 GPU。預設情況下， @distributed 會用到全部八張 GPU。設定 num_accelerators 為使用較少。這個 accelerator_type 參數也是可選的，會根據你的筆記型電腦所連接的加速器自動偵測。

每次呼叫 .distributed() 時，都會建立一個 MLflow 執行；如果已有作用中的執行，則會建立巢狀子執行，並在儲存格輸出中顯示執行連結。完整且可執行的攻略，請參見完整範例。

支援的架構

該 @distributed API 整合了主要的分散式訓練函式庫：

PyTorch 分散式資料平行（DDP）：標準多 GPU 資料平行處理。
全分片資料平行（FSDP）：適用於大型模型的可節省記憶體訓練。
DeepSpeed：Microsoft 用於大型模型訓練的優化函式庫。

關於使用各函式庫的真實訓練情境，請參見分散式訓練範例。

裝飾師的工作原理`@distributed`

當你呼叫以 .distributed() 裝飾的函式時，AI Runtime 會處理原本需要使用分散式啟動器手動設定的相關細節：

序列化與扇出：該函式會序列化並在每個選定的 GPU 上啟動。每台 GPU 都會執行一個包含相同參數的函式副本。
環境同步：Python 環境及其相依性在所有階級中複製，因此每個程序執行相同的程式碼。
Rank 環境變數：系統會針對每個程序自動設定標準變數，例如 LOCAL_RANK。在你的函式中讀取它們，將模型和資料放到正確的裝置上。
結果收集：從所有階級收集回傳值並返回給呼叫者。
MLflow 追蹤：每次 .distributed() 呼叫都會產生一個 MLflow 執行，或如果已有子執行，則會產生巢狀子執行，這樣你函式中記錄的指標會落在同一執行中。
生命週期與逾時：分散式執行會在筆記本的生命週期內執行。終止筆記本即終止該運行。裝飾器的預設逾時時間為 3 小時。在幾秒內傳入 timeout 即可變更，或傳入 timeout=None 將其停用。自訂逾時設定需要使用 GPU 環境 v5 以上版本。

API 建立在標準的 PyTorch 函式庫之上：分散式資料平行（Distributed Data Parallel，DDP ）、全分片資料平行（Fully Sharded Data Parallel，FSDP ）以及 DeepSpeed。

遷移自 `serverless_gpu`

該 serverless_gpu 套件已棄用，不再獲得新功能。從 GPU 環境第 6 版開始，Databricks AI Runtime SDK（databricks-sdk-air）取代了它，並在 databricks.air 命名空間下提供等效的 API，且引數已重新命名。你使用的套件取決於 GPU 環境版本：

版本 5 及以下：僅包含 serverless_gpu。在這些版本上繼續使用 serverless_gpu。
版本 6：包含兩個套件，但只有 AI Runtime SDK 會收到更新。版本 6 是最後一個包含 serverless_gpu的 GPU 環境版本。

AI 執行階段 SDK 是與適用於 Python 的 Databricks SDK（databricks-sdk）分開的套件。

	`serverless_gpu` (已取代)	AI 執行階段 SDK（`databricks.air`）
GPU 環境版本	6 以下	6及以上
Import	`from serverless_gpu import distributed, ray_init`	`from databricks.air import distributed, ray_init`
GPU 數量參數	`gpus`	`num_accelerators`
GPU 類型參數	`gpu_type` （`GPUType`）	`accelerator_type` （`AcceleratorType`）
GPU 類型的值	`"A10"`、`"H100"`、`"B300"`	`"GPU_1xA10"`、`"GPU_1xH100"`、`"GPU_8xH100"`、`"GPU_8xB300"`
資料載入與檢查點模組	`serverless_gpu.data`	`databricks.air.data`。請參見 AI 執行時的載入資料。

若要遷移 Notebook：

將 Notebook 連接到 GPU 環境 6 或以上版本，其中包含 AI Runtime SDK。
將匯入從 serverless_gpu 更改為 databricks.air。
將 gpus 重新命名為 num_accelerators，並將 gpu_type 重新命名為 accelerator_type，並使用新的加速器類型值。舊的參數名稱和 GPU 類型值仍然能在 databricks.air 中運作，但會發出 FutureWarning。對 GPU 環境第 7 版的支援將終止。

Tip

請告訴 Genie Code 幫你做這件事：

Migrate this notebook from the serverless_gpu SDK to the AI Runtime SDK.

對於使用 databricks.air.data 的 Databricks CLI 工作負載，將 environment.dependencies 加入工作負載的 databricks-sdk-air[data]。

來自 TorchDistributor

如果你在 Spark 叢集上使用 TorchDistributor 執行分散式訓練，且工作負載可在單一節點上執行，@distributed API 是 AI 執行環境中的建議替代方案。你的 DDP、FSDP 或 DeepSpeed 模型程式碼不會改變。

Feature	`@distributed` 應用程式介面	火炬分配器
基礎結構	完全無伺服器，無需叢集管理	需要一個帶有 GPU 工作者的 Spark 叢集
設定	單一裝飾器，極簡配置	需要 Spark 叢集與 TorchDistributor 設定
架構支援	PyTorch DDP、FSDP、DeepSpeed	主要是 PyTorch DDP
數據載入	在裝飾器中，使用 Unity Catalog 磁碟區（`UCVolumeDataset` 用於串流檔案資料）	透過 Spark 或檔案系統

要遷移單一節點工作負載：

將 train_fn.distributed(...) 上的 @distributed 呼叫替換為 TorchDistributor(...).run(train_fn, ...) 裝飾器，然後使用 train_fn 啟動。
移除 Spark 叢集和 GPU 工作節點設定。改用你的notebook接到8xH100加速器上。 @distributed 預設使用全部八顆 GPU。
在裝飾函式中移動資料載入。請參見資料載入。
保留你現有的 DDP、FSDP 或 DeepSpeed 模型程式碼。裝飾器支援這三者。

@distributed Decorator 目前處於 Beta 階段，運行於單一節點（詳見限制），因此不會取代所有 TorchDistributor 的工作負載。將依賴 Spark 整合的工作負載保留在 TorchDistributor 上。若要從本地機器或跨多個節點執行分散式訓練，請使用 Databricks CLI 的 AI 執行時指令。請參閱搭配 AI 執行階段使用 Databricks CLI。

完整範例

以下範例是在筆記型電腦中訓練一個多層感知器（MLP）模型，使用8顆H100 GPU。

建立你的模型並定義效用函數。


# Define the model
import os
import torch
import torch.distributed as dist
import torch.nn as nn

def setup():
    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
    dist.init_process_group("nccl")

def cleanup():
    dist.destroy_process_group()

class SimpleMLP(nn.Module):
    def __init__(self, input_dim=10, hidden_dim=64, output_dim=1):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim, output_dim)
        )

    def forward(self, x):
        return self.net(x)

從 AI Runtime SDK 匯入 distributed 裝飾器。
```
from databricks.air import distributed
```

把模型訓練程式碼包裝成一個函式，然後用 @distributed 裝飾器裝飾這個函式。裝飾過的函式是分散式執行的入口，因此在其中定義所有訓練邏輯、資料載入與模型初始化。

@distributed()
def run_train(num_epochs: int, batch_size: int) -> None:
    import mlflow
    import torch.optim as optim
    from torch.nn.parallel import DistributedDataParallel as DDP
    from torch.utils.data import DataLoader, DistributedSampler, TensorDataset

    # 1. Set up multi-GPU environment
    setup()
    device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}")

    # 2. Apply the Torch distributed data parallel (DDP) library for data-parallel training.
    model = SimpleMLP().to(device)
    model = DDP(model, device_ids=[device])

    # 3. Create and load dataset.
    x = torch.randn(5000, 10)
    y = torch.randn(5000, 1)

    dataset = TensorDataset(x, y)
    sampler = DistributedSampler(dataset)
    dataloader = DataLoader(dataset, sampler=sampler, batch_size=batch_size)

    # 4. Define the training loop.
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    loss_fn = nn.MSELoss()

    for epoch in range(num_epochs):
        sampler.set_epoch(epoch)
        model.train()
        total_loss = 0.0
        for step, (xb, yb) in enumerate(dataloader):
            xb, yb = xb.to(device), yb.to(device)
            optimizer.zero_grad()
            loss = loss_fn(model(xb), yb)
            # Log loss to MLflow metric
            mlflow.log_metric("loss", loss.item(), step=step)

            loss.backward()
            optimizer.step()
            total_loss += loss.item() * xb.size(0)

        mlflow.log_metric("total_loss", total_loss)
        print(f"Total loss for epoch {epoch}: {total_loss}")

    cleanup()

透過呼叫分散式函式並使用使用者定義的參數來執行分散式訓練。
```
run_train.distributed(num_epochs=3, batch_size=1)
```
執行時，會在筆記本儲存格輸出中產生 MLflow 執行連結。點擊 MLflow 跑動連結，或在實驗面板中找到它，查看跑動結果。關於如何自訂實驗名稱、追蹤指標及恢復運行的細節，請參見實驗追蹤與可觀察性。

資料載入

把資料載入程式碼放在 @distributed 函式裡。資料集可能會超過 pickle 允許的大小上限，因此在裝飾器內產生或載入資料集可避免序列化錯誤：

from databricks.air import distributed

# Don't do this: loading the dataset outside the function means it's captured
# and pickled with the function, which can raise a pickle error.
# dataset = get_dataset(file_path)

@distributed()
def run_train():
    # Load the dataset inside the decorated function instead.
    dataset = get_dataset(file_path)
    ...

對於儲存在 Unity Catalog 磁碟區中的檔案型資料，請使用來自 UCVolumeDataset 的 databricks.air.data，其會以本機快取串流處理檔案，並自動在各個 rank 與 worker 之間分割檔案。若要將分散式訓練的檢查點儲存到磁碟區，請使用 UCVolumeWriter 和 UCVolumeReader。請參見在 AI 執行階段載入資料和使用分散式檢查點（DCP）的檢查點。

Limitations

@distributed 只支援單一節點：它會跨越你的筆記本所連線的節點上的 GPU。若要跨多個節點訓練，請使用 AI Runtime Databricks CLI。請參閱搭配 AI 執行階段使用 Databricks CLI。
如果你設定 accelerator_type，它必須與你的 Notebook 所連接的加速器相符（例如，"GPU_8xH100"）。不相符會導致工作負載執行失敗。該參數為可選，省略時會自動偵測。
AI 執行時支援 GPU 環境 v4 及以上版本。自訂逾時設定（timeout 參數）需要 v5 或以上版本的 GPU 環境。
裝飾器預設會在 3 小時後逾時。在幾秒內傳入 timeout 即可變更，或傳入 timeout=None 將其停用。
執行過程在筆記本的生命週期內執行。終止筆記本即終止該運行。

其他資源

關於 @distributed 裝飾器、AcceleratorType 和 Ray API，請參見 Databricks AI Runtime SDK API 參考。
若想了解如何讓你的訓練流程更有效率且更具韌性，請參閱績效與韌性指南。
關於端對端訓練情境，請參見分散式訓練範例。

意見反應

此頁面對您有幫助嗎？

上次更新時間： 2026-10-22

在筆記本中進行分散式訓練

快速入門

支援的架構

裝飾師的工作原理@distributed

遷移自 serverless_gpu

來自 TorchDistributor

完整範例

資料載入

Limitations

其他資源

意見反應

其他資源

裝飾師的工作原理`@distributed`

遷移自 `serverless_gpu`