noao-chat-7-nonochat 多卡训练指南

12 分钟阅读时长

发布时间： December 25, 2025

LLM 多卡训练完全指南

nonochat

本文档全面解析 nanochat 项目中的分布式训练（DDP），从环境准备到代码实现，让小白也能轻松理解和使用多卡训练。

1. 什么是多卡训练

1.1 基本概念

多卡训练 = 数据并行 (Data Parallel)

简单来说：用多张 GPU 同时训练一个模型。

单卡训练:
GPU 0: [模型副本] → 处理 batch 1, 2, 3, 4, 5, 6, 7, 8
       时间: ████████

多卡训练 (8卡):
GPU 0: [模型副本] → 处理 batch 1
GPU 1: [模型副本] → 处理 batch 2
GPU 2: [模型副本] → 处理 batch 3
...
GPU 7: [模型副本] → 处理 batch 8
       时间: █
       速度提升: 8x (理论)

1.2 关键术语

术语	英文	解释	示例
进程组	Process Group	所有参与训练的进程	8 卡 = 8 个进程
世界大小	World Size	总共有多少个进程	8 张卡 → world_size=8
全局排名	Rank	当前进程的全局编号	0, 1, 2, …, 7
本地排名	Local Rank	当前节点内的进程编号	单节点 = Rank
主进程	Master Process	Rank 0，负责日志和保存	ddp_rank == 0
DDP	DistributedDataParallel	PyTorch 分布式训练框架	torch.nn.parallel.DDP

1.3 工作流程

初始化阶段:
1. 每个 GPU 加载相同的模型副本
2. 建立进程间通信（NCCL）
3. 同步所有进程

训练阶段:
每个步骤:
  1. 每个 GPU 处理不同的 mini-batch
  2. 各自计算梯度
  3. 通信：所有 GPU 的梯度取平均
  4. 所有 GPU 用相同的梯度更新模型
  5. 模型保持同步

结束阶段:
主进程保存模型检查点

2. 为什么需要多卡训练

2.1 训练加速

实际效果（nanochat 项目）：

配置	训练时间（Base, d=12）	加速比
单卡 A100	~16 小时	1x
8卡 A100	~2 小时	~8x

为什么不是完美的 8x？

通信开销（梯度同步）
数据加载开销
通常能达到 85-95% 的线性扩展

2.2 批次大小限制

问题： 单卡显存不够大

# 单卡场景
device_batch_size = 32  # 受限于单卡显存
total_batch_size = 32 * 2048 = 65,536 tokens

# 8卡场景
device_batch_size = 32  # 每卡相同
total_batch_size = 32 * 2048 * 8 = 524,288 tokens
# 更大的批次 → 更稳定的梯度 → 更好的收敛

2.3 实验效率

对比：

# 单卡训练 - 需要 16 小时
python -m scripts.base_train --depth=12

# 8卡训练 - 只需 2 小时
torchrun --nproc_per_node=8 -m scripts.base_train --depth=12

# 节省时间 = 可以做更多实验！

3. 环境准备

3.1 硬件要求

最低配置：

2+ 张 NVIDIA GPU
GPU 之间有 NVLink 或 PCIe 连接
推荐：相同型号的 GPU

本项目测试配置：

8x NVIDIA A100 80GB
NVLink 互联

3.2 软件环境

必需组件：

# 1. PyTorch (支持分布式)
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

# 2. 验证 NCCL 后端
python -c "import torch; print(torch.distributed.is_nccl_available())"
# 应输出: True

# 3. 验证多卡可见
python -c "import torch; print(torch.cuda.device_count())"
# 应输出: 8 (如果你有 8 张卡)

3.3 环境变量

自动设置（torchrun 会处理）：

RANK          # 全局进程编号: 0, 1, 2, ..., 7
LOCAL_RANK    # 本地进程编号: 0, 1, 2, ..., 7
WORLD_SIZE    # 总进程数: 8
MASTER_ADDR   # 主节点地址: localhost (单节点)
MASTER_PORT   # 主节点端口: 29500 (默认)

手动检查：

# 查看当前环境
echo $CUDA_VISIBLE_DEVICES  # 可见的 GPU: 0,1,2,3,4,5,6,7

