torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.04 GiB. GPU

<think>好的，用户遇到了PyTorch的CUDA内存不足问题，需要解决。首先，我得回忆一下常见的解决方案。根据引用[1]，释放缓存分配器持有的未占用缓存内存可能有帮助。用户提到分配2.04GiB失败，可能显存已经被占用较多。 首先应该检查当前显存使用情况，用torch.cuda.memory_summary()看看哪些部分占用了内存。然后考虑减少批量大小，因为批量大显存需求也大。另外，混合精度训练可以节省显存，使用torch.cuda.amp可能有效。 如果模型太大，可能需要梯度检查点技术，比如torch.utils.checkpoint，这样可以在反向传播时重新计算中间结果，节省内存。模型并行或数据并行也是选项，但用户可能没有多卡，所以数据并行可能需要调整。 清理缓存也很重要，手动调用torch.cuda.empty_cache()释放未使用的缓存。但要注意，这不会释放正在使用的显存。另外，检查是否有不必要的张量留在内存中，比如中间变量，及时用del释放。 根据引用[2]，使用8位量化(--load-8bit)可以在11G显存下运行，用户的情况可能需要模型量化。比如使用bitsandbytes库进行8位优化，或者更激进的4位量化。同时，模型转换到更高效的框架如vLLM可能有效，但用户提到vLLM不支持，可能需要调整。 另外，检查代码是否有内存泄漏，比如循环中不断累积张量。确保在不需要的时候将张量移到CPU或删除。使用内存分析工具，比如memory_profiler，帮助定位问题。 如果这些方法都不行，可能需要升级硬件，或者使用云实例。但用户可能希望软件层面的解决。综合这些步骤，给出用户具体的解决方案，从简单到复杂排列，确保他们能一步步尝试。</think>### PyTorch CUDA内存不足（OutOfMemoryError）解决方案 当PyTorch报错`CUDA out of memory`时，通常是因为显存被模型、数据或缓存占用过多。以下是逐步解决方案： --- #### **1. 检查显存占用** 在代码中添加显存监控语句，定位问题点： ```python import torch print(torch.cuda.memory_summary()) # 输出显存使用详情 ``` 通过此命令可查看哪些操作占用了显存[^1]。 --- #### **2. 减少批量大小（Batch Size）** 降低`batch_size`是最直接的优化方式： ```python # 修改前 dataloader = DataLoader(dataset, batch_size=64) # 修改后 dataloader = DataLoader(dataset, batch_size=32) ``` 批量大小与显存占用呈线性关系，降低后可立即缓解压力。 --- #### **3. 使用混合精度训练** 通过`torch.cuda.amp`自动转换精度，减少显存占用： ```python from torch.cuda.amp import autocast, GradScaler scaler = GradScaler() for data, target in dataloader: with autocast(): output = model(data) loss = loss_fn(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() ``` 此方法可节省约30%显存。 --- #### **4. 释放缓存与无用变量** 手动清理未使用的缓存和变量： ```python torch.cuda.empty_cache() # 释放未使用的显存缓存 del intermediate_tensor # 删除中间变量（如不需要反向传播的中间结果） ``` --- #### **5. 梯度累积（Gradient Accumulation）** 在显存不足时，通过累积梯度模拟大批量训练： ```python optimizer.zero_grad() for i, (data, target) in enumerate(dataloader): output = model(data) loss = loss_fn(output, target) loss.backward() if (i+1) % 4 == 0: # 每4个batch更新一次参数 optimizer.step() optimizer.zero_grad() ``` --- #### **6. 模型量化与8位优化** 使用`bitsandbytes`库进行8位量化（适用于推理和小规模训练）： ```python from transformers import AutoModelForCausalLM model = AutoModelForCausalLM.from_pretrained("model_name", load_in_8bit=True) ``` 此方法可将显存占用降低至原大小的1/3~1/2[^2]。 --- #### **7. 分布式训练与模型并行** 多卡训练显存优化： ```python # 数据并行（单机多卡） model = torch.nn.DataParallel(model) # 模型并行（拆分模型到不同GPU） model.layer1 = model.layer1.to('cuda:0') model.layer2 = model.layer2.to('cuda:1') ``` --- #### **8. 其他高级优化** - **梯度检查点（Gradient Checkpointing）**：牺牲时间换空间，适用于大模型： ```python from torch.utils.checkpoint import checkpoint def forward_with_checkpoint(x): return checkpoint(custom_forward, x) output = forward_with_checkpoint(input) ``` - **卸载到CPU**：将中间结果暂存到CPU内存（需手动管理）。 --- ### 典型场景示例 若显存总容量为12GB，运行`chatglm3-6b`模型时： - 原始模型需约12GB显存，可能触发OOM。 - 使用`load_in_8bit`后，显存占用降至约7GB[^2]。 - 进一步降低`batch_size`可适配更低显存设备。 --- ###