在 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 上进行分布式数据并行性的工作原理是：

1. 将当前全局训练批次拆分为每个 GPU 上的小局部批次。例如，如果你有 8 个 GPU，且全局批次设置为 32 个样本，则每个 8 个 GPU 将具有 4 个样本的局部批次大小。

2. 将模型复制到每个设备，以便每个设备可以独立处理其局部批次。

3. 运行前向传播，然后是反向传播，并输出模型权重相对于局部批次损失的渐变。这会在多个设备上并发发生。

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架构。其组件包括：

• 自注意力机制：允许解码器在生成输出时权衡输入序列的不同部分。

• 掩码自注意力：通过掩盖后续位置，防止解码器在训练期间看到未来的令牌。

• 解码器层：由多层子层组成，如多头注意力和前馈神经网络，促进信息处理和生成。

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.shapek = 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 outclass 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 outclass 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_headself.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 xclass GPTLanguageModel(nn.Module):def __init__(self):super().__init__()# 每个令牌直接从查找表中读取下一个令牌的logitsself.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 normself.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 = Noneelse:B, T, C = logits.shapelogits = logits.view(B*T, C)targets = targets.view(B*T)loss = F.cross_entropy(logits, targets)return logits, lossdef 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_dataix = 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），每个训练/验证损失是不同的。这是因为每个进程有不同的种子，以确保它们不会在相同的数据批次上进行训练。