在 AMD GPU 上构建解码器 Transformer 模型

Building a decoder transformer model on AMD GPU(s) — ROCm Blogs

2024年3月12日作者 Phillip Dang.

在这篇博客中，我们展示了如何使用 PyTorch 2.0 和 ROCm 在单个节点上的单个和多个 AMD GPU 上运行Andrej Karpathy’s beautiful PyTorch re-implementation of GPTGPT PyTorch 重新实现Andrej Karpathy’s beautiful PyTorch re-implementation of GPT。

GPT 模型的核心是仅解码器的 transformer 架构。这种架构以自回归的方式一次生成一个输出标记（在我们的情况下是字符），这意味着每个生成的标记都依赖于之前生成的标记。如需深入了解该模型的工作原理，我们强烈推荐查看以下内容：

Andrej Karpathy的 Let’s build GPT: 从头开始，用代码详细阐述视频
Attention Is All You Need 论文。

我们首先使用 PyTorch 和 ROCm 在单个 GPU 上训练我们的模型，然后稍微修改代码以使用 to run on two GPUs using PyTorch 的分布式数据并行（Distributed Data Parallel, DDP）在两个 GPU 上运行。

在多 GPU 上进行分布式数据并行性的工作原理是：

将当前全局训练批次拆分为每个 GPU 上的小局部批次。例如，如果你有 8 个 GPU，且全局批次设置为 32 个样本，则每个 8 个 GPU 将具有 4 个样本的局部批次大小。
将模型复制到每个设备，以便每个设备可以独立处理其局部批次。
运行前向传播，然后是反向传播，并输出模型权重相对于局部批次损失的渐变。这会在多个设备上并发发生。
同步每个设备计算的局部梯度并将它们结合起来以更新模型权重。更新后的权重会重新分发到每个设备。

如需深入了解 PyTorch 中的分布式训练，请参阅使用 DDP 进行多 GPU 训练。

先决条件

要跟随本博客的内容进行操作，您需要具备以下软件：

ROCm
PyTorch
Linux OS

接下来，通过运行以下代码确保您的系统能够识别两块 AMD GPU：

! rocm-smi --showproductname

输出结果应如下所示：

================= ROCm System Management Interface ================
========================= Product Info ============================
GPU[0] : Card series: Instinct MI210
GPU[0] : Card model: 0x0c34
GPU[0] : Card vendor: Advanced Micro Devices, Inc. [AMD/ATI]
GPU[0] : Card SKU: D67301
GPU[1] : Card series: Instinct MI210
GPU[1] : Card model: 0x0c34
GPU[1] : Card vendor: Advanced Micro Devices, Inc. [AMD/ATI]
GPU[1] : Card SKU: D67301
===================================================================
===================== End of ROCm SMI Log =========================

确保 PyTorch 也能识别这些 GPU：

import torch
import torch.nn as nn
from torch.nn import functional as F
print(f"number of GPUs: {torch.cuda.device_count()}")
print([torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())])

输出结果应如下所示：

number of GPUs: 2
['AMD Radeon Graphics', 'AMD Radeon Graphics']

加载数据

我们使用 tiny_shakespeare 数据集，该数据集包含来自莎士比亚各种剧作的40,000行文本。让我们加载它并查看前200个字符。

! wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt
with open('input.txt', 'r', encoding='utf-8') as f:
    text = f.read()
print(text[:200])

输出：

First Citizen:
Before we proceed any further, hear me speak.

All:
Speak, speak.

First Citizen:
You are all resolved rather to die than to famish?

All:
Resolved. resolved.

First Citizen:
First, you

接下来，我们从文本中的唯一字符创建我们的词汇表，并建立字符与整数之间的映射。

chars = sorted(list(set(text)))
vocab_size = len(chars)
# create a mapping from characters to integers
stoi = { ch:i for i,ch in enumerate(chars) }
itos = { i:ch for i,ch in enumerate(chars) }
encode = lambda s: [stoi[c] for c in s] # encoder: take a string, output a list of integers
decode = lambda l: ''.join([itos[i] for i in l]) # decoder: take a list of integers, output a string

现在，我们将数据在编码字符之后分成训练集和测试集（即将每个字符转换为一个整数）。

data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9*len(data)) # first 90% will be train, rest val
train_data = data[:n]
val_data = data[n:]

