ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

PyTorch Lightning GPU 分布式训练实战(中级):DDP、DDP Spawn、DDP Notebook 与 TorchRun 多机扩展指南

PyTorch Lightning GPU 分布式训练实战(中级):DDP、DDP Spawn、DDP Notebook 与 TorchRun 多机扩展指南 PyTorch Lightning GPU 分布式训练实战中级DDP、DDP Spawn、DDP Notebook 与 TorchRun 多机扩展指南【免费下载链接】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本文聚焦 PyTorch Lightning 在 GPU 场景下的分布式训练策略面向已能单卡训练、希望在多卡或多机间横向扩展的开发者。你将完整掌握DDP、DDP Spawn、DDP Notebook/Fork三种策略的原理、适用场景与取舍学会用torchrun启动容错的多机训练并能显式控制nccl等进程组后端以优化多机通信。文中所有结论均可对照仓库源码验证。原文出处docs/source-pytorch/accelerators/gpu_intermediate.rst本文在此基础上结合src/lightning/pytorch/strategies源码、Trainer 连接器 与 torchrun 集群文档 进行了纵深扩充。三种分布式训练策略速览PyTorch Lightning 的 Trainer 支持多种分布式训练方式按启动进程的机制不同主要分为三类Regular标准 DDPstrategyddp以子进程脚本方式启动是生产环境的首选。Spawnstrategyddp_spawn基于torch.multiprocessing.spawn()启动仅建议用于调试或迁移旧代码。Notebook / Forkstrategyddp_notebook别名ddp_fork基于进程fork机制专为 Jupyter Notebook、Google Colab、Kaggle 等交互式环境设计。关键行为如果你请求了多个 GPU 或多个节点但没有显式设置 strategyLightning 会自动使用 DDP。该逻辑体现在 accelerator_connector.py 的_choose_strategy方法中num_nodes 1时返回ddp多设备且处于交互环境时返回ddp_fork其余情况返回ddp。Distributed Data Parallel标准 DDP 的工作原理与用法DDPtorch.nn.parallel.DistributedDataParallel在 Lightning 中的执行流程如下每个节点上的每张 GPU 各自拥有一个独立进程每张 GPU 只看到整个数据集的一个子集且只处理这个子集由分布式采样器保证划分每个进程各自初始化模型每个进程并行地执行完整的前向与反向传播所有进程的梯度被同步并求平均每个进程用平均后的梯度更新自己的优化器。对应的 Trainer 配置非常简洁# 单机 8 卡 trainer Trainer(acceleratorgpu, devices8, strategyddp) # 4 节点共 32 卡 trainer Trainer(acceleratorgpu, devices8, strategyddp, num_nodes4)注意 Lightning 2.x 中使用acceleratorgpu与devicesN的组合等价于早期版本的gpusN。num_nodes必须为正整数否则连接器会抛出ValueError见 accelerator_connector.py。Lightning 是如何启动这些进程的与用户手动spawn不同Lightning 的 DDP 实现会在底层用正确的环境变量多次调用你的脚本。以 3 卡 DDP 为例实际效果等价于MASTER_ADDRlocalhost MASTER_PORTrandom() WORLD_SIZE3 NODE_RANK0 LOCAL_RANK0 python my_file.py --accelerator gpu --devices 3 --etc MASTER_ADDRlocalhost MASTER_PORTrandom() WORLD_SIZE3 NODE_RANK0 LOCAL_RANK1 python my_file.py --accelerator gpu --devices 3 --etc MASTER_ADDRlocalhost MASTER_PORTrandom() WORLD_SIZE3 NODE_RANK0 LOCAL_RANK2 python my_file.py --accelerator gpu --devices 3 --etc该机制由 _SubprocessScriptLauncher 实现主进程LOCAL_RANK0通过subprocess.Popen为其余设备各启动一个子进程并逐一设置MASTER_ADDR主节点 IP、MASTER_PORT进程通信端口、NODE_RANK节点索引0 到num_nodes-1、LOCAL_RANK节点内进程索引与WORLD_SIZE所有节点的进程总数即num_processes * num_nodes。源码中的 docstring 以python train.py --devices 4为例说明它会额外创建LOCAL_RANK1/2/3三个子进程。这个启动方式与torch.distributed.run的进程模型非常相似。相比torch.multiprocessing.spawn()的优势使用这种重复调用脚本的 DDP 方式相比手动spawn有几个明确优点所有进程包括主进程都参与训练主进程持有最新的模型与 Trainer 状态不存在 multiprocessing pickle 错误因为模型无需通过队列回传天然支持多节点扩展只要节点间网络互通即可。重要限制不能用于交互式环境标准 DDP 依赖重复执行整个脚本因此无法在 Jupyter Notebook、Google Colab、Kaggle 等交互式环境中使用。Lightning 会在检测到交互环境时给出明确提示若当前策略不兼容交互环境会抛出MisconfigurationException并建议改用Trainer(strategyddp_notebook)见 accelerator_connector.py。交互环境的判定来自sys.ps1是否存在或sys.flags.interactive是否为真见 src/lightning/fabric/utilities/imports.py。DDP Spawn调试与迁移场景的备选ddp_spawn与标准 DDP 的唯一区别在于它使用torch.multiprocessing.spawn()来启动训练进程。# 单机 8 卡 trainer Trainer(acceleratorgpu, devices8, strategyddp_spawn)文档明确警告强烈建议优先使用 DDP 以获得更快的速度与更好的性能。ddp_spawn只应在调试时使用或用于迁移那些原本依赖 spawn 机制的旧代码库。不推荐它的原因受 Python 与 PyTorch 机制限制.fit()结束后只有模型的权重会被回传到主进程Trainer 的其他状态优化器状态、调度器状态、epoch 计数等不会同步不支持多节点训练总体而言比 DDP 更慢。从源码看spawn 与标准 DDP 共用同一个DDPStrategy类只是start_method不同ddp.py 中的register_strategies将ddp映射到start_methodpopen、将ddp_spawn映射到start_methodspawn并在_configure_launcher中根据 start method 选择 _SubprocessScriptLauncher 或_MultiProcessingLauncher。因此ddp_spawn本质上就是用 spawn 方式启动的 DDP其功能边界由 Python 多进程的 spawn 语义决定。DDP Notebook / Fork交互式环境的多卡方案DDP Notebook/Fork 是 Spawn 的替代方案可在交互式 Python、Jupyter Notebook、Google Colab、Kaggle 等环境中使用。Trainer 在检测到此类环境时会默认启用它对应_choose_strategy中len(self._parallel_devices) 1 and _IS_INTERACTIVE时返回ddp_fork的分支。# Jupyter notebook 中训练 8 卡自动启用 fork trainer Trainer(acceleratorgpu, devices8) # 也可以显式指定 trainer Trainer(acceleratorgpu, devices8, strategyddp_notebook) # 非交互环境也可以显式使用 trainer Trainer(acceleratorgpu, devices8, strategyddp_fork)从注册表看ddp_notebook与ddp_fork的start_method均为fork二者本质是同一个机制notebook只是语义化的别名见 ddp.py 与 ddp.py。另外注意 accelerator_connector.py 的校验如果平台不支持fork启动方法如 Windows选择 fork 别名会抛出ValueError并建议改用ddp_spawn。在原生分布式策略中标准 DDPstrategyddp仍然是速度与稳定性俱佳的首选策略只是它只能用于脚本形式运行。Fork/Notebook 在交互场景下还有一条额外限制在调用Trainer.fit之前不允许执行把张量移到 GPU或调用torch.cuda函数等 GPU 操作见下文的对比表。DDP 三种变体的对比与取舍特性DDPDDP SpawnDDP Notebook/Fork可在 Jupyter / IPython 环境中工作否否是支持多节点是是是支持的平台Linux、Mac、WinLinux、Mac、WinLinux、Mac要求所有对象可 pickle否是否主进程中的限制无返回主进程后对象状态不是最新的Trainer.fit()等只有模型参数被回传不允许在调用Trainer.fit之前执行 GPU 操作如将张量移到 GPU、调用torch.cuda函数进程创建时间慢慢快需要说明的是表中进程创建时间的慢/快指的是子进程的启动开销popen/spawn都需要重新初始化 Python 解释器而fork通过复制当前进程内存实现因此创建更快这也是它适合 notebook 场景的原因之一。用 TorchRunTorchElastic实现弹性多机调度Lightning 支持 TorchRun原 TorchElastic用于实现容错fault-tolerant与弹性elastic的分布式任务调度。用法分两步第一步在 Trainer 中指定 DDP 策略与 GPU 数量Trainer(acceleratorgpu, devices8, strategyddp)第二步用torchrun命令启动脚本。详细的参数说明见 集群文档其核心命令模板为torchrun \ --nproc_per_nodeGPUS_PER_NODE \ --nnodesNUM_NODES \ --node_rank NODE_RANK \ --master_addr MASTER_ADDR \ --master_port MASTER_PORT \ train.py --arg1 --arg2各参数含义与约束--nproc_per_node每个节点启动的进程数默认 1必须与Trainer(devices...)中设置的数值一致--nnodes参与训练的节点/机器数默认 1必须与Trainer(num_nodes...)中设置的数值一致--node_rank当前节点的索引从 0 开始--master_addrrank 0 主节点的 IP 地址--master_port节点间通信端口必须在每个节点的防火墙上开放 TCP 流量。两节点各 8 卡的实际示例假设主节点 IP 为10.10.10.16、通信端口为50000。第一个节点node rank 0上运行torchrun \ --nproc_per_node8 --nnodes2 --node_rank 0 \ --master_addr 10.10.10.16 --master_port 50000 \ train.py第二个节点node rank 1上运行torchrun \ --nproc_per_node8 --nnodes2 --node_rank 1 \ --master_addr 10.10.10.16 --master_port 50000 \ train.py注意两条命令唯一的区别就是--node_rank。除 Lightning 侧的配置外你还需自行保证节点间的网络连通性防火墙放行MASTER_PORT、确定主节点MASTER_ADDR以及各节点的 rank。torchrun会从这些参数自动生成 PyTorch 分布式初始化所需的MASTER_ADDR、MASTER_PORT、RANK、WORLD_SIZE等环境变量。关于弹性、容错等更高级的配置可查阅 torchrun 官方文档。从源码角度印证_choose_and_init_cluster_environment见 accelerator_connector.py会按优先级依次检测TorchElasticEnvironment、SLURMEnvironment、LSFEnvironment、MPIEnvironment其中 TorchElastic 优先级最高因为它也能在 SLURM 内部使用当这些外部集群环境都不存在时才回退到LightningEnvironment。因此在torchrun启动时Lightning 会自动识别 TorchElastic 提供的环境变量_SubprocessScriptLauncher也会因creates_processes_externally为真而跳过重复创建子进程直接把所有进程交给 torchrun 管理。优化多机通信显式指定进程组后端默认情况下在 GPU 上运行时 Lightning 会选择nccl后端而非gloo。PyTorch 的分布式包支持多种后端各有适用的硬件与场景。Lightning 允许你通过策略类Strategy class的构造参数process_group_backend显式指定后端未指定时Lightning 会根据当前硬件自动选择合适的后端from lightning.pytorch.strategies import DDPStrategy # 显式指定进程组后端 ddp DDPStrategy(process_group_backendnccl) # 将策略配置到 Trainer trainer Trainer(strategyddp, acceleratorgpu, devices8)源码级原理默认后端是如何决定的在 DDPStrategy.setup_distributed 中Lightning 调用_get_process_group_backend()若用户显式传入了process_group_backend则直接使用否则调用_get_default_process_group_backend_for_device按设备类型查表见 src/lightning/fabric/utilities/distributed.py该函数读取torch.distributed.Backend.default_device_backend_mapCUDA 设备对应nccl其他未登记的设备类型回退到gloo。也就是说文档所说的GPU 上默认 nccl实际上是委托给 PyTorch 的设备-后端映射表完成的。确定后端后setup_distributed会完成reset_seed()、通过set_world_ranks()计算global_rank node_rank * num_processes local_rank与world_size num_nodes * num_processes最后调用_init_dist_connection初始化进程组。若你想调整进程组超时时间还可以在构造DDPStrategy时传入timeout默认来自default_pg_timeout。DDP 构造参数速查DDPStrategy的完整构造签名见 ddp.py常用参数包括process_group_backend进程组后端如nccl、gloo默认None表示按硬件自动选择start_method进程启动方式popen标准 DDP、spawnDDP Spawn、fork/forkserverNotebook/Fork默认popentimeout进程组通信超时类型为datetime.timedeltaddp_comm_state/ddp_comm_hook/ddp_comm_wrapperDDP 通信钩子相关用于自定义梯度通信如 post-localSGD 模型平均化model_averaging_period与 post-localSGD 配合的模型平均周期其他**kwargs会直接透传给torch.nn.parallel.DistributedDataParallel例如find_unused_parametersTrue。此外注册表还提供了若干快捷字符串策略均可在Trainer(strategy...)中直接使用见 ddp.pyddp_find_unused_parameters_true/ddp_find_unused_parameters_false显式控制find_unused_parametersddp_spawn_find_unused_parameters_true/..._falsespawn 版本ddp_fork_find_unused_parameters_true/..._false与ddp_notebook_find_unused_parameters_true/..._falsefork 版本。当模型存在未参与 loss 计算的参数时DDP 的梯度归约会报错。此时若该行为是有意为之应通过strategyddp_find_unused_parameters_true或DDPStrategy(find_unused_parametersTrue)显式启用未使用参数检测——这是DDPStrategy.on_exception中对这类错误的标准处置建议见 ddp.py。上述注册表条目均有对应的单元测试覆盖可参考 tests/tests_pytorch/strategies/test_registry.py 中的参数化测试。实战建议与排错清单脚本训练选 DDP只要以python train.py方式运行就用strategyddp或直接省略 strategy兼顾速度与稳定性notebook 训练用 ddp_notebook交互环境下 Lightning 会自动切换也可显式指定fork 启动前不要先执行 GPU 相关操作调试/迁移旧代码才用 ddp_spawn接受只有权重回传、不支持多节点、更慢这三条限制多机训练确认三件事--nproc_per_node与devices一致、--nnodes与num_nodes一致、--master_port在节点防火墙放行遇到梯度归约报错确认模型是否有未参与 loss 计算的参数按需开启find_unused_parameters希望自定义通信用DDPStrategy(process_group_backend...)显式指定后端与超时。通过本文的组合配置你可以从单机单卡平滑扩展到单机多卡、再到多机多卡的分布式训练同时兼顾交互式环境与弹性调度的实际需求。【免费下载链接】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),仅供参考
返回列表