fairseq 中的 Fully Sharded Data Parallel(FSDP)实战指南:单卡与多卡训练 130 亿参数模型 📅 发布时间:2026/9/14 10:06:23 👁 浏览次数: fairseq 中的 Fully Sharded Data ParallelFSDP实战指南单卡与多卡训练 130 亿参数模型【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilm导读本指南以kosmos-2/fairseq/examples/fully_sharded_data_parallel/README.md为骨架系统讲解 fairseq本仓库kosmos-2/fairseq内置的版本对 fairscaleFullyShardedDataParallelFSDP的完整支持包括开启方式、关键参数、适用限制、工作原理以及两个可直接运行的真实示例——在 1 张 V100 GPU 上借助 CPU offload 训练 13B 参数 GPT-3 模型和在 8 张 V100 GPU 上通过全量参数分片训练同一模型。读完本文你将掌握如何用--ddp-backend fully_sharded突破单卡显存瓶颈并理解 FSDP 与 PyTorch DDP 的差异及其在 fairseq 训练管线中的底层实现。概述FSDP 是什么微软ZeRO与 GoogleGShard 相关工作的研究表明将模型参数与优化器状态在数据并行 worker 之间进行分片sharding可以显著提升数据并行训练的效率。这些思想被封装在 fairscale 提供的全新FullyShardedDataParallelFSDP包装器中。fairseq 在其分布式训练框架中对该包装器做了 fairseq 专属的封装checkpoint 保存/加载逻辑相关实现位于 fairseq/distributed/fully_sharded_data_parallel.py。与 PyTorch DDP 的对比原文档明确指出 FSDP 与 fairseq 默认的 PyTorch DDP 后端相比具有以下特性维度说明训练结果与 PyTorch DDP 完全一致仍然是同步数据并行训练内存占用在数据并行 GPU 间分片参数FP16 FP32与优化器状态训练速度快于 PyTorch DDP因为优化器更新步骤本身被分片且通信可与前向传播重叠可扩展性支持在 8 张 GPU 上训练 13B 参数模型在 128 张 GPU 上训练 175B 参数模型fairseq 中的启用方式FSDP 在 fairseq 中通过以下新增命令行参数获得完整支持--ddp-backendfully_sharded通过 FSDP 启用全量分片--cpu-offload将优化器状态与 FP32 模型副本卸载到 CPU需配合--optimizercpu_adam使用--no-reshard-after-forward提升大模型10 亿参数以上的训练速度行为类似 ZeRO stage 2其他常用选项--fp16、--update-freq、--checkpoint-activations、--offload-activations等继续正常工作。参数定义与底层实现在 fairseq 的分布式训练配置类DistributedTrainingConfig中这些 FSDP 专属参数都有明确定义见 fairseq/dataclass/configs.py参数默认值作用--no-reshard-after-forwardFalse前向传播后不再对参数执行 reshard--fp32-reduce-scatterFalse以 FP32 精度执行梯度 reduce-scatter--cpu-offloadFalse将 FP32 参数卸载到 CPU--use-sharded-stateFalse使用分片 checkpoint 文件--not-fsdp-flatten-parametersFalse不为 FSDP 扁平化参数这些参数最终会被映射为 fairscale FSDP 的构造配置。在 fully_sharded_data_parallel.py 的fsdp_enable_wrap上下文中可以看到完整的映射关系fsdp_config { process_group: group, reshard_after_forward: not cfg.no_reshard_after_forward, mixed_precision: cfg.fp16 and not cfg.memory_efficient_fp16, fp32_reduce_scatter: cfg.fp32_reduce_scatter, flatten_parameters: not cfg.not_fsdp_flatten_parameters, cpu_offload: cfg.cpu_offload, compute_dtype: torch.float16 if cfg.fp16 else torch.float32, bucket_cap_mb: cfg.bucket_cap_mb, state_dict_device: torch.device(cpu), # reduce GPU mem usage }其中几个关键点值得注意reshard_after_forward直接由--no-reshard-after-forward取反得到。默认开启 reshard下前向传播结束后参数会被立即释放回分片状态以节省显存对大模型关闭该行为可以省去反复 gather/reshard 的通信开销从而提速对应 ZeRO stage 2 的语义。mixed_precision仅在--fp16且未开启--memory-efficient-fp16时启用若开启memory_efficient_fp16则要求同时开启--fp16源码中有对应 assert。state_dict_device固定为 CPU以减少 GPU 显存占用这是 fairseq 对 FSDP 的定制行为。fsdp_wrap辅助函数fully_sharded_data_parallel.py还支持按最小参数量决定是否包装某个子模块若 fairscale 未安装则自动退化为 no-op。从源码结构看--ddp-backendfully_sharded分支在 fairseq/models/distributed_fairseq_model.py 中处理它要求模型已经被 FSDP 包装否则 assert 报错并且仅在未开启--cpu-offload时才将模型搬到 GPU 设备——这正是 CPU offload 模式下模型常驻 CPU 的机制。限制与兼容性说明与 fairseq 默认的 PyTorch DDP 后端相比FSDP 目前存在以下限制FSDP 与逐元素pointwise优化器如 Adam、AdamW、Adadelta、Adamax、SGD 等完全兼容但当前不兼容非逐元素优化器如 Adagrad、Adafactor、LAMB 等FSDP 依赖参数扁平化flattening因此当前需要--fp16-no-flatten-grads的模型可能不受支持。这些限制在 fairseq 训练器中也有对应的硬性校验在 fairseq/trainer.py 附近源码会拒绝--ddp-backendfully_sharded与 BMUF--use-bmuf的组合并校验优化器必须是 pointwise 类型。此外--cpu-offload必须配合--ddp-backendfully_sharded使用否则训练器会直接抛出ValueError: --cpu-offload requires --ddp-backendfully_sharded见 trainer.py当--cpu-offload与 FSDP 组合时use_distributed_wrapper属性恒为真trainer.py保证模型包装器在单卡 CPU offload 场景下也被正确启用。工作原理FSDP 的核心思想是在数据并行进程中分片存储模型参数FP16 与 FP32 两个副本和优化器状态训练时按需**全收集all-gather**完整参数完成前向与反向传播梯度计算完成后以reduce-scatter的方式聚合梯度并更新各自分片内的优化器状态从而让每一步的显存占用随数据并行规模线性下降。同时由于优化器步骤只在自己的分片上执行通信可以与前向传播重叠这也解释了为何 FSDP 可以比 DDP 更快。在 checkpoint 方面fairseq 的 FSDP 封装fully_sharded_data_parallel.py提供两种行为默认模式state_dict只在数据并行 rank 0 上返回完整模型权重其他 rank 返回空load_state_dict时从 rank 0 广播权重到所有 rank——always_call_state_dict_during_save_checkpoint恒为真trainer.py因为合并 checkpoint 需要通信集体操作--use-sharded-state模式直接使用local_state_dict/load_local_state_dict每个 rank 保存自己的分片checkpoint 文件名带-shard{rank}后缀trainer.py且所有 rank 都会保存 checkpointshould_save_checkpoint_on_current_rank返回 True。实战示例一13B 参数在单张 V100 GPU 上训练CPU offload以下命令在单张 V100 GPU 上训练 13B 参数 GPT-3 模型使用--cpu-offload将参数与优化器状态卸载到 CPU。此时优化器步骤Adam在 CPU 上执行同时使用--checkpoint-activations即梯度检查点进一步节省显存代价是少量计算增加。运行前提安装最新 master 版 fairscalepip install githttps://github.com/facebookresearch/fairscale.gitmaster训练 13B 参数模型需要 32GB GPU 显存与约 256GB 系统内存若系统内存不足可改用 6.7B 参数模型仅需约 128GB 系统内存将架构参数替换为--arch transformer_lm_gpt3_6_7即可示例使用 DeepSpeed 提供的 CPU Adam 优化器运行前需pip install deepspeed。注意事项命令启动后约 5 分钟内看起来像卡住是正常的因为随机初始化 13B 权重本身较慢--cpu-offload要求混合精度训练--fp16可通过调整OMP_NUM_THREADS环境变量优化 CPU offload 下的性能以下示例仅训练 10 步--max-update 10且不保存 checkpoint--no-save。OMP_NUM_THREADS20 CUDA_VISIBLE_DEVICES0 \ fairseq-train>OMP_NUM_THREADS20 CUDA_VISIBLE_DEVICES0,1,2,3,4,5,6,7 \ fairseq-train contenteditable="false">【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilm创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考