笔记本中的分布式训练

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 会使用节点上的所有显卡。

Note

本页介绍了使用 Databricks AI 运行时 SDK 从 Databricks 笔记本进行分布式训练。要从本地机器提交分布式训练工作负载，请使用 Databricks 的 CLI 命令用于 AI 运行时。参见将 Databricks CLI 与 AI Runtime 配合使用。

快速入门

本页示例使用 Databricks AI 运行时 SDK，该 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 它会使用全部八块显卡。设置 num_accelerators 以使用更少。 accelerator_type 参数也是可选的，并从你的笔记本连接的加速器自动检测。

每次 .distributed() 调用都会创建一个 MLflow 运行（如果已有活动运行，则创建一个嵌套的子运行），并在单元格输出中显示该运行的链接。有关完整且可运行的攻略，请参见完整示例。

支持的框架

API @distributed 与主要的分布式训练库集成：

PyTorch 分布式数据并行（DDP）：标准多 GPU 数据并行度。
完全分片数据并行（FSDP）：面向大型模型的节省内存训练。
DeepSpeed：Microsoft用于大型模型训练的优化库。

关于使用每个库的真实训练场景，请参见分布式训练示例。

装饰师的工作原理`@distributed`

当你调用带有 .distributed()装饰的函数时，AI 运行时会处理你本来需要用分布式启动器手动配置的机制：

序列化与扇出：该函数被序列化并在每个选定的GPU上启动。每块GPU都运行一个带有相同参数的函数副本。
环境同步：Python 环境和依赖在所有等级中复制，因此每个进程运行相同的代码。
排名环境变量：会为每个进程填充标准变量，例如 LOCAL_RANK。在函数中读取它们，将模型和数据放置在正确的设备上。
结果收集：从所有等级收集返回值并返回给呼叫者。
MLflow 跟踪：每次 .distributed() 调用都会创建一次 MLflow 运行，或者如果已有子运行，则创建一个嵌套子运行，这样你函数记录的指标会回到同一运行中。
生命周期与超时：分布式执行在笔记本生命周期内运行。终止笔记本即终止运行。装饰师默认有3小时的超时时间。传入以秒为单位的 timeout 以更改此设置，或传入 timeout=None 以将其禁用。自定义超时需要 GPU 环境 v5 及以上版本。

该API基于标准PyTorch库：分布式数据并行（DDP）、全分片数据并行（FSDP）和 DeepSpeed。

从 `serverless_gpu` 迁移

该 serverless_gpu 软件包已被弃用，不再获得新功能。从 GPU 环境版本 6 开始，Databricks AI 运行时 SDK（databricks-sdk-air）取代了它，并在 databricks.air 命名空间下提供了等效的 API，这些 API 的参数名称有所更改。你使用的软件包取决于 GPU 环境版本：

版本5及以下：仅包含 serverless_gpu。在这些版本上继续使用 serverless_gpu。
版本6：包含两个软件包，但只有 AI 运行时 SDK 会收到更新。第6版是最后一个包含 serverless_gpu的GPU环境版本。

AI 运行时 SDK 是与 Python 的 Databricks SDK 分开的软件包（databricks-sdk）。

	`serverless_gpu`（已弃用）	AI 运行时 SDK （`databricks.air`）
GPU 环境版本	6岁及以下	6岁及以上
进口	`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 运行时上加载数据。

迁移笔记本文件：

将笔记本连接到 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 Runtime 上推荐的替代方案。你的DDP、FSDP或DeepSpeed模型代码不会改变。

Feature	`@distributed` API（应用程序接口）	TorchDistributor
基础结构	完全无服务器，无群集管理	需要具有 GPU 工作节点的 Spark 群集
Setup	单个修饰器，最小配置	需要 Spark 群集和 TorchDistributor 设置
框架支持	PyTorch DDP、FSDP、DeepSpeed	主要是 PyTorch DDP
数据加载	在该装饰器中，使用 Unity Catalog 卷（`UCVolumeDataset` 用于流式处理文件数据）	通过 Spark 或文件系统

迁移单节点工作负载：

将 train_fn 上的 TorchDistributor(...).run(train_fn, ...) 调用替换为 @distributed 装饰器，然后使用 train_fn.distributed(...) 启动。
移除 Spark 集群和 GPU 工作节点配置。改为将你的笔记本连接到 8xH100 加速器。 @distributed 默认使用全部八块显卡。
将数据加载移到被装饰的函数内部。参见数据加载。
保留你现有的DDP、FSDP或DeepSpeed模型代码。装饰师支持这三者。

@distributed装饰器处于Beta阶段，运行在单一节点上（见限制），因此它不会取代所有TorchDistributor的工作负载。依赖Spark集成的工作负载保留在TorchDistributor上。要从本地机器或跨多个节点运行分布式训练，可以使用适用于 AI Runtime 的 Databricks CLI 命令。参见将 Databricks CLI 与 AI Runtime 配合使用。

完整示例

以下示例在笔记本电脑上的8个H100 GPU上训练一个多层感知器（MLP）模型。

设置模型并定义实用工具函数。


# 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 卷中的基于文件的数据，使用 databricks.air.data 中的 UCVolumeDataset，它会流式传输文件并进行本地缓存，并自动跨秩和辅助角色分区。要将分布式训练的检查点保存到卷，请使用 UCVolumeWriter 和 UCVolumeReader。参见在 AI Runtime 上加载数据和使用 Distributed Checkpoint (DCP) 进行检查点保存。

Limitations

@distributed 只支持单个节点：它会在你的笔记本所连接节点上的各个 GPU 上运行。要跨多个节点训练，可以使用 AI 运行时 Databricks CLI。参见将 Databricks CLI 与 AI Runtime 配合使用。
如果你设置 accelerator_type，它必须与 Notebook 实例连接的加速器匹配（例如，"GPU_8xH100"）。不匹配会导致工作负载失败。该参数为可选，省略时自动检测。
AI 运行时支持 GPU 环境 v4 及以上版本。自定义超时（参数 timeout ）需要GPU环境v5及以上。
装饰师默认在3小时后超时。传入以秒为单位的 timeout 以更改此设置，或传入 timeout=None 以将其禁用。
执行过程在笔记本的生命周期内运行。终止笔记本即终止运行。

其他资源

关于 @distributed 装饰器、 AcceleratorType和 Ray API，请参见 Databricks AI 运行时 SDK API 参考资料。
关于让你的培训流程更高效、更有韧性的模式，请参阅绩效与韧性指南。
关于端到端训练场景，请参见分布式训练示例。

反馈

此页面是否有帮助？

上次更新时间为 2026-10-22

笔记本中的分布式训练

快速入门

支持的框架

装饰师的工作原理@distributed

从 serverless_gpu 迁移

来自 TorchDistributor

完整示例

数据加载

Limitations

其他资源

反馈

其他资源

装饰师的工作原理`@distributed`

从 `serverless_gpu` 迁移