本地 PyTorch 训练 OOM,第一次租 RTX 4090 云 GPU 怎么迁移项目?从环境检查到 100 Step 跑通

很多小型 AI 团队都会遇到类似问题:

项目在本地已经能运行,但 12GB 或 16GB 显卡频繁 OOM;想临时租一张 RTX 4090,又担心 CUDA、PyTorch、数据路径和依赖全部需要重配。第一次使用云 GPU,应该从哪里开始?

把 PyTorch 项目迁移到云 GPU,真正容易出错的并不是"上传代码",而是下面四件事:

  1. 选错基础镜像,导致 PyTorch 与 CUDA 环境不匹配;
  2. 代码、数据、模型权重混在一起,上传和备份效率很低;
  3. 没做小规模验证,直接启动完整训练;
  4. 只保存模型权重,没有验证断点续训。

本文以一个已经能在本地运行的 PyTorch 项目为例,演示如何迁移到 RTX 4090 云 GPU,并完成从环境检查到 100 Step 试跑的全过程。


一、为什么选择 RTX 4090,而不是直接上更贵的卡?

RTX 4090 配备 24GB GDDR6X 显存,适合以下常见任务:

  • 中小型视觉模型训练;
  • Stable Diffusion、LoRA 等生成模型微调;
  • 7B 级模型的 LoRA 或 QLoRA 实验;
  • 本地 12GB、16GB 显卡放不下的较大 batch;
  • 需要在一两周内快速完成训练验证的项目。

RTX 4090 的官方标准显存容量为 24GB,但"24GB"并不等于所有模型都能直接训练。模型参数量、优化器、序列长度、激活值和 batch size 都会占用显存。NVIDIA RTX 4090 官方规格

本文面向的是有 Linux、Python、PyTorch 操作基础,但没有太多云 GPU 租赁经验的小团队。因此,以算家云专业版 RTX 4090 为例操作演示。

算家云官网将专业版定位为"长期生产、推理训练、稳定运行",并提供项目镜像、网络存储、配置调整等功能。对于需要连续调试和反复训练的小团队,这类环境比一次性实验环境更容易保存工作状态。具体 GPU 库存、区域和实例配置,应以创建实例页面的实时信息为准。算家云专业版


二、迁移前先整理项目,不要把整个电脑环境搬上去

假设本地项目结构如下:

text 复制代码
my_project/
├── configs/
├── datasets/
├── checkpoints/
├── src/
├── train.py
├── requirements.txt
└── README.md

建议把内容分成三类:

内容 是否上传 建议保存位置
Python 代码、配置文件 是 项目代码目录
数据集 是 数据盘或项目网络存储
模型权重、Checkpoint 是 持久化存储目录
.git、__pycache__ 否 不上传
本地虚拟环境 .venv 否 云端重新创建
pip、模型下载缓存 通常否 云端重新下载或单独迁移

不要直接复制本地的 Conda 或 .venv 目录。虚拟环境中可能包含与本机系统、Python 路径和 CUDA 环境绑定的二进制文件,上传后往往不能直接使用。

迁移前,先在本地记录关键版本:

bash 复制代码
python --version

python -c "import torch; print('torch:', torch.__version__); print('cuda build:', torch.version.cuda)"

pip list

如果项目依赖不多,建议手动整理一份精简的 requirements.txt,不要无条件把本地环境中所有包全部冻结进去。

例如:

text 复制代码
transformers==4.x.x
datasets==3.x.x
accelerate==1.x.x
peft==0.x.x
sentencepiece
tensorboard

PyTorch 可以先不写入这个文件,等云端基础镜像启动后,根据已有环境决定是否需要安装或调整。


三、以算家云为例操作演示:创建专业版 RTX 4090 实例

进入实例创建页面后,重点检查以下信息:

  1. 选择专业版资源;
  2. 查看当前区域是否有 RTX 4090 可用;
  3. 确认显存、CPU、内存和磁盘配置;
  4. 选择与项目接近的 PyTorch 基础镜像;
  5. 检查 Python、PyTorch 和 CUDA 版本;
  6. 再创建实例。

不要只看"CUDA 版本越新越好"。

