linux pytorch怎么省内存

在 Linux 上跑 PyTorch 时，“省内存”主要从模型、数据、训练策略、系统层面几个方向入手。下面按实用程度给你总结一套可操作的方法。

一、模型层面省内存

1. 用 torch.float16 / bfloat16

model.half()  # float16
# 或
model.to(torch.bfloat16)

训练/推理都可用（推理最明显）。

2. 用更小的模型 / 结构

• 用 torch.nn.Linear 替代大矩阵
• 用 Depthwise / Group Conv
• 减少 embedding 维度

3. 使用 torch.compile（PyTorch 2.x）

model = torch.compile(model)

可减少中间变量，加速+省显存。

二、训练时最关键的技巧（强烈推荐）

4. 梯度检查点（最常用）

from torch.utils.checkpoint import checkpoint

out = checkpoint(model.layer, x)

或 HuggingFace：

model.gradient_checkpointing_enable()

✅ 用时间换显存，非常有效

5. 减小 Batch Size + 梯度累积

accumulation_steps = 4

for i, (x, y) in enumerate(dataloader):
    loss = model(x, y) / accumulation
    loss.backward()

    if (i+1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

✅ 等效大 batch，但省内存

6. optimizer.zero_grad(set_to_none=True)

optimizer.zero_grad(set_to_none=True)

比默认更省内存。

三、显存 / 内存管理技巧

7. 及时释放变量

del x, y, loss
torch.cuda.empty_cache()

⚠️ empty_cache() 不要频繁调用，训练循环外用。

8. 关闭不必要的梯度

with torch.no_grad():
    out = model(x)

9. 冻结部分参数

for p in model.encoder.parameters():
    p.requires_grad = False

四、数据加载省内存（Linux 很关键）

10. 使用 DataLoader 正确配置

DataLoader(
    dataset,
    batch_size=8,
    num_workers=4,
    pin_memory=True,
    prefetch_factor=2
)

• num_workers 不要太多（吃 RAM）
• 数据尽量 on-the-fly 读取，不要全加载进内存

11. 用内存映射（大文件）

np.load("data.npy", mmap_mode="r")

五、分布式 / 多卡方案

12. 模型并行 / 流水线并行

• torch.nn.DataParallel（简单但不推荐）
• torch.nn.parallel.DistributedDataParallel（推荐）

13. ZeRO（大模型必备）

• DeepSpeed
• Accelerate
• FSDP（PyTorch 原生）

六、Linux 系统层面

14. 查看内存/显存

nvidia-smi
free -h
htop

15. 限制进程内存（防止 OOM）

ulimit -v 20000000

16. 用 swap（临时救急）

sudo fallocate -l 16G /swapfile
sudo mkswap /swapfile
sudo swapon /swapfile

七、总结（最推荐组合）

✅ 推理

float16 + torch.compile + 小 batch

✅ 训练

gradient checkpointing + 梯度累积 + zero_grad(set_to_none=True)

✅ 大模型

DeepSpeed / FSDP + CPU offload

如果你愿意，可以告诉我：

• 是 训练还是推理
• 模型大小（如 7B / ViT / CNN）
• 是 CPU 内存不够 还是 GPU 显存不够

我可以直接给你可抄的配置。