网站建设 6万元,北京知名的品牌设计公司,网站开发动态,rp网站做多大Building a decoder transformer model on AMD GPU(s) — ROCm Blogs 2024年3月12日 作者 Phillip Dang.
在这篇博客中#xff0c;我们展示了如何使用 PyTorch 2.0 和 ROCm 在单个节点上的单个和多个 AMD GPU 上运行Andrej Karpathy’s beautiful PyTorch re-implementation …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: D67301End of ROCm SMI Log 确保 PyTorch 也能识别这些 GPU
import torch
import torch.nn as nn
from torch.nn import functional as F
print(fnumber 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, encodingutf-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), dtypetorch.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_size1]
for t in range(block_size):context x[:t1]target y[t]print(fwhen 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, biasFalse)self.query nn.Linear(n_embd, head_size, biasFalse)self.value nn.Linear(n_embd, head_size, biasFalse)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_headn_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, mean0.0, std0.02)if module.bias is not None:torch.nn.init.zeros_(module.bias)elif isinstance(module, nn.Embedding):torch.nn.init.normal_(module.weight, mean0.0, std0.02)def forward(self, idx, targetsNone):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, devicedevice)) # (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_samples1) # (B, 1)# 将采样得到的索引追加到运行的序列中idx torch.cat((idx, idx_next), dim1) # (B, T1)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:iblock_size] for i in ix])y torch.stack([data[i1:iblock_size1] 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(), lrlearning_rate)for iter in range(max_iters):# 定期评估训练集和验证集的损失if iter % eval_interval 0 or iter max_iters - 1:losses estimate_loss(model)print(fstep {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_noneTrue)loss.backward()optimizer.step()# 使用模型进行推理以生成文本context torch.zeros((1, 1), dtypetorch.long, devicedevice)print(decode(model.generate(context, max_new_tokens2000)[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 fellowdst 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 - 1os.environ[MASTER_ADDR] localhostos.environ[MASTER_PORT] 12355# 初始化进程组backend ncclinit_process_group(backendbackend, rankrank, world_sizeworld_size)torch.cuda.set_device(rank)在实例化模型之前我们需要调用这个设置函数。 接下来我们对主函数进行些许修改使其能够在多个 GPU 上运行。 请注意主函数现在接收两个参数rankGPU 的 ID和 world_sizeGPU 的数量。
def main(rank:int, world_size:int):print(fTraining 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), dtypetorch.long, devicedevice)# 因为模型现在是一个分布式模型我们需要通过添加 module 来解包它print(decode(model.module.generate(context, max_new_tokens500)[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.输出来自两个不同的GPUrank 0和1每个训练/验证损失是不同的。这是因为每个进程有不同的种子以确保它们不会在相同的数据批次上进行训练。
