My-Notes

002.FSDP

5分钟阅读

Fully Sharded Data Parallel (FSDP) 是一种数据并行方法，最早是在 2021 年由 FairScale-FSDP 提出的，并在后续被集成到了 PyTorch 1.11 版本中

FSDP 可以看作是微软 Deepspeed 框架中提出的三种级别的 ZERO 算法中的 ZERO-3 的实现。它通过将模型的梯度、优化器状态和参数进行分片操作，使得每个 GPU 只存储部分参数信息，从而优化了资源的利用和提高了训练效率。此外，FSDP 也与包括 Tensor 实现、调度器系统和 CUDA 内存缓存分配器在内的几个关键 PyTorch 核心组件紧密协同设计，以提供非侵入式用户体验和高训练效率

ZeRO 有三种级别的算法，分别是 ZERO-1、ZERO-2 和 ZERO-3。ZERO-3 是最高级别的算法，它将模型的梯度、优化器状态和参数进行分片操作，使得每个 GPU 只存储部分参数信息，从而优化了资源的利用和提高了训练效率。FSDP 是 ZeRO-3 的实现

FSDF PyTorch

在 PyTorch 中使用 FSDP 可以有效地训练大型模型，特别是在显存或内存受限的情况下。FSDP 是一种数据并行技术，它将模型的参数、梯度和优化器状态跨多个设备进行分片。以下是基本步骤：

  1. 初始化分布式环境：首先，需要初始化分布式环境以帮助进程间通信。这通常通过 torch.distributed.init_process_group 函数完成
  2. 设置本地排名：每个进程需要根据其 local_rank 设置应该使用的 GPU，这可以通过环境变量或命令行参数来获取
  3. 创建 FSDP 模型：使用 FullyShardedDataParallel 类来包装你的模型，这将允许模型参数在多个 GPU 上进行分片。例如：
from torch.distributed.fsdp import FullyShardedDataParallel
 
model = MyModel()
model = model.to(device)  # 将模型移动到GPU
fsdp_model = FullyShardedDataParallel(model, ...其他参数...)
  1. 配置 FSDP 参数：FSDP提供了多种参数来配置其行为，例如 cpu_offload 用于决定是否将参数卸载到 CPU，以及 sharding_strategy 用于指定分片策略
  2. 训练模型：在训练循环中，FSDP 会自动处理参数的分片和梯度的聚合。你只需要像往常一样进行前向和反向传播
  3. 保存和加载模型：当使用 FSDP 时，保存和加载模型可能需要一些特殊的处理，以确保分片的参数被正确处理

下面是一个更详细的示例代码，展示了如何使用 FSDP 来训练一个简单的模型：

import torch
import torch.nn as nn
from torch.distributed.fsdp import FullyShardedDataParallel, CPUOffload
 
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.layer1 = nn.Linear(8, 4)
        self.layer2 = nn.Linear(4, 16)
        self.layer3 = nn.Linear(16, 4)
 
    def forward(self, x):
        x = torch.relu(self.layer1(x))
        x = torch.relu(self.layer2(x))
        x = self.layer3(x)
        return x
 
# 初始化分布式环境
torch.distributed.init_process_group(backend='nccl')
 
# 设置本地排名和设备
local_rank = torch.distributed.get_rank()
world_size = torch.distributed.get_world_size()
torch.cuda.set_device(local_rank)
 
# 创建模型并移动到对应的GPU
model = MyModel().to(local_rank)
 
# 使用FSDP包装模型
fsdp_model = FullyShardedDataParallel(
    model,
    cpu_offload=CPUOffload(offload_params=True),
    # 其他FSDP参数
)
 
# 定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(fsdp_model.parameters(), lr=0.001)
 
# 训练循环
for epoch in range(num_epochs):
    for data, target in dataloader:
        data, target = data.to(local_rank), target.to(local_rank)
        optimizer.zero_grad()
        output = fsdp_model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

FSDP Huggingface/Accelerate

作为一个高级的深度学习库，Huggingface 提供了一个名为 Accelerate 的库，它可以帮助用户更容易地使用分布式训练技术，包括 FSDP。Accelerate 提供了一个简单的 API，可以在几行代码中将模型转换为 FSDP 模型，并自动处理分布式训练的细节

compute_environment: LOCAL_MACHINE
debug: false
distributed_type: FSDP # 使用FSDP的配置
downcast_bf16: 'no'
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch_policy: BACKWARD_PRE
  fsdp_forward_prefetch: false
  fsdp_cpu_ram_efficient_loading: true
  fsdp_offload_params: false
  fsdp_sharding_strategy: FULL_SHARD
  fsdp_state_dict_type: SHARDED_STATE_DICT
  fsdp_sync_module_states: true
  fsdp_transformer_layer_cls_to_wrap: BertLayer
  fsdp_use_orig_params: true
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 2
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

Tips：2024 年 9 月 13 号，在 Accelerate 开发趋于稳定将近一年后的，正式发布了 Accelerate 1.0.0 —— Accelerate 的第一个发布候选版本

关系图谱