使用完全分片数据并行（FSDP）进行高级模型训练
创建于：2024年10月31日 | 最后更新：2024年10月31日 | 最后验证：2024年11月5日
作者: Hamid Shojanazeri, Less Wright, Rohan Varma, Yanli Zhao
PyTorch的全分片数据并行模块：一个用于跨分片模块参数的包装器
数据并行工作者。
PyTorch 1.12 或更高版本
阅读有关FSDP API的内容。
本教程介绍了作为PyTorch 1.12版本一部分的完全分片数据并行（FSDP）的更高级功能。要熟悉FSDP，请参考FSDP入门教程。
在本教程中，我们以文本摘要为例，使用FSDP微调了一个HuggingFace (HF) T5模型。
该示例使用Wikihow，为了简单起见，我们将在单个节点P4dn实例上展示训练，该实例配备8个A100 GPU。我们现在有几篇博客文章（ (link1), (link2)） 和一篇关于在多节点集群上进行大规模FSDP训练的论文。
FSDP 是一个生产就绪的软件包，专注于易用性、性能和长期支持。FSDP 的主要优势之一是减少每个 GPU 上的内存占用。这使得与 DDP 相比，可以用更少的总内存训练更大的模型，并利用计算和通信的重叠来高效地训练模型。这种减少的内存压力可以用来训练更大的模型或增加批量大小，可能有助于提高整体训练吞吐量。你可以在这里阅读更多关于 PyTorch FSDP 的信息。
FSDP 工作原理回顾
在高层次上，FDSP 的工作方式如下：
在构造函数中
分片模型参数，每个等级只保留自己的分片
在前向传播过程中
运行 all_gather 以从所有等级收集所有分片，以恢复此 FSDP 单元的完整参数并运行前向计算
丢弃刚刚收集的非自有参数分片以释放内存
在反向传播过程中
运行 all_gather 以从所有等级收集所有分片，以恢复此 FSDP 单元中的完整参数并运行反向计算
丢弃非自有参数以释放内存。
运行reduce_scatter以同步梯度
微调HF T5
HF T5 预训练模型有四种不同大小，从6000万参数的小型模型到110亿参数的XXL模型。在本教程中，我们展示了如何使用FSDP对T5 3B模型进行微调，以使用WikiHow数据集进行文本摘要。本教程的主要重点是突出FSDP中可用的不同功能，这些功能有助于训练超过30亿参数的大规模模型。此外，我们还涵盖了基于Transformer模型的特定功能。本教程的代码可在Pytorch示例中找到。
设置
1.1 安装最新的 PyTorch
pip3 install torch torchvision torchaudio
1.2 数据集设置
请创建一个data文件夹，从wikihowAll.csv和wikihowSep.cs下载WikiHow数据集，并将它们放入data文件夹中。我们将使用来自summarization_dataset的wikihow数据集。
接下来，我们将以下代码片段添加到Python脚本“T5_training.py”中。
注意
本教程的完整源代码可在PyTorch examples中找到。
1.3 导入必要的包：
import os
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from transformers import AutoTokenizer, GPT2TokenizerFast
from transformers import T5Tokenizer, T5ForConditionalGeneration
import functools
from torch.optim.lr_scheduler import StepLR
import torch.nn.functional as F
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
from transformers.models.t5.modeling_t5 import T5Block
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
checkpoint_wrapper,
CheckpointImpl,
apply_activation_checkpointing_wrapper)
from torch.distributed.fsdp import (
FullyShardedDataParallel as FSDP,
MixedPrecision,
BackwardPrefetch,
ShardingStrategy,
FullStateDictConfig,
StateDictType,
)
from torch.distributed.fsdp.wrap import (
transformer_auto_wrap_policy,
enable_wrap,
wrap,
)
from functools import partial
from torch.utils.data import DataLoader
from pathlib import Path
from summarization_dataset import *
from transformers.models.t5.modeling_t5 import T5Block
from typing import Type
import time
import tqdm
from datetime import datetime
1.4 分布式训练设置。 在这里，我们使用两个辅助函数来初始化分布式训练的进程，然后在训练完成后进行清理。在本教程中，我们将使用torch elastic，使用torchrun，它将自动设置工作者的RANK和WORLD_SIZE。
def setup():
# initialize the process group
dist.init_process_group("nccl")
def cleanup():
dist.destroy_process_group()
2.1 设置HuggingFace T5模型：
def setup_model(model_name):
model = T5ForConditionalGeneration.from_pretrained(model_name)
tokenizer = T5Tokenizer.from_pretrained(model_name)
return model, tokenizer
我们还在这里添加了一些辅助函数，用于日期和格式化内存指标。
def get_date_of_run():
"""create date and time for file save uniqueness
example: 2022-05-07-08:31:12_PM'
"""
date_of_run = datetime.now().strftime("%Y-%m-%d-%I:%M:%S_%p")
print(f"--> current date and time of run = {date_of_run}")
return date_of_run
def format_metrics_to_gb(item):
"""quick function to format numbers to gigabyte and round to 4 digit precision"""
metric_num = item / g_gigabyte
metric_num = round(metric_num, ndigits=4)
return metric_num
2.2 定义一个训练函数：
def train(args, model, rank, world_size, train_loader, optimizer, epoch, sampler=None):
model.train()
local_rank = int(os.environ['LOCAL_RANK'])
fsdp_loss = torch.zeros(2).to(local_rank)
if sampler:
sampler.set_epoch(epoch)
if rank==0:
inner_pbar = tqdm.tqdm(
range(len(train_loader)), colour="blue", desc="r0 Training Epoch"
)
for batch in train_loader:
for key in batch.keys():
batch[key] = batch[key].to(local_rank)
optimizer.zero_grad()
output = model(input_ids=batch["source_ids"],attention_mask=batch["source_mask"],labels=batch["target_ids"] )
loss = output["loss"]
loss.backward()
optimizer.step()
fsdp_loss[0] += loss.item()
fsdp_loss[1] += len(batch)
if rank==0:
inner_pbar.update(1)
dist.all_reduce(fsdp_loss, op=dist.ReduceOp.SUM)
train_accuracy = fsdp_loss[0] / fsdp_loss[1]
if rank == 0:
inner_pbar.close()
print(
f"Train Epoch: \t{epoch}, Loss: \t{train_accuracy:.4f}"
)
return train_accuracy
2.3 定义一个验证函数：
def validation(model, rank, world_size, val_loader):
model.eval()
correct = 0
local_rank = int(os.environ['LOCAL_RANK'])
fsdp_loss = torch.zeros(3).to(local_rank)
if rank == 0:
inner_pbar = tqdm.tqdm(
range(len(val_loader)), colour="green", desc="Validation Epoch"
)
with torch.no_grad():
for batch in val_loader:
for key in batch.keys():
batch[key] = batch[key].to(local_rank)
output = model(input_ids=batch["source_ids"],attention_mask=batch["source_mask"],labels=batch["target_ids"])
fsdp_loss[0] += output["loss"].item() # sum up batch loss
fsdp_loss[1] += len(batch)
if rank==0:
inner_pbar.update(1)
dist.all_reduce(fsdp_loss, op=dist.ReduceOp.SUM)
val_loss = fsdp_loss[0] / fsdp_loss[1]
if rank == 0:
inner_pbar.close()
print(f"Validation Loss: {val_loss:.4f}")
return val_loss
2.4 定义一个分布式训练函数，将模型封装在FSDP中：
def fsdp_main(args):
model, tokenizer = setup_model("t5-base")
local_rank = int(os.environ['LOCAL_RANK'])
rank = int(os.environ['RANK'])
world_size = int(os.environ['WORLD_SIZE'])
dataset = load_dataset('wikihow', 'all', data_dir='data/')
print(dataset.keys())
print("Size of train dataset: ", dataset['train'].shape)
print("Size of Validation dataset: ", dataset['validation'].shape)
#wikihow(tokenizer, type_path, num_samples, input_length, output_length, print_text=False)
train_dataset = wikihow(tokenizer, 'train', 1500, 512, 150, False)
val_dataset = wikihow(tokenizer, 'validation', 300, 512, 150, False)
sampler1 = DistributedSampler(train_dataset, rank=rank, num_replicas=world_size, shuffle=True)
sampler2 = DistributedSampler(val_dataset, rank=rank, num_replicas=world_size)
setup()
train_kwargs = {'batch_size': args.batch_size, 'sampler': sampler1}
test_kwargs = {'batch_size': args.test_batch_size, 'sampler': sampler2}
cuda_kwargs = {'num_workers': 2,
'pin_memory': True,
'shuffle': False}
train_kwargs.update(cuda_kwargs)
test_kwargs.update(cuda_kwargs)
train_loader = torch.utils.data.DataLoader(train_dataset,**train_kwargs)
val_loader = torch.utils.data.DataLoader(val_dataset, **test_kwargs)
t5_auto_wrap_policy = functools.partial(
transformer_auto_wrap_policy,
transformer_layer_cls={
T5Block,
},
)
sharding_strategy: ShardingStrategy = ShardingStrategy.SHARD_GRAD_OP #for Zero2 and FULL_SHARD for Zero3
torch.cuda.set_device(local_rank)
#init_start_event = torch.cuda.Event(enable_timing=True)
#init_end_event = torch.cuda.Event(enable_timing=True)
#init_start_event.record()
bf16_ready = (
torch.version.cuda
and torch.cuda.is_bf16_supported()
and LooseVersion(torch.version.cuda) >= "11.0"
and dist.is_nccl_available()
and nccl.version() >= (2, 10)
)
if bf16_ready:
mp_policy = bfSixteen
else:
mp_policy = None # defaults to fp32
# model is on CPU before input to FSDP
model = FSDP(model,
auto_wrap_policy=t5_auto_wrap_policy,
mixed_precision=mp_policy,
#sharding_strategy=sharding_strategy,
device_id=torch.cuda.current_device())
optimizer = optim.AdamW(model.parameters(), lr=args.lr)
scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
best_val_loss = float("inf")
curr_val_loss = float("inf")
file_save_name = "T5-model-"
if rank == 0:
time_of_run = get_date_of_run()
dur = []
train_acc_tracking = []
val_acc_tracking = []
training_start_time = time.time()
if rank == 0 and args.track_memory:
mem_alloc_tracker = []
mem_reserved_tracker = []
for epoch in range(1, args.epochs + 1):
t0 = time.time()
train_accuracy = train(args, model, rank, world_size, train_loader, optimizer, epoch, sampler=sampler1)
if args.run_validation:
curr_val_loss = validation(model, rank, world_size, val_loader)
scheduler.step()
if rank == 0:
print(f"--> epoch {epoch} completed...entering save and stats zone")
dur.append(time.time() - t0)
train_acc_tracking.append(train_accuracy.item())
if args.run_validation:
val_acc_tracking.append(curr_val_loss.item())
if args.track_memory:
mem_alloc_tracker.append(
format_metrics_to_gb(torch.cuda.memory_allocated())
)
mem_reserved_tracker.append(
format_metrics_to_gb(torch.cuda.memory_reserved())
)
print(f"completed save and stats zone...")
if args.save_model and curr_val_loss < best_val_loss:
# save
if rank == 0:
print(f"--> entering save model state")
save_policy = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
with FSDP.state_dict_type(
model, StateDictType.FULL_STATE_DICT, save_policy
):
cpu_state = model.state_dict()
#print(f"saving process: rank {rank} done w state_dict")
if rank == 0:
print(f"--> saving model ...")
currEpoch = (
"-" + str(epoch) + "-" + str(round(curr_val_loss.item(), 4)) + ".pt"
)
print(f"--> attempting to save model prefix {currEpoch}")
save_name = file_save_name + "-" + time_of_run + "-" + currEpoch
print(f"--> saving as model name {save_name}")
torch.save(cpu_state, save_name)
if curr_val_loss < best_val_loss:
best_val_loss = curr_val_loss
if rank==0:
print(f"-->>>> New Val Loss Record: {best_val_loss}")
dist.barrier()
cleanup()
2.5 解析参数并设置主函数：
if __name__ == '__main__':
# Training settings
parser = argparse.ArgumentParser(description='PyTorch T5 FSDP Example')
parser.add_argument('--batch-size', type=int, default=4, metavar='N',
help='input batch size for training (default: 64)')
parser.add_argument('--test-batch-size', type=int, default=4, metavar='N',
help='input batch size for testing (default: 1000)')
parser.add_argument('--epochs', type=int, default=2, metavar='N',
help='number of epochs to train (default: 3)')
parser.add_argument('--lr', type=float, default=.002, metavar='LR',
help='learning rate (default: .002)')
parser.add_argument('--gamma', type=float, default=0.7, metavar='M',
help='Learning rate step gamma (default: 0.7)')
parser.add_argument('--no-cuda', action='store_true', default=False,
help='disables CUDA training')
parser.add_argument('--seed', type=int, default=1, metavar='S',
help='random seed (default: 1)')
parser.add_argument('--track_memory', action='store_false', default=True,
help='track the gpu memory')
parser.add_argument('--run_validation', action='store_false', default=True,
help='running the validation')
parser.add_argument('--save-model', action='store_false', default=True,
help='For Saving the current Model')
args = parser.parse_args()
torch.manual_seed(args.seed)
fsdp_main(args)
要使用torchrun运行训练：
torchrun --nnodes 1 --nproc_per_node 4 T5_training.py
变压器包装政策
正如在之前的教程中所讨论的， auto_wrap_policy 是 FSDP 功能之一，使得自动分片给定模型并将模型、优化器和梯度分片放入不同的 FSDP 单元变得容易。
对于一些架构，如Transformer编码器-解码器，模型的某些部分（如嵌入表）在编码器和解码器之间共享。在这种情况下，我们需要将嵌入表放置在外部FSDP单元中，以便编码器和解码器都可以访问它。此外，通过为Transformer注册层类，可以使分片计划在通信效率上大大提高。在PyTorch 1.12中，FSDP增加了这一支持，现在我们有了针对Transformer的包装策略。
它可以如下创建，其中T5Block代表T5变压器层类（包含MHSA和FFN）。
t5_auto_wrap_policy = functools.partial(
transformer_auto_wrap_policy,
transformer_layer_cls={
T5Block,
},
)
torch.cuda.set_device(local_rank)
model = FSDP(model,
auto_wrap_policy=t5_auto_wrap_policy)
要查看包装后的模型，您可以轻松地打印模型并直观地检查分片和FSDP单元。
混合精度
FSDP 支持灵活的混合精度训练，允许使用任意降低精度的类型（如 fp16 或 bfloat16）。目前，BFloat16 仅在 Ampere GPU 上可用，因此在使用之前需要确认本地支持。例如，在 V100 上，BFloat16 仍然可以运行，但由于它是非本地运行的，可能会导致显著的减速。
要检查是否原生支持BFloat16，您可以使用以下方法：
bf16_ready = (
torch.version.cuda
and torch.cuda.is_bf16_supported()
and LooseVersion(torch.version.cuda) >= "11.0"
and dist.is_nccl_available()
and nccl.version() >= (2, 10)
)
FSDP中混合精度的优势之一是提供对参数、梯度和缓冲区的不同精度级别的细粒度控制，如下所示：
fpSixteen = MixedPrecision(
param_dtype=torch.float16,
# Gradient communication precision.
reduce_dtype=torch.float16,
# Buffer precision.
buffer_dtype=torch.float16,
)
bfSixteen = MixedPrecision(
param_dtype=torch.bfloat16,
# Gradient communication precision.
reduce_dtype=torch.bfloat16,
# Buffer precision.
buffer_dtype=torch.bfloat16,
)
fp32_policy = MixedPrecision(
param_dtype=torch.float32,
# Gradient communication precision.
reduce_dtype=torch.float32,
# Buffer precision.
buffer_dtype=torch.float32,
)
请注意，如果未指定某种类型（参数、reduce、buffer），它们将不会被转换。
这种灵活性允许用户进行细粒度的控制，例如仅设置梯度通信以降低精度进行，而所有参数/缓冲区计算则以全精度完成。这在节点内通信是主要瓶颈且参数/缓冲区必须保持全精度以避免精度问题的情况下可能非常有用。可以通过以下策略实现：
grad_bf16 = MixedPrecision(reduce_dtype=torch.bfloat16)
在2.4版本中，我们只需将相关的混合精度策略添加到FSDP包装器中：
model = FSDP(model,
auto_wrap_policy=t5_auto_wrap_policy,
mixed_precision=bfSixteen)
在我们的实验中，我们观察到使用BFloat16进行训练时，速度提升了多达4倍，并且在某些实验中内存减少了约30%，这可以用于增加批量大小。
在设备上初始化FSDP模型
在1.12版本中，FSDP支持一个device_id参数，用于在由device_id指定的设备上初始化输入CPU模块。当整个模型无法适应单个GPU，但可以适应主机的CPU内存时，这非常有用。当指定device_id时，FSDP将以每个FSDP单元为基础将模型移动到指定设备，避免GPU内存不足问题，同时初始化速度比基于CPU的初始化快几倍：
torch.cuda.set_device(local_rank)
model = FSDP(model,
auto_wrap_policy=t5_auto_wrap_policy,
mixed_precision=bfSixteen,
device_id=torch.cuda.current_device())
分片策略
FSDP 分片策略默认设置为完全分片模型参数，梯度和优化器状态在所有等级之间分片。（也称为 Zero3 分片）。如果您有兴趣使用 Zero2 分片策略，其中仅优化器状态和梯度被分片，FSDP 支持此功能，通过使用“ShardingStrategy.SHARD_GRAD_OP”而不是“ShardingStrategy.FULL_SHARD”传递给 FSDP 初始化，如下所示：
torch.cuda.set_device(local_rank)
model = FSDP(model,
auto_wrap_policy=t5_auto_wrap_policy,
mixed_precision=bfSixteen,
device_id=torch.cuda.current_device(),
sharding_strategy=ShardingStrategy.SHARD_GRAD_OP # ZERO2)
这将减少FSDP中的通信开销，在这种情况下，它在正向传播后和反向传播过程中保持完整的参数。
这节省了反向传播期间的all_gather操作，从而减少了通信量，但代价是更高的内存占用。请注意，完整的模型参数在反向传播结束时会被释放，并且all_gather操作将在下一次前向传播时进行。
向后预取
向后预取设置控制下一个FSDP单元的参数何时应被请求。通过将其设置为BACKWARD_PRE，可以在当前单元的计算开始之前更早地请求并接收下一个FSDP单元的参数。这可以重叠all_gather通信和梯度计算，从而以稍微增加内存消耗为代价提高训练速度。它可以在2.4版本的FSDP包装器中如下使用：
torch.cuda.set_device(local_rank)
model = FSDP(model,
auto_wrap_policy=t5_auto_wrap_policy,
mixed_precision=bfSixteen,
device_id=torch.cuda.current_device(),
backward_prefetch = BackwardPrefetch.BACKWARD_PRE)
backward_prefetch 有两种模式，BACKWARD_PRE 和 BACKWARD_POST。 BACKWARD_POST 意味着下一个 FSDP 单元的参数在当前 FSDP 单元处理完成之前不会被请求， 从而最小化内存开销。在某些情况下，使用 BACKWARD_PRE 可以将模型训练速度提高 2-10%， 对于更大的模型，速度提升更为显著。
模型检查点保存，通过流式传输到Rank0 CPU
为了使用FULL_STATE_DICT保存模型检查点，这种方式保存模型的方式与本地模型相同，PyTorch 1.12提供了一些工具来支持保存更大的模型。
首先，可以指定一个FullStateDictConfig，允许仅在rank 0上填充state_dict并将其卸载到CPU。
使用此配置时，FSDP将allgather模型参数，将它们逐个卸载到CPU上，仅在rank 0上执行。当最终保存state_dict时，它仅在rank 0上填充并包含CPU张量。这避免了模型大于单个GPU内存时可能出现的OOM，并允许用户检查点模型，其大小大致为用户机器上可用的CPU RAM。
此功能可以按如下方式运行：
save_policy = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
with FSDP.state_dict_type(
model, StateDictType.FULL_STATE_DICT, save_policy
):
cpu_state = model.state_dict()
if rank == 0:
save_name = file_save_name + "-" + time_of_run + "-" + currEpoch
torch.save(cpu_state, save_name)
总结
在本教程中，我们介绍了Pytorch 1.12中FSDP的许多新功能，并使用HF T5作为运行示例。使用适当的包装策略，特别是对于transformer模型，以及混合精度和反向预取，应该可以加快您的训练运行。此外，诸如在设备上初始化模型和通过流式传输到CPU保存检查点等功能，应该有助于避免在处理大型模型时出现OOM错误。
我们正在积极努力为下一个版本的FSDP添加新功能。如果您有反馈、功能请求、问题或在使用FSDP时遇到问题，请随时通过在PyTorch Github仓库中提交问题与我们联系。
