Ubuntu上PyTorch资源如何管理

Ubuntu上PyTorch资源管理实战指南

一 基础监控与诊断

• 使用 nvidia-smi 实时查看 GPU显存、利用率、功耗与温度，配合 watch 做周期性采样：watch -n 1 nvidia-smi。多卡训练时关注各卡的显存占用是否均衡。
• 启用 PyTorch 内存分配器历史记录 并生成 火焰图，定位占用热点与泄漏：
  1. 训练前开启记录：torch.cuda.memory._record_memory_history(True)
  2. 触发一次快照：s = torch.cuda.memory._snapshot()
  3. 使用官方提供的 _memory_viz.py 生成交互式火焰图，直观查看各分配栈与块生命周期。
• 在训练循环中定期打印显存统计：torch.cuda.memory_allocated()、torch.cuda.memory_reserved()、torch.cuda.max_memory_allocated()，帮助发现异常增长与缓存未释放问题。
• 若使用 AMD ROCm，升级到 ROCm 7 可获得更好的框架兼容与可观测性，并配套 资源管理器仪表板 与 AI 工作台，便于多 GPU/多实例 的资源分配与监控。

二 显存优化与批大小策略

• 采用 自动混合精度 AMP：前向用 autocast，反向用 GradScaler，通常可显著减少显存并提速；示例：
  scaler = torch.cuda.amp.GradScaler()
  with torch.cuda.amp.autocast(): loss = model(x)
  scaler.scale(loss).backward(); scaler.step(optimizer); scaler.update()
• 使用 梯度累积 在单卡上模拟大 batch：累积 N 步后再 optimizer.step()，等价于 batch_size * N 的更新效果，代价是更新频率降低。
• 开启 激活检查点（Checkpointing）：以计算换显存，适合 Transformer/LLM 等深层模型。
• 选择更省显存的 优化器（如 AdamW），并尽量使用 融合/高效实现（如 APEX 的融合 Adam，视环境可用性）。
• 减少 CPU↔GPU 频繁拷贝：创建张量时直接指定 device，传输时使用 non_blocking=True，避免在热路径上 .cpu()/.numpy()。
• 打开 cuDNN benchmark（输入尺寸固定时）：torch.backends.cudnn.benchmark = True，加速卷积等算子选择。
• 小技巧：用 zero_grad(set_to_none=True) 降低梯度清零开销；验证阶段用 torch.no_grad()。

三 多GPU与分布式资源管理

• 单机多卡优先使用 DistributedDataParallel（DDP），比 DataParallel 更高效、可扩展。
• 结合 张量分片/模型并行 将大矩阵拆分到多卡，降低单卡显存峰值。
• 使用 FSDP（Fully Sharded Data Parallel） 做参数/梯度/优化器状态分片；必要时开启 CPU Offload 将部分状态转移到 CPU 内存，以时间换显存。
• 在 AMD ROCm 环境下，ROCm 7 原生支持 PyTorch 2.7/2.9，并改进多 GPU 兼容与可观测性，适合在 Ubuntu 24.04.3 LTS 等平台部署大规模训练/推理。

四 进程与作业调度

• 使用 CUDA_VISIBLE_DEVICES 精确绑定进程到指定 GPU，避免误占用：CUDA_VISIBLE_DEVICES=0,2 python train.py。
• 通过 cgroups + systemd 或 nvidia-container-runtime 设置 显存/算力限额，防止单个作业耗尽整卡资源。
• 多用户/多任务场景建议采用 队列系统（如 Slurm） 或 Kubernetes + GPU Operator，结合节点标签与资源配额实现公平调度与隔离。
• 容器化训练时，优先选择带 GPU 驱动与 ROCm/CUDA 的基础镜像，减少环境不一致导致的资源浪费与调度冲突。

五 快速排障清单

• 显存“爆掉”或 OOM：先降 batch size，开启 AMP 与 梯度累积，对关键模块启用 Checkpointing，再考虑 FSDP/CPU Offload。
• 训练变慢但显存有余：检查是否频繁 CPU↔GPU 拷贝、是否关闭了 cudnn.benchmark、是否使用了低效优化器或数据加载瓶颈。
• 多卡不均衡：确认 DDP 正确初始化、各卡 batch 一致、数据增强无随机种子导致的负载漂移；必要时改用 FSDP 做更细粒度分片。
• 疑似泄漏：开启 _record_memory_history 生成 火焰图，定位持续增长分配栈；定期打印 max_memory_allocated 观察趋势。
• AMD 平台：确认 ROCm 7 与框架版本匹配，并使用其 仪表板 观察集群/设备利用率与配额。