第二课：环境准备——把工具箱收拾利索

说实话，做RLHF训练，最让人头疼的往往不是算法本身，而是环境配置。我见过太多同学在环境上卡了一整天，最后发现是CUDA版本不对。这节课，咱们就把这些基础工作一次性搞定。

Python虚拟环境：给你的项目一个“隔离区”

为什么要用虚拟环境？说白了，就是防止不同项目之间的依赖打架。我有个血的教训——曾经同时做两个项目，一个要PyTorch 1.12，另一个要2.0，结果pip install的时候互相覆盖，最后两个都跑不起来。

我个人习惯用 conda 来管理环境，当然 venv 也行。这里我推荐 conda，因为它在处理CUDA相关依赖时更省心。

# 创建环境，指定Python版本
conda create -n rlhf_env python=3.10

# 激活环境
conda activate rlhf_env

# 如果你用venv，也可以这样
python -m venv rlhf_env
source rlhf_env/bin/activate  # Linux/Mac
# rlhf_env\Scripts\activate  # Windows

PyTorch安装：选对版本是关键

PyTorch安装看似简单，但坑不少。你想想看，如果装了个CPU版本的PyTorch，后面跑RLHF训练，那速度简直让人崩溃。

先确认你的CUDA版本：

nvidia-smi  # 查看CUDA版本，比如12.1

然后去PyTorch官网选对应的命令。我个人习惯用pip安装，因为conda有时候解析依赖太慢。

# CUDA 12.1 版本
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

# 或者用conda（慢但稳定）
conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia

装完后验证一下：

import torch
print(torch.__version__)        # 应该显示2.x
print(torch.cuda.is_available()) # 应该返回True
print(torch.cuda.device_count()) # GPU数量

Hugging Face Transformers与Datasets

这两个库是RLHF的基石。Transformers提供模型，Datasets提供数据。安装很简单：

pip install transformers datasets

但要注意版本兼容性。我记得有一次，transformers 4.30和datasets 2.14有个小bug，加载某些数据集时会报错。所以建议指定版本：

pip install transformers==4.44.0 datasets==2.21.0

为什么要指定版本？因为RLHF训练中，我们经常要用到 AutoModelForCausalLM 和 Trainer 类，新版本可能会改API。我一般会锁定一个经过验证的版本组合。

DeepSpeed与TRL：RLHF的“加速器”和“脚手架”

DeepSpeed是微软出的分布式训练框架，能让你在有限的GPU上跑更大的模型。TRL是Hugging Face出的强化学习库，专门为RLHF设计。

# 安装DeepSpeed
pip install deepspeed

# 安装TRL
pip install trl

这里有个坑：DeepSpeed对CUDA和PyTorch版本很敏感。我遇到过最离谱的一次，DeepSpeed 0.12和PyTorch 2.1不兼容，训练时直接报 RuntimeError: CUDA error: invalid device function。解决方案是降级DeepSpeed或升级PyTorch。

硬件需求概览：你的GPU够用吗？

RLHF训练对GPU显存的要求，说白了就是“有多少吃多少”。我们来估算一下：

为什么会差这么多？因为RLHF需要同时加载四个模型：Actor（策略模型）、Reference（参考模型）、Reward（奖励模型）、Critic（价值模型）。每个模型都占一份显存。再加上优化器状态、梯度等，显存消耗是普通微调的3-4倍。

知识体系总览

下面这张图，把咱们这节课的核心逻辑串起来了：

验证环境：跑一个最小示例

所有东西装完后，我建议你跑一个最小示例来验证。别急着上大模型，先确认每个组件都能正常工作：

# test_env.py
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
import deepspeed
from trl import PPOTrainer

print("PyTorch version:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())

# 加载一个小模型测试
model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")
print("Model loaded successfully")

# 测试数据集
dataset = load_dataset("imdb", split="train[:100]")
print(f"Dataset loaded: {len(dataset)} samples")

# 测试DeepSpeed
deepspeed.ops.op_builder.CPUAdamBuilder().load()
print("DeepSpeed ops loaded")

print("All components ready!")

如果这个脚本能顺利跑完，恭喜你，环境准备就绪。如果报错，别慌——99%的问题都是版本不匹配。去查一下官方文档的兼容性表格，或者直接搜错误信息，一般都能解决。