4. 多卡启动方式

4.1 torchrun 命令详解

基本语法：

torchrun [torchrun参数] -m [模块名] [脚本参数]

常用参数：

torchrun \
    --standalone \              # 单节点模式
    --nproc_per_node=8 \       # 每个节点的进程数（GPU数）
    -m scripts.base_train \    # 要运行的 Python 模块
    --depth=12 \               # 脚本参数
    --device_batch_size=32

4.2 完整示例

训练 Base 模型

# 单卡（不使用 torchrun）
python -m scripts.base_train \
    --depth=12 \
    --device_batch_size=8 \
    --run=base_d12_single

# 8卡（使用 torchrun）
torchrun --standalone --nproc_per_node=8 \
    -m scripts.base_train \
    --depth=12 \
    --device_batch_size=32 \
    --run=base_d12_multi

训练 Mid 模型

torchrun --standalone --nproc_per_node=8 \
    -m scripts.mid_train \
    --source=base \
    --device_batch_size=16

训练 SFT 模型

torchrun --standalone --nproc_per_node=8 \
    -m scripts.chat_sft \
    --source=mid \
    --device_batch_size=8

训练 RL 模型

torchrun --standalone --nproc_per_node=8 \
    -m scripts.chat_rl \
    --source=sft \
    --device_batch_size=8 \
    --num_samples=16

4.3 参数对比

参数	单卡	多卡	说明
启动方式	`python -m`	`torchrun`	多卡需要 torchrun
device_batch_size	8	32	多卡可以更大
梯度累积	多	少	多卡减少累积步数
训练时间	长	短	多卡显著加速

5. 代码实现详解

5.1 分布式初始化

位置： nanochat/common.py

def compute_init(device_type="cuda"):
    """基础初始化，包括分布式设置"""
    
    # 1. 获取分布式信息
    ddp, ddp_rank, ddp_local_rank, ddp_world_size = get_dist_info()
    
    # 2. 如果是分布式 + CUDA
    if ddp and device_type == "cuda":
        # 为每个进程指定不同的 GPU
        device = torch.device("cuda", ddp_local_rank)
        torch.cuda.set_device(device)
        
        # 初始化进程组（使用 NCCL 后端）
        dist.init_process_group(
            backend="nccl",      # NVIDIA GPU 专用
            device_id=device
        )
        
        # 同步所有进程
        dist.barrier()
    else:
        device = torch.device(device_type)
    
    return ddp, ddp_rank, ddp_local_rank, ddp_world_size, device

关键点：

进程绑定 GPU：

# Rank 0 → GPU 0
# Rank 1 → GPU 1
# ...
# Rank 7 → GPU 7
device = torch.device("cuda", ddp_local_rank)

初始化通信：

dist.init_process_group(backend="nccl")
# 建立所有进程之间的通信通道

同步屏障：

dist.barrier()
# 等待所有进程都到达这里

5.2 判断是否启用 DDP

def is_ddp():
    """检查是否在分布式环境中"""
    return int(os.environ.get('RANK', -1)) != -1

def get_dist_info():
    """获取分布式信息"""
    if is_ddp():
        # torchrun 会设置这些环境变量
        ddp_rank = int(os.environ['RANK'])
        ddp_local_rank = int(os.environ['LOCAL_RANK'])
        ddp_world_size = int(os.environ['WORLD_SIZE'])
        return True, ddp_rank, ddp_local_rank, ddp_world_size
    else:
        # 单卡模式
        return False, 0, 0, 1

5.3 主进程判断

为什么需要主进程？

避免重复日志（8 个进程打印 8 遍）
避免重复保存（8 个进程保存 8 次）
避免重复评估（8 个进程评估 8 次）

# 方法 1: 直接判断
if ddp_rank == 0:
    print("只有主进程打印这条消息")
    save_checkpoint(model)

# 方法 2: 使用 print0 工具
print0("自动判断主进程的打印")

print0 实现：

def print0(s="", **kwargs):
    ddp_rank = int(os.environ.get('RANK', 0))
    if ddp_rank == 0:
        print(s, **kwargs)

5.4 模型包装

不需要显式包装！

