【llm对话系统】大模型 Llama 源码分析之并行训练方案

1. 引言

训练大型语言模型 (LLM) 需要巨大的计算资源和内存。为了高效地训练这些模型，我们需要采用各种并行策略，将计算和数据分布到多个 GPU 或设备上。Llama 作为当前最流行的开源大模型之一，其训练代码中采用了多种并行技术。本文将深入 Llama 的训练代码，分析其并行训练方案，主要关注参数并行和部分结构参数共享。

2. 并行训练策略概述

常见的并行训练策略包括：

• 数据并行 (Data Parallelism, DP)：将数据分成多个 batch，每个 GPU 处理一个 batch，所有 GPU 使用相同的模型副本。
• 模型并行 (Model Parallelism, MP)：将模型分成多个部分，每个 GPU 负责模型的一部分。
• 流水线并行 (Pipeline Parallelism, PP)：将模型的不同层分配到不同的 GPU 上，形成一个流水线。
• 张量并行 (Tensor Parallelism, TP)：将模型的张量 (例如，权重矩阵) 分片到多个 GPU 上。
• 序列并行 (Sequence Parallelism, SP): 将序列长度分片到多个 GPU 上。

Llama 主要采用了数据并行、张量并行，以及一些结构参数共享的优化。

3. Llama 中的参数并行方案

Llama 使用了 ZeRO (Zero Redundancy Optimizer) 技术，这是一种强大的内存优化方法，它结合了数据并行和模型并行。ZeRO 的核心思想是将模型状态 (权重、梯度和优化器状态) 分片到多个 GPU 上，从而减少每个 GPU 的内存占用。

ZeRO 有三个阶段：

• ZeRO-1 (Optimizer State Partitioning)：将优化器状态 (例如，Adam 的动量和方差) 分片。
• ZeRO-2 (Gradient Partitioning)：在 ZeRO-1 的基础上，将梯度也分片。
• ZeRO-3 (Parameter Partitioning)：在 ZeRO-2 的基础上，将模型参数也分片。

Llama 主要使用了 ZeRO-3，将模型参数、梯度和优化器状态都分片到多个 GPU 上。

3.1 参数分片计算

在 Llama 的训练代码中, 以 torch.distributed.fsdp 库为例 (Fully Sharded Data Parallel, FSDP)，它实现了 ZeRO-3 的功能。

以下是一个简化的 FSDP 参数分片示例：

import torch
import torch.distributed as dist
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import (
    transformer_auto_wrap_policy,
    enable_wrap,
    wrap,
)
import functools

# 假设我们有一个简单的 Transformer 模型
class TransformerLayer(torch.nn.Module):
    def __init__(self, hidden_dim):
        super().__init__()
        self.linear1 = torch.nn.Linear(hidden_dim, 4 * hidden_dim)
        self.linear2 = torch.nn.Linear(4 * hidden_dim, hidden_dim)

    def forward(self, x):
        x = self.linear1(x)
        x = torch.nn.functional.relu(x)
        x = self.linear2(x)
        return x

# 初始化分布式环境
dist.init_process_group("nccl")
rank = dist.get_rank()
world_size = dist.get_world_size()
device = torch.device(f"cuda:{rank}")

# 模型和优化器
hidden_dim = 768
model = TransformerLayer(hidden_dim).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# 使用 FSDP 包装模型
# 使用自动包装策略
auto_wrap_policy = functools.partial(
    transformer_auto_wrap_policy,
    transformer_layer_cls={TransformerLayer,},
)
model = FSDP(model, fsdp_auto_wrap_policy=auto_wrap_policy, device_id=device)

# 模拟训练数据
x = torch.randn(1, 10, hidden_dim, device=device)

# 前向传播
y = model(x)

# 反向传播
loss = y.sum()
loss.backward()

# 优化器更新
optimizer.step()

# 清空梯度
optimizer.zero_grad()

print(f"Rank {rank}: 训练完成")

代码解释：

1. 初始化分布式环境：dist.init_process_group("nccl") 初始化分布式训练环境，使用 NCCL 后端。
2. 模型和优化器：定义一个简单的 Transformer 层和 Adam 优化器。
3. FSDP 包装：使用 FullyShardedDataParallel 包装模型。这会将模型参数分片到多个 GPU 上。
   • transformer_auto_wrap_policy 会自动将模型的每一层都用 FSDP 包装起来。
4. 前向和反向传播：执行模型的前向和反向传播。
5. 优化器更新：执行优化器的 step 方法。
6. 清空梯度：清空梯度以进行下一次迭代。

运行方式：

你需要使用 torchrun（或 torch.distributed.launch）来启动这个脚本，例如：

torchrun --nproc_per_node=4 fsdp_example.py