如果项目原来基于某个稳定的 PyTorch 版本运行,优先选择接近该环境的基础镜像。例如,本地项目使用 PyTorch 2.x,就没有必要为了追求最新版本,直接切换到一个完全不同的开发环境。

实例启动后,可以通过 SSH、JupyterLab 或 VS Code 进入环境。第一次迁移项目时,SSH 更适合批量上传文件,JupyterLab 更适合检查数据和逐步运行代码。


四、第一步不是安装依赖,而是检查 GPU 是否正常

进入实例后,先执行:

bash 复制代码
nvidia-smi

确认输出中能够看到:

  • GPU 型号;
  • 显存容量;
  • 驱动版本;
  • 当前显存占用;
  • 正在运行的 GPU 进程。

然后检查 PyTorch:

bash 复制代码
python - <<'PY'
import torch

print("PyTorch:", torch.__version__)
print("PyTorch CUDA build:", torch.version.cuda)
print("CUDA available:", torch.cuda.is_available())
print("GPU count:", torch.cuda.device_count())

if torch.cuda.is_available():
    print("GPU:", torch.cuda.get_device_name(0))
PY

PyTorch 官方说明中,torch.cuda.is_available() 用于判断当前 PyTorch 环境是否可以使用 CUDA。PyTorch 官方文档

正常情况下,应该看到类似结果:

text 复制代码
PyTorch: 2.x.x
PyTorch CUDA build: 12.x
CUDA available: True
GPU count: 1
GPU: NVIDIA GeForce RTX 4090

常见异常怎么判断?

现象 优先排查方向
nvidia-smi 无法执行 GPU、驱动或实例状态异常
nvidia-smi 正常,但 torch.cuda.is_available() 为 False 当前 Python 环境可能安装了 CPU 版 PyTorch,或运行了错误的虚拟环境
GPU 型号不是预期的 RTX 4090 检查实例规格和当前连接的实例
torch.version.cuda 与本地不同 不一定是故障,先运行项目验证,不要仅凭版本号重装
导入 PyTorch 时报动态库错误 检查是否混装了多个 Conda、pip 或系统环境

如果基础镜像已经能够正常识别 GPU,不要一上来就重新安装 PyTorch,更不要自行重装 NVIDIA 驱动。很多环境问题正是由覆盖原有可用环境引起的。


五、上传代码时排除无用文件

如果实例提供 SSH 连接信息,可以在本地使用 rsync:

bash 复制代码
rsync -av \
  --exclude '.git' \
  --exclude '.venv' \
  --exclude '__pycache__' \
  --exclude 'datasets' \
  --exclude 'checkpoints' \
  ./my_project/ 用户名@实例地址:/目标目录/my_project/

上面的用户名、实例地址和目标目录,需要替换为控制台显示的实际信息。

代码通常体积不大,可以先上传。数据集和模型权重建议单独处理,避免每次更新代码都重新传输几十 GB 数据。

进入云端项目目录后,先检查:

bash 复制代码
cd /目标目录/my_project

find . -maxdepth 2 -type f | head -50

然后确认训练配置中的路径没有写死成本地路径,例如:

python 复制代码
data_path = "D:/datasets/train"

建议改成命令行参数或环境变量:

python 复制代码
import os

data_path = os.environ.get(
    "DATA_DIR",
    "./datasets"
)

启动训练前设置:

bash 复制代码
export DATA_DIR=/实际数据目录
export OUTPUT_DIR=/实际输出目录

这样更换实例或目录时,不需要修改业务代码。


六、先安装项目依赖,不要覆盖能用的 PyTorch

先确认当前 Python:

bash 复制代码
which python
which pip
python --version

然后安装业务依赖:

bash 复制代码
python -m pip install -r requirements.txt

如果 requirements.txt 中锁定了另一个 PyTorch 版本,建议先删除或注释对应行,避免覆盖镜像已有的 GPU 版 PyTorch。

安装完成后执行一次导入检查:

bash 复制代码
python - <<'PY'
import torch
import transformers
import datasets

print("imports ok")
print("cuda:", torch.cuda.is_available())
PY

如果这里失败,应先修复依赖,再运行训练。不要把依赖错误、数据错误和显存问题混在一次完整训练中排查。