nanochat 项目不使用 torch.nn.parallel.DistributedDataParallel 包装模型。

# 传统 DDP 做法（其他项目）:
if ddp:
    model = DDP(model, device_ids=[local_rank])

# nanochat 做法:
# 不包装！直接使用原始模型
# 梯度同步在优化器中处理

为什么？

使用自定义的分布式优化器（DistMuon, DistAdamW）
优化器内部处理梯度通信
更灵活、更高效

6. 数据并行机制

6.1 数据分配策略

核心思想： 每个进程处理不同的数据

# 循环分配（Round-Robin）
for example_idx in range(ddp_rank, len(dataset), ddp_world_size):
    process_example(example_idx)

# 示例：8 卡处理 24 个样本
# Rank 0: 处理 0, 8, 16
# Rank 1: 处理 1, 9, 17
# Rank 2: 处理 2, 10, 18
# ...
# Rank 7: 处理 7, 15, 23

6.2 Base 训练的数据加载

位置： nanochat/dataloader.py

def tokenizing_distributed_data_loader(B, T, split, device="cuda"):
    """分布式数据加载器"""
    
    ddp, ddp_rank, ddp_local_rank, ddp_world_size = get_dist_info()
    
    # 读取 Parquet 文件
    parquet_paths = list_parquet_files()
    
    # 每个 rank 从不同的 row group 开始
    rg_idx = ddp_rank  # 起始索引 = rank
    
    while rg_idx < pf.num_row_groups:
        # 读取数据
        rg = pf.read_row_group(rg_idx)
        batch = rg.column('text').to_pylist()
        
        # ... tokenize ...
        
        # 跳到下一个 row group（间隔 = world_size）
        rg_idx += ddp_world_size  # 0, 8, 16, ...

可视化：

Row Groups (Parquet 文件):
[0] [1] [2] [3] [4] [5] [6] [7] [8] [9] [10] ...

8 卡分配:
Rank 0: [0]             [8]              [16] ...
Rank 1:     [1]             [9]              ...
Rank 2:         [2]             [10]         ...
...
Rank 7:                     [7]              ...

结果: 
- 所有 rank 处理不同的数据
- 不重复、不遗漏
- 负载均衡

6.3 SFT/RL 的数据加载

对话数据的分配：

# SFT 训练
conversations = load_conversations()
for i in range(ddp_rank, len(conversations), ddp_world_size):
    conversation = conversations[i]
    # 处理这个对话...

# RL 训练
questions = load_questions()
for idx in range(ddp_rank, len(questions), ddp_world_size):
    question = questions[idx]
    # 采样 16 个答案...

7. 分布式优化器

7.1 为什么需要分布式优化器

问题： 每个 GPU 计算的梯度不同

GPU 0: grad_0 = [-0.1, 0.2, -0.3]
GPU 1: grad_1 = [-0.2, 0.1, -0.4]
...
GPU 7: grad_7 = [-0.15, 0.25, -0.35]

需要: 平均梯度
avg_grad = mean([grad_0, ..., grad_7])

解决方案： 分布式优化器自动同步梯度

7.2 DistAdamW 实现

位置： nanochat/adamw.py

class DistAdamW(torch.optim.Optimizer):
    """分布式 AdamW"""
    
    @torch.no_grad()
    def step(self):
        rank = dist.get_rank()
        world_size = dist.get_world_size()
        
        # 1. Reduce-Scatter: 梯度切片并平均
        for param in params:
            grad = param.grad
            rank_size = grad.shape[0] // world_size
            grad_slice = torch.empty_like(grad[:rank_size])
            
            # 分发梯度的一部分给每个 rank
            dist.reduce_scatter_tensor(
                grad_slice,      # 输出: 当前 rank 的切片
                grad,            # 输入: 完整梯度
                op=dist.ReduceOp.AVG  # 操作: 平均
            )
        
        # 2. 更新参数（每个 rank 只更新自己的切片）
        for param in params:
            p_slice = param[rank * rank_size:(rank + 1) * rank_size]
            # ... AdamW 更新逻辑 ...
            p_slice.add_(update, alpha=-lr)
        
        # 3. All-Gather: 收集所有 rank 的参数
        for param in params:
            dist.all_gather_into_tensor(
                param,      # 输出: 完整参数
                p_slice     # 输入: 当前 rank 的切片
            )

