PyTorch Scala 高校计算机硕士研一课程

全分片数据并行（FSDP）

DistributedDataParallel (DDP) 通过复制模型和平均梯度有效地在多个GPU上扩展训练。然而，它根本上要求每个GPU都持有完整的模型、其梯度和优化器状态。这在处理包含数十亿参数的模型时成为一个限制因素，因为这些模型可能轻易超出甚至高端加速器的内存容量。

全分片数据并行（FSDP）通过扩展数据并行理念，同时大幅减少每个GPU的内存占用，提供了一种解决办法。FSDP不是复制整个模型，而是将模型的参数、梯度和优化器状态分片或分区到数据并行工作器（GPU）上。

FSDP运行机制

FSDP的核心是确保数据并行组中的每个GPU在任何给定时间点只持有模型参数、梯度和优化器状态的一部分（一个“分片”）。完整张量仅在计算需要时才临时重建。

以下是训练过程的详细步骤：

1. 初始化： 模型使用 FullyShardedDataParallel 模块进行封装。在初始化期间，参数、梯度和优化器状态被划分到参与进程组的GPU上。每个GPU负责管理其分配到的分片。
2. 前向传播：
   • 当FSDP封装的层需要执行计算时，每个GPU使用 all_gather 集合通信操作从组中所有其他GPU收集该特定层所需的完整参数。
   • 该层的计算在当前本地的完整参数上进行。
   • 该层的前向计算完成后，除了GPU自己的分片外，完整参数会立即被丢弃，从而释放内存。这个过程逐层（或逐块，取决于封装策略）重复进行。
3. 反向传播：
   • 当给定层的参数在本地计算出梯度时，它们不会像DDP那样立即在所有GPU上求平均。
   • 相反，会执行 reduce_scatter 操作。此操作计算所有GPU上梯度的平均值，并同时对平均结果进行分片，仅向每个GPU发送与其管理参数分片对应的梯度部分（分片）。
   • 这个分片后的梯度被存储起来，再次保持了较低的内存使用。reduce_scatter 后，完整梯度被丢弃。
4. 优化器步骤：
   • 每个GPU的优化器只需要更新其负责的参数分片。
   • 由于优化器状态（如Adam中的动量缓冲区）也与参数一同分片，因此优化器步骤可以在每个GPU上仅使用其梯度和优化器状态的分片进行本地处理。

这种方法大幅减少了每个GPU所需的峰值内存，因为只有当前执行层的参数以及完整模型、梯度和优化器状态的分片才会被持久存储。

每GPU内存使用比较DDP内存布局FSDP内存布局GPU 0完整模型完整梯度完整优化器状态GPU 1完整模型完整梯度完整优化器状态GPU 0参数分片0梯度分片0优化器状态分片0GPU 1参数分片1梯度分片1优化器状态分片1cluster_ddpcluster_fsdp

分布式数据并行（DDP）和全分片数据并行（FSDP）每GPU内存分配比较。DDP复制所有组件，而FSDP将其分片。

在PyTorch中实现FSDP

PyTorch通过 torch.distributed.fsdp.FullyShardedDataParallel 类提供了对FSDP的原生支持。集成它通常涉及封装您的模型定义。

import torch
import torch.nn as nn
import torch.distributed as dist
import torch.distributed.fsdp.FullyShardedDataParallel as FSDP
import torch.distributed.fsdp.wrap.size_based_auto_wrap_policy
import functools

// 假设分布式环境已初始化（rank, world_size 等）
// dist.init_process_group(backend="nccl")
// torch.cuda.set_device(local_rank) // local_rank 通常获取

class LargeTransformerBlock extends nn.Module:
    // 子模块定义示例
    def __init__(self, dim, ff_dim):
        super().__init__()
        val layer_norm = nn.LayerNorm(dim)
        val attention = nn.MultiheadAttention(dim, num_heads=8) // Simplified
        val ffn = nn.Sequential(
            nn.Linear(dim, ff_dim),
            nn.ReLU(),
            nn.Linear(ff_dim, dim)
        )

    def forward(x: torch.Tensor):
        x = layer_norm(x + attention(x, x, x)[0])
        x = x + ffn(x)
        return x

class BigModel extends nn.Module:
    // 子模块定义示例
    def __init__(self, num_layers, dim, ff_dim, vocab_size):
        super().__init__()
        val embedding = nn.Embedding(vocab_size, dim)
        val layers = nn.ModuleList(
            [LargeTransformerBlock(dim, ff_dim) for _ in range(num_layers)]
        )
        val output_head = nn.Linear(dim, vocab_size)

    def forward(x: torch.Tensor):
        x = embedding(x)
        for layer in layers:
            x = layer(x)
        x = output_head(x)
        return x

// --- FSDP 设置 ---
val model = BigModel(num_layers=48, dim=2048, ff_dim=8192, vocab_size=50000).to(torch.cuda.current_device())

// 定义一个自动封装策略（可选但推荐用于大型模型）
// 这会根据大小封装子模块（例如 LargeTransformerBlock）
val auto_wrap_policy = functools.partial(
    size_based_auto_wrap_policy, min_num_params=1_000_000 # 示例阈值
)

// 用 FSDP 封装模型
val fsdp_model = FSDP(
    model,
    auto_wrap_policy=auto_wrap_policy,
    // 其他配置选项可在此处添加
    // 例如，cpu_offload=CPUOffload(offload_params=True)
    // 例如，mixed_precision=MixedPrecision(...)
    // 例如，sharding_strategy=ShardingStrategy.SHARD_GRAD_OP
)