让我们来看一个输入和目标的例子。

block_size = 8
x = train_data[:block_size]
y = train_data[1:block_size+1]
for t in range(block_size):
    context = x[:t+1]
    target = y[t]
    print(f"when input is {context} the target: {target}")

输出：

when input is tensor([18]) the target: 47
when input is tensor([18, 47]) the target: 56
when input is tensor([18, 47, 56]) the target: 57
when input is tensor([18, 47, 56, 57]) the target: 58
when input is tensor([18, 47, 56, 57, 58]) the target: 1
when input is tensor([18, 47, 56, 57, 58,  1]) the target: 15
when input is tensor([18, 47, 56, 57, 58,  1, 15]) the target: 47
when input is tensor([18, 47, 56, 57, 58,  1, 15, 47]) the target: 58

创建解码器transformer模型

让我们设置一些超参数，用于构建和训练模型。因为我们能够使用强大的AMD GPU，可以扩展我们的网络，并设置与Andrej在其视频教程中使用的相同的超参数。这包括增加我们的批量大小、块大小、层数和头数以及嵌入大小。我们希望这将使我们的损失值更低。

# 超参数
batch_size = 64 # 我们将并行处理多少个独立序列？
block_size = 256 # 预测的最大上下文长度是多少？
max_iters = 5000
eval_interval = 100
learning_rate = 3e-4
device = 'cuda' if torch.cuda.is_available() else 'cpu'
eval_iters = 200
n_embd = 384
n_head = 6
n_layer = 6
dropout = 0.2
# ------------

torch.manual_seed(1337) # 设置手动种子以保证可重复性

解码器Transformer架构

以下是我们的主要模型类，它创建了一个解码器Transformer架构。其组件包括：

自注意力机制：允许解码器在生成输出时权衡输入序列的不同部分。
掩码自注意力：通过掩盖后续位置，防止解码器在训练期间看到未来的令牌。
解码器层：由多层子层组成，如多头注意力和前馈神经网络，促进信息处理和生成。

Transformer decoder process

class Head(nn.Module):
    """ 一个自注意力头  """

    def __init__(self, head_size):
        super().__init__()
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))

        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # 输入大小 (batch, time-step, channels)
        # 输出大小 (batch, time-step, head size)
        B,T,C = x.shape
        k = self.key(x)   # (B,T,hs)
        q = self.query(x) # (B,T,hs)
        # 计算注意力分数 ("affinities")
        wei = q @ k.transpose(-2,-1) * k.shape[-1]**-0.5 # (B, T, hs) @ (B, hs, T) -> (B, T, T)
        # 使用上三角矩阵技巧创建掩码用于自注意力机制
        # 该掩码确保在训练期间，解码器只能关注到当前生成的令牌之前的位置，防止看到未来的令牌，即仅从左到右注意
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # (B, T, T)
        wei = F.softmax(wei, dim=-1) # (B, T, T)
        wei = self.dropout(wei)
        #  执行值的加权聚合
        v = self.value(x) # (B,T,hs)
        out = wei @ v # (B, T, T) @ (B, T, hs) -> (B, T, hs)
        return out

class MultiHeadAttention(nn.Module):
    """ 并行的多头自注意力机制 """

    def __init__(self, num_heads, head_size):
        super().__init__()
        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
        self.proj = nn.Linear(head_size * num_heads, n_embd)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        out = torch.cat([h(x) for h in self.heads], dim=-1)
        out = self.dropout(self.proj(out))
        return out

class FeedFoward(nn.Module):
    """ 一个简单的线性层及其非线性激活 """

    def __init__(self, n_embd):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),
            nn.ReLU(),
            nn.Linear(4 * n_embd, n_embd),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        return self.net(x)

class Block(nn.Module):
    """ Transformer块: 通信后是计算  """

    def __init__(self, n_embd, n_head):
        # n_embd: 嵌入维度，n_head: 多头数量
        super().__init__()
        head_size = n_embd // n_head
        self.sa = MultiHeadAttention(n_head, head_size)
        self.ffwd = FeedFoward(n_embd)
        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)

    def forward(self, x):
        x = x + self.sa(self.ln1(x))
        x = x + self.ffwd(self.ln2(x))
        return x

