从零搭建大模型训练环境：CUDA、PyTorch、Transformers、Accelerate与DeepSpeed

2026-03-19 Ai大模型训练教程,CUDA,PyTorch,Transformers,Accelerate,DeepSpeed,大模型训练环境搭建,ZeRO优化,分布式训练

原始问题：

本篇Ai大模型训练教程详解从零搭建大模型训练环境的实操流程：检查NVIDIA驱动与GPU、安装带CUDA runtime的PyTorch、部署Transformers与datasets、配置Accelerate实现单机多卡启动，并引入DeepSpeed进行ZeRO显存优化，附验收脚本与常见故障排查清单。

重新提问

写在前面：你要搭建的是“能稳定复现”的训练环境

在 Ai大模型训练教程 这个系列里，训练环境搭建不是“装上能跑就行”，而是要做到：

版本可控：CUDA / PyTorch / 驱动 / Transformers 等能对齐，避免“今天能跑明天报错”。
可复现：同一套步骤在新机器、同事电脑、服务器上都能复刻。
可扩展：从单卡到多卡，从简单微调到 DeepSpeed ZeRO、混合精度都能覆盖。

本文将从零开始，按“从基础到进阶”的顺序，完成 CUDA（或 CUDA Runtime）+ PyTorch + Transformers + Accelerate + DeepSpeed 的环境搭建与验收，并给出常见坑位排查方法。

默认读者使用 Linux（推荐 Ubuntu 20.04/22.04）。Windows 也能做，但大模型训练的稳定性、生态与排错成本通常不如 Linux。

环境规划：先定标准，再装软件

你需要明确的 5 个关键点

GPU 型号：例如 A100 / 3090 / 4090 / L40 等，不同显卡对 CUDA 版本和性能差异有影响。
驱动版本（NVIDIA Driver）：驱动决定你能使用的最高 CUDA 运行时能力（兼容层面）。
CUDA 选择策略：
- 推荐：不手动安装完整 CUDA Toolkit，直接用 PyTorch 自带的 CUDA runtime（如 cu121）。
- 仅在需要编译自定义算子、特定 CUDA 工具链时才装完整 Toolkit。
Python 与虚拟环境：推荐使用 conda/mamba 或 uv/venv，避免污染系统 Python。
训练框架组合：Transformers + Accelerate 是主流“标配”，DeepSpeed 用于大模型分布式/显存优化。

第一步：检查 GPU 与驱动是否正常

在服务器上执行：

1) 查看 GPU 与驱动

nvidia-smi

你需要关注：

GPU 是否能被识别
Driver Version（驱动版本）
CUDA Version（驱动支持的最高 CUDA 版本提示）

2) 查看更详细信息（可选）

nvidia-smi -L
nvidia-smi --query-gpu=name,driver_version,memory.total --format=csv

如果 nvidia-smi 都无法运行：优先解决驱动安装/权限问题（例如宿主机未安装驱动、容器未挂载 GPU、或 Secure Boot 导致驱动未加载）。

第二步：创建隔离的 Python 环境（强烈建议）

以下给出 conda 与 venv 两种方式，选其一。

方案 A：使用 conda/mamba（更省心）

conda create -n llm-train python=3.10 -y
conda activate llm-train

方案 B：使用 venv（更轻量）

python3.10 -m venv .venv
source .venv/bin/activate
python -m pip install -U pip setuptools wheel

建议额外安装一些常用工具：

pip install -U packaging ninja psutil

第三步：安装 PyTorch（优先使用官方带 CUDA runtime 的 wheel）

为什么推荐“PyTorch 自带 CUDA”

避免系统 CUDA Toolkit 版本混乱
减少环境变量配置
对初学者更容易“装完就能用”

安装命令（示例：CUDA 12.1 runtime）

请以 PyTorch 官方指引为准，但常见安装方式如下：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

安装后立刻验收：

python - << 'PY'
import torch
print('torch:', torch.__version__)
print('cuda available:', torch.cuda.is_available())
print('cuda version (torch):', torch.version.cuda)
print('gpu count:', torch.cuda.device_count())
if torch.cuda.is_available():
    print('gpu name:', torch.cuda.get_device_name(0))