三个关键步骤：

Reduce-Scatter（梯度平均）

Before:
GPU 0: [grad_full]
GPU 1: [grad_full]
...
   
After:
GPU 0: [avg_grad_slice_0]
GPU 1: [avg_grad_slice_1]
...

Update（本地更新）

GPU 0: 更新 param_slice_0
GPU 1: 更新 param_slice_1
...

All-Gather（参数同步）

Before:
GPU 0: [param_slice_0]
GPU 1: [param_slice_1]
...
   
After:
GPU 0: [param_full]
GPU 1: [param_full]
...

7.3 DistMuon 实现

位置： nanochat/muon.py

class DistMuon(torch.optim.Optimizer):
    """分布式 Muon 优化器"""
    
    @torch.no_grad()
    def step(self):
        # 1. Reduce-Scatter: 梯度平均
        for group in self.param_groups:
            for param in group["params"]:
                # 将 world_size 个梯度收集并平均
                dist.reduce_scatter(
                    rs_output,
                    rs_input,
                    op=dist.ReduceOp.AVG
                )
        
        # 2. 正交化 + 更新（每个 rank 的切片）
        for param in params:
            g = averaged_grad
            # ... Muon 逻辑 ...
            g = zeropower_via_newtonschulz5(g)
            param.add_(g, alpha=-lr)
        
        # 3. All-Gather: 同步参数
        for param in params:
            dist.all_gather(ag_output, ag_input)

与 DistAdamW 的区别：

梯度平均方式相同
更新算法不同（正交化 vs 自适应学习率）

7.4 通信开销分析

通信量：

# 假设模型参数量: 100M
# 梯度大小: 100M * 4 bytes (fp32) = 400MB

# Reduce-Scatter: 400MB / 8 GPUs = 50MB per GPU
# All-Gather: 50MB * 8 = 400MB per GPU

# 总通信: 450MB per GPU per step

通信时间：

NVLink 带宽: ~300 GB/s
通信时间: 450MB / 300GB/s ≈ 1.5ms

前向+反向: ~50ms
通信占比: 1.5 / 50 = 3%

效率: 97% (非常高！)

8. 实战案例

8.1 完整训练流程（8 卡）

步骤 1: Base 训练

# 启动 Base 训练
torchrun --standalone --nproc_per_node=8 \
    -m scripts.base_train \
    --depth=12 \
    --device_batch_size=32 \
    --num_iterations=50000 \
    --run=base_d12_8gpu

# 预期输出（每个 GPU）:
# Step 0: loss=4.234 (8 个进程同时打印)
# Step 1: loss=4.123 (但只有 rank 0 保存日志)

步骤 2: 监控训练

# 查看 GPU 使用
watch -n 1 nvidia-smi

# 应该看到 8 张 GPU 都在运行
# GPU 0: Python (28GB / 80GB)
# GPU 1: Python (28GB / 80GB)
# ...
# GPU 7: Python (28GB / 80GB)

步骤 3: Mid 训练

torchrun --standalone --nproc_per_node=8 \
    -m scripts.mid_train \
    --source=base \
    --device_batch_size=16

步骤 4: SFT 训练

torchrun --standalone --nproc_per_node=8 \
    -m scripts.chat_sft \
    --source=mid \
    --device_batch_size=8

步骤 5: RL 训练

torchrun --standalone --nproc_per_node=8 \
    -m scripts.chat_rl \
    --source=sft \
    --device_batch_size=8 \
    --num_samples=16

8.2 单卡 vs 多卡对比

相同的总批次大小

# 目标: total_batch_size = 524,288 tokens

# ===== 单卡 =====
python -m scripts.base_train \
    --device_batch_size=8 \
    --max_seq_len=2048
# 计算: 8 * 2048 = 16,384 tokens/step
# 需要梯度累积: 524,288 / 16,384 = 32 steps
# 训练时间: ~16 小时

# ===== 8 卡 =====
torchrun --nproc_per_node=8 -m scripts.base_train \
    --device_batch_size=32 \
    --max_seq_len=2048