// --- 训练循环 ---
// 优化器必须在用 FSDP 封装模型后构建
val optimizer = torch.optim.AdamW(fsdp_model.parameters(), lr=1e-4)

// 训练步骤示例（简化）
// for batch in dataloader:
//     inputs = batch['input_ids'].to(torch.cuda.current_device())
//     labels = batch['labels'].to(torch.cuda.current_device())
//
//     optimizer.zero_grad()
//     outputs = fsdp_model(inputs)
//     loss = criterion(outputs.view(-1, vocab_size), labels.view(-1))
//     loss.backward()
//     optimizer.step()

实现中的要点：

1. 模型定义： 像往常一样使用 nn.Module 定义您的模型。
2. 封装： 使用 FSDP 封装模型实例。请注意，模型应在封装 之前 移至目标设备。
3. 优化器： 在使用 FSDP 封装模型 之后 构建优化器，并将 fsdp_model.parameters() 传递给它。这确保优化器了解分片参数和状态。
4. 自动封装策略： 对于复杂模型，定义 auto_wrap_policy 非常重要。此策略告知FSDP如何递归地封装主模型内的子模块。封装单个块（例如Transformer层）可以实现更细粒度的分片以及更好的通信和计算重叠。size_based_auto_wrap_policy 是一种常见选择，用于封装参数数量超出特定阈值的模块。

配置选项

FSDP提供了多种配置选项来调整其行为：

• sharding_strategy：控制参数、梯度和优化器状态的分片程度。
  • ShardingStrategy.FULL_SHARD：（默认）分片参数、梯度和优化器状态。提供最大的内存节省，但通信开销可能更高。
  • ShardingStrategy.SHARD_GRAD_OP：仅分片梯度和优化器状态。参数被复制（类似于ZeRO 阶段2）。内存节省少于 FULL_SHARD，但通信开销可能更低。
  • ShardingStrategy.NO_SHARD：等同于DDP（复制所有内容）。有助于调试或基准比较。
  • ShardingStrategy.HYBRID_SHARD：在节点内结合完全分片，跨节点复制。在多节点场景中有用。
• cpu_offload：通过 CPUOffload(offload_params=True/False) 配置。当参数和梯度的分片未主动用于计算时，允许将其卸载到CPU内存。这以CPU和GPU之间显著的通信开销为代价，进一步增加了可行的模型大小。当GPU内存是绝对瓶颈时使用此选项。
• mixed_precision：通过 MixedPrecision(param_dtype=torch.float16, reduce_dtype=torch.float16, buffer_dtype=torch.float16) 配置。将混合精度训练直接集成到FSDP封装器中，自动处理类型转换和梯度缩放。通常建议使用FSDP内置的混合精度，而不是在外部应用 torch.cuda.amp.autocast。
• auto_wrap_policy：如前所述，它定义了嵌套模块如何封装。size_based_auto_wrap_policy 的替代方案包括基于模块类型（transformer_auto_wrap_policy）的封装或手动封装。
• backward_prefetch：控制反向传播的参数预取，以实现通信和计算重叠。像 BackwardPrefetch.BACKWARD_PRE （在当前层的反向传播期间预取下一层的参数）这样的选项可以提高性能。

权衡与考量

尽管FSDP能够训练显著更大的模型，但它也引入了一些权衡：

• 增加的通信量：all_gather（前向）和 reduce_scatter（反向）操作相比DDP在反向传播中的单个 all_reduce 引入了更大的通信量。性能影响严重依赖于GPU/节点之间的互连速度。更快的互连（例如NVLink，InfiniBand）能更有效地减轻此开销。
• 计算/通信重叠：FSDP旨在使通信（为下一层收集参数）与计算（当前层的执行）重叠。使用 auto_wrap_policy 的有效封装策略对最大化这种重叠很重要。
• 激活检查点：为了进一步减少正向传播期间存储的激活所带来的内存使用，FSDP通常与激活检查点（也称为梯度检查点）一起使用。PyTorch提供了 torch.utils.checkpoint.checkpoint，FSDP有特定的实用工具（fsdp_checkpointing）可以高效地将其应用于封装模块。
• 复杂性：设置和调优FSDP可能比基本的DDP更复杂，特别是在最佳封装策略和特定硬件设置的配置方面。

总而言之，FSDP是一种强大的技术，用于训练不适合单个GPU内存的超大型模型。通过将参数、梯度和优化器状态分片到数据并行工作器上，它显著降低了每GPU的内存需求。然而，这可能会增加通信开销，使得快速互连和仔细的配置对于获得良好的训练性能非常重要。这代表了PyTorch在大规模模型训练能力方面的一个重要进展。

使用 torch.distributed 通信原语

收藏

高级抽象，例如 DistributedDataParallel (DDP)，会自动处理分布式训练的许多细节。然而，了解 torch.distributed 包提供的底层通信原语可以提供全面的理解，并能够实现自定义的并行化策略。这些原语是用于在分布式环境中协调不同进程之间通信的基本组成部分。

在使用任何通信原语之前，必须先初始化分布式环境，通常使用 torch.distributed.init_process_group 函数。这会建立通信后端（例如 NCCL 或 Gloo），并为总 world_size 中的每个进程分配一个唯一的 rank。一旦初始化完成，默认组（或自定义创建的组）中的进程就可以使用集体和点对点操作进行协调。