class GPTLanguageModel(nn.Module):

    def __init__(self):
        super().__init__()
        # 每个令牌直接从查找表中读取下一个令牌的logits
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
        self.position_embedding_table = nn.Embedding(block_size, n_embd)
        self.blocks = nn.Sequential(*[Block(n_embd, n_head=n_head) for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd) # final layer norm
        self.lm_head = nn.Linear(n_embd, vocab_size)

        # 更好的权重初始化，原始GPT视频中未涵盖，但很重要，将在后续视频中介绍
        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.shape

        # idx 和 targets 都是 (B, T) 的整数张量
        tok_emb = self.token_embedding_table(idx) # (B,T,C)
        pos_emb = self.position_embedding_table(torch.arange(T, device=device)) # (T,C)
        x = tok_emb + pos_emb # (B,T,C)
        x = self.blocks(x) # (B,T,C)
        x = self.ln_f(x) # (B,T,C)
        logits = self.lm_head(x) # (B,T,vocab_size)

        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            logits = logits.view(B*T, C)
            targets = targets.view(B*T)
            loss = F.cross_entropy(logits, targets)

        return logits, loss

    def generate(self, idx, max_new_tokens):
        # idx 是 (B, T) 当前上下文的索引数组
        for _ in range(max_new_tokens):
            # 将 idx 裁剪为最后 block_size 个令牌
            idx_cond = idx[:, -block_size:]
            # 获取预测结果
            logits, loss = self(idx_cond)
            # 只关注最后一个时间步
            logits = logits[:, -1, :] # becomes (B, C)
            # 应用 softmax 获得概率
            probs = F.softmax(logits, dim=-1) # (B, C)
            # 从概率分布中采样
            idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)
            # 将采样得到的索引追加到运行的序列中
            idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)
        return idx

实用函数

在训练我们的模型之前，我们需要两个实用函数：

在训练我们的模型之前，我们需要两个实用函数：
一个是获取随机批量数据的函数

为了估计损失：

@torch.no_grad()
def estimate_loss(model):
    out = {}
    model.eval()
    for split in ['train', 'val']:
        losses = torch.zeros(eval_iters)
        for k in range(eval_iters):
            X, Y = get_batch(split)
            logits, loss = model(X, Y)
            losses[k] = loss.item()
        out[split] = losses.mean()
    model.train()
    return out

为了获取一个小批量的数据：

# 数据加载
def get_batch(split):
    # 数据加载
    data = train_data if split == 'train' else val_data
    ix = torch.randint(len(data) - block_size, (batch_size,))
    x = torch.stack([data[i:i+block_size] for i in ix])
    y = torch.stack([data[i+1:i+block_size+1] for i in ix])
    x, y = x.to(device), y.to(device)
    return x, y

训练和推理

现在我们已经准备好所有的部分，让我们实例化我们的模型，进行训练，并运行一些推理来生成我们（希望）类似莎士比亚风格的文本。

这是我们的主函数：

def main():
    model = GPTLanguageModel()
    model = model.to(device)

    # 输出模型的参数数量
    print(sum(p.numel() for p in model.parameters())/1e6, 'M parameters')

    # 创建一个 PyTorch 优化器
    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

    for iter in range(max_iters):

        # 定期评估训练集和验证集的损失
        if iter % eval_interval == 0 or iter == max_iters - 1:
            losses = estimate_loss(model)
            print(f"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")

        # 抽取一个数据批次
        xb, yb = get_batch('train')

        # 计算损失
        logits, loss = model(xb, yb)
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        optimizer.step()

    # 使用模型进行推理以生成文本
    context = torch.zeros((1, 1), dtype=torch.long, device=device)
    print(decode(model.generate(context, max_new_tokens=2000)[0].tolist()))

我们创建了一个名为 gpt_single_gpu.py 的脚本，包含所有必要的命令。要运行它，请使用以下代码：

python3 gpt_single_gpu.py

以下是脚本的预期输出：

10.788929 M parameters
step 0: train loss 4.2221, val loss 4.2306
step 100: train loss 2.4966, val loss 2.5012
step 200: train loss 2.4029, val loss 2.4295
...
step 4900: train loss 0.8676, val loss 1.5644
step 4999: train loss 0.8598, val loss 1.5677

以下是我们生成的文本的前几行：

