笔记本中的分布式训练

重要

此功能在 Beta 版中。 工作区管理员可以从 预览 页控制对此功能的访问。 请参阅 Manage Azure Databricks 预览版。

来自 Serverless GPU Python API 的 @distributed 装饰器是在 Databricks 笔记本中运行分布式训练的最便捷方式。 为你的训练函数添加装饰器并调用它，AI Runtime 就会在你的笔记本所连接节点的所有 GPU 上运行该训练函数。 同样的代码可以从单GPU扩展到多GPU，无需集群配置，也无需分布式启动器配置。

Tip

• @distributed 装饰器从笔记本内部在节点上的每个 GPU 上运行训练函数。
• 它支持 PyTorch DDP、FSDP 和 DeepSpeed，并将单 GPU 代码以最小改动转换为多 GPU。
• 将笔记本连接到 8xH100 加速器，并将 gpus=8 设置为启用完整的多 GPU 训练。

Note

本页介绍使用无服务器 GPU Python API 从 Databricks 笔记本进行分布式训练。 要从本地机器提交分布式训练工作负载，可以使用 Data Bricks 的 CLI 命令用于 AI 运行时，这些命令目前处于 公开预览阶段。 参见 将 Databricks CLI 与 AI Runtime 配合使用。

快速入门

当你的笔记本连接到无服务器显卡时，该 serverless_gpu 软件包是预装的。 用 @distributed 装饰训练函数，然后使用 .distributed() 调用它：

from serverless_gpu import distributed

# gpus is the number of GPUs on the node. gpu_type is optional and
# auto-detected from the accelerator your notebook is connected to.
@distributed(gpus=8, gpu_type="H100")
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 修饰器时，请设置 gpus=8。 该 gpu_type 参数是可选的，并且从笔记本连接到的加速器中自动检测到该参数。

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

支持的框架

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

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

关于使用每个库的真实训练场景，请参见 笔记本示例。

装饰师的工作原理@distributed

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

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

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

来自 TorchDistributor

如果你目前在 Spark 上使用 TorchDistributor 运行分布式 PyTorch，并且你的工作负载可在单个节点上运行，那么 serverless_gpu@distributed API 是针对新的深度学习工作负载推荐使用的替代方案。 它无需 Spark 集群，并让你从单 GPU 到多 GPU 使用相同的代码路径。

功能	serverless_gpu @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 加速器，并改为设置 gpus=8。
• 将数据加载移到被装饰的函数内部。 参见数据加载。
• 保留你现有的DDP、FSDP或DeepSpeed模型代码。 装饰师支持这三者。

@distributed 运行在单一节点上（参见 限制），因此不会取代所有 TorchDistributor 工作负载。 依赖Spark集成的工作负载保留在TorchDistributor上。 要从本地机器或跨多个节点运行分布式训练，请使用公开 预览版中的 AI 运行时 CLI。 参见 将 Databricks CLI 与 AI Runtime 配合使用。

完整示例

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

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

   
   # 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)
   

2. 导入 serverless_gpu 库和 distributed 模块。

   import serverless_gpu
   from serverless_gpu import distributed
   

3. 将模型训练代码包装在函数中，并使用修饰器修饰函数 @distributed 。 装饰函数是分布式执行的入口，因此在其中定义所有训练逻辑、数据加载和模型初始化。

   @distributed(gpus=8, gpu_type='H100')
   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-parellel 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()
   

4. 通过调用带有用户自定义参数的分布式函数来运行分布式训练。

   run_train.distributed(num_epochs=3, batch_size=1)
   

5. 执行时，会在笔记本单元输出中生成 MLflow 运行链接。 单击 MLflow 运行链接或在 “试验 ”面板中找到它以查看运行结果。 有关自定义试验名称、跟踪指标和恢复运行的详细信息，请参阅 试验跟踪和可观测性。

数据加载

将数据加载代码放在 @distributed 函数中。 数据集可以超过由 pickle允许的最大大小，因此在装饰器中生成或加载它可以避免序列化错误：

from serverless_gpu import distributed

# This may cause a pickle error because the dataset is captured by the function.
dataset = get_dataset(file_path)

@distributed(gpus=8, gpu_type='H100')
def run_train():
    # Load the dataset inside the decorated function instead.
    dataset = get_dataset(file_path)
    ...

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

Limitations

• 分布式训练在你的笔记本所连接的单个节点上的各个 GPU 之间运行。 为了实现完整的多GPU训练，连接一个8xH100加速器，该加速器为一个节点配置8个GPU，然后设置 gpus=8。
• 加速器类型必须匹配。 如果你在 @distributed 中设置 gpu_type，其必须与笔记本所连接的加速器匹配（"H100" 或 "A10"）。 不匹配会导致工作负载失败。 该参数为可选，省略时自动检测。
• AI Runtime推荐GPU环境v4及以上版本。 自定义超时（参数 timeout ）需要GPU环境v5及以上。
• 装饰师默认在3小时后超时。 传入以秒为单位的 timeout 以更改此设置，或传入 timeout=None 以将其禁用。
• 执行过程在笔记本的生命周期内运行。 终止笔记本即终止运行。

了解详细信息

• 关于 @distributed 装饰器、GPUType 以及 Ray API，请参见 Serverless GPU Python API 参考文档。
• 关于让你的培训流程更高效、更有韧性的模式，请参阅 绩效与韧性指南。
• 关于端到端的培训场景，请参见 笔记本示例。