# 计算: 32 * 2048 * 8 = 524,288 tokens/step
# 需要梯度累积: 1 step (无累积！)
# 训练时间: ~2 小时

对比表格：

指标	单卡	8卡	提升
device_batch_size	8	32	4x
梯度累积步数	32	1	32x 更快
训练时间	16h	2h	8x
最终 Loss	3.45	3.45	相同

8.3 调试技巧

检查分布式状态

# 在训练脚本中添加
print(f"Rank {ddp_rank}/{ddp_world_size} on GPU {ddp_local_rank}")

# 应该看到:
# Rank 0/8 on GPU 0
# Rank 1/8 on GPU 1
# ...
# Rank 7/8 on GPU 7

验证数据不重复

# 在数据加载器中
print(f"Rank {ddp_rank} processing examples: {list(range(ddp_rank, 100, ddp_world_size))}")

# 输出:
# Rank 0: [0, 8, 16, 24, ...]
# Rank 1: [1, 9, 17, 25, ...]
# ...

验证梯度同步

# 在训练循环中
if step % 100 == 0 and ddp:
    # 检查所有 rank 的第一个参数是否相同
    param = list(model.parameters())[0]
    local_sum = param.sum().item()
    print(f"Rank {ddp_rank}: param sum = {local_sum}")
    
    # 所有 rank 应该打印相同的值!

9. 常见问题

9.1 环境问题

Q: `NCCL error: unhandled system error`

原因： 网络配置问题

解决：

# 方法 1: 设置网络接口
export NCCL_SOCKET_IFNAME=eth0  # 或 ens3, ib0 等

# 方法 2: 使用 IB 网络（如果有）
export NCCL_IB_DISABLE=0

# 方法 3: 禁用 IB（如果不需要）
export NCCL_IB_DISABLE=1

Q: `RuntimeError: Address already in use`

原因： 端口被占用

解决：

# 更换端口
torchrun --master_port=29501 --nproc_per_node=8 ...

# 或者杀死之前的进程
pkill -9 python

9.2 性能问题

Q: 多卡训练没有加速

可能原因 1：通信瓶颈

# 检查 GPU 连接方式
nvidia-smi topo -m

# 理想情况: NVLink
# GPU0 <-> GPU1: NV12 (12条 NVLink)

# 不好: PCIe
# GPU0 <-> GPU1: PHB (通过 PCIe Host Bridge)

可能原因 2：数据加载慢

# 增加数据加载线程
loader = DataLoader(
    dataset,
    num_workers=4,  # 增加到 4-8
    pin_memory=True
)

可能原因 3：批次太小

# 批次太小，通信开销占比高
device_batch_size = 1  # 太小 ❌
device_batch_size = 32  # 合适 ✅

Q: 显存占用不均衡

现象：

GPU 0: 60GB / 80GB
GPU 1: 30GB / 80GB

原因： 数据不均衡或主进程额外操作

解决：

# 确保数据均匀分配
assert len(dataset) % world_size == 0

# 主进程避免额外的显存操作
if ddp_rank == 0:
    # 在 CPU 上评估，不在 GPU 上
    model.cpu().eval()

9.3 训练问题

Q: Loss 在不同 rank 不一致

原因： 模型没有正确同步

检查：

# 1. 确保使用分布式优化器
optimizer = DistAdamW(...)  # ✅
optimizer = torch.optim.AdamW(...)  # ❌ 不会同步

# 2. 确保调用 barrier
if ddp:
    dist.barrier()

Q: 检查点保存/加载问题

最佳实践：

# 保存: 只有主进程保存
if ddp_rank == 0:
    torch.save(model.state_dict(), path)

# 加载: 所有进程都加载
model.load_state_dict(torch.load(path))

# 同步
if ddp:
    dist.barrier()

Q: 如何从单卡切换到多卡？

非常简单！

# 1. 单卡训练
python -m scripts.base_train --depth=12

# 2. 多卡训练（只需添加 torchrun）
torchrun --nproc_per_node=8 -m scripts.base_train --depth=12

# 代码完全不需要修改！
# nanochat 已经处理好了所有细节

10. 高级话题

10.1 多节点训练

如果你有多台机器（例如 2 台，每台 8 卡）：