七、不要直接跑完整训练,先做四级验证

第一次迁移最有效的方式,是逐级扩大运行范围。

第一级:语法和导入检查

bash 复制代码
python -m compileall src train.py

第二级:配置和数据检查

确认训练脚本能读取配置、找到数据集,并取出一个样本。

如果项目有独立的数据测试脚本,可以执行:

bash 复制代码
python check_dataset.py

如果没有,至少打印一条样本的字段和张量形状。

第三级:单个 batch 前向与反向传播

让模型只运行一个 batch,确认:

  • 数据能够移动到 GPU;
  • 模型能够执行前向传播;
  • loss 可以正常计算;
  • backward 不会立即 OOM。

伪代码如下:

python 复制代码
batch = next(iter(train_loader))
batch = {
    key: value.cuda() if hasattr(value, "cuda") else value
    for key, value in batch.items()
}

optimizer.zero_grad()
outputs = model(**batch)
loss = outputs.loss
loss.backward()
optimizer.step()

print("single batch passed:", loss.item())

具体代码需要根据项目的数据格式调整。

第四级:运行 100 Step

如果训练脚本支持最大步数参数,可以执行:

bash 复制代码
python train.py \
  --config configs/train.yaml \
  --max_steps 100 \
  --output_dir "$OUTPUT_DIR/smoke_test"

--max_steps 和 --output_dir 只是示例,实际参数名称以项目代码为准。

这 100 Step 需要观察:

  • loss 是否正常变化;
  • GPU 利用率是否长期接近 0;
  • 显存是否持续增长;
  • DataLoader 是否卡住;
  • 日志和 Checkpoint 是否写入正确目录。

另开一个终端观察 GPU:

bash 复制代码
watch -n 1 nvidia-smi

100 Step 正常完成后,再启动完整训练。这样可以避免运行数小时后才发现路径或保存逻辑有问题。


八、如何检查 RTX 4090 的真实显存占用?

在一次迭代开始前清空峰值统计:

python 复制代码
torch.cuda.reset_peak_memory_stats()

完成前向、反向和参数更新后输出:

python 复制代码
allocated = torch.cuda.max_memory_allocated() / 1024**3
reserved = torch.cuda.max_memory_reserved() / 1024**3

print(f"峰值已分配显存: {allocated:.2f} GB")
print(f"峰值保留显存: {reserved:.2f} GB")

需要区分:

  • allocated:张量实际占用的显存;
  • reserved:PyTorch 缓存分配器保留的显存。

因此,nvidia-smi 显示的显存占用通常不会与 max_memory_allocated() 完全相同。

如果 24GB 仍然 OOM,可以按以下顺序调整:

  1. 减小 batch_size;
  2. 减小图片分辨率或序列长度;
  3. 使用梯度累积;
  4. 开启混合精度;
  5. 使用 gradient checkpointing;
  6. 大模型微调改用 LoRA 或 QLoRA;
  7. 确认没有意外保存计算图;
  8. 再考虑更大显存的 GPU。

不要通过反复调用 torch.cuda.empty_cache() 掩盖代码中的显存泄漏。它不能释放仍被 Python 对象引用的张量。


九、必须验证一次断点续训

云端训练最容易忽略的不是"保存模型",而是"保存后能不能继续训练"。

用于恢复训练的 Checkpoint,通常至少需要包含:

python 复制代码
torch.save(
    {
        "step": global_step,
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "scheduler_state_dict": scheduler.state_dict(),
        "scaler_state_dict": scaler.state_dict()
            if scaler is not None else None,
    },
    checkpoint_path,
)

然后实际停止 100 Step 测试,再从 Checkpoint 恢复几步:

bash 复制代码
python train.py \
  --config configs/train.yaml \
  --resume_from_checkpoint /检查点目录/checkpoint.pt \
  --max_steps 110

PyTorch 官方也建议,在用于恢复训练的通用 Checkpoint 中,同时保存模型和优化器状态,而不只是模型参数。PyTorch 模型保存与加载教程

需要确认恢复后:

  • step 没有从 0 重新开始;
  • optimizer 状态已加载;
  • learning rate 没有意外重置;
  • loss 没有明显异常;
  • 新 Checkpoint 能正常写入。