Thou fellow'dst idst the game of his names;
And yet since was Menenius, one would thrident again
That Anne. But where shall do become me injuries?

JULIET:
O though often thee cortainted matter,--
A better with him he gone hath
A colder-balm equal-deniving,
Of what a peril the people, when he did make me
Disobedition, become him to see
That conceive on earth fitting his finger,

在验证损失为1.5677时，生成的文本看起来几乎像英语，并且模型能够学习莎士比亚输入的对话风格。

每次运行推理时，我们期望的结果会有所不同。这是因为模型从所有可能的标记分布中进行采样，其中每个标记的概率由 softmax 函数给出。

分布式训练在多 GPU 上

为了在 PyTorch 中使用多个 GPU 在单个节点上训练我们的模型，我们将使用 PyTorch 的分布式数据并行 (Distributed Data Parallel)。为此，我们只需要对当前代码进行一点点修改。

首先，让我们从 PyTorch 导入一些所需的函数：

import torch.multiprocessing as mp
from torch.distributed import init_process_group, destroy_process_group
from torch.nn.parallel import DistributedDataParallel as DDP

接下来，让我们设置分布式数据并行组。通常，每个 GPU 运行一个进程，因此我们需要设置一个组，以便所有进程和 GPU 之间可以相互通信。让我们创建一个执行此操作的小函数。

def ddp_setup(rank, world_size):
    """
    world_size: GPU 的数量
    rank: GPU 的 ID，从 0 到 world_size - 1
    """

    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    # 初始化进程组
    backend = "nccl"
    init_process_group(backend=backend, rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

在实例化模型之前，我们需要调用这个设置函数。接下来，我们对主函数进行些许修改，使其能够在多个 GPU 上运行。请注意，主函数现在接收两个参数：rank（GPU 的 ID）和 world_size（GPU 的数量）。

def main(rank:int, world_size:int):
    print(f"Training DDP model on rank/gpu {rank}.")
    ddp_setup(rank, world_size)

    # 每个 gpu/进程获得不同的种子
    torch.manual_seed(1337 + rank)

    model = GPTLanguageModel()
    model.to(rank)
    model = DDP(model, device_ids=[rank])
    ... # 剩下的训练过程与单进程主函数的训练过程相同

    # 运行推理
    if rank == 0: # 仅在主进程上运行推理。没有这个 if 语句，每个进程都会运行自己的预测
        print('generating text')
        context = torch.zeros((1, 1), dtype=torch.long, device=device)
        # 因为模型现在是一个分布式模型，我们需要通过添加 "module" 来解包它
        print(decode(model.module.generate(context, max_new_tokens=500)[0].tolist()))

    #  一旦模型训练完成，销毁进程以干净退出
    destroy_process_group()

完整代码

我们创建了一个脚本 gpt_multiple_gpus.py ，其中包含所有必需的命令。要运行该脚本，请使用以下代码：

python3 gpt_multiple_gpus.py

以下是该脚本的预期输出。

We have 2 GPUs! Using 2 GPUs
Training DDP model on rank/gpu 1.
Training DDP model on rank/gpu 0.
10.788929 M parameters
10.788929 M parameters
GPU/rank 0 step 0: train loss 4.2221, val loss 4.2306
GPU/rank 1 step 0: train loss 4.2228, val loss 4.2304
GPU/rank 0 step 500: train loss 1.6010, val loss 1.7904
GPU/rank 1 step 500: train loss 1.5984, val loss 1.7871
...
GPU/rank 1 step 4999: train loss 0.5810, val loss 1.7733
GPU/rank 0 step 4999: train loss 0.5807, val loss 1.7723

以下是生成的文本：

HENRY BOLINGBROKE:
Warwick, It say; and he is safe, and whose
With unmorable slaves, they have stafd too:
So say the tidings you for Richmond, with ride?

BUSHY:
Marry, my Lord Clarence to my noble lady.

GLOUCESTER:
Go, to thee to thy daughter as I may break;
And what do now thy will I, I do me say
My name; it is the king.

BUCKINGHAM:
'Twas every stranger:--
Nay, my good son.

输出来自两个不同的GPU（rank 0和1），每个训练/验证损失是不同的。这是因为每个进程有不同的种子，以确保它们不会在相同的数据批次上进行训练。