# 节点 0（主节点）
torchrun \
    --nnodes=2 \
    --node_rank=0 \
    --master_addr=192.168.1.100 \
    --master_port=29500 \
    --nproc_per_node=8 \
    -m scripts.base_train

# 节点 1
torchrun \
    --nnodes=2 \
    --node_rank=1 \
    --master_addr=192.168.1.100 \
    --master_port=29500 \
    --nproc_per_node=8 \
    -m scripts.base_train

10.2 混合精度训练

nanochat 已经使用 bfloat16：

# 自动混合精度
autocast_ctx = torch.amp.autocast(
    device_type="cuda",
    dtype=torch.bfloat16  # 训练更稳定
)

with autocast_ctx:
    loss = model(inputs, targets)

10.3 梯度检查点

对于更大的模型（depth > 20）：

# 在模型中启用
model = torch.compile(model)
# PyTorch 2.0 自动优化内存使用

总结

多卡训练的核心要点

启动方式

torchrun --nproc_per_node=8 -m scripts.xxx

数据分配

for i in range(ddp_rank, len(data), ddp_world_size):
    process(data[i])

梯度同步

# 自动由 DistAdamW / DistMuon 处理
optimizer.step()

主进程检查

if ddp_rank == 0:
    save_checkpoint()

为什么 nanochat 的多卡训练这么简单？

自动检测：代码自动检测是否在分布式环境
透明处理：分布式优化器自动同步梯度
无需修改：单卡代码直接支持多卡
工具函数：print0(), get_dist_info() 等简化开发

下一步

实际运行多卡训练
监控 GPU 使用和训练速度
尝试不同的 batch size 配置
理解通信模式和优化策略

参考资源：

分享到

Bluesky Facebook LinkedIn Mastodon X (formerly Twitter)

Jinyang Zhang

LLM 多卡训练完全指南

目录

1. 什么是多卡训练

1.1 基本概念

1.2 关键术语

1.3 工作流程

2. 为什么需要多卡训练

2.1 训练加速

2.2 批次大小限制

2.3 实验效率

3. 环境准备

3.1 硬件要求

3.2 软件环境

3.3 环境变量

4. 多卡启动方式

4.1 torchrun 命令详解

4.2 完整示例

训练 Base 模型

训练 Mid 模型

训练 SFT 模型

训练 RL 模型

4.3 参数对比

5. 代码实现详解

5.1 分布式初始化

5.2 判断是否启用 DDP

5.3 主进程判断

5.4 模型包装

6. 数据并行机制

6.1 数据分配策略

6.2 Base 训练的数据加载

6.3 SFT/RL 的数据加载

7. 分布式优化器

7.1 为什么需要分布式优化器

7.2 DistAdamW 实现

7.3 DistMuon 实现

7.4 通信开销分析

8. 实战案例

8.1 完整训练流程（8 卡）

步骤 1: Base 训练

步骤 2: 监控训练

步骤 3: Mid 训练

步骤 4: SFT 训练

步骤 5: RL 训练

8.2 单卡 vs 多卡对比

相同的总批次大小

8.3 调试技巧

检查分布式状态

验证数据不重复

验证梯度同步

9. 常见问题

9.1 环境问题

Q: NCCL error: unhandled system error

Q: RuntimeError: Address already in use

9.2 性能问题

Q: 多卡训练没有加速

Q: 显存占用不均衡

9.3 训练问题

Q: Loss 在不同 rank 不一致

Q: 检查点保存/加载问题

Q: 如何从单卡切换到多卡？

10. 高级话题

10.1 多节点训练

10.2 混合精度训练

10.3 梯度检查点

总结

多卡训练的核心要点

为什么 nanochat 的多卡训练这么简单？

下一步

分享到

你可能感兴趣的

noao-chat-6-训练评估指南

LLM 模型评估验证完全指南

noao-chat-5-训练四阶段数据报告

nanochat 项目四阶段训练数据完全报告

noao-chat-4-rl阶段训练

LLM RL 训练完整解析

noao-chat-3-sft阶段训练

LLM SFT 训练完整解析

Q: `NCCL error: unhandled system error`

Q: `RuntimeError: Address already in use`