PY

期望输出：

cuda available: True
能看到 GPU 名称

如果 torch.cuda.is_available() 为 False：

先确认你是否装了 CPU 版 torch（index-url 是否正确）
再确认当前环境能访问 GPU（容器是否 --gpus all，或是否在无 GPU 的节点上）

第四步：安装 Transformers 与常用训练依赖

基础安装

pip install -U transformers datasets tokenizers safetensors

训练与日志常用组件（建议装）

pip install -U accelerate evaluate tqdm sentencepiece
pip install -U tensorboard wandb

说明：

datasets：标准数据集加载与处理
safetensors：更安全的权重格式
sentencepiece：一些模型分词器依赖（如 LLaMA 系）
wandb/tensorboard：实验记录与可视化

快速验收：加载一个小模型做前向

python - << 'PY'
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

model_name = 'gpt2'
tok = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

text = 'Hello, I am learning LLM training.'
inputs = tok(text, return_tensors='pt')
with torch.no_grad():
    out = model(**inputs)
print('logits shape:', out.logits.shape)
PY

这一步主要验证：Transformers 下载、缓存、依赖是否正常。

第五步：安装 Accelerate 并完成配置（单机多卡的基础）

Accelerate 的价值在于：

统一分布式启动方式
简化 FP16/BF16、DDP、多 GPU 等配置
与 Transformers Trainer / 自定义训练脚本协同很好

安装

pip install -U accelerate

交互式配置

accelerate config

常见选择建议（以单机多卡为例）：

Compute environment：This machine
Machine type：No distributed training（单卡）或 Multi-GPU（多卡）
Mixed precision：
- A100/H100 等支持 BF16 的卡：优先 bf16
- 其他卡：可用 fp16
DeepSpeed：先选 no，跑通基础流程后再引入（本文后面会讲）

配置完成后会生成类似：

~/.cache/huggingface/accelerate/default_config.yaml

快速验收：Accelerate 启动一个小脚本

创建 acc_test.py：

import torch
from accelerate import Accelerator

acc = Accelerator()
device = acc.device
x = torch.randn(2, 3).to(device)
y = x @ torch.randn(3, 4).to(device)
acc.print('device:', device)
acc.print('y shape:', y.shape)

运行：

accelerate launch acc_test.py

如果你配置了多卡，会自动按配置做进程启动。

第六步：安装 DeepSpeed（进阶：更大的模型、更省显存）

DeepSpeed 主要用来解决：

模型/优化器状态过大导致显存不够
ZeRO 分片（Stage 1/2/3）
CPU/NVMe Offload
更强的吞吐优化

先确认你的需求

只是做中小模型（几百 M 到数 B）微调：Accelerate + FSDP 或 DDP 可能已足够。
要冲更大模型、显存压力明显：引入 DeepSpeed 会更合适。

安装方式 1：直接 pip（最简单，但可能踩编译坑）

pip install -U deepspeed

如果安装过程中编译失败，常见原因：

缺少 gcc/g++、python-dev、ninja
CUDA 工具链或头文件缺失（某些算子需要）
PyTorch 与 CUDA 组合不匹配

Ubuntu 上可先准备基础编译工具：

sudo apt-get update
sudo apt-get install -y build-essential

安装方式 2：带环境变量控制（更稳一点）

如果你不需要编译某些可选算子，可尝试关闭以减少失败概率（不同版本变量略有差异，以下为常见思路）：

DS_BUILD_OPS=0 pip install -U deepspeed

说明：不同 DeepSpeed 版本对 DS_BUILD_* 的支持不同。若失败，请查看安装日志或官方文档对应版本的构建参数。

验收：DeepSpeed 是否可用

python - << 'PY'
import deepspeed
print('deepspeed:', deepspeed.__version__)
PY

以及检查 deepspeed 命令：

deepSpeed --help 2>/dev/null || deepspeed --help

第七步：把 Accelerate + DeepSpeed 串起来（推荐的落地方式）

在实际训练中，更常见的是：

