Gpu Pytorch Setup

GPU PyTorch 环境搭建与验证

全局 skill，对所有项目生效。本文以一台 8 × NVIDIA GeForce RTX 3090（24GB） 服务器为例（驱动 595.71.05，最高支持 CUDA 13.2， conda 位于 /opt/anaconda3），固化「检测 GPU → 建环境 → 装匹配的 GPU torch → 非沙箱验证」流程。 使用前请按自己机器的实际情况修改路径、GPU 数量与 torch 版本。

触发场景

• 【首要】每遇到一个需要 GPU 的新项目：无论是训练还是推理，只要新拉取 / 新建的项目要在 GPU 上跑，就按本 skill 为它建独立 conda 环境并装匹配的 GPU torch，不复用其他项目环境（避免依赖版本互相污染）
• 用户要求「创建环境跑通代码」「给新项目配 GPU 环境」「新项目装 torch」
• 需要安装 / 重装 GPU 版 PyTorch，或从 CPU 版切换到 GPU 版
• 报错 CUDA initialization ... Error 304: OS call failed、torch.cuda.is_available() 返回 False
• 需要确认模型 / 张量是否真正运行在 GPU 上，或需要挑选空闲 GPU 卡

默认约定：一个 GPU 项目 ↔ 一个独立 conda 环境（环境名建议与项目同名），统一走「检测 → 建环境 → 装 cu121 torch → 非沙箱验证」。

关键前置事实（本机）

• conda：source /opt/anaconda3/etc/profile.d/conda.sh && conda activate <env>
• GPU：8× RTX 3090（Ampere, sm_86），兼容 CUDA 11.x / 12.x 全系列 wheel
• 推荐 torch wheel：cu121（torch 2.5.1+cu121，生态最稳）
• 共享服务器上部分 GPU 常被他人任务（如 vLLM 推理服务）长期占用，优先用 nvidia-smi 确认空闲的卡
• pip 默认走公网 pypi.org 可通但偶发超时；清华镜像 pypi.tuna.tsinghua.edu.cn 在本机不可达，不要用

执行步骤

1. 检测 GPU、驱动、占用情况

nvidia-smi

记录：驱动版本、CUDA Version（驱动上限）、各卡显存占用与进程。 选卡原则：挑 memory.used 最小、无他人 compute 进程的卡。

2. 创建 conda 环境

source /opt/anaconda3/etc/profile.d/conda.sh
conda create -n <env_name> python=3.10 -y
conda activate <env_name>

3. 安装项目依赖（先装非 torch 依赖）

pip install --timeout 120 <其他依赖，如 numpy fastapi sentence-transformers ...>

4. 安装 GPU 版 torch

若当前已是 CPU 版（torch+x.x.x+cpu），先卸载：

pip uninstall -y torch

安装 cu121 GPU 版（torch 2.5.1）：

pip install --timeout 600 torch --index-url https://download.pytorch.org/whl/cu121

torch 2.5.1 会自动带上 nvidia-cublas/cudnn/nccl 等 cu12 运行时库，无需本机装 CUDA toolkit（不依赖 nvcc）。

5. 验证 CUDA —— 必须禁用 sandbox

这是本机最容易踩的坑：在 TRAE 默认 sandbox 内执行任何 CUDA 调用都会报：

CUDA initialization ... Error 304: OS call failed or operation not supported on this OS
torch.cuda.is_available() -> False

这不是 torch / 驱动问题，而是沙箱拦截了 GPU 设备文件访问。所有 GPU 验证和实际训练 / 推理命令，Shell 工具必须设 dangerouslyDisableSandbox: true。

验证脚本：

source /opt/anaconda3/etc/profile.d/conda.sh && conda activate <env_name> && python -c "
import torch
print('torch:', torch.__version__)
print('cuda available:', torch.cuda.is_available())
print('cuda version:', torch.version.cuda)
print('gpu count:', torch.cuda.device_count())
x = torch.randn(3,3, device='cuda:0'); print('matmul:', (x@x).device)
"

期望输出：cuda available: True、gpu count: 8、matmul: cuda:0。

6. 指定空闲卡并验证项目代码

export CUDA_VISIBLE_DEVICES=0   # 选第 1 步确认的空闲卡
python app.py                    # 或项目入口

确认模型真在 GPU（三选一）：

1. 服务进程在该卡有显存占用： nvidia-smi --query-compute-apps=pid,used_memory --format=csv,noheader -i 0
2. 代码内检查：next(model.parameters()).device → cuda:0
3. 推理器一般有自动选择逻辑，确认类似 torch.device("cuda" if torch.cuda.is_available() else "cpu") 解析为 cuda

7. 收尾

后台跑的服务验证完务必停止并确认显存释放，避免残留：

pkill -9 -f "python app.py"
nvidia-smi --query-compute-apps=pid,used_memory --format=csv,noheader -i 0

注意：pkill 可能因进程属于其他用户而报「不允许的操作」，用 ps -p <pid> -o pid,user,cmd 辨认，只清理自己启动的进程，不要误杀共享服务器上他人的训练任务。

常见问题速查