集体通信操作

集体操作涉及组内所有进程之间的通信。它们对于同步梯度或分发模型参数等任务非常重要。以下是一些最常用的集体操作：

广播 (dist.broadcast)

此操作将张量从一个源进程 (src) 发送到组中的所有其他进程。它通常用于确保所有进程都使用相同的初始模型参数开始。

import torch
import torch.distributed as dist
import os

def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    # 初始化进程组
    dist.init_process_group("gloo", rank=rank, world_size=world_size)

def run_broadcast(rank, world_size):
    setup(rank, world_size)
    tensor = torch.zeros(1)
    if rank == 0:
        //源进程创建数据
        tensor += 1

    // Rank 0 将 'tensor' 广播到所有其他进程
    dist.broadcast(tensor=tensor, src=0)

    println(f"Rank {rank} has data: {tensor[0]}")
    dist.destroy_process_group()

// 假设 world_size = 4 进行演示
// 在实际脚本中，这将通过 torchrun 或类似方式启动
// run_broadcast(0, 4)
// run_broadcast(1, 4)
// run_broadcast(2, 4)
// run_broadcast(3, 4)

此操作后，Rank 1、2 和 3 上的 tensor 将从 0 更新为 1。

Rank 0张量=1Rank 1张量=0广播Rank 2张量=0Rank 3张量=0

dist.broadcast 操作中数据流的示意图，从 Rank 0 到 4 进程组中的所有其他 Rank。

全局归约 (dist.all_reduce)

此操作使用指定的归约操作（op，例如 dist.ReduceOp.SUM、dist.ReduceOp.AVG）合并所有进程的张量，并将最终结果分发回所有进程。这是 DDP 中梯度同步的根本。每个进程贡献其局部梯度，这些梯度在所有进程中求和（或求平均），然后每个进程都收到合并后的梯度。

import torch
import torch.distributed as dist
import os

// 假设 setup 函数已如上定义

def run_all_reduce(rank, world_size):
    setup(rank, world_size)
    // 每个 Rank 根据其 Rank 创建数据
    val tensor = torch.tensor([rank + 1], dtype=torch.float32) 

    println(f"Rank {rank} initial tensor: {tensor[0]}")

    // 执行 SUM 操作的全归约
    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)

    // 结果 (1+2+3+4 = 10 的和) 在所有 Rank 上都可用
    println(f"Rank {rank} final tensor: {tensor[0]}") 
    dist.destroy_process_group()

// world_size = 4 的执行示例
// run_all_reduce(0, 4) // 初始值: 1, 最终值: 10
// run_all_reduce(1, 4) // 初始值: 2, 最终值: 10
// run_all_reduce(2, 4) // 初始值: 3, 最终值: 10
// run_all_reduce(3, 4) // 初始值: 4, 最终值: 10

全归约前全归约后 (求和)Rank 0值=A合并 (求和)S = A+B+C+DRank 1值=BRank 2值=CRank 3值=DRank 0值=SRank 1值=SRank 2值=SRank 3值=S

dist.all_reduce 进行求和操作的数据流示意图。所有 Rank 贡献数据，数据被聚合，结果再分发回所有 Rank。

归约 (dist.reduce)

与 all_reduce 类似，reduce 使用归约操作合并所有进程的张量。但其结果只存储在目标进程 (dst) 上。其他进程不会收到结果。

分散 (dist.scatter)

此操作获取单个源进程 (src) 上的一个张量列表 (scatter_list)，并将列表中的一个张量分发给组中的每个进程，包括其自身。scatter_list 中第 ii 个张量会发送给 Rank 为 ii 的进程。这有助于在进程间分发数据批次。

import torch
import torch.distributed as dist
import os

// 假设 setup 函数已如上定义

def run_scatter(rank, world_size):
    setup(rank, world_size)
    val my_tensor = torch.zeros(1)
    val scatter_list = None

    if rank == 0:
        // 源 Rank 准备要分散的张量列表
        val scatter_list = [torch.tensor([i + 1.0]) for i in range(world_size)]
        println(f"Rank 0 scatter list: {[t.item() for t in scatter_list]}")

    // Rank 0 分散列表。每个 Rank 都会收到一个张量到 my_tensor 中。
    dist.scatter(tensor=my_tensor, scatter_list=scatter_list, src=0)

    print(f"Rank {rank} received tensor: {my_tensor.item()}")
    dist.destroy_process_group()

// world_size = 4 的执行示例
// run_scatter(0, 4) // 收到: 1.0
// run_scatter(1, 4) // 收到: 2.0
// run_scatter(2, 4) // 收到: 3.0
// run_scatter(3, 4) // 收到: 4.0

Rank 0 (源)ABCDRank 0接收 A分散 ARank 1接收 B分散 BRank 2接收 C分散 CRank 3接收 D分散 D

dist.scatter 的数据流示意图。Rank 0 持有张量列表 [A, B, C, D]，并将 A 发送给 Rank 0，B 发送给 Rank 1，C 发送给 Rank 2，D 发送给 Rank 3。

收集 (dist.gather)

scatter 的反向操作。每个进程将其张量发送到一个目标进程 (dst)。目标进程接收这些张量并将它们存储在一个列表 (gather_list) 中。gather_list 中的顺序与发送进程的 Rank 对应。

import torch
import torch.distributed as dist
import os

