PyTorch Lightning Fabric 策略体系全解析从单设备到大规模分布式训练【免费下载链接】pytorch-lightningPretrain, finetune ANY AI model of ANY size on 1 or 10,000 GPUs with zero code changes.项目地址: https://gitcode.com/gh_mirrors/py/pytorch-lightning导读lightning.fabric.strategies是 PyTorch Lightning Fabric 分布式训练的核心抽象层它定义了如何在多个设备上组织模型、数据与通信的全部细节。本文以 docs/source-fabric/api/strategies.rst 中列出的 10 个策略类为主线结合 src/lightning/fabric/strategies 目录下的源码实现逐一讲解每个策略的定位、适用场景、核心参数与底层工作机制。读完本文你将理解 Fabric 如何通过strategy参数一行切换训练范式并能根据硬件与模型规模正确选择ddp、fsdp、deepspeed、dp、xla等策略。策略Strategy是什么API 文档的骨架官方 API 参考文档 strategies.rst 是一个 autosummary 索引页它在lightning.fabric.strategies命名空间下挂载了如下 10 个公开类策略类对应注册名定位Strategy—抽象基类所有策略的公共契约ParallelStrategy—抽象基类多进程并行策略的公共基类SingleDeviceStrategysingle_device单设备训练SingleDeviceXLAStrategysingle_xla单 XLATPU设备训练DDPStrategyddp/ddp_spawn/ddp_fork/ddp_notebook/ddp_find_unused_parameters_true多进程数据并行DataParallelStrategydp单进程内多卡数据并行FSDPStrategyfsdp/fsdp_cpu_offload全分片数据并行ZeRO-3 风格XLAStrategyxla多 TPU 训练XLAFSDPStrategyxla_fsdpTPU 上的 FSDPDeepSpeedStrategydeepspeed基于 DeepSpeed 的 ZeRO 优化ModelParallelStrategy—实验性需显式实例化2D 并行FSDP2 张量并行这些类全部实现在 src/lightning/fabric/strategies 目录下并在init.py 中集中导出。理解这一族类的继承关系是掌握 Fabric 并行能力的钥匙。Strategy 基类所有策略的公共契约Strategy 是一个抽象基类它定义了训练循环、验证循环与测试循环中所有可能因并行方式而不同的行为。从源码看其职责可以划分为六组。生命周期方法环境与模型的接管setup_environment()在每个进程开始任何分布式通信之前被调用负责调用加速器的setup_device完成设备初始化strategy.py。setup_module(module)/setup_optimizer(optimizer)将模型与优化器包装成并行所需的形态例如 DDP 会把模型包进DistributedDataParallel。默认实现是原样返回。setup_module_and_optimizers(module, optimizers, scheduler)联合设置模型与多个优化器。默认实现会依次调用setup_module与setup_optimizerstrategy.py。module_to_device(module)/batch_to_device(batch)把模型与 batch 移动到正确的设备。teardown()训练结束时释放内存与资源依次调用 precision、accelerator 与 checkpoint_io 的 teardownstrategy.py。初始化上下文控制张量与模型参数的创建方式tensor_init_context()返回一个上下文管理器在创建张量时控制设备与 dtypestrategy.py。module_init_context(empty_initNone)包裹模型实例化过程empty_init为True时用未初始化内存创建参数meta device从而显著降低峰值内存并加速初始化None时由策略自行决定strategy.py。FSDP 正是通过覆写该方法在参数创建时直接分片。集体通信抽象all_gather、all_reduce、barrier、broadcast是四个抽象方法每个具体策略都必须实现。reduce_boolean_decision提供跨进程的布尔归约默认按all语义执行。反向传播与优化器步骤backward(tensor, module, *args, **kwargs)会把反向传播转发给 precision 插件依次执行pre_backward、backward、post_backwardstrategy.pyoptimizer_step同样委托给precision.optimizer_step。这意味着 AMP 缩放、梯度反缩放等细节都被策略统一接管。梯度裁剪基类提供两种实现strategy.pyclip_gradients_norm(module, optimizer, max_norm, norm_type2.0, error_if_nonfiniteTrue)按范数裁剪会先unscale_gradients再取main_params。clip_gradients_value(module, optimizer, clip_val)按数值裁剪。FSDP 策略覆写了clip_gradients_norm改用 FSDP 包装器上的clip_grad_norm_见下文。检查点读写基类save_checkpointstrategy.py会把 state 中的 Module/Optimizer 自动转成 state_dict并且只在is_global_zero进程上实际落盘load_checkpoint支持四种形态完整 dict、Module、Optimizer 或空 dict直接返回原始 checkpoint并支持strict与weights_only参数strategy.py。filter参数允许按 key 过滤需要保存/加载的状态。策略注册表用字符串切换策略Fabric 允许用字符串而非类实例来指定策略背后的机制是注册表 registry.py 中的_StrategyRegistry。它的register方法把一个名字映射到策略类与初始化参数registry.pyget方法则用存储的init_params实例化策略对象registry.py。在 strategies/init.py 中导入时即调用_register_classes(STRATEGY_REGISTRY, register_strategies, ...)逐个收集各策略类上的register_strategies类方法完成注册。例如 ddp.py 一次性注册了ddp、ddp_spawn、ddp_fork、ddp_notebook与ddp_find_unused_parameters_true五个名字其中前四者只是start_method不同。在 connector.py 中Fabric(strategy...)接受的策略可以是字符串或策略实例当传入字符串时会校验其是否存在于注册表非法名字直接报错。auto模式下则由_choose_strategy依据加速器与设备数推导默认策略。实际用法根据 docs/source-fabric/api/fabric_args.rst 的说明可用的策略字符串为dp、ddp、ddp_spawn、ddp_find_unused_parameters_true、xla、deepspeed、fsdpfrom lightning.fabric import Fabric # 4 卡 GPUDistributedDataParallel fabric Fabric(strategyddp, acceleratorgpu, devices4) # DDP 允许存在未使用参数 fabric Fabric(strategyddp_find_unused_parameters_true, acceleratorgpu, devices4) # 4 个 CPU 进程spawn 方式启动 fabric Fabric(strategyddp_spawn, acceleratorcpu, devices4)单设备族SingleDeviceStrategy 与 SingleDeviceXLAStrategySingleDeviceStrategySingleDeviceStrategy 处理单个设备上的训练构造函数签名SingleDeviceStrategy(devicecpu, acceleratorNone, checkpoint_ioNone, precisionNone)它把global_rank、local_rank固定为 0world_size固定为 1single_device.py因此is_global_zero恒为True。其all_reduce、all_gather都是恒等操作barrier为空操作broadcast直接返回原对象single_device.py——这是单进程无需通信的最直观体现。module_to_device直接执行module.to(root_device)。SingleDeviceXLAStrategySingleDeviceXLAStrategy 继承自SingleDeviceStrategy用于单个 XLATPU设备。它要求torch_xla可用否则在构造时抛出ModuleNotFoundError其checkpoint_io强制为XLACheckpointIOsingle_xla.py。ParallelStrategy多进程并行的公共基类ParallelStrategy 是所有多进程并行策略DDP、DP、FSDP、XLA、ModelParallel的抽象基类。它新增了parallel_devices与cluster_environment两个成员并提供了四个关键属性parallel.pyglobal_rank全局进程编号local_rank本节点内进程编号node_rank节点编号world_size总进程数is_global_zero即global_rank 0。此外distributed_sampler_kwargs默认返回{num_replicas: world_size, rank: global_rank}这正是 DataLoader 中DistributedSampler的参数来源。reduce_boolean_decision通过all_reduce的 SUM 归约实现all/any语义parallel.py。DDPStrategy多进程单设备数据并行DDPStrategy 是一台或多台节点上做多进程单设备训练的标准策略也是 Fabric 默认的分布式方案。构造参数包括process_group_backend默认按设备自动选择如 NCCL、timeout默认default_pg_timeout、start_methodpopen/spawn/fork/forkserver默认popen其余**kwargs直接透传给torch.nn.parallel.DistributedDataParallel。关键机制启动方式_configure_launcher中start_methodpopen时使用_SubprocessScriptLauncher重新以子进程脚本方式启动其余方式使用_MultiProcessingLauncherddp.py。模型包装setup_module将模型包进DistributedDataParallel并依据设备类型决定device_idsCPU 时为NoneGPU 时为[root_device.index]见 ddp.py。源码中还处理了 CUDA Graph 捕获时的侧流side-stream初始化细节。梯度同步控制通过_DDPBackwardSyncControl实现no_backward_sync上下文管理器即module.no_sync()用于梯度累积时跳过冗余的梯度同步ddp.py。世界排名_set_world_ranks中global_rank node_rank * num_processes local_rankworld_size num_nodes * num_processesddp.py。状态读写get_module_state_dict/load_module_state_dict会先解开DistributedDataParallel包装再读写底层 module 的 state_dictddp.py。梯度累积示例fabric Fabric(strategyddp, acceleratorgpu, devices8) model, optimizer fabric.setup(model, optimizer) for epoch in range(num_epochs): for i, batch in enumerate(dataloader): # 累积 4 个 batch 的梯度只同步一次 is_last_batch (i 1) % 4 0 with fabric.no_backward_sync(model, enablednot is_last_batch): output model(batch) loss loss_fn(output, target) fabric.backward(loss) if is_last_batch: optimizer.step() optimizer.zero_grad()DataParallelStrategy单进程内多卡数据并行DataParallelStrategy 实现单进程内的数据并行模型被复制到每张卡每个卡处理数据的一个分片。它的setup_module把模型包进torch.nn.DataParalleldp.py并把 batch 的搬运交给 DataParallel 处理batch_to_device直接返回原 batchdp.py。其all_reduce是对 collection 中的每个张量求均值。需要说明的是由于nn.DataParallel使用单进程多线程barrier与broadcast均为空实现reduce_boolean_decision直接返回本进程决策。该类在注册表中注册的名字是dpdp.py。在dp下distributed_sampler_kwargs返回None即不会自动套用DistributedSampler。FSDPStrategy全分片数据并行FSDPStrategy 基于torch.distributed.fsdp.FullyShardedDataParallel把整个模型参数、梯度、优化器状态分片到所有 GPU 上从而在不改动代码的前提下训练远大于单卡显存的模型其思路与 ZeRO-Stage 3 类似。它是_Sharded接口的实现者通过module_sharded_context底层是 FSDP 的enable_wrap见 fsdp.py在模型实例化阶段即完成逐层分片。核心参数参数默认值说明sharding_strategyFULL_SHARDFULL_SHARD参数梯度优化器状态全分片、SHARD_GRAD_OP仅梯度与优化器状态、NO_SHARD等价普通 DDP、HYBRID_SHARD节点内分片、节点间复制也接受ShardingStrategy枚举fsdp.pydevice_meshNone(replication_size, sharding_size)元组或DeviceMesh乘积必须等于 world size仅配合HYBRID_SHARD使用cpu_offloadNone参数 CPU offloadTrue/CPUOffload实例mixed_precisionNoneFSDP 的MixedPrecision配置auto_wrap_policyNoneFSDP 自动切分策略也接受一组需要包装的层类会转换为ModuleWrapPolicyactivation_checkpointingNone已弃用改用activation_checkpointing_policyactivation_checkpointing_policyNone为指定层开启激活检查点显著省显存但反向需重算激活值state_dict_typeshardedsharded每个 rank 保存自己的分片到目录full在 rank 0 汇总成单文件use_orig_paramsTrue框架强制默认保持优化器中原始参数引用支持多参数组与torch.compile与基类的差异精度插件固定只能搭配FSDPPrecision否则抛TypeErrorfsdp.py。检查点不走 CheckpointIOcheckpoint_io属性直接抛出NotImplementedErrorfsdp.py其save_checkpoint/load_checkpoint自行实现分片或全量两种格式底层使用torch.distributed.checkpoint的save/loadfsdp.py。sharded格式下用户元数据非 Module/Optimizer 对象会由 rank 0 单独保存到meta.pt分片 checkpoint 的过滤保存/加载filter当前不支持。梯度裁剪要求模型已用FullyShardedDataParallel包装否则报错裁剪本身调用module.clip_grad_norm_fsdp.py。激活检查点在包装完成后通过apply_activation_checkpointing统一设置fsdp.py若模型已含CheckpointWrapper层会忽略配置并告警。共享参数检测源码额外实现了_warn_if_shared_params_across_fsdp_units当 auto_wrap_policy 可能把绑定权重如 tied embedding分到不同 FSDP 单元时提前发出警告避免前向/反向时报出晦涩的 size mismatch。注册表中fsdp与fsdp_cpu_offload即cpu_offloadTrue两个名字都指向该类fsdp.py。DeepSpeedStrategy亿级参数训练的 ZeRO 优化DeepSpeedStrategy 继承自DDPStrategy并实现_Sharded通过 DeepSpeed 库提供训练优化默认已开启 ZeRO-Offload 相关配置。它既可以直接通过构造参数配置也可以通过config传入 DeepSpeed JSON 配置文件路径或字典。主要构造参数源码中的默认值ZeRO 基础zero_optimizationTrue、stage2默认 ZeRO Stage 2。参数 offloadoffload_parametersFalse、offload_params_devicecpu、nvme_path/local_nvme、params_buffer_count5、params_buffer_size100_000_000、max_in_cpu1_000_000_000。优化器 offloadoffload_optimizerFalse、offload_optimizer_devicecpu、optimizer_buffer_count4。通信与显存contiguous_gradientsTrue、overlap_commTrue、allgather_partitionsTrue、reduce_scatterTrue、allgather_bucket_size200_000_000、reduce_bucket_size200_000_000。动态 loss scalingloss_scale0、initial_scale_power16、loss_scale_window1000、hysteresis2、min_loss_scale1。其他partition_activationsFalse、cpu_checkpointingFalse、load_full_weightsFalse、logging_levellogging.WARN、exclude_frozen_parametersFalse。这些默认值在 deepspeed.py 的构造函数签名中全部可见可按模型规模自行调整。它同样暴露了DEEPSPEED_ENV_VAR PL_DEEPSPEED_CONFIG_PATH环境变量常量。在 docs/source-pytorch/advanced/model_parallel/deepspeed.rst 与 docs/source-fabric/advanced/model_parallel 系列文档中可找到更完整的 DeepSpeed 用法说明。XLA 族XLAStrategy 与 XLAFSDPStrategyXLAStrategyXLAStrategy 面向多 TPU 设备使用torch_xla.distributed.xla_multiprocessing.spawn方式启动多进程训练。其checkpoint_io固定为XLACheckpointIOprecision固定为XLAPrecisioncluster_environment固定为XLAEnvironment()xla.py。值得注意的是它的_backward_sync_control为None因为 XLA 在optimizer.step()时统一同步梯度不提供no_backward_sync接口。root_device属性在进程尚未 spawn 时访问会抛出RuntimeError提示Accessing the XLA device before processes have spawned is not allowedxla.py。XLAFSDPStrategyXLAFSDPStrategy 继承ParallelStrategy并实现_Sharded在 XLA 设备上使用torch_xla.distributed.xla_fully_sharded_data_parallel.XlaFullyShardedDataParallel做全分片训练属于实验性 API。它的register_strategies注册名为xla_fsdpxla_fsdp.py。对应测试位于 tests/tests_fabric/strategies/test_xla_fsdp.py。ModelParallelStrategy2D 并行实验性ModelParallelStrategy 支持用户自定义的模型并行方式当前支持到 2D 并行即 FSDP2torch.distributed.fsdp.fully_sharded_data_parallel2与基于 DTensor 的张量并行组合。从源码结构看它复用了 FSDP 策略中的分布式 checkpoint 保存/加载工具_distributed_checkpoint_save、_get_full_state_dict_context等因此分片与全量检查点逻辑与 FSDP 保持一致。该类在文档中标注为实验性 API建议仅在需要同时进行张量并行与全分片时使用。如何选择策略Fabric 的自动推导当Fabric(strategyauto)默认时connector.py 中的_choose_strategy会根据加速器与设备数量选择默认策略多设备时优先 DDP单设备时退化为单设备策略。同时存在一层回退校验_check_strategy_and_fallback例如在 MPS 加速器上使用 DDP 家族策略会被明确拒绝并提示改为acceleratorcpu或更换策略connector.py。选择建议基于源码行为单 GPU / CPU无需指定默认单设备策略。多 GPU、模型可放进单卡strategyddp通信开销低、实现成熟。模型超过单卡显存、或追求极致显存利用率strategyfsdp或fsdp_cpu_offload并配合auto_wrap_policy、sharding_strategy、activation_checkpointing_policy调优。需要 ZeRO-Offload、NVMe offload 或已有的 DeepSpeed 配置strategydeepspeed通过config传配置文件。TPU 环境strategyxla单/多 TPU或strategyxla_fsdp需要分片。单进程内快速多卡实验strategydp注意nn.DataParallel的线程模型限制。张量并行 全分片组合实例化ModelParallelStrategy实验性。测试佐证仓库在 tests/tests_fabric/strategies 下为每个策略提供了单元与集成测试例如 test_ddp.py、test_dp.py、test_fsdp.py、test_deepspeed.py、test_model_parallel.py、test_strategy.py 与 test_registry.py。阅读这些测试可以快速验证各策略的注册名、默认参数与容错行为。总结lightning.fabric.strategies是一个精心设计的策略族Strategy基类把所有可能随并行方式变化的行为抽象为统一的钩子ParallelStrategy在多进程场景下提供排名与世界规模信息具体的 DDP、DP、FSDP、DeepSpeed、XLA 与 ModelParallel 策略则各自落实启动方式、模型包装、通信原语与检查点格式。配合字符串注册表开发者只需改动Fabric(strategy...)一行即可在单机调试、多卡 DDP、全分片 FSDP、DeepSpeed ZeRO 与 TPU 训练之间自由切换这正是 Fabric 面向任意规模、任意硬件训练的核心设计。【免费下载链接】pytorch-lightningPretrain, finetune ANY AI model of ANY size on 1 or 10,000 GPUs with zero code changes.项目地址: https://gitcode.com/gh_mirrors/py/pytorch-lightning创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
