LitGPT 显存溢出CUDA OOM问题排查与解决从微批大小到多卡分片的完整实践指南【免费下载链接】litgpt20 high-performance LLMs with recipes to pretrain, finetune and deploy at scale.项目地址: https://gitcode.com/GitHub_Trending/li/litgpt导读在 LitGPT 中预训练pretrain或微调finetune大语言模型时OutOfMemoryError: CUDA out of memory是最常见的拦路虎之一。本文以仓库中的 tutorials/oom.md 为核心结合 pretrain.py、finetune/lora.py 等源码与 TrainArgs 等配置实现系统讲解 OOM 的成因判断与六类行之有效的缓解手段缩小微批大小、缩减上下文长度、降低训练精度、多 GPU 分片FSDP、启用 CPU 卸载以及更换轻量优化器。读完本文你将能独立判断脚本配置与硬件显存之间的失衡点并熟练运用--train.micro_batch_size、--train.max_seq_length、--precision、--devices等核心参数把模型训练塞进现有 GPU。OOM 错误的识别先读懂报错再动手在运行脚本时如果看到如下报错说明你的 GPU 显存大小不足以承载当前的模型与脚本配置OutOfMemoryError: CUDA out of memory. Tried to allocate 2.22 GiB. GPU 0 has a total capacity of 79.15 GiB of which 228.38 MiB is free. Including non-PyTorch memory, this process has 78.93 GiB memory in use. Of the allocated memory 76.28 GiB is allocated by PyTorch, and 2.14 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF这份报错信息本身极具诊断价值它告诉你尝试分配了多少显存2.22 GiB、GPU 总容量79.15 GiB、已被 PyTorch 分配/预留的量以及一个可选的碎片化优化方向max_split_size_mb/PYTORCH_CUDA_ALLOC_CONF。当显存中reserved but unallocated已预留但未分配部分过大时可以尝试调整PYTORCH_CUDA_ALLOC_CONF环境变量如设置max_split_size_mb来减少碎片。从根本上看OOM 意味着模型权重 优化器状态 激活值activations 训练批次的总需求超过了单卡显存。因此解决思路无外乎两条要么减少某一项的内存占用要么把内存需求分摊到多张卡上。下面按成本从低到高、改动从易到难的顺序逐一展开。缩小微批大小micro batch size最直接的第一招参数含义与用法微批大小由--train.micro_batch_size参数控制它决定了一次迭代iteration中同时加载的样本数量。在 TrainArgs 中其默认值为4注释明确写道Number of samples per>litgpt finetune lora --train.micro_batch_size 1 litgpt pretrain --train.micro_batch_size 2取更小的值最小为1意味着同时加载的样本更少单次前向/反向传播的激活值占用显著下降这是所有手段中改动最小、效果最直接的一个。与全局批大小的关系需要留意的是micro_batch_size只是整体训练配置的一部分。TrainArgs 中还有global_batch_size默认64即所有数据并行 rank 在一次优化器步骤间处理的样本总数。在 pretrain.py 与各微调脚本中梯度累积次数由gradient_accumulation_iters()计算global_batch_size // (devices * num_nodes) // micro_batch_size。也就是说单纯调小微批大小不会改变全局批大小而是通过增加梯度累积步数来保持整体训练语义不变——这正是它可以在不损失模型效果相同 global batch的前提下降低显存的原因。权衡与实验建议更小的微批大小占用显存更少但由于每步处理的样本少、梯度累积步数增多训练收敛速度wall-clock 时间可能变慢更大的微批大小吞吐更高但对显存要求更高官方建议多尝试不同的微批大小在显存占用与计算效率之间找到平衡点。对微调任务micro_batch_size1通常是保底选择finetune/lora.py 中 LoRA 微调的默认值即为1。缩减上下文长度context length针对注意力机制显存的主要杠杆为什么 context length 如此关键上下文长度源码中的block_size对带注意力机制的模型影响极大。注意力矩阵的大小与序列长度的平方成正比对每个样本而言因此序列越长激活值显存增长得越剧烈。在 Config 中block_size默认值为4096各模型配置会覆盖该值例如 Llama 3 系列为8192Llama 3.1/3.2 系列为131072。预训练与微调脚本的默认行为差异仓库对两类脚本采用了不同的默认策略见 tutorials/oom.md预训练脚本默认使用模型的完整上下文长度进行训练即直接使用Config(block_size...)配置值微调脚本默认使用训练数据中最长样本的长度以避免为不存在的长序列分配多余显存。这一逻辑在 finetune/lora.py 的fit()中实现longest_seq_length, longest_seq_ix get_longest_seq_length(...) model.max_seq_length min(longest_seq_length, train.max_seq_length or float(inf))同时微调脚本在validate_args()finetune/lora.py等校验逻辑中会检查如果样本长度超过模型上下文长度则报错如果尝试运行的批次超过max_seq_length同样会报错。硬件受限时的两个对策预训练直接调小配置中的block_size。例如在 Python 中from litgpt.config import Config config Config.from_name(Llama-2-7b-hf, block_size2048)或在 yaml 配置文件中覆盖block_size字段后加载Config.from_file(...)见 litgpt/config.py。微调截断数据集中的样本长度。所有微调脚本都暴露了--data.max_seq_length...参数例如litgpt finetune lora --data.max_seq_length 256这在样本长度高度不均衡的数据集上尤其有效——因为只要存在一条特别长的样本批次内所有其他较短样本都会被 padding 到同样的长度导致整批内存膨胀。一个可复现的量化收益tutorials/prepare_dataset.md 的Truncating datasets章节给出了具体数据Alpaca 数据集中样本长度的中位数约为 110 个 token将 Alpaca 数据集截断到 256 个 token 上限后Falcon 7B 模型在 LoRA 微调micro batch size 为 1、bf16 精度下的显存需求从23.52 GB 降至 15.73 GB节省约三分之一。需要始终牢记的是缩减上下文长度会影响模型对超长文本序列的建模能力。如果你的下游任务包含长文档请评估截断对效果的潜在影响而不是无脑调小。使用更低精度用数值稳定性换显存精度选项一览所有脚本都暴露了--precision参数它直接决定模型权重与计算的内存占用。仓库支持的常见取值包括32-true、16-true、bf16-true、16-mixed、bf16-mixed各脚本 docstring 中有明确声明见 finetune/lora.py。核心原则真正的低精度16-true、bf16-true权重与计算均以半精度fp16/bf16进行相比32-true显存占用直接减半代价是表示范围受限模型可能因溢出/舍入产生 NaN混合精度16-mixed、bf16-mixed主权重保存在 fp32仅在计算时使用半精度训练稳定性更好但显存节省有限优化器与主权重仍在 fp32。仓库的默认精度策略如果你不显式传入--precision脚本会调用 get_default_supported_precision() 自动选择当前硬件支持的精度训练时默认返回bf16-mixed若 GPU 支持 bf16否则返回16-mixed推理时相应返回-true版本。这意味着默认配置已经是混合精度想进一步压缩显存时可以向-true方向调整但务必留意数值稳定性。与量化quantization的边界需要注意的是--precision与--quantize如bnb.nf4是两条独立的路径且存在互斥约束在 finetune/lora.py 等脚本中若同时使用bnb.*量化与mixed精度会直接抛出ValueError(Quantization and mixed precision is not supported.)。量化相关内容详见 tutorials/quantize.md。多 GPU 分片用 FSDP 把显存摊到多张卡上当模型规模过大、上述手段都不够用时可以用通信换显存。方法很简单——把脚本中的--devices 1改为大于 1 的值litgpt pretrain --devices 4 litgpt finetune full --devices 4启用该选项后会激活一种并行技术——FSDPFully Sharded Data Parallelism把模型参数、梯度和优化器状态分片存储到不同 GPU 上从而降低单卡显存峰值。源码中的 FSDP 配置预训练pretrain.py 在多卡时使用strategy FSDPStrategy(auto_wrap_policy{Block}, state_dict_typefull, sharding_strategyHYBRID_SHARD)微调adapterfinetune/adapter.py 使用strategy FSDPStrategy( auto_wrap_policy{Block}, activation_checkpointing_policy{Block}, state_dict_typefull, limit_all_gathersTrue, cpu_offloadFalse, )其中auto_wrap_policy{Block}表示以 Transformer Block 为单位进行分片微调脚本还额外设置了activation_checkpointing_policy{Block}即默认已启用激活检查点activation checkpointing——它通过不保存中间激活值、在反向传播时重新计算来换取显存是默认配置已开启的省显存机制。CPU 卸载CPU offloading用 CPU 内存换显存如果多卡分片后依然吃紧可以进一步启用 CPU 卸载将cpu_offloadFalse改为True见 finetune/adapter.py。该选项把部分参数状态卸载到 CPU 内存显著降低 GPU 显存峰值但会引入 CPU-GPU 之间的数据传输开销训练速度随之下降。它适合GPU 显存极小但 CPU 内存充足的机器。实践提示多卡训练前脚本会通过check_nvlink_connectivity()见 finetune/lora.py检查 GPU 间的 NVLink 连接提示通信拓扑质量量化--quantize目前不支持多卡训练相关脚本会抛出NotImplementedError详见 finetune/lora.py更多分布式与策略细节可参考 tutorials/developer-docs/python-api.md。更换优化器针对优化器状态的二次方内存为什么优化器也是显存大户仓库脚本默认使用AdamW优化器见 tutorials/oom.md。AdamW 为每个可训练参数维护两个状态一阶动量与二阶动量因此优化器内存约为模型参数本身的 2 倍相比之下SGD等优化器维护的状态更少甚至为零显存开销显著更低。替换方式与注意事项你可以通过脚本的--optimizer参数替换为任意内存更轻的优化器。例如litgpt pretrain --optimizer SGD不过要清醒认识到不同优化器的收敛行为差异很大替换后必须评估对训练过程与最终模型效果的影响。文档中还举例了两个近期发布的研究型优化器Sophia论文 arXiv:2305.14342二阶信息预条件与 Lion论文 arXiv:2302.06675符号梯度更新它们都以更低的显存开销为目标设计。何时最值得动优化器这一建议对预训练尤其相关预训练时模型全部参数都可训练优化器状态占大头而 LoRA、Adapter 等微调场景下可训练参数只占模型总参数的很小一部分优化器状态的绝对增量有限更换优化器的收益相对较小。组合拳一份 OOM 排查路线图把上述手段按干预成本与收益整理成建议的排查顺序优先级手段参数/配置成本与收益1缩小微批大小--train.micro_batch_size最小 1改动最小直接降低激活值显存2缩减上下文长度预训练block_size微调--data.max_seq_length激活值显存随序列长度平方级下降注意长文本能力损失3使用更低精度--precision16-true/bf16-true减半显存-mixed更稳定但省得少注意 NaN 风险4多 GPU 分片--devices N启用 FSDP默认已开激活检查点用通信换显存多卡均摊参数/梯度/优化器状态5CPU 卸载cpu_offloadFalse→True用 CPU 内存换显存训练变慢6更换优化器--optimizer如 SGD、Sophia、Lion预训练收益大微调收益有限额外提示显存碎片若报错信息中 reserved but unallocated 部分较大可尝试设置PYTORCH_CUDA_ALLOC_CONF如max_split_size_mb缓解碎片化验证参数合法性微调脚本的validate_args()finetune/lora.py会检查--train.max_seq_length与模型上下文长度的关系超限会直接报错而非静默截断精度默认值不传--precision时get_default_supported_precision() 按硬件自动选择bf16-mixed/16-mixed可据此判断当前默认显存基线。小结CUDA OOM 的本质是模型配置的显存需求与硬件显存容量的错配。LitGPT 通过 TrainArgs、Config 与 FSDPStrategy 等组件为每个可调参数都预留了清晰的 CLI 入口。实践中最有效的路径通常是先降 micro batch size 到 1再按需截断序列长度微调场景或调小 block_size预训练场景必要时切换-true精度最后考虑多卡 FSDP 分片与 CPU 卸载。每调整一步都建议结合报错信息中的显存分配明细观察变化找到最适合你硬件配置的组合。【免费下载链接】litgpt20 high-performance LLMs with recipes to pretrain, finetune and deploy at scale.项目地址: https://gitcode.com/GitHub_Trending/li/litgpt创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考