1. 环境准备

任务目标

本案例完成开源大模型本地部署与训练的前置环境搭建：安装Python 3.10+与CUDA Toolkit，并验证两者可被PyTorch正常调用。完成本案例后，你的机器将具备运行PyTorch模型训练与推理的基础条件。

环境需求

• 操作系统：Ubuntu 20.04/22.04（Windows/WSL2亦可，命令略有差异）
• GPU：NVIDIA显卡（需支持CUDA），已安装NVIDIA驱动（nvidia-smi可正常输出）
• 磁盘空间：至少10GB可用空间

安装步骤

1. 安装Python 3.10

使用apt安装Python 3.10及pip与venv模块：

sudo apt update
sudo apt install -y python3.10 python3.10-venv python3.10-dev
# 安装pip
curl -sS https://bootstrap.pypa.io/get-pip.py | python3.10
# 验证
python3.10 --version
python3.10 -m pip --version

2. 安装CUDA Toolkit

首先确认驱动支持的CUDA版本：

nvidia-smi  # 查看右上角 "CUDA Version"

以下以CUDA 12.1为例（若驱动版本较旧，请到NVIDIA官网选择对应版本）：

wget https://developer.download.nvidia.com/compute/cuda/12.1.1/local_installers/cuda_12.1.1_530.30.02_linux.run
sudo sh cuda_12.1.1_530.30.02_linux.run --toolkit --silent --override

安装完成后配置环境变量，编辑 ~/.bashrc 追加：

export PATH=/usr/local/cuda-12.1/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64:$LD_LIBRARY_PATH

source ~/.bashrc
nvcc --version  # 验证CUDA Toolkit安装

3. 创建虚拟环境并安装PyTorch

python3.10 -m venv ~/llm-env
source ~/llm-env/bin/activate
# 安装与CUDA 12.1匹配的PyTorch
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

验证安装

运行以下Python脚本，确认PyTorch能正确调用CUDA：

import torch
print("PyTorch版本:", torch.__version__)
print("CUDA是否可用:", torch.cuda.is_available())
if torch.cuda.is_available():
    print("GPU型号:", torch.cuda.get_device_name(0))
    # 实际执行一次GPU上的张量运算
    x = torch.rand(1000, 1000).cuda()
    y = torch.mm(x, x)
    print("GPU矩阵乘法验证通过，结果shape:", y.shape)
else:
    print("CUDA不可用，请检查驱动与CUDA安装")

通过条件：脚本输出中 CUDA是否可用: True，且能看到GPU型号名称与矩阵乘法结果shape为 torch.Size([1000, 1000])。若输出 False，请依次检查：NVIDIA驱动是否正常（nvidia-smi）、CUDA Toolkit版本是否与驱动匹配、PyTorch安装的CUDA版本是否与Toolkit一致。