DistributedDataParallel (DDP) 通过复制模型和平均梯度有效地在多个GPU上扩展训练。然而,它根本上要求每个GPU都持有完整的模型、其梯度和优化器状态。这在处理包含数十亿参数的模型时成为一个限制因素,因为这些模型可能轻易超出甚至高端加速器的内存容量。全分片数据并行(FSDP)通过扩展数据并行理念,同时大幅减少每个GPU的内存占用,提供了一种解决办法。FSDP不是复制整个模型,而是将模型的参数、梯度和优化器状态分片或分区到数据并行工作器(GPU)上。FSDP运行机制FSDP的核心是确保数据并行组中的每个GPU在任何给定时间点只持有模型参数、梯度和优化器状态的一部分(一个“分片”)。完整张量仅在计算需要时才临时重建。以下是训练过程的详细步骤:初始化: 模型使用 FullyShardedDataParallel 模块进行封装。在初始化期间,参数、梯度和优化器状态被划分到参与进程组的GPU上。每个GPU负责管理其分配到的分片。前向传播:当FSDP封装的层需要执行计算时,每个GPU使用 all_gather 集合通信操作从组中所有其他GPU收集该特定层所需的完整参数。该层的计算在当前本地的完整参数上进行。该层的前向计算完成后,除了GPU自己的分片外,完整参数会立即被丢弃,从而释放内存。这个过程逐层(或逐块,取决于封装策略)重复进行。反向传播:当给定层的参数在本地计算出梯度时,它们不会像DDP那样立即在所有GPU上求平均。相反,会执行 reduce_scatter 操作。此操作计算所有GPU上梯度的平均值,并同时对平均结果进行分片,仅向每个GPU发送与其管理参数分片对应的梯度部分(分片)。这个分片后的梯度被存储起来,再次保持了较低的内存使用。reduce_scatter 后,完整梯度被丢弃。优化器步骤:每个GPU的优化器只需要更新其负责的参数分片。由于优化器状态(如Adam中的动量缓冲区)也与参数一同分片,因此优化器步骤可以在每个GPU上仅使用其梯度和优化器状态的分片进行本地处理。这种方法大幅减少了每个GPU所需的峰值内存,因为只有当前执行层的参数以及完整模型、梯度和优化器状态的分片才会被持久存储。digraph FSDP_vs_DDP { rankdir=LR; node [shape=record, style=filled, fillcolor="#e9ecef", fontname="helvetica"]; edge [fontname="helvetica"]; subgraph cluster_ddp { label = "DDP内存布局"; bgcolor="#f8f9fa"; style=dashed; node [fillcolor="#a5d8ff"]; ddp_gpu0 [label="{ GPU 0 | { 完整模型 | 完整梯度 | 完整优化器状态 } }"]; ddp_gpu1 [label="{ GPU 1 | { 完整模型 | 完整梯度 | 完整优化器状态 } }"]; ddp_gpu0 -> ddp_gpu1 [style=invis]; // 如有需要,确保垂直布局 } subgraph cluster_fsdp { label = "FSDP内存布局"; bgcolor="#f8f9fa"; style=dashed; node [fillcolor="#96f2d7"]; fsdp_gpu0 [label="{ GPU 0 | { 参数分片 0 | 梯度分片 0 | 优化器状态分片 0 } }"]; fsdp_gpu1 [label="{ GPU 1 | { 参数分片 1 | 梯度分片 1 | 优化器状态分片 1 } }"]; fsdp_gpu0 -> fsdp_gpu1 [style=invis]; // 如有需要,确保垂直布局 } invisible_node [style=invis]; invisible_node -> cluster_ddp [style=invis]; invisible_node -> cluster_fsdp [style=invis]; label="每GPU内存使用比较"; fontsize=12; fontcolor="#495057"; }分布式数据并行(DDP)和全分片数据并行(FSDP)每GPU内存分配比较。DDP复制所有组件,而FSDP将其分片。在PyTorch中实现FSDPPyTorch通过 torch.distributed.fsdp.FullyShardedDataParallel 类提供了对FSDP的原生支持。集成它通常涉及封装您的模型定义。import torch import torch.nn as nn import torch.distributed as dist from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.distributed.fsdp.wrap import 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(nn.Module): # 子模块定义示例 def __init__(self, dim, ff_dim): super().__init__() self.layer_norm = nn.LayerNorm(dim) self.attention = nn.MultiheadAttention(dim, num_heads=8) # Simplified self.ffn = nn.Sequential( nn.Linear(dim, ff_dim), nn.ReLU(), nn.Linear(ff_dim, dim) ) def forward(self, x): x = self.layer_norm(x + self.attention(x, x, x)[0]) x = x + self.ffn(x) return x class BigModel(nn.Module): def __init__(self, num_layers, dim, ff_dim, vocab_size): super().__init__() self.embedding = nn.Embedding(vocab_size, dim) self.layers = nn.ModuleList( [LargeTransformerBlock(dim, ff_dim) for _ in range(num_layers)] ) self.output_head = nn.Linear(dim, vocab_size) def forward(self, x): x = self.embedding(x) for layer in self.layers: x = layer(x) x = self.output_head(x) return x # --- FSDP 设置 --- model = BigModel(num_layers=48, dim=2048, ff_dim=8192, vocab_size=50000).to(torch.cuda.current_device()) # 定义一个自动封装策略(可选但推荐用于大型模型) # 这会根据大小封装子模块(例如 LargeTransformerBlock) auto_wrap_policy = functools.partial( size_based_auto_wrap_policy, min_num_params=1_000_000 # 示例阈值 ) # 用 FSDP 封装模型 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 封装模型后构建 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()实现中的要点:模型定义: 像往常一样使用 nn.Module 定义您的模型。封装: 使用 FSDP 封装模型实例。请注意,模型应在封装 之前 移至目标设备。优化器: 在使用 FSDP 封装模型 之后 构建优化器,并将 fsdp_model.parameters() 传递给它。这确保优化器了解分片参数和状态。自动封装策略: 对于复杂模型,定义 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在大规模模型训练能力方面的一个重要进展。