这将使用 4 个 GPU 来训练模型。

原理说明：

• 参数分片：在 model = FSDP(model, ...) 这一步，模型参数被分片到 4 个 GPU 上。每个 GPU 只存储一部分参数。
• All-gather：在前向传播过程中，当需要完整的参数进行计算时（例如，矩阵乘法），FSDP 会自动执行 all-gather 操作，将所有 GPU 上的参数片段收集起来，组成完整的参数。
• Reduce-scatter：在反向传播过程中，梯度也是分片的。计算完梯度后，FSDP 会执行 reduce-scatter 操作，将每个参数的梯度片段 reduce 到对应的 GPU 上。
• 优化器更新：每个 GPU 使用自己分片到的参数和梯度来更新优化器状态。

通过这种方式，FSDP 显著减少了每个 GPU 的内存占用，使得训练更大的模型成为可能。

4. Llama 中的部分结构参数共享

除了参数分片，Llama 还采用了一些结构参数共享的优化，以进一步减少内存占用和提高训练效率。

例如在 Transformer 的多头注意力 (Multi-Head Attention) 机制中，不同 head 的 query, key, value 矩阵的计算通常是独立的。Llama 通过共享 key 和 value 矩阵，减少了参数量和计算量。更具体地说，llama使用了分组注意力机制(Grouped-Query Attention)。

4.1 分组注意力 (Grouped-Query Attention, GQA)

GQA 介于标准的多头注意力 (MHA) 和 Multi-Query Attention (MQA) 之间。

• MHA: 每个 head 都有独立的 Q, K, V 矩阵。
• MQA: 所有 head 共享 K, V 矩阵，只有 Q 矩阵是独立的。
• GQA: 将 head 分成多个组，每个组内的 head 共享 K, V 矩阵。

例如：

假设我们有 8 个 head，可以将它们分成 4 个组，每个组 2 个 head。这样，我们就只有 4 个 K 矩阵和 4 个 V 矩阵，而不是 8 个。

代码示例 (简化版)

import torch
import torch.nn as nn

class GroupedQueryAttention(nn.Module):
    def __init__(self, embed_dim, num_heads, num_groups):
        super().__init__()
        self.num_heads = num_heads
        self.num_groups = num_groups
        self.head_dim = embed_dim // num_heads
        self.group_size = num_heads // num_groups

        self.q_proj = nn.Linear(embed_dim, embed_dim)
        # 共享 K, V 矩阵
        self.k_proj = nn.Linear(embed_dim, num_groups * self.head_dim)
        self.v_proj = nn.Linear(embed_dim, num_groups * self.head_dim)
        self.out_proj = nn.Linear(embed_dim, embed_dim)

    def forward(self, x):
        batch_size, seq_len, embed_dim = x.shape

        # 计算 Q, K, V
        q = self.q_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim)
        k = self.k_proj(x).view(batch_size, seq_len, self.num_groups, self.head_dim)
        v = self.v_proj(x).view(batch_size, seq_len, self.num_groups, self.head_dim)

        # 将 head 分组
        q = q.view(batch_size, seq_len, self.num_groups, self.group_size, self.head_dim)

        # 计算注意力
        attn_scores = torch.einsum("bqlgd,bklhd->bqgkh", q, k) / (self.head_dim ** 0.5)
        attn_probs = attn_scores.softmax(dim=-1)
        attn_output = torch.einsum("bqgkh,bklhd->bqlgd", attn_probs, v)

        # 拼接 head
        attn_output = attn_output.reshape(batch_size, seq_len, embed_dim)

        # 输出投影
        output = self.out_proj(attn_output)

        return output

# 示例
embed_dim = 768
num_heads = 8
num_groups = 4
model = GroupedQueryAttention(embed_dim, num_heads, num_groups)
x = torch.randn(1, 10, embed_dim)
y = model(x)
print(y.shape)

代码解释：

1. num_groups：将 head 分成多少个组。
2. k_proj 和 v_proj：只输出 num_groups 个 head 的 K 和 V 矩阵。
3. 将 q 分成 num_groups 个组，每个组 group_size 个 head。
4. 计算注意力时，每个组内的 head 共享 K 和 V 矩阵。

GQA 的优势：

• 减少参数量：K 和 V 矩阵的数量减少了。
• 减少计算量：计算注意力时，每个 head 需要处理的 K 和 V 数量减少了。
• 性能接近 MHA：实验表明，GQA 的性能接近 MHA，明显优于 MQA。

llama的GQA实现在llama/model.py文件中，class Attention(nn.Module) 类下的forward函数中，更具体地，体现在self.num_heads和self.num_kv_heads的参数上, 分别控制query和kv的head数量，num_kv_heads小于num_heads