// 假设 setup 函数已如上定义

def run_gather(rank, world_size):
    setup(rank, world_size)
    // 每个 Rank 创建自己的张量
    val my_tensor = torch.tensor([rank + 1.0]) 
    val gather_list = None

    if rank == 0 then
        // 目标 Rank 准备一个列表来存储收集到的张量
        val gather_list = [torch.zeros(1) for _ in range(world_size)]

    // 所有 Rank 将其张量发送给 Rank 0
    dist.gather(tensor=my_tensor, gather_list=gather_list, dst=0)

    if rank == 0:
        print(f"Rank 0 gathered list: {[t.item() for t in gather_list]}")
    else:
        print(f"Rank {rank} sent tensor: {my_tensor.item()}")

    dist.destroy_process_group()

// world_size = 4 的执行示例
// run_gather(0, 4) // 收集到: [1.0, 2.0, 3.0, 4.0]
// run_gather(1, 4) // 已发送: 2.0
// run_gather(2, 4) // 已发送: 3.0
// run_gather(3, 4) // 已发送: 4.0

Rank 0ARank 0 (目标)ABCD收集Rank 1BRank 2CRank 3D

dist.gather 的数据流示意图。Rank 0、1、2、3 分别将其张量 A、B、C、D 发送给 Rank 0，Rank 0 将它们收集到列表 [A, B, C, D] 中。

全局收集 (dist.all_gather)

与 gather 类似，但从所有进程收集到的张量列表结果会分发回组中的所有进程。每个进程都会收到相同的最终列表。

点对点通信操作

这些操作涉及两个特定进程之间的通信，通过它们的 Rank 进行识别。

• dist.send(tensor, dst): 将张量从当前进程发送到目标进程 (dst)。这是发送方的阻塞操作。
• dist.recv(tensor, src): 从源进程 (src) 接收张量到提供的 tensor 缓冲区中。这是接收方的阻塞操作，直到张量被接收。

虽然功能强大，但点对点操作需要仔细管理以避免死锁（例如，两个进程在发送之前互相等待接收）。与集体操作相比，它们在标准数据并行训练中不那么常用，但对于模型并行或自定义算法等更复杂的通信模式很重要。

阻塞与非阻塞操作

大多数集体操作（broadcast、all_reduce、scatter、gather 等）默认是阻塞（同步）的。这意味着进程上的程序执行会暂停，直到该进程完成了其在集体通信中的部分。

PyTorch 还提供了许多操作的非阻塞（异步）版本，通常以 i 作为前缀（例如 dist.isend、dist.irecv、dist.all_reduce(..., async_op=True)）。这些调用会启动通信并立即返回一个 Work 对象（或类似的句柄）。在通信在后台进行的同时，程序可以继续执行其他任务。您可以稍后使用返回句柄上的 wait() 等方法检查完成情况或等待操作结束。

// 非阻塞全归约示例
val tensor = torch.ones(1) * rank
// ... 其他设置 ...

// 启动非阻塞全归约
val work_handle = dist.all_reduce(tensor, op=dist.ReduceOp.SUM, async_op=True)

// 在通信进行时执行其他计算...
// result = compute_something_else() 

// 等待 all_reduce 操作完成
work_handle.wait() 

// 现在 'tensor' 包含了归约后的结果
println(f"Rank {rank} async all_reduce result: {tensor[0]}")

使用非阻塞操作可以将计算与通信重叠，从而显著提升性能，尤其是在具有快速互连的系统上。但是，这需要仔细管理依赖关系和同步点。

了解这些 torch.distributed 原语为实现复杂的分布式训练流程奠定了根基。它们允许对进程间通信进行细粒度控制，这对于流水线并行、自定义梯度聚合方案或与专用硬件通信库交互等技术是必要的。

设置分布式环境

建立进程间的连接是分布式训练任务中重要第一步。PyTorch 的 torch.distributed 包提供管理这种通信设置所需的工具。主要思路是创建一个 进程组，它包含所有参与任务的进程。每个进程被分配一个唯一 rank，它们共同了解进程总数，称为 size。

初始化进程组

设置分布式环境的主要函数是 torch.distributed.init_process_group。此函数初始化默认进程组，且必须由每个参与分布式任务的进程调用。

import torch
import torch.distributed as dist
import os

def setup_distributed(backend='nccl'):
    """初始化分布式环境。"""
    if not dist.is_available():
        print("分布式训练不可用。")
        return

    if not dist.is_initialized():
        // 从环境变量获取配置。这些变量通常由启动工具设置
        // (例如，torchrun, Slurm)
        val rank = int(os.environ.get("RANK", "0"))
        val world_size = int(os.environ.get("WORLD_SIZE", "1"))
        val master_addr = os.environ.get("MASTER_ADDR", "localhost")
        val master_port = os.environ.get("MASTER_PORT", "29500") // 默认端口

        println(f"正在初始化进程组: Rank {rank}/{world_size}")
        dist.init_process_group(
            backend=backend,
            init_method=f'tcp://{master_addr}:{master_port}',
            rank=rank,
            world_size=world_size
        )
        println(f"进程组已初始化 ({backend})。")

    // 为当前进程设置设备。这很重要！
    // 假设每个GPU一个进程。
    if backend == 'nccl' && torch.cuda.is_available():
        val local_rank = int(os.environ.get("LOCAL_RANK", "0"))
        torch.cuda.set_device(local_rank)
        println(f"Rank {dist.get_rank()} 正在使用 GPU {local_rank}")