用 Transformers 写训练脚本
用 Accelerate 负责启动与设备管理
用 DeepSpeed 负责 ZeRO/Offload 等

1）准备一个最小 DeepSpeed 配置文件

创建 ds_config_zero2.json（示例：ZeRO Stage 2，偏通用）

{
  "train_micro_batch_size_per_gpu": 1,
  "gradient_accumulation_steps": 8,
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "allgather_bucket_size": 200000000,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 200000000,
    "contiguous_gradients": true
  },
  "bf16": {
    "enabled": true
  },
  "fp16": {
    "enabled": false
  },
  "steps_per_print": 50,
  "wall_clock_breakdown": false
}

说明与建议：

train_micro_batch_size_per_gpu：每张卡一次前向的 batch size，显存紧张时设小。
gradient_accumulation_steps：梯度累积，配合 micro batch 形成等效大 batch。
BF16 优先（若硬件支持），否则用 FP16。

2）用 accelerate config 启用 DeepSpeed

执行：

accelerate config

当询问 DeepSpeed 时选择 yes，并填入 ds_config_zero2.json 的路径。

3）启动训练（示例命令形态）

当你有训练脚本 train.py 时：

accelerate launch train.py \
  --model_name_or_path <your_model> \
  --train_file <your_data> \
  --output_dir outputs \
  --per_device_train_batch_size 1 \
  --gradient_accumulation_steps 8 \
  --learning_rate 2e-5 \
  --num_train_epochs 1

注意：

如果你在 DeepSpeed 配置里指定了 batch/accumulation，尽量与脚本参数保持一致，避免混淆。
Transformers Trainer 与自定义训练循环对参数读取方式不同，以你实际框架为准。

常见坑与排查清单（非常实用）

1）`torch.cuda.is_available() == False`

排查顺序：

nvidia-smi 是否正常
是否装成了 CPU 版 torch（重新用 cu121 index-url 安装）
是否在容器里未挂载 GPU（Docker 需要 --gpus all，并安装 nvidia-container-toolkit）

2）Transformers 下载很慢或失败

建议：

设置 Hugging Face 镜像或使用公司内部缓存
预先下载到共享缓存目录

常用环境变量（按需）：

export HF_HOME=~/.cache/huggingface
export TRANSFORMERS_CACHE=~/.cache/huggingface/transformers

3）DeepSpeed 安装编译失败

优先策略：

先 pip install deepspeed 失败的话，看日志是否缺少 build-essential/gcc
尝试 DS_BUILD_OPS=0 降低编译复杂度
确认 PyTorch + CUDA runtime 组合是常见组合（如 cu121）

4）多卡启动后卡住或 NCCL 报错

常见原因：

多网卡/IB 配置不正确
端口冲突
容器网络隔离

建议先做单机多卡的 NCCL 基础排查：

export NCCL_DEBUG=INFO
export NCCL_P2P_DISABLE=0
export NCCL_IB_DISABLE=1

如果禁用 IB 后正常，说明问题多半在 IB/RDMA 网络路径上。

5）显存爆炸：OOM

处理顺序建议（从最容易到更强力）：

降 per_device_train_batch_size
增大 gradient_accumulation_steps
开启 bf16/fp16
开启 gradient checkpointing（模型侧）
引入 DeepSpeed ZeRO（Stage 2/3）或 Offload

建议的“最终验收”标准（你搭完环境后要做到这些）

你可以用以下 checklist 判断环境是否合格：

nvidia-smi 正常，能看到 GPU
Python 环境隔离完成（conda/venv），pip list 不污染系统
PyTorch 能识别 GPU：torch.cuda.is_available() == True
Transformers 能正常加载模型并前向
Accelerate 能 launch 脚本，单卡/多卡均可启动
DeepSpeed 能导入、命令可用；配置文件可被 Accelerate 引用

做到以上 6 点，你就拥有了一个可用于后续微调（SFT）、对齐（DPO/ORPO）、以及更大规模训练的可靠底座。

下一步你可以做什么

在 Ai大模型训练教程 的后续实战里，这套环境会直接用于：

