3. 实验环境搭建：PyTorch版本、CUDA版本、硬件配置

好，咱们直接进入正题。做实验嘛，环境搭不对，后面全白费。我见过太多同学在环境上翻车，折腾一整天结果发现是CUDA版本不匹配。嗯，咱们今天就把这事一次性说清楚。

3.1 PyTorch版本选择

我个人习惯用PyTorch 2.0以上版本。为什么？因为torch.compile这个功能太香了。你想想看，同样的SwiGLU代码，加上一行编译指令，训练速度能提升20%-30%。

我建议你用PyTorch 2.1.0或2.2.0。这两个版本我都在项目中实测过，稳定性不错。别追最新版，有时候新版本会有一些奇怪的bug。我曾经在PyTorch 2.3.0刚发布时直接升级，结果发现某个算子有内存泄漏问题，折腾了两天才回退。

安装命令很简单：

# PyTorch 2.1.0 + CUDA 11.8
pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu118

# PyTorch 2.2.0 + CUDA 12.1
pip install torch==2.2.0 torchvision==0.17.0 torchaudio==2.2.0 --index-url https://download.pytorch.org/whl/cu121

3.2 CUDA版本与驱动

这里有个坑，我得重点说一下。CUDA版本不是越高越好，关键看你的显卡驱动支持到什么程度。

怎么查？打开终端，输入：

nvidia-smi

看右上角的CUDA Version，那个数字就是你的驱动支持的最高CUDA版本。比如显示CUDA Version: 12.1，那你装CUDA 12.1或更低的版本都行。

我个人建议用CUDA 11.8。为什么？兼容性好。大部分开源项目、预训练模型都是用这个版本测试的。你想想看，如果遇到一个只支持CUDA 11.8的库，你装的是12.1，那就得重新配环境，多麻烦。

3.3 硬件配置要求

做SwiGLU实验，GPU是核心。我直接说结论：

GPU型号	显存	能否跑SwiGLU实验	备注
RTX 3090	24GB	✅ 完全够用	性价比之王，我主力卡
RTX 4090	24GB	✅ 非常流畅	速度快，但贵
A100	40GB/80GB	✅ 绰绰有余	云端常用，按需租用
RTX 3060	12GB	⚠️ 勉强能跑	batch size要调小
RTX 2080 Ti	11GB	⚠️ 有点吃力	建议用混合精度

如果你用的是RTX 3090或4090，那基本不用担心。24GB显存跑SwiGLU对比实验绰绰有余。我自己的主力机就是RTX 3090，batch size设到64，模型参数量在1B以内，完全没问题。

显存不够怎么办？两个办法：

1. 减小batch size——从64降到32甚至16，效果差别不大
2. 用混合精度训练——torch.cuda.amp，显存占用直接减半

3.4 环境验证脚本

环境搭好之后，别急着跑实验。先跑个验证脚本，确保一切正常。这是我每次搭环境必做的步骤：

import torch
import torch.cuda as cuda

print(f"PyTorch版本: {torch.__version__}")
print(f"CUDA是否可用: {cuda.is_available()}")
print(f"CUDA版本: {cuda.get_device_capability()}")
print(f"GPU型号: {cuda.get_device_name(0)}")
print(f"显存总量: {cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB")

# 简单测试SwiGLU
x = torch.randn(4, 1024, 768).cuda()
w1 = torch.randn(768, 2048).cuda()
w2 = torch.randn(768, 2048).cuda()

def swiglu(x, w1, w2):
    gate = torch.matmul(x, w1)
    up = torch.matmul(x, w2)
    return gate * torch.sigmoid(gate) * up

y = swiglu(x, w1, w2)
print(f"SwiGLU前向传播成功，输出形状: {y.shape}")

如果这个脚本能顺利跑完，说明环境没问题。如果报错，多半是CUDA版本或PyTorch版本不匹配。别慌，按我前面说的检查一下。

3.5 知识体系总览

下面这张图把咱们这章的核心逻辑串起来了。你一看就明白：

说白了，环境搭建就三步：选对PyTorch版本、配好CUDA、确认硬件够用。然后跑个验证脚本，没问题就开干。别在这上面花太多时间，咱们的重点是SwiGLU对收敛速度的影响，不是折腾环境。