// 示例用法（通常在脚本开头调用）
// setup_distributed(backend='nccl') # 或 'gloo'

我们来分析一下 init_process_group 的参数：

• backend: 指定使用的通信库。此选择取决于您的硬件和要求。
• init_method: 定义进程如何互相发现。最常用方法是使用 TCP，并需要 rank 0 进程（主进程）的地址和端口。基于环境变量的初始化 ('env://') 也常使用，它依赖于预设的特定环境变量。
• rank: 当前进程的唯一标识符，范围从 0 到 world_size - 1。
• world_size: 参与训练任务的进程总数。

通信后端

PyTorch 支持多种后端来处理进程间的底层通信：

1. NCCL (NVIDIA Collective Communication Library): 这是 NVIDIA 硬件上基于 GPU 的分布式训练的推荐后端。它为 CUDA 张量提供了集体通信操作（如 all_reduce、broadcast）的优化实现，提供高带宽和低延迟。当您的所有进程都在配备 NVIDIA GPU 并通过 NVLink 或 InfiniBand 等高速互连连接的机器上运行时，请使用 NCCL。
2. Gloo: Gloo 是一个更通用的后端，适用于 CPU 和 GPU 通信。对于异构环境、仅限 CPU 的训练，或者当 NCCL 不可用或不适用时（例如，某些云环境或旧硬件），它是一个不错的选择。虽然在 GPU 到 GPU 通信方面通常比 NCCL 慢，但它提供更广泛的兼容性。
3. MPI (Message Passing Interface): MPI 是一个高性能计算标准。如果您的集群环境已配置 MPI 实现（如 Open MPI），您可以使用 MPI 后端。PyTorch 需要编译时包含 MPI 支持才能使用此选项。它常用于学术或研究集群。

后端选择会显著影响性能。对于 NVIDIA 硬件上的典型多 GPU 训练，NCCL 几乎总是最佳选择。

