基于 fairseq FSDP 的 13B 大模型训练实战全参数分片与 CPU 卸载指南【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilm导读本文以 edgelm/examples/fully_sharded_data_parallel/README.md 为核心指南系统讲解如何在 fairseqEdgeLM 内置版本中通过 fairscale 的FullyShardedDataParallelFSDP后端训练超大语言模型包括 FSDP 与 PyTorch DDP 的差异、--ddp-backendfully_sharded/--cpu-offload/--no-reshard-after-forward等核心参数的含义以及两套可直接运行的实战命令——在 1 张 V100 GPU 上通过 CPU 卸载训练 13B GPT-3 模型与在 8 张 V100 GPU 上通过全参数分片叠加 CPU 卸载将吞吐从约 310 词/秒提升到约 3200 词/秒。读完后你将具备使用 fairseq 单机多卡分片训练百亿级参数模型的实际能力并理解其底层实现原理与已知限制。FSDP把 ZeRO 风格的分片思想引入 fairseq从数据并行到参数分片微软的 ZeRO 系列工作论文《ZeRO: Memory Optimizations Toward Training Trillion Parameter Models》与 Google 的 Mesh-TensorFlow 工作表明数据并行训练的效率可以通过把模型参数与优化器状态在数据并行 worker 之间进行**分片sharding**而显著提升。这些思想被封装进 fairscale 提供的FullyShardedDataParallelFSDP包装器中并作为 fairseq 的一种分布式训练后端接入。与 PyTorch DDP 相比FSDP 的核心差异体现在训练结果等价FSDP 仍然是同步的数据并行训练因此理论上与 PyTorch DDP 产生一致的结果显存占用大幅下降FSDP 把 FP16 与 FP32 的参数以及优化器状态分片分布在多张数据并行 GPU 上而非每张卡都保留完整副本训练速度更快优化器 step 本身被分片执行且通信可以与前向传播重叠overlap从而降低等待时间模型规模上限显著提升按该文档的实践经验FSDP 可以在 8 张 GPU 上训练 13B 参数模型在 128 张 GPU 上训练 175B 参数模型。fairseq 中启用 FSDP 的入口在 fairseq 中FSDP 通过--ddp-backend参数选择。在 edgelm/fairseq/dataclass/constants.py 中定义了可选后端列表其中fully_sharded即对应 fairscale 的FullyShardedDataParallelDDP_BACKEND_CHOICES ChoiceEnum([ c10d, # alias for pytorch_ddp fully_sharded, # FullyShardedDataParallel from fairscale legacy_ddp, no_c10d, # alias for legacy_ddp pytorch_ddp, slow_mo, ])底层实现位于 edgelm/fairseq/distributed/fully_sharded_data_parallel.py。该文件对 fairscale 的 FSDP 做了两件关键事情包一层 fairseq 专属的 checkpoint 保存/加载逻辑FullyShardedDataParallel子类第 23-76 行支持use_sharded_state两种模式——为 True 时使用local_state_dict/load_local_state_dict读写分片 checkpoint为 False 时只在 rank 0 上导出完整权重、其余 rank 返回空 dict加载时通过dist_utils.broadcast_object从 rank 0 广播保证任意进程都能恢复完整模型。提供fsdp_enable_wrap上下文管理器第 79-111 行把 fairseq 的分布式训练配置翻译为 fairscale 的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: True, 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取反得到即该参数关闭的是前向后重新分片行为mixed_precision只有在--fp16且未开启--memory-efficient-fp16时才启用说明 FSDP 的混合精度由 fairseq 的fp16开关驱动flatten_parameters恒为 True这正是文档限制一节所说的FSDP 依赖参数扁平化的代码依据state_dict_device固定为 CPU以降低 GPU 显存占用。在模型装配层面edgelm/fairseq/models/distributed_fairseq_model.py 中的DistributedFairseqModel分支会断言模型已被 FSDP 包装并在未开启 CPU 卸载时把模型移到 GPUelif args.ddp_backend fully_sharded: try: from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP except ImportError: raise ImportError( Cannot find FullyShardedDataParallel. Please install fairscale with: pip install fairscale ) assert isinstance(model, FSDP), expected model to already be wrapped in FSDP wrapped_model model if args.memory_efficient_fp16: wrapped_model wrapped_model.half() if not args.cpu_offload: wrapped_model wrapped_model.to(devicedevice)核心参数速查启用 FSDP 的参数组在 edgelm/fairseq/dataclass/configs.py 中DistributedTrainingConfig集中定义了 FSDP 专属配置参数默认值含义--ddp-backend fully_shardedpytorch_ddp启用 FSDP 后端对参数与优化器状态做全量分片--cpu-offloadFalse把 FP32 参数副本卸载到 CPU需配合--optimizercpu_adam--no-reshard-after-forwardFalse前向之后不再重新分片参数行为接近 ZeRO stage 2可提升 1B 大模型训练速度--fp32-reduce-scatterFalse以 FP32 精度做 reduce-scatter 梯度归约--use-sharded-stateFalse保存/加载分片 checkpoint 文件此外bucket_cap_mb默认 25控制参数归约桶大小同样会透传给 FSDP--memory-efficient-fp16会改变 FSDP 的混合精度路径mixed_precision置为 False 并断言必须同时开启--fp16。文档明确指出其他常用选项在 FSDP 下继续正常工作--fp16、--update-freq、--checkpoint-activations梯度检查点、--offload-activations等。Trainer 内置的兼容性校验在 edgelm/fairseq/trainer.py 中fairseq 的Trainer初始化时会针对 FSDP 做一系列防御性校验这些校验本身就是参数间约束关系的最佳说明与--bf16/--memory-efficient-bf16互斥FSDP 当前不兼容 bf16与--zero-sharding互斥FSDP 已内建参数分片能力无需再叠加 ZeRO 配置组合--update-freq时需要 fairscale 版本 ≥ 0.4.0非fully_sharded后端一旦传入--cpu-offload会直接报错--cpu-offload requires --ddp-backendfully_sharded。另外从 trainer.py 中可以看到is_fsdp属性即判断ddp_backend fully_sharded第 366 行并据此决定是否跳过模型的常规 FP16 转换与设备搬移这些工作交由 FSDP 包装器完成这也是理解 FSDP 训练初始化流程的关键。已知限制与适用边界文档在 Limitations 一节给出了两条明确的限制结合源码可以进一步确认成因优化器兼容性FSDP 完全兼容逐点pointwise优化器Adam、AdamW、Adadelta、Adamax、SGD 等但不兼容非逐点优化器Adagrad、Adafactor、LAMB 等。原因是分片后的优化器 step 需要逐参数独立计算非逐点优化器依赖跨参数的全局统计量无法在分片视图下保持一致语义。参数扁平化依赖FSDP 依赖把参数扁平化flatten为单一连续张量因此凡是需要--fp16-no-flatten-grads的模型可能无法被 FSDP 支持。这一限制在 fully_sharded_data_parallel.py 的flatten_parameters: True配置中得到了印证。fairscale 官方文档对这些限制及其它细节有更详细的解释建议在落地前通读。实战一单卡 V100 训练 13B 模型CPU 卸载模式适用场景与硬件要求该模式面向只有 1 张 GPU、但机器内存充裕的场景把参数和优化器状态卸载到 CPU让 Adam 优化器 step 在 CPU 上执行从而在单卡上训练 13B 参数模型。文档给出的硬件前提为GPU 显存 32GB系统内存约 256GB若系统内存不足可将架构改为--arch transformer_lm_gpt3_6_76.7B 参数约需 128GB 系统内存由于使用 DeepSpeed 的 CPU Adam 优化器需先执行pip install deepspeed需安装 fairscale 最新 master 版本pip install githttps://github.com/facebookresearch/fairscale.gitmaster。文档还特别提醒三点使用细节命令启动后约 5 分钟内看似卡住属正常现象——随机初始化 13B 权重本身很慢--cpu-offload必须配合混合精度训练--fp16使用通过OMP_NUM_THREADS环境变量调节线程数可获得更好的 CPU 卸载性能。完整命令以下命令使用 WikiText-103 数据集需按 RoBERTa/GPT-2 词表预处理为data-bin/wikitext-103-roberta-bpe-bin演示训练流程仅跑 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>CUDA_VISIBLE_DEVICES0,1,2,3 fairseq-train /path/to/stories_mmap \ --ddp-backend fully_sharded --fp16 --fp16-init-scale 4 \ --cpu-offload --checkpoint-activations \ --task language_modeling --tokens-per-sample 256 --batch-size 8 \ --arch transformer_lm_gpt2_tiny \ --optimizer cpu_adam --adam-betas (0.9,0.98) \ --lr 0.0001 --lr-scheduler polynomial_decay --warmup-updates 5 --total-num-update 10 \ --max-update 5 --log-format json --log-interval 1 \ --save-interval-updates 5 --save-dir fsdp_dummy --disable-validation总结与选型建议fairseq 的 FSDP 后端为超大模型训练提供了零模型改动的分片能力其选型路径可以概括为单卡 大内存机器--ddp-backend fully_sharded --cpu-offload --optimizer cpu_adam在 1 张 32GB V100 256GB 内存上训练 13B 模型多卡分片8 × 32GB V100 上对 13B 模型做全参数分片若每卡 batch size 受限再叠加--cpu-offload换取吞吐追求极致速度对 1B 参数模型可加--no-reshard-after-forward接近 ZeRO stage 2 的行为内存再省一层叠加--checkpoint-activations与--offload-activations进一步压低激活显存。同时需要牢记两条硬性约束优化器必须是逐点式Adam/AdamW/SGD 等模型不能被--fp16-no-flatten-grads依赖所阻塞训练必须使用--fp16混合精度且不可与--bf16、--zero-sharding混用。结合 fully_sharded_data_parallel.py、configs.py 与 test_fsdp.sh 中的实现与验证脚本你可以快速在自己的集群上复现 13B 模型训练并将这套组合迁移到更大规模的语言模型训练任务中。【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilm创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考