只有通过这一步,才算具备完整训练的基本条件。


十、环境、数据和 Checkpoint 应该怎么保存?

这三类内容不要混为一谈:

1. 项目镜像保存环境

依赖已经调通后,可以把实例系统环境保存为项目镜像,减少下次重新安装依赖的时间。

需要注意:算家云项目镜像保存的是系统盘环境,不包含数据盘内容。因此,保存镜像不能替代数据和 Checkpoint 备份。项目镜像说明

2. 数据集放在持久化存储

数据集不应只保存在可能随实例释放而清除的位置。具体目录和存储类型,应以控制台及官方帮助文档为准。

3. Checkpoint 单独保存

建议按实验划分目录:

text 复制代码
outputs/
└── exp_20260909_001/
    ├── config.yaml
    ├── train.log
    ├── checkpoint-100.pt
    ├── checkpoint-1000.pt
    └── metrics.json

配置文件、日志、代码版本和 Checkpoint 应对应保存,否则几周后很难复现实验。


十一、训练结束后,关机和释放不是一回事

按量实例训练完成后,应及时检查实例状态。

通常情况下,关机后实例计算资源停止计费,但磁盘、网络存储等资源可能仍然独立计费。释放实例则可能清除系统盘和本地数据盘,而且操作不可逆。

因此,释放前至少确认:

text 复制代码
[ ] 最终模型已经保存
[ ] Checkpoint 已复制到持久化位置
[ ] 训练日志已经保存
[ ] 配置文件已经保存
[ ] 断点续训已经验证
[ ] 需要复用的环境已保存为项目镜像
[ ] 已确认实例释放后的数据影响

不要把"实例已关机"理解为"所有资源均已停止计费",也不要在没有备份的情况下直接释放实例。


十二、第一次迁移云 GPU,真正的完成标准是什么?

不是看到 nvidia-smi,也不是代码开始打印 loss。

一套可以交付给团队继续使用的迁移结果,至少应满足:

text 复制代码
[ ] PyTorch 能识别 RTX 4090
[ ] 代码和数据路径已改成可配置形式
[ ] 单个 batch 可以完成前向和反向传播
[ ] 100 Step 测试能够正常结束
[ ] 显存峰值已经记录
[ ] 日志和 Checkpoint 写入持久化目录
[ ] Checkpoint 能够恢复训练
[ ] 环境版本已经记录或保存为项目镜像

对于第一次租用云 GPU 的小团队,最稳妥的第一步不是立即跑几天完整训练,而是先在专业版 RTX 4090 上完成一次 100 Step 冒烟测试。

这一步跑通后,再增加 batch、训练轮数和数据规模。出现问题时,也能快速判断它来自环境、数据、显存还是训练代码。

相关推荐
2501_938140078 小时前
做内容很省心,2026好用AI笔记,支持抖音小红书视频无水印保存
人工智能
sell-cloud11 小时前
Prompt Engineering要凉了?Fable 5.5能“猜“意图,但银行正在封杀AI Agent——企业AI的真正护城河是什么
人工智能·prompt
Navicat中国12 小时前
Navicat 18 内置 AI 助手让数据库开发从繁琐操作迈向智能协作
大数据·人工智能·数据库开发
4SAPI13 小时前
多模型请求路由怎么设计?成本、速度和能力如何平衡
人工智能·python·gateway·php
魔乐社区13 小时前
翼芯・智算测试与适配优化专区,让国产算力“选得准、跑得通、用得好“
人工智能·开源·大模型
XCDoris15 小时前
郴州logo设计机构
大数据·人工智能·python
Goboy15 小时前
一个人带三个Agent,我用全新TRAE从0到1搭了个资讯周刊站
人工智能·trae
俊哥V15 小时前
每日 AI 研究简报 · 2026-10-19
人工智能·ai
AI_AGENT_DEV_AI15 小时前
AI英语教培智能体的开发
人工智能
物联网软硬件开发-轨物科技15 小时前
【技术白皮书】光伏清检一体化机器人系统解决方案技术白皮书
大数据·运维·网络·人工智能·机器人