分布式进程 (WORLD_SIZE=N)进程 0Rank 0(主)进程 1Rank 1通信后端(NCCL / Gloo / MPI)init_process_group…init_process_group进程 N-1Rank N-1init_process_group主地址和端口(例如，tcp://10.1.1.2:29500)使用初始化方法

使用通信后端和发现方法（如带主地址/端口的 TCP）进行进程组初始化。每个进程都调用 init_process_group。

环境变量

分布式训练设置通常依赖环境变量进行配置，尤其是在使用启动工具时。最重要的包括：

• MASTER_ADDR: 托管 rank 0 进程的机器的 IP 地址或主机名。所有其他进程都需要连接到此地址。
• MASTER_PORT: rank 0 机器上的一个开放网络端口，供进程协调初始化。选择一个不太可能被占用的端口（例如，10000 以上）。
• RANK: 当前进程的唯一 rank。
• WORLD_SIZE: 参与任务的进程总数。
• LOCAL_RANK (常用): 当在同一节点上运行多个进程时（例如，每个 GPU 一个进程），LOCAL_RANK 通常标识进程在该节点内的索引（通常从 0 到 num_gpus_on_node - 1）。这常用于通过 torch.cuda.set_device(local_rank) 为每个进程分配一个特定 GPU。

这些变量通常由启动脚本或集群调度器自动设置。

启动分布式脚本

手动设置环境变量并在每台机器上启动 Python 脚本可能繁琐且易错。PyTorch 提供 torchrun（以前是 torch.distributed.launch）工具来简化此过程。

torchrun 负责设置必要的环境变量（RANK、WORLD_SIZE、MASTER_ADDR、MASTER_PORT、LOCAL_RANK）并启动每个节点指定数量的进程。

在单机 2 GPU 上使用 torchrun 的示例：

# 假设您的训练脚本名为 train.py
# 并且它包含前面所示的 setup_distributed() 函数。

# --nproc_per_node 指定在此机器上启动多少个进程。
# 通常设置为可用 GPU 的数量。
# --nnodes=1 表示我们在单机上运行。
# train.py 是您的脚本，后面跟着它的参数。

torchrun --nproc_per_node=2 --nnodes=1 train.py --arg1 value1 --arg2 value2

对于多节点训练，您通常在每个节点上运行 torchrun，指定节点总数（--nnodes）、当前节点的 rank（--node_rank）以及主节点的地址和端口（--rdzv_endpoint）：

# 在主节点 (Rank 0):
torchrun \
    --nproc_per_node=<gpus_on_master> \
    --nnodes=<total_nodes> \
    --node_rank=0 \
    --rdzv_id=<job_id> \
    --rdzv_backend=c10d \
    --rdzv_endpoint="<master_node_ip>:<port>" \
    train.py --args...

# 在工作节点 1 (Rank 1):
torchrun \
    --nproc_per_node=<gpus_on_worker1> \
    --nnodes=<total_nodes> \
    --node_rank=1 \
    --rdzv_id=<job_id> \
    --rdzv_backend=c10d \
    --rdzv_endpoint="<master_node_ip>:<port>" \
    train.py --args...

# ... 其他工作节点以此类推

这里，rdzv 是 rendezvous 的缩写。torchrun 使用此机制（通常构建在 torch.distributed 的 c10d 后端之上）来协调跨节点启动。<job_id> 对于此特定训练运行应是唯一的。

清理

训练结束后，清理分布式环境资源是良好实践：

def cleanup_distributed():
    """销毁默认进程组。"""
    if dist.is_initialized():
        dist.destroy_process_group()
        print("进程组已销毁。")

// 示例用法（通常在脚本末尾或 finally 块中）
// try:
//     // setup_distributed(...)
//     // ... training loop ...
// finally:
//     cleanup_distributed()

调用 dist.destroy_process_group() 会释放与进程组关联的资源。虽然 Python 的退出通常会处理此问题，但明确清理更安全，尤其是在复杂应用或长期运行的服务中。

正确配置分布式环境是构建所有并行训练策略的根本。了解进程组、rank、size、通信后端以及 torchrun 等启动工具对于有效扩展您的 PyTorch 训练任务非常重要。

实践操作：设置DDP训练脚本

演示了如何将一个标准的单GPU PyTorch训练脚本转换为使用分布式数据并行（DDP）进行单机多GPU训练。这主要说明启用DDP所需的具体修改，而不是构建一个完整的先进模型训练器。我们假设您已有一个可用的单GPU脚本。

DDP 的主要思想很简单：在每个可用GPU上复制模型，为每个副本提供输入数据批次的不同切片，在每个GPU上独立计算梯度，然后在更新模型参数之前将这些梯度在所有GPU上平均。这保证了所有模型副本保持同步。

前提条件

• 一个适用于单GPU的PyTorch训练脚本。
• 一台PyTorch可检测到的、具有多个CUDA功能的GPU的机器。
• 已安装具有分布式支持的PyTorch（通常在CUDA构建中默认包含）。

DDP转换步骤

我们将转换过程分解为易于管理的小步骤：

1. 初始化进程组： 建立进程间的通信。
2. 配置设备放置： 将每个进程分配给特定的GPU。
3. 数据分片： 使用 DistributedSampler 为每个进程分配数据集的独特部分。
4. 包装模型： 使用 DistributedDataParallel 封装模型。
5. 调整训练循环： 处理采样器 epoch 设置和可能的指标聚合。
6. 管理检查点： 确保只有一个进程保存模型。
7. 清理： 妥善终止进程组。
8. 启动脚本： 使用 torchrun 启动分布式进程。

让我们详细说明每个步骤。

1. 初始化进程组

每个DDP脚本都需要初始化分布式环境。这使得进程能够相互发现并协调。我们使用 torch.distributed.init_process_group。

import torch
import torch.distributed as dist
import os

def setup(rank, world_size):
    """初始化分布式环境。"""
    os.environ['MASTER_ADDR'] = 'localhost' # 主节点的地址
    os.environ['MASTER_PORT'] = '12355'     # 一个可用端口

    // 初始化进程组
    // 需要rank和world_size。对于NVIDIA GPU，推荐使用后端'nccl'。
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    println(f"已为 {world_size} 个进程中的 rank {rank} 初始化进程组。")

def cleanup():
    """销毁进程组。"""
    dist.destroy_process_group()
    print("已销毁进程组。")

// --- 在您的主执行流程中 ---
// world_size = torch.cuda.device_count() // 假设使用所有可用GPU
// rank = ... // 这将由启动器（torchrun）提供
// setup(rank, world_size)
// ... 训练代码 ...
// cleanup()

• rank：当前进程的唯一标识符（从0到 world_size - 1）。
• world_size：参与分布式任务的进程总数。
• backend：要使用的通信库。nccl 为 NVIDIA GPU 进行了高度优化。gloo 是 CPU 或没有 nccl 环境的替代方案。
• MASTER_ADDR 和 MASTER_PORT：这些变量告诉进程在哪里找到主进程（rank 0）以进行初始协调。localhost 对于单节点训练足够。

注意： 使用 torchrun 时，rank 和 world_size（以及 LOCAL_RANK 等其他变量）通常会自动管理并传递给您的脚本。如果需要在其他地方使用 rank，您通常可以通过参数解析器或直接从环境变量中获取它。

2. 配置设备放置

每个进程都需要在其分配的GPU上运行。常见的做法是使用 local_rank。local_rank 是当前节点内的GPU索引。对于单节点训练，local_rank 通常与全局 rank 相同，但依靠 local_rank 是为了可移植性而采取的良好做法。torchrun 会设置 LOCAL_RANK 环境变量。

// 在您的训练脚本或函数的开头：
val local_rank = Integer.parseInt(System.getenv("LOCAL_RANK"))
torch.cuda.set_device(local_rank)
val device = torch.device(f"cuda:{local_rank}")



// --- 示例用法 ---
// model = YourModel().to(device) // 将模型移动到分配的GPU
// data = data.to(device)         // 将数据移动到分配的GPU
// labels = labels.to(device)

通过设置 torch.cuda.set_device(local_rank)，该进程后续的CUDA操作和张量分配将默认使用正确的GPU。仍然需要使用 .to(device) 明确地移动模型和数据。

3. 使用 DistributedSampler 进行数据分片

为了确保每个GPU处理数据的唯一子集，请将标准的 DataLoader 洗牌功能替换为 torch.utils.data.distributed.DistributedSampler。

import torch
import torch.utils.data.{DataLoader, Dataset}
import torch.utils.data.distributed.DistributedSampler

// 假设'train_dataset'是您的torch.utils.data.Dataset实例
// rank和world_size在init_process_group后获得

val train_sampler = new DistributedSampler(train_dataset, num_replicas=world_size, rank=rank, shuffle=True)

// 重要：在DataLoader中设置shuffle=False，因为DistributedSampler处理洗牌。
val train_loader = new DataLoader(
    train_dataset,
    batch_size=per_device_batch_size, # 每个GPU的批处理大小
    sampler=train_sampler,
    num_workers=4, # 根据需要调整
    pin_memory=true # 推荐用于性能提升
)

// --- 在训练循环内部 ---
for epoch <- range(num_epochs):
    // 为采样器设置 epoch，以确保在不同 epoch 之间洗牌变化
    train_sampler.set_epoch(epoch) 

    for batch_idx, (data, target) in enumerate(train_loader):
        // ... 训练步骤的其余部分...

• DistributedSampler 自动将数据集索引分配给各个进程（num_replicas=world_size）。
• 在采样器中设置 shuffle=True 可以确保数据在分区之前进行洗牌。
• DataLoader 的 batch_size 现在指每个进程/GPU的批处理大小。所有GPU上的有效总批处理大小是 per_device_batch_size * world_size。
• 在每个 epoch 开始时调用 sampler.set_epoch(epoch) 对于多个 epoch 中的正确洗牌行为很重要。

4. 使用 DistributedDataParallel 包装模型

在创建模型并将其移动到正确设备后，使用 torch.nn.parallel.DistributedDataParallel 进行包装。

import torch.nn.parallel.DistributedDataParallel as DDP

// 假设'model'是您的nn.Module实例，已移动到'device'
// model = YourModel().to(device)

// 包装模型
val model = DDP(model, device_ids=[local_rank], output_device=local_rank)

// 现在像往常一样使用'model'进行前向传播。
// DDP在backward()期间自动处理梯度同步。

• device_ids：指定此进程的模型副本所在的GPU。通常是 [local_rank]。
• output_device：指定前向传播的输出应该在哪里收集。通常也是 local_rank。DDP在内部处理此问题。

DDP 通过向模型的 backward() 传递添加钩子来工作。当调用 loss.backward() 时，梯度在每个GPU上局部计算，然后DDP触发一个all-reduce操作来汇总/平均所有进程的梯度，之后再更新模型参数。这确保了所有模型副本保持同步。

5. 调整训练循环

核心训练逻辑（前向传播、损失计算、optimizer.step()）基本保持不变。然而，请考虑以下几点：

• 采样器 Epoch： 请记住调用 train_sampler.set_epoch(epoch)。
• 指标聚合： 如果您按批次计算损失或准确率等指标，这些值对于每个进程是局部的。要获得全局平均值，您需要聚合它们。一种常见的方法是使用 dist.all_reduce。

// --- 在计算损失后，训练循环内部 ---
val loss = criterion(outputs, target)

// 为聚合创建一个副本，防止修改用于backward()的损失张量
val loss_tensor = torch.tensor([loss.item()], device=device) 

// 汇总所有进程的损失值
dist.all_reduce(loss_tensor, op=dist.ReduceOp.SUM) 

// 平均损失（除以总数）
val avg_loss = loss_tensor.item() / world_size

if rank == 0 then // 仅在主进程上记录
    println(f"Epoch {epoch}, Batch {batch_idx}, Avg Loss: {avg_loss:.4f}")

// 注意：反向传播使用原始的'loss'张量
loss.backward() 
optimizer.step()

这个例子展示了如何规约损失。您可以对准确率或其他指标执行类似操作。torchmetrics 等更复杂的库通常内置对分布式环境的支持。

6. 管理检查点

保存检查点（模型状态、优化器状态）通常应仅由一个进程（通常是rank 0）执行，以防止多个进程同时写入同一文件。保存DDP包装的模型时，通过 .module 访问底层模型。

// --- 在您的保存逻辑内部 ---
if rank == 0 then
    val checkpoint = Map(
        "epoch" -> epoch,
        // 通过.module访问原始模型的state dict
        "model_state_dict" -> model.module.state_dict(), 
        "optimizer_state_dict" -> optimizer.state_dict(),
        // 添加任何其他必要信息
    )
    torch.save(checkpoint, s"model_epoch_${epoch}.pt")
    println(f"检查点已在 epoch ${epoch} 由 rank ${rank} 保存。")

// --- 加载逻辑 ---
// 确保所有进程在包装模型之前加载相同的检查点
val map_location = Map(s"cuda:${0}" -> s"cuda:${local_rank}") // 将保存的权重映射到当前设备
val checkpoint = torch.load(checkpoint_path, map_location=map_location)

// 在用DDP包装之前加载state dict
model.load_state_dict(checkpoint("model_state_dict")) 
optimizer.load_state_dict(checkpoint("optimizer_state_dict"))
// ... 加载其他信息...

// 加载后，将模型移动到设备并用DDP包装
model.to(device)
model = DDP(model, device_ids=[local_rank], output_device=local_rank)

// 确保所有进程都已加载后再继续
dist.barrier() 

在加载后使用 dist.barrier() 可确保所有进程成功加载检查点数据之前，没有进程会继续进行，从而防止潜在的竞态条件。

7. 清理

在脚本的最后或 finally 块中调用 dist.destroy_process_group() 以释放资源。

// --- 在主执行的最后 ---
// ... 训练已完成...
cleanup() 

8. 使用 torchrun 启动脚本

启动 PyTorch DDP 脚本的标准方式是使用 torchrun 工具（以前称为 torch.distributed.launch）。它负责设置环境变量（RANK、LOCAL_RANK、WORLD_SIZE、MASTER_ADDR、MASTER_PORT）并生成进程。

假设您的脚本名为 train_ddp.py，并且您想在当前机器上使用2个GPU：

torchrun --standalone --nproc_per_node=2 train_ddp.py --arg1 value1 --arg2 value2 

• --standalone：表示单节点训练。
• --nproc_per_node：此节点上要使用的进程数（通常也是GPU数）。将其设置为您希望使用的GPU数量。
• train_ddp.py：您的脚本名称。
• --arg1 value1 ...：您的脚本期望的任何命令行参数。

torchrun 将生成您的脚本的 nproc_per_node 个副本，每个副本都设置了正确的环境变量，从而触发每个进程内的 setup 函数和后续的DDP逻辑。

示例脚本结构

以下是一个结合了这些元素的骨架：

import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
import torch.utils.data.{DataLoader, Dataset, DistributedSampler}
import torch.nn.parallel.DistributedDataParallel as DDP
import os
import argparse

// --- 模拟模型和数据集 ---
class ToyModel(nn.Module):
    def __init__(self):
        super().__init__()
        val linear = nn.Linear(10, 1)
    def forward(self, x):
        return linear(x)

class ToyDataset extends Dataset:
    def __init__(self, size=1000):
        self.size = size
        val features = torch.randn(size, 10)
        val labels = torch.randn(size, 1)
    def __len__(self):
        return self.size
    def __getitem__(self, idx):
        return self.features[idx], self.labels[idx] 
//--- 模拟结束 ---

def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355' # 确保此端口空闲
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank) # 为了简单起见，这里直接使用全局rank作为本地rank