用 Transformers + Accelerate 写一个可复现的 SFT 训练脚本
引入 DeepSpeed ZeRO2/ZeRO3 做显存优化
加入数据清洗、打包、评估与日志体系，形成可落地的训练工程结构

从零搭建大模型训练环境：CUDA、PyTorch、Transformers、Accelerate与DeepSpeed

https://aissn.com/90.html

我要提问

写在前面：你要搭建的是“能稳定复现”的训练环境

环境规划：先定标准，再装软件

你需要明确的 5 个关键点

推荐版本组合（2026 年初仍普遍稳定的思路）

第一步：检查 GPU 与驱动是否正常

第二步：创建隔离的 Python 环境（强烈建议）

方案 A：使用 conda/mamba（更省心）

方案 B：使用 venv（更轻量）

第三步：安装 PyTorch（优先使用官方带 CUDA runtime 的 wheel）

为什么推荐“PyTorch 自带 CUDA”

安装命令（示例：CUDA 12.1 runtime）

第四步：安装 Transformers 与常用训练依赖

基础安装

训练与日志常用组件（建议装）

快速验收：加载一个小模型做前向

第五步：安装 Accelerate 并完成配置（单机多卡的基础）

安装

交互式配置

快速验收：Accelerate 启动一个小脚本

第六步：安装 DeepSpeed（进阶：更大的模型、更省显存）

先确认你的需求

安装方式 1：直接 pip（最简单，但可能踩编译坑）

安装方式 2：带环境变量控制（更稳一点）

验收：DeepSpeed 是否可用

第七步：把 Accelerate + DeepSpeed 串起来（推荐的落地方式）

1）准备一个最小 DeepSpeed 配置文件

2）用 accelerate config 启用 DeepSpeed

3）启动训练（示例命令形态）

常见坑与排查清单（非常实用）

1）`torch.cuda.is_available() == False`

2）Transformers 下载很慢或失败

3）DeepSpeed 安装编译失败

4）多卡启动后卡住或 NCCL 报错

5）显存爆炸：OOM

建议的“最终验收”标准（你搭完环境后要做到这些）

下一步你可以做什么

Prev：训练一套大模型要多少显存与多少钱：按参数量与序列长度估算成本

写在前面：你要搭建的是“能稳定复现”的训练环境

环境规划：先定标准，再装软件

你需要明确的 5 个关键点

推荐版本组合（2026 年初仍普遍稳定的思路）

第一步：检查 GPU 与驱动是否正常

第二步：创建隔离的 Python 环境（强烈建议）

方案 A：使用 conda/mamba（更省心）

方案 B：使用 venv（更轻量）

第三步：安装 PyTorch（优先使用官方带 CUDA runtime 的 wheel）

为什么推荐“PyTorch 自带 CUDA”

安装命令（示例：CUDA 12.1 runtime）

第四步：安装 Transformers 与常用训练依赖

基础安装

训练与日志常用组件（建议装）

快速验收：加载一个小模型做前向

第五步：安装 Accelerate 并完成配置（单机多卡的基础）

安装

交互式配置

快速验收：Accelerate 启动一个小脚本

第六步：安装 DeepSpeed（进阶：更大的模型、更省显存）

先确认你的需求

安装方式 1：直接 pip（最简单，但可能踩编译坑）

安装方式 2：带环境变量控制（更稳一点）

验收：DeepSpeed 是否可用

第七步：把 Accelerate + DeepSpeed 串起来（推荐的落地方式）

1）准备一个最小 DeepSpeed 配置文件

2）用 accelerate config 启用 DeepSpeed

3）启动训练（示例命令形态）

常见坑与排查清单（非常实用）

1）torch.cuda.is_available() == False

2）Transformers 下载很慢或失败

3）DeepSpeed 安装编译失败

4）多卡启动后卡住或 NCCL 报错

5）显存爆炸：OOM

建议的“最终验收”标准（你搭完环境后要做到这些）

下一步你可以做什么

Prev：训练一套大模型要多少显存与多少钱：按参数量与序列长度估算成本

1）`torch.cuda.is_available() == False`