在第 7 章完成训练数据的采集、清洗、Token 化和质量分级之后，预训练的核心挑战就从“数据够不够干净”转向“机器能不能稳定地吃下这些数据”。很多团队在这一步犯的致命错误，是拿单机跑 Demo 的思维来应付预训练集群——结果往往是训练启动三天，两天花在排查“某张卡掉驱动”“某节点 NCCL 超时”“checkpoint 写到一半存储掉线”上。
本节讲的是工业化预训练的“开工仪式”：选一套能 Scale 的框架，配一个能复现的环境，把一架子服务器初始化成“一台逻辑上的大机器”。这是 8.2 节模型初始化和 8.3 节训练执行的前提，也是第 9 章分布式训练技术的工程底座。
一、框架选型：没有银弹，只有场景匹配
预训练框架不是选一个“最好”的，而是选一套“在你的集群规模和模型尺寸下能跑得起来且调得动”的组合。当前工业界的基座几乎统一在 PyTorch 之上，但在分布式策略层存在明确的分工。
1. 基座：PyTorch
今天做大模型预训练，TensorFlow/JAX 仍有特定团队使用（如 Google 的 TPU 生态偏好 JAX），但PyTorch 是绝对主流。原因很现实：
- 开源模型权重（Llama、Qwen、Mistral 等）的 release 代码几乎全是 PyTorch；
- HuggingFace
transformers生态的接口已成为事实标准； - 动态图调试友好，分布式策略封装成熟。
选型建议：除非你的团队有历史包袱或 TPU 独占资源，否则直接锁定 PyTorch。
2. 分布式训练框架：必选项
原生 PyTorch 的 DistributedDataParallel（DDP）只解决数据并行，对大模型预训练远远不够。你需要专门的大模型分布式框架：
| 框架 | 核心优势 | 适用场景 | 学习曲线 | 备注 |
|------|---------|---------|----------|------|
| DeepSpeed（微软） | ZeRO 显存优化极致；Fused Adam、Offload 成熟 | 7B–70B 模型，百卡级集群，中小团队首选 | 平缓 | 与 HuggingFace 集成极好，社区踩坑文档丰富 |
| Megatron-LM（NVIDIA） | Tensor Parallel + Pipeline Parallel 性能极致；千卡级线性扩展优秀 | 65B+ 超大模型，千卡集群，追求 MFU 极致 | 陡峭 | 代码侵入性强，需要改模型定义，NVIDIA 官方优化 |
| FSDP（PyTorch 官方） | 原生集成，无需额外安装 | 7B–13B 小模型快速实验 | 低 | 显存优化力度弱于 DeepSpeed，大模型易 OOM |
| Colossal-AI | 统一封装 DP/TP/PP/Sequence Parallel | 学术实验、特定优化场景 | 中等 | 工业级超大集群验证案例少于前两者 |
真实决策逻辑：
- 如果你训练 7B–13B 模型，集群规模 ≤ 64 张卡：DeepSpeed（ZeRO-2/3）是性价比最高的选择。一行配置就能启动，且能直接复用 HuggingFace 的模型定义。
- 如果你训练 70B+ 模型，或拥有千卡 A100/H100 集群：必须上 Megatron-LM，或 Megatron-DeepSpeed 混合方案（Megatron 负责 TP/PP，DeepSpeed 负责 ZeRO 和优化器状态分片）。此时模型本身需要按 Megatron 的接口改写，不再是“拿来就用”。
- 不要同时学两套框架。预训练周期紧张，选定一套后深度定制，比左右横跳更省时间。
3. 周边加速库（必装）
- FlashAttention / FlashAttention-2：显存节省 + 计算加速的标配，安装前需确认 CUDA 版本匹配（见下文）。
- Apex（NVIDIA）：提供 FusedAdam、FusedLayerNorm、混合精度训练支持。Megatron-LM 依赖它。
- TransformerEngine（Hopper/Ada 架构）：若使用 H100，开启 FP8 训练可显著提速，但需 TransformerEngine ≥ 0.10 且 PyTorch ≥ 2.1。
二、环境配置：版本地狱与容器化救赎
大模型预训练对环境一致性有近乎苛刻的要求。一台节点上 CUDA 驱动 535、另一台 525，就可能导致 NCCL 握手失败或 FlashAttention 编译报错。
1. 版本对齐：一条铁律
驱动 ≥ CUDA Toolkit ≥ cuDNN ≥ NCCL ≥ PyTorch，必须形成兼容链。以 NVIDIA A100/H100 为例，2024 年主流稳定组合：
| 组件 | 推荐版本（示例） | 检查命令 |
|------|------------------|----------|
| NVIDIA Driver | 535.104 或更高 | nvidia-smi |
| CUDA Toolkit | 12.1 或 12.2 | nvcc -V |
| cuDNN | 8.9.x | cat /usr/local/cuda/include/cudnn_version.h | grep CUDNN_MAJOR |
| NCCL | 2.18.3+ | dpkg -l | grep nccl |
| PyTorch | 2.1.2 / 2.2.0（对应 CUDA 12.1） | python -c "import torch; print(torch.version.cuda)" |
踩坑提醒：
nvidia-smi显示的 CUDA Version 是驱动兼容上限，不等于实际安装的 CUDA Toolkit 版本。很多人被这个误导。- NCCL 必须使用 NVIDIA 官方编译的版本，且与 CUDA 大版本严格匹配。自行编译 NCCL 是预训练延迟的“经典陷阱”。
- PyTorch 的 CUDA 版本必须与系统 CUDA 一致。若系统装 12.2，却装了
torch+cu118，会出现各种诡异的 CUDA kernel 错误。
2. Python 环境：放弃 Conda，拥抱容器
在 8 卡、16 卡的实验机上，Conda 或 venv 尚可应付。但在 64 节点 × 8 卡的集群上，任何非容器化的环境管理都是灾难。
推荐方案：
- Docker + NVIDIA Container Toolkit：将驱动之外的整个软件栈（CUDA、cuDNN、PyTorch、训练框架、依赖包）打成一个镜像。
- 每个节点只需装好 NVIDIA Driver 和 Container Runtime，其余全部从镜像加载。
- 镜像仓库建议用私有 Harbor 或云厂商镜像服务，避免 DockerHub 拉取超时。
基础镜像 Dockerfile 片段示意：
FROM nvcr.io/nvidia/pytorch:24.02-py3
RUN pip install flash-attn --no-build-isolation
RUN pip install deepspeed>=0.12.0 transformers>=4.36.0
# 将 Megatron-LM 或自定义代码 COPY 进去关键原则：所有节点跑同一个镜像 ID。如果发现某节点行为异常，先 docker images | grep <tag> 确认镜像哈希是否一致。
3. 编译依赖与系统调优
- GCC 版本：编译 FlashAttention、Apex 时需要 GCC ≥ 9.4。若系统默认 GCC 4.8，必须手动切换。
- NUMA 与 CPU 绑核：数据加载和预处理会吃满 CPU，建议安装
numactl，并在启动脚本中绑定 NUMA 节点，避免跨 NUMA 访存拖慢数据流水线。 - 文件句柄数与栈大小：分布式训练会打开大量 socket 连接，需修改
/etc/security/limits.conf，将nofile提高到 65535 以上。
三、集群初始化：让 64 台机器表现得像 1 台
预训练集群不是“64 台各自为政的服务器”，而是一台逻辑上的超级计算机。初始化阶段要解决四个问题：网络通、存储通、时间同步、环境一致。
1. 网络层：计算网络与存储网络分离（与 19.2 节呼应）
- 计算网络：InfiniBand（IB）或 RoCE v2。预训练启动前，必须用
ibstat或ib_write_bw验证每张网卡是否处于 Active 状态，且速率符合预期（NDR 400Gbps、HDR 200Gbps 等）。 - 多网卡绑定：若单节点有多张 IB 网卡，需在 NCCL 环境变量中显式指定
NCCL_IB_HCA=mlx5_0,mlx5_1，否则 NCCL 可能只走单卡，带宽直接腰斩。 - NCCL 测试：正式训练前，全集群运行
all_reduce_perf（NCCL Tests 套件）。所有节点两两之间必须在预期带宽的 80% 以上，且没有明显掉速的节点。这是排查“慢节点”的黄金标准。
2. 存储挂载：数据在哪里，瓶颈就在哪里
第 7 章清洗好的数据通常存放在并行文件系统（Lustre/GPFS/JuiceFS，详见 19.3 节）或对象存储上。预训练前必须确认：
- 所有节点挂载同一路径：例如
/data/pretrain必须在每个节点上指向同一个命名空间。莫明其妙的“某节点没挂载”会导致数据加载进程 hang 住。 - 本地 NVMe SSD 缓存：若共享存储带宽不足，建议先将数据
rsync或cp到各节点的本地 SSD，再从本地读取。这对小文件或随机读取场景（如 WebDataset）尤为关键。 - Checkpoint 目录的容量与性能：以 70B 模型、DeepSpeed ZeRO-3 为例，单份 checkpoint 包含模型参数 + 优化器状态 + 调度器状态，FP16 下约 1.5–2 TB。如果 checkpoint 路径放在慢速 NAS 上，每保存一次 checkpoint 训练就要停摆 10 分钟以上。务必预留高速存储空间给 checkpoint。
3. 节点互通：SSH 免密、Hosts 与 NTP
- SSH 免密：分布式训练启动器（如
pdsh、mpirun、torchrun）依赖节点间无密码 SSH。需配置好~/.ssh/authorized_keys，并确保 严格模式权限正确（chmod 600 ~/.ssh/authorized_keys，权限不对会导致静默失败）。 /etc/hosts同步**：所有节点需能互相解析主机名。建议使用内部 DNS 或统一的 hosts 文件，避免“hostname 解析到 127.0.1.1”这类陷阱。- NTP 时间同步：训练日志、监控指标、checkpoint 时间戳必须在全集群对齐。安装
chrony或ntpd，以统一的 NTP Server 为源。时间不同步会导致 TensorBoard 曲线乱跳、日志顺序错乱，甚至某些分布式锁机制异常。
4. 环境一致性：Ansible 批量运维
手动登录每台节点装环境是不可接受的。推荐使用 Ansible 或 PSSH 做批量操作：
- 批量挂载存储、批量拉取 Docker 镜像；
- 批量检查
nvidia-smi、nvcc -V、ibstat； - 批量修改系统参数（
ulimit、sysctl网络缓冲区大小）。
一个实用的初始化 Checklist（建议在正式训练前跑一遍）：
| 检查项 | 通过标准 | 常用命令 |
|--------|---------|----------|
| GPU 驱动 | 所有卡可见，无 ECC Error | nvidia-smi -q | grep ECC |
| CUDA 可用 | PyTorch 能调起 CUDA | python -c "import torch; print(torch.cuda.is_available())" |
| NCCL 全互联带宽 | ≥ 理论值的 80% | all_reduce_perf -b 8M -e 1G -f 2 -g 8 |
| 共享存储挂载 | 所有节点同一路径可读写 | touch /data/pretrain/.write_test_$(hostname) |
| 时间同步 | 节点间时差 < 1 秒 | date; ssh node2 date |
| 容器环境 | 所有节点镜像 ID 一致 | docker images | grep training-image |
四、预训练特有的环境容量规划
在环境准备阶段，有几个容易算错账的资源需求：
- 显存：不是只看模型参数。以 2.1 节的估算为基础，训练时还要加上激活值（Activations）、梯度和优化器状态。DeepSpeed ZeRO-3 会将优化器状态切分到所有卡上，因此单卡显存压力降低，但集群通信量增加。需在启动前用
deepspeed --num_gpus 8 --num_nodes 1 pretrain.py做小规模显存压测。 - 内存：若开启 CPU Offload（ZeRO-Infinity），CPU 内存需要容纳部分优化器状态，建议单节点配备 1–2 TB DRAM。
- 临时存储：日志、tensorboard 事件文件、频繁的中途 checkpoint 会迅速占满磁盘。建议单独挂载一块大容量 NVMe 分区给
/tmp或训练工作目录。
五、小结
预训练环境准备的核心，是把“算法论文里的理想条件”翻译成“机房里能 7×24 小时稳定燃烧的工程条件”。
关键 takeaway：
- 框架：7B–70B 优选 DeepSpeed，70B+ 或千卡级优选 Megatron-LM，PyTorch 是通用基座；
- 环境：用容器固化 CUDA/cuDNN/NCCL/PyTorch 版本链，拒绝节点间差异；
- 集群：计算网络（IB/RoCE）、共享存储、SSH 免密、NTP 同步、NCCL 带宽测试，缺一不可；
- 校验：正式烧卡前，先用小规模 all_reduce_perf 和单节点显存压测暴露问题。
环境就绪后，下一节（8.2）将进入模型初始化——如何设置随机种子、初始化权重分布、配置模型结构，让这台上百亿美元的“超级乐器”发出第一个准确的音。