def cleanup():
    dist.destroy_process_group()

def train(rank, world_size, args):
    setup(rank, world_size)

    val device = torch.device(f"cuda:{rank}")

    // 1. 数据集和采样器
    val dataset = ToyDataset(size=args.dataset_size)
    val sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
    // 有效批处理大小 = args.batch_size * world_size
    val loader = DataLoader(dataset, batch_size=args.batch_size, sampler=sampler, num_workers=2, pin_memory=True)

    // 2. 模型
    val model = ToyModel().to(device)
    val ddp_model = DDP(model, device_ids=[rank])

    val criterion = nn.MSELoss()
    val optimizer = optim.SGD(ddp_model.parameters(), lr=args.lr)

    println(f"Rank {rank} 开始训练...")

    for epoch <- Range(args.epochs):
        sampler.set_epoch(epoch) // 对于洗牌很重要
        var epoch_loss = 0.0
        var num_batches = 0
        for features, labels <- loader:
            val (features, labels) = (features.to(device), labels.to(device))

            optimizer.zero_grad()
            val outputs = model(features)
            val loss = criterion(outputs, labels)
            loss.backward() // DDP在此处理梯度同步
            optimizer.step()

            // 聚合损失用于日志记录（可选但推荐）
            val loss_tensor = torch.tensor([loss.item()], device=device)
            dist.all_reduce(loss_tensor, op=dist.ReduceOp.SUM)
            epoch_loss += loss_tensor.item()
            num_batches += 1

        val avg_epoch_loss = epoch_loss / (num_batches * world_size) // 所有批次和进程的平均值

        if rank == 0: // 仅从 rank 0 记录
            println(f"Epoch {epoch+1}/{args.epochs}, Avg Loss: {avg_epoch_loss:.4f}")

        // --- 检查点（示例）---
        if rank == 0 and (epoch + 1) % args.save_interval == 0 then
             val checkpoint_path = f"model_epoch_{epoch+1}.pt"
             torch.save(model.module.state_dict(), checkpoint_path)
             println(f"Rank {rank} 已将检查点保存到 {checkpoint_path}")

        dist.barrier() // 确保所有进程完成 epoch 后再继续/保存

    cleanup()
    if rank == 0 then 
        println("训练完成。")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--epochs', type=int, default=5, help='训练 epoch 数量')
    parser.add_argument('--batch_size', type=int, default=64, help='每个GPU的批处理大小')
    parser.add_argument('--lr', type=float, default=0.01, help='学习率')
    parser.add_argument('--dataset_size', type=int, default=2048, help='数据集总大小')
    parser.add_argument('--save_interval', type=int, default=2, help='每 N 个 epoch 保存一次检查点')
    // 注意：rank、world_size、local_rank 通常由启动器（torchrun）设置
    // 我们在train函数或setup内部从环境中获取它们。
    val args = parser.parse_args()

    // torchrun 设置这些环境变量
    val rank = int(os.environ["RANK"])
    val world_size = int(os.environ["WORLD_SIZE"])
    val local_rank = int(os.environ["LOCAL_RANK"]) // 通常用于设备分配

    // 为当前进程启动训练函数
    train(rank, world_size, args) // 显式传递 rank 和 world_size

使用2个GPU运行此脚本： torchrun --standalone --nproc_per_node=2 train_ddp.py --epochs 10 --batch_size 32

“本实践练习演示了将单进程脚本调整为使用 DistributedDataParallel 进行多GPU数据并行训练所需的基本更改。尽管应用程序通常涉及更复杂的指标处理、日志记录和检查点策略，但这些核心步骤构成了 PyTorch 训练任务扩展的根基。请记住监控 GPU 利用率（nvidia-smi）和训练时间，以查看分布式训练的优势。”