尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

PyTorch Lightning 梯度累积(Gradient Accumulation)完全指南:原理、配置与动态调度

PyTorch Lightning 梯度累积(Gradient Accumulation)完全指南:原理、配置与动态调度 PyTorch Lightning 梯度累积Gradient Accumulation完全指南原理、配置与动态调度【免费下载链接】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 中在显存受限时扩大有效批次大小的核心手段它通过连续运行 K 个小批次再执行一次参数更新将有效批次大小放大 K 倍同时几乎不增加显存开销。本文以 docs/source-pytorch/common/gradient_accumulation.rst 为主线结合仓库源码训练循环、优化器闭包、回调实现与测试用例深入讲解Trainer(accumulate_grad_batches...)的完整用法、分布式场景下的真实有效批次计算、以及用GradientAccumulationScheduler按 epoch 动态调整累积因子的实战方案。读完本文你将能够在 DDP/DP 等策略下准确预估有效批次大小并安全地在训练中切换累积窗口。一、核心概念什么是梯度累积为什么它能“免费”增大批次梯度累积Accumulated Gradients的基本思想是在真正执行backward更新之前先连续运行K个大小为N的小批次效果等同于使用大小为K×N的大批次进行训练这里 N 为单次批大小。关键点在于其内部实现方式原文说明它并不是把 K 个批次堆叠起来做一次前向传播而是逐个累积 K 个批次的梯度最后再调用optimizer.step()从而在增大有效批次大小的同时不产生额外的显存开销。也就是说显存中始终只保存一个小批次的数据与中间激活而梯度本身在参数上累积。这使得在单卡显存无法容纳大 batch 时依然可以模拟大 batch 的统计特性进行稳定训练。在 PyTorch Lightning 中该功能由Trainer的accumulate_grad_batches参数控制默认值为1即不做累积定义见 src/lightning/pytorch/trainer/trainer.pyaccumulate_grad_batches: int 1,其文档注释为“Accumulates gradients over k batches before stepping the optimizer.”并在构造时保存为self.accumulate_grad_batches见 trainer.py。二、基础用法一行代码开启梯度累积最简单的用法是直接给Trainer传入一个整数表示每累积多少个 batch 执行一次优化器 step# DEFAULT (即不做梯度累积) trainer Trainer(accumulate_grad_batches1) # 每 7 个 batch 累积一次梯度等效有效批次为 7*N trainer Trainer(accumulate_grad_batches7)上例中accumulate_grad_batches7意味着在连续的 7 个批次上执行loss.backward()累积梯度到第 7 个 batch 结束时才调用optimizer.step()。源码视角累积窗口内到底发生了什么从源码看这一过程由 src/lightning/pytorch/loops/optimization/automatic.py 中的_AutomaticOptimization循环与Closure协作完成主要行为包括损失归一化ClosureResult.from_training_step_output在累积模式下会对 loss 除以累积因子normalize即accumulate_grad_batches代码注释明确写道# accumulate the loss. If accumulate_grad_batches 1, no effect # note: avoid in-place operation x / y here on purpose closure_loss closure_loss / normalize这一步与 PyTorch 官方推荐的梯度累积写法loss loss / accumulation_steps; loss.backward()一致用于将多次 backward 的梯度平均到等效大 batch 的尺度避免学习率被隐式放大。只在前向阶段阻止梯度同步在累积窗口内_AutomaticOptimization.run中通过_block_parallel_sync_behavior(self.trainer.strategy, blockTrue)包裹 closure 执行保证累积期间的每次loss.backward()不会触发跨设备梯度同步只有在真正的optimizer_step时才同步。zero_grad只在窗口起点触发_make_zero_grad_fn中以batch_idx % self.trainer.accumulate_grad_batches 0判断是否为累积窗口的第一个 batch只有第一个 batch 才执行optimizer_zero_grad其余 batch 直接返回None跳过 zero_grad从而让梯度在小批次间持续累积。optimizer.step()只在窗口结束触发训练 epoch 循环通过_accumulated_batches_reached()training_epoch_loop.py判断batch_progress.current.ready % accumulate_grad_batches 0决定当前 batch 结束后是继续累积还是执行参数更新。测试用例也验证了这一行为在 tests/tests_pytorch/callbacks/test_gradient_accumulation_scheduler.py 中limit_train_batches20、accumulate_grad_batchesk时断言zero_grad的调用次数恰好为math.ceil(20 / k)即每个窗口只清一次梯度。三、分布式训练下的有效批次DDP 与 DP 的区别重要警告原文档特别警告了分布式场景下有效批次大小的计算差异这是最容易踩坑的地方。DDP如 DDP 策略当使用 DDP 且设备数为P时每个设备各自独立累积梯度——即每个设备在各自的loss.backward()后保存本地梯度直到调用optimizer.step()才进行跨设备梯度同步。因此对于每个设备自身而言一个累积窗口内的有效批次大小是N×K但在optimizer.step()之前的梯度同步all-reduce会把 P 个设备的梯度求和/求平均因此全局有效批次大小实际上是 P×N×K。DPDataParallelDP 策略下一个批次的数据会被切分到多个设备上每个设备只处理其中的一部分因此最终的有效批次大小保持为N×K不会乘上设备数 P。关键推论当你用 DDP 且accumulate_grad_batchesK时若想精确复现单卡上batch_sizeN的统计行为需要把“每步梯度下降所见样本数”按P×N×K来估算这直接影响 BatchNorm 的统计、学习率缩放等超参数设定。累积期间的梯度不会跨设备同步见上文_block_parallel_sync_behavior的实现因此累积窗口内部的通信开销被显著降低这也是梯度累积在分布式训练中兼具省显存与省通信的原因之一。四、累积窗口与 epoch 结束边界最后一个 batch 的行为一个容易被忽视的细节是当累积因子不能整除一个 epoch 的 batch 数时Lightning 会在 epoch 的最后一个 batch 上强制执行一次optimizer.step()。在 src/lightning/pytorch/loops/training_epoch_loop.py 的_should_accumulate()中可以看到其判定逻辑def _should_accumulate(self) - bool: Checks if the optimizer step should be performed or gradients should be accumulated for the current step. accumulation_done self._accumulated_batches_reached() # Lightning steps on the final batch is_final_batch self._num_ready_batches_reached() # but the strategy might not strategy_accumulates_on_final_batch self.trainer.strategy.handles_gradient_accumulation or not is_final_batch return not accumulation_done and strategy_accumulates_on_final_batch若当前 batch 已满足“累积窗口结束”或“已是 epoch 最后一个 batch”_should_accumulate()返回False即执行优化器 step若策略自身声明handles_gradient_accumulation如 DeepSpeed 这类由策略内部处理累积的策略则交给策略处理Lightning 始终调用 step 入口以路由到策略。同时当 epoch 内 batch 总数小于accumulate_grad_batches时src/lightning/pytorch/loops/fit_loop.py 会发出警告由于 Lightning 总是在 epoch 最后一个 batch 上执行 step本 epoch 实际累积的 batch 数将小于设定值例如设定accumulate_grad_batches7但该 epoch 只有 4 个 batch则实际按 4 个 batch 累积。这也提醒我们当每个 epoch 的 batch 数不固定如 IterableDataset或过小时梯度累积的实际效果可能与预期不符。此外epoch 循环中update_lr_schedulers也会配合_should_accumulate()工作——在累积窗口内跳过按 step 频率更新的 LR scheduler见 training_epoch_loop.py确保学习率调度与真实参数更新步数对齐。五、进阶用法用 GradientAccumulationScheduler 动态调整累积因子在实际训练中累积因子未必需要全程固定。PyTorch Lightning 提供了GradientAccumulationScheduler回调允许按 epoch 动态改变累积窗口大小。原文档给出的示例为from lightning.pytorch.callbacks import GradientAccumulationScheduler # 直到第 5 个 epoch每个累积窗口累积 8 个 batch # 从第 5 个 epoch 到第 9 个 epoch每个窗口累积 4 个 batch # 之后不再累积。 # 注意epoch 键是 0 起始zero-indexed的 accumulator GradientAccumulationScheduler(scheduling{0: 8, 4: 4, 8: 1}) trainer Trainer(callbacksaccumulator)调度字典的键表示“从该 epoch 开始生效”值表示“该 epoch 起的累积因子”。上述字典的含义是epoch 03累积 8 个 batchepoch 47累积 4 个 batchepoch 8 及以后累积 1 个 batch即不再累积。这正是原文注释“till 5th epoch … From 5th epoch till 9th epoch … after that no accumulation”对应的 0 起始索引表述。回调实现细节GradientAccumulationScheduler的实现位于 src/lightning/pytorch/callbacks/gradient_accumulation_scheduler.py其关键机制包括构造时校验输入空字典直接抛出TypeErrorepoch 键必须是非负整数否则抛出MisconfigurationException(Epoch should be an int greater than or equal to 0...)累积因子值必须是大于 0 的整数否则抛出MisconfigurationException(Accumulation factor should be an int greater than 0...)若最小 epoch 键不为 0用户未定义第 0 个 epoch 的因子会自动补上{0: 1}即默认从 epoch 0 起不累积。按 epoch 查询因子get_accumulate_grad_batches(epoch)从大到小遍历排序后的 epoch 键返回“最后一个满足epoch iter_epoch的调度值”实现区间取值。在on_train_epoch_start中生效每个 epoch 开始时将trainer.accumulate_grad_batches更新为当前 epoch 的累积因子从而让训练循环在下个 epoch 使用新窗口。调度器的限制与校验务必注意on_train_start中会对兼容性做严格校验见 gradient_accumulation_scheduler.py手动优化manual optimization不支持若pl_module.automatic_optimization为False直接抛出RuntimeError。因为自动梯度累积与GradientAccumulationScheduler仅支持自动优化模式DeepSpeed 策略不支持窗口动态变化DeepSpeedStrategy要求累积因子固定使用调度器会抛出RuntimeError不能与Trainer(accumulate_grad_batches...)同时使用一旦在Trainer中设置了非 1 的固定累积因子又挂载该回调会抛出ValueError要求二者只选其一若你重写了LightningModule.optimizer_step或optimizer_zero_grad且累积因子大于 1会收到警告这两个钩子将不再每个 batch 调用而是每个优化 step 调用一次。这些约束与测试用例一一对应test_gradient_accumulation_scheduler.py非法 epoch 键、非法累积值均会抛出MisconfigurationExceptionDeepSpeedStrategy会被判定不支持。六、手动优化模式为何不能自动累积需要特别强调的是自动梯度累积仅适用于自动优化模式。在 src/lightning/pytorch/trainer/configuration_validator.py 的__verify_manual_optimization_support中if trainer.accumulate_grad_batches ! 1: raise MisconfigurationException( Automatic gradient accumulation is not supported for manual optimization. f Remove Trainer(accumulate_grad_batches{trainer.accumulate_grad_batches}) or switch to automatic optimization. )如果你在LightningModule中设置了self.automatic_optimization False并自行管理optimizer.zero_grad()/backward()/optimizer.step()那么需要手动实现梯度累积逻辑例如经典的“每 K 步才 step”写法而不能依赖Trainer(accumulate_grad_batches...)。七、最佳实践与注意事项总结显存受限优先考虑梯度累积它扩大有效批次但不增加单次前向/反向的峰值显存是比“直接调大 batch_size”更稳妥的替代方案。注意分布式下的有效批次DDP 下全局有效批次为P×N×KDP 下为N×K据此校准学习率与归一化层行为。损失会自动归一化Lightning 在累积模式下自动将 loss 除以 K无需在training_step中手动loss / K避免二次归一化导致学习率失真。epoch 末尾的强制 step累积因子无法整除 epoch batch 数时最后一批会提前 step属预期行为若 epoch 内 batch 数小于累积因子请留意 fit_loop.py 的警告。动态调度注意兼容性GradientAccumulationScheduler不支持手动优化、不支持 DeepSpeed 策略且不能与固定accumulate_grad_batches同时使用原文档末尾的提示“并非所有策略与加速器都支持可变累积窗口”即指此类约束。回调钩子调用频率变化累积窗口内optimizer_step/optimizer_zero_grad只按窗口调用依赖逐 batch 执行这些钩子的自定义逻辑需要相应调整。八、相关资源导航本文主文档docs/source-pytorch/common/gradient_accumulation.rstTrainer.accumulate_grad_batches参数定义与文档src/lightning/pytorch/trainer/trainer.py自动优化循环与损失归一化实现src/lightning/pytorch/loops/optimization/automatic.pyepoch 循环中的累积判定_should_accumulatesrc/lightning/pytorch/loops/training_epoch_loop.py动态调度回调实现src/lightning/pytorch/callbacks/gradient_accumulation_scheduler.py手动优化模式校验src/lightning/pytorch/trainer/configuration_validator.py相关测试用例tests/tests_pytorch/callbacks/test_gradient_accumulation_scheduler.py【免费下载链接】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),仅供参考
返回列表