经典神经网络(13)GPT-1、GPT-2原理及nanoGPT源码分析(GPT-2)

2022 年 11 月，ChatGPT 成功面世，成为历史上用户增长最快的消费者应用。与 Google、FaceBook等公司不同，OpenAI 从初代模型 GPT-1 开始，始终贯彻只有解码器（Decoder-only）的技术路径，不断迭代升级。

模型	发布日期
GPT	2018-11-14
GPT-2	2019-11-27
GPT-3	2020-6-11
InstructGPT	2022-3-4
ChatGPT	2022-11-30
GPT-4	2023-3-14
ChatGPT Plugin	2023-5-12

可以看到，2022 年是 GPT 系列模型围绕 GPT-3、GPT-3.5 加速版本迭代的年份；
- 2022 年 3 月，基于 GPT-3 微调的 InstructGPT 发布，验证了人类反馈强化学习RLHF对模型输出对齐（alignment）的重要作用；
- 2022年4-6月，基于Codex、InstructGPT，OpenAI 加速迭代形成 GPT-3.5 模型；
- 2022 年 11 月，基于 GPT-3.5 微调的 ChatGPT 发布，成为 Instruction-tuning、RLHF、思维链等 LLM 相关技术的集大成者。
  - ChatGPT与InstructGPT的训练方法基本一致，区别在于InstructGPT、ChatGPT分别基于GPT-3、GPT-3.5进行模型微调。
  - InstructGPT具体可分为有监督微调、奖励模型训练、PPO 强化学习三个步骤。
  - ChatGPT技术原理解析可参考：https://blog.csdn.net/v_JULY_v/article/details/128579457
- 2023年3月中旬，OpenAI正式对外发布GPT-4，增加了多模态(支持图片的输入形式)，且ChatGPT底层的语言模型直接从GPT3.5升级到了GPT4。
ChatGPT 的发展不仅得益于 GPT 模型参数、训练数据的持续优化，也得益于各类 LLM 新技术的融会贯通，OpenAI 博采众长，加速 Instruction-tuning、RLHF、思维链等新技术在 GPT 系列模型中的深度应用，ChatGPT 是现有技术的集大成者。
我们今天主要了解下GPT-2。

1 GPT-1简介

2017年，Google推出了Transformer模型，这一架构因其在性能上的显著优势迅速吸引了OpenAI团队的注意。OpenAI随后将研发重点转移到Transformer架构，并在2018年发布了GPT-1模型。
OpenAI在2018年提出了GPT（Generative Pre-training）生成式预训练模型，采用了仅有解码器的Transformer模型，专注于预测下一个词元。
GPT采用了transformer的Decoder作为框架，并采用了两阶段的训练方式。首先，在大量的无标记数据集中，进行生成式训练（Generative Pre-training）；然后，在在特定任务进行微调（fine-tuning）。
论文：language_understanding_paper (openai.com)

1.1 GPT-1网络结构

下图左半部分是transformer架构图，右半部分是GPT的架构图。与GPT相比，transformer的Decoder除了中间被隐去的cross-attention，其余部分结构相同，GPT由12个Decoder串联而成。
- GPT的输入部分与transformer的Encoder的输入部分相同，与transformer不同的是位置向量采用随机初始化，并在训练中进行更新；
- GPT的Decoder可以看做是 transformer 的 Decoder 去掉中间Cross-Attention后的结构。
- Decoder主要有三个子模块组成：Masked Multi-Head Attention、残差网络&LayerNorm 以及 Feed Forward。
GPT的输出部分对应于下游任务，不同的任务使用不同的全连接层作为输出。

在这里插入图片描述

1.2 两阶段训练

GPT的训练过程分为两个阶段，第一阶段是采用大量文本预料进行无监督训练，第二个阶段是采用少量有标注的数据进行有监督的微调。

1.2.1 无监督预训练

GPT采用标准的语言模型进行无监督训练，即通过上文前 $k$ 个词来预测当前词
预训练时只有Text Prediction，没有Task Classifier

在这里插入图片描述

预训练模型的参数设置

GPT-1使用BooksCorpus数据集(约5GB)来训练语言模型。BooksCorpus有大约7000本未出版的书籍，这些书籍帮助在不可见的数据上训练语言模型。

采用L=12层decoder，每个自注意层有A=12个注意头
使用了带有40,000个合并的字节对编码(BPE)词汇表
分词嵌入维度H为768，位置编码嵌入维度为768，通过模型学习获得位置编码
位置前馈层采用3072维
dropout为0.1
GELU函数作为激活函数
batch_size为64，序列长度为512，epoch为100，参数共117M，即1.17 亿的参数量。

1.2.2 有监督微调

无监督训练完毕后，并不能直接用于下游任务，需要在具体任务的标注数据集上进行微调(如下图右半部分)。
对于大多数下游任务，监督的微调只需要3个epoch。
对比预训练阶段，只是多增加了“Task Classifier”模块。

当得到无监督的预训练模型后，可以将该模型直接应用到有监督任务中。每个实例有m个输入 $x^1,x^2,..,x^m$ ，以及标签y组成。首先将这些token输入到预训练的GPT1中，得到最终的特征向量 $h_l^m$ ,然后再通过一个全连接层得到预测结果y：
$P(y|x^1,...,x^m)=softmax(h_l^mW_y)$
有监督的目标是最大化以下损失：
$L_2(C)=∑logP(y|x^1,...,x^m)$
GPT的实验中发现，加入语言模型学习目标作为辅助任务，也就是损失函数中加入 $L_1(u)$ 能带来两点好处：不仅能提升监督模型的泛化能力；还能加快收敛；

因此，作者并没有直接使用L2，而是使用联合损失函数来进行微调(Text Prediction【L1】 + Task Classifier【L2】)，并使用 $λ$ 进行两个任务权值的调整， $λ$ 的值一般为0.5：
$L_3(C)=L_2(C)+λ*L_1(C)$
在这里插入图片描述

文本分类：可以看到，有两个特殊符号（Start和Extract）通过Transformer后，添加了一个线性层。
蕴含理解：给一段话，提出一个假设，看看假设是否成立。
- 将前提（premise）和假设（hypothesis）通过分隔符（Delimiter）隔开
- 两端加上起始和终止token
- 再依次通过Transformer和全连接得到预测结果；
文本相似：断两段文字是不是相似。相似是一个对称关系，A和B相似，那么B和A也是相似的
- 所以先有Text1+分隔符+Text2，再有Text2+分隔符+Text1；
- 两个序列分别经过Transformer后，各自得到输出的向量；
- 我们把它按元素加到一起，然后送给一个线性层。这也是一个二分类问题。
多项选择：多个序列，每个序列都由相同的问题Context和不同的Answer构成。
- 如果有N个答案，就构造N个序列；
- 每个QA序列都各自经过Transformers和线性层，对每个答案都计算出一个标量；
- 最后经过softmax生成一个各个答案的概率密度分布。这是一个N分类问题。

2 GPT-2简介

Bert论文：https://arxiv.org/pdf/1810.04805
GPT-2论文：Language Models are Unsupervised Multitask Learners (openai.com)
同年10月，谷歌发布了BERT模型。BERT用了Transformer中的Encoder部分，更类似完形填空，根据上下文来确定中间词：
- 和GPT相比，BERT所使用的掩码语言模型任务（Masked Language Model）虽然让它失去了直接生成文本的能力，但换来的是双向编码的能力，这让模型拥有了更强的文本编码性能，直接的体现则是下游任务效果的大幅提升。
- 而GPT为了保留生成文本的能力，只能采用单向编码。
如下图所示，为了和GPT进行对比Bert论文作者设计了Base模型，模型参数大小和GPT大致相等。但是BERT数据集的数据量大概是GPT的四倍。
对比结果如下图所示，BERT-Large模型在多个NLU任务上取得了显著的性能提升，成为当时自然语言处理领域的明星模型，引领了一波研究热潮。

在这里插入图片描述

以当年的眼光来看，BERT绝对是一个更加优秀的模型。因为既然BERT和GPT两者都是采用「预训练+微调」的范式，并且下游任务依然是分类、匹配、序列标注等等「经典」的NLP任务形式，那么像BERT模型这种更注重特征编码的质量，下游任务选一个合适的损失函数去配合任务做微调，显然比GPT这种以文本生成的方式去「迂回地」完成这些任务更加直接。
如果当时OpenAI放弃生成式预训练这条路，也许我们要等更长的时间才能见到ChatGPT这样的模型。

2.1 GPT-2核心思想

Q：当模型被别人用更大的数据集、参数量(Bert)打败时，应该怎么做？

A：自己也增大数据集和模型参数量。

Q：但是当文本数据集增大，模型参数量增大的基础上，模型的优势没有那么高的情况下应该怎么办？

A：找到另一个卖点（另辟蹊径）——zero shot

如果这篇文章单纯讲结果比Bert好一些，其实大家一看没啥意思，太工程化了，但是换一个角度，做一个更难的问题，但是得到一个好一点的结果，文章新颖性一下子有了。做工程和做研究的区别就是，做工程可以一直死盯着精度，但是做研究需要另辟蹊径找到一个创新点。

来自《李沐GPT、GPT-2、GPT-3论文精读》

GPT-2的核心思想就是，当模型的容量非常大且数据量足够丰富时，仅仅靠语言模型的学习便可以完成其他有监督学习的任务，不需要在下游任务微调。

GPT-2和GPT的区别在于GPT-2使用了更多的网络参数和更大的数据集，以此来训练一个泛化能力更强的词向量模型。GPT-2相比于GPT有如下几点区别：

主推zero-shot，而GPT-1为pre-train+fine-tuning；
模型更大，参数量达到了15亿个，而GPT-1只有1.17亿个；
数据集更大，WebText数据集包含了40GB的文本数据，而GPT-1只有5GB；
- 没有选择Common Crawl这种具有很多冗余无用信息的项目
- 选用的是reddit里面已经被人工筛选出的有意义的，并且具有至少3karma值的网页进行数据处理，大概有800万个文本，40GB的文字。
训练参数变化，batch_size从64增加到 512，上文窗口大小从512增加到1024；

在这里插入图片描述

2.2 GPT-2模型结构

GPT-2 提供了四种规模的模型：

在这里插入图片描述

在模型结构方面，整个GPT-2的模型框架与GPT相同，调整更多的是被当作训练时的trick，而不作为GPT-2的创新，具体为以下几点：

后置层归一化(post-norm)改为前置层归一化(pre-norm)
在模型最后一个自注意力层之后，额外增加一个层归一化
调整参数的初始化方式，按残差层个数进行缩放，缩放比例为 1 : sqrt(n)
输入序列的最大长度从512扩充到1024

在这里插入图片描述

2.3 GPT-2的贡献

2.3.1 预训练和zero-shot

GPT-2的预训练和GPT基本没什么区别，但是对下游任务用了zero-shot。
GPT-2可以在zero-shot设定下实现下游任务，即不需要用有标签的数据进行微调。
- 为了实现zero-shot，下游任务的输入就不能像GPT那样在构造输入时加入开始、中间和结束的特殊字符，这些是模型在预训练时没有见过的，而是应该和预训练模型看到的文本一样，更像一个自然语言。
- 可以通过做prompt的方式来zero-shot。
  - 例如：在做句子翻译任务时，训练的句子可以被写为： (translate to french, english text, french text). 其中translate to french在后文叫做prompt也叫做提示，相当于做了一个特殊的提示词。
  - 如果要做阅读理解任务时：可以写作(answer the question, document（阅读的文本）, question, answer)，answer the question相当于任务提示
- 为何zero-shot这种方式是有效的呢?从一个尽可能大且多样化的数据集中一定能收集到不同领域不同任务相关的自然语言描述示例，数据集里就存在展示了这些prompt示例，所以训练出来就自然而然有一定zero-shot的能力了。

2.3.2 总结

GPT-2的最大贡献是验证了通过海量数据和大量参数训练出来的词向量模型有迁移到其它类别任务中而不需要额外的训练，即zero-shot learning的能力。但是效果其实很一般。
GPT-2表明随着模型容量和数据量的增大，其潜能还有进一步开发的空间，基于这个思想，促使了GPT3的出现。
GPT-2虽然提出zero-shot，比Bert有新意，但是有效性方面不佳。GPT-3考虑few-shot，用少量文本提升有效性。

3 nanoGPT源码分析(GPT-2)

nanoGPT是李飞飞教授的学生，前特斯拉AI总监Andrej Karpathy基于OpenWebText重现 GPT-2 (124M)的开源项目。
nanoGPT删繁就简，代码非常简单，cpu也可以跑。无论是从头开始训练新模型，或者微调，都能很容易满足需求。

在这里插入图片描述

GPT-2原理解释可参考：https://jalammar.github.io/illustrated-gpt2/
nanoGPT源码仓库：https://github.com/karpathy/nanoGPT
首先我们clone下整个仓库，并按照仓库的Install的部分安装需要的package即可
如果我们想使用torch.compile的技术则需要安装2.x版本的PyTorch【torch.compile是PyTorch 2.x版本提出的一个新的技术，可以有效降低我们训练模型的时间】

pip install torch numpy transformers datasets tiktoken wandb tqdm

nanoGPT提供了两个案例：
- 第一个案例是在shakespeare上的构建的字符级别（character-level）的GPT2。这个数据很小，并且训练很快，很适合理解原理。
- 第二个就是根据原始论文在OpenWebText数据上重现GPT2，但是这个需要在8块A100 40GB机器上训练4天。

3.1 GPT-2模型

3.1.1 LayerNorm

import math
import inspect
from dataclasses import dataclass

import torch
import torch.nn as nn
from torch.nn import functional as F

class LayerNorm(nn.Module):
    """ LayerNorm but with an optional bias. PyTorch doesn't support simply bias=False """
    def __init__(self, ndim, bias):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(ndim))
        self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None

    def forward(self, input):
        return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)

3.1.2 CausalSelfAttention

class CausalSelfAttention(nn.Module):

    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        # key, query, value projections for all heads, but in a batch
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
        # output projection
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
        # regularization
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)
        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.dropout = config.dropout
        # flash attention make GPU go brrrrr but support is only in PyTorch >= 2.0
        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
        if not self.flash:
            print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
            # causal mask to ensure that attention is only applied to the left in the input sequence
            # 目的：使用下三角矩阵屏蔽未来词汇
            # torch.ones(config.block_size, config.block_size) 创建一个 block_size * block_size的矩阵
            # torch.tril() 将上三角元素设置为0，下三角仍为1
            self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
                                        .view(1, 1, config.block_size, config.block_size))

    def forward(self, x):
        B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)

        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
        q, k, v  = self.c_attn(x).split(self.n_embd, dim=2)
        # 多头注意力机制，需要拆头
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)

        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
        if self.flash:
            # efficient attention using Flash Attention CUDA kernels
            y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 0, is_causal=True)
        else:
            # manual implementation of attention
            # QK^T / sqrt(d_k)
            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
            # 将矩阵的为0的地方填为-inf（即上三角填充为-inf、屏蔽未来词汇）
            att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float('-inf'))
            # 得到的softmax以后，矩阵上三角部分的注意力权重 -> 0
            att = F.softmax(att, dim=-1)
            # 如果提供了dropout，对注意力权重att进行dropout操作
            att = self.attn_dropout(att)
            # 对value进行加权求和
            y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
        # 维度转换
        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side

        # output projection
        y = self.resid_dropout(self.c_proj(y))
        return y

3.1.3 FFN

class MLP(nn.Module):
    """前馈神经网络FFN"""
    def __init__(self, config):
        super().__init__()
        # 注意：输出维度为n_embd的4倍
        self.c_fc    = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
        self.gelu    = nn.GELU()
        self.c_proj  = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        x = self.c_fc(x)
        x = self.gelu(x)
        x = self.c_proj(x)
        x = self.dropout(x)
        return x

3.1.4 组装Block

包含两个子层，一个为Masked Multi-Head Attention，一个为FFN
后置层归一化(post-norm)改为前置层归一化(pre-norm)

class Block(nn.Module):

    def __init__(self, config):
        super().__init__()
        self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)
        self.mlp = MLP(config)

    def forward(self, x):
        # 1、第一个子层
        # Sublayer(LayerNorm(x)) + x，其中Sublayer = Masked Multi-Head Attention
        x = x + self.attn(self.ln_1(x))

        # 2、第二个子层
        # Sublayer(LayerNorm(x)) + x，其中Sublayer = FFN
        x = x + self.mlp(self.ln_2(x))
        return x

3.1.5 封装GPT-2模型

这里主要看下初始化、forward函数、以及generate函数。当然作者还提供了从transformers库加载GPT2模型、配置优化器等函数，可以自己看下。
这里作者把word embedding和lm_head权重进行共享，减少了模型参数
c_proj的权重初始化按n_layer层数进行缩放

class GPT(nn.Module):

    def __init__(self, config):
        super().__init__()
        assert config.vocab_size is not None
        assert config.block_size is not None
        self.config = config
        self.transformer = nn.ModuleDict(dict(
            wte = nn.Embedding(config.vocab_size, config.n_embd), # word table embedding
            wpe = nn.Embedding(config.block_size, config.n_embd), # word position embedding
            drop = nn.Dropout(config.dropout),
            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
            ln_f = LayerNorm(config.n_embd, bias=config.bias),
        ))

        # 定义一个线性层，将模型的输出维度映射到词汇表大小
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)

        # with weight tying when using torch.compile() some warnings get generated:
        # "UserWarning: functional_call was passed multiple values for tied weights.
        # This behavior is deprecated and will be an error in future versions"
        # not 100% sure what this is, so far seems to be harmless. TODO investigate

        # word embedding和lm_head权重共享
        self.transformer.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying

        # init all weights
        self.apply(self._init_weights)
        # apply special scaled init to the residual projections, per GPT-2 paper
        for pn, p in self.named_parameters():
            if pn.endswith('c_proj.weight'):
                # 初始化方式按n_layer层数进行缩放
                torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))

        # report number of parameters(参数共享后的参数量)
        print("number of parameters: %.2fM" % (self.get_num_params()/1e6,))

在模型最后一个自注意力层之后，额外增加一个层归一化
推理时只用最后一个时刻的hidden state去预测logits

    ......
    def forward(self, idx, targets=None):
        device = idx.device
        b, t = idx.size()
        assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
        pos = torch.arange(0, t, dtype=torch.long, device=device) # shape (t)

        # forward the GPT model itself
        tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)
        pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd)
        x = self.transformer.drop(tok_emb + pos_emb)
        for block in self.transformer.h:
            x = block(x)
        # 在模型最后一个自注意力层之后，额外增加一个层归一化
        x = self.transformer.ln_f(x)

        if targets is not None:
            # if we are given some desired targets also calculate the loss
            logits = self.lm_head(x)
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
        else:
            # inference-time mini-optimization: only forward the lm_head on the very last position
            # 小优化: 只用最后一个时刻的hidden state去预测logits
            logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim
            loss = None

        return logits, loss

推理过程中，如果上下文长度大于block_size(1024)，就裁剪到block_size长度，再输入到模型中推理
temperature的是作用在logits上，通过temperature的大小来调整模型输出概率分布的"尖锐程度"。
- 当temperature值较高（大于1）时，将使得概率分布变得更加均匀，模型的预测结果将更加多样化，但可能不太准确。换句话说，模型会有更大的概率去尝试预测不太可能的结果。
- 相反，当temperature值较低（小于1）时，将使得概率分布变得更加尖锐，模型的预测结果将更加聚焦于最有可能的结果，但可能牺牲多样性。换句话说，模型会更倾向于预测最可能的结果。
top_k在这段代码中的作用就是在生成预测结果后，只保留得分最高的前k个选项，以减少计算量，并且可能提高模型的生成质量。

    ......
	@torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        """
        Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete
        the sequence max_new_tokens times, feeding the predictions back into the model each time.
        Most likely you'll want to make sure to be in model.eval() mode of operation for this.
        """
        for _ in range(max_new_tokens):
            # if the sequence context is growing too long we must crop it at block_size
            # 如果上下文长度大于block_size，就裁剪到block_size长度，再输入到模型中推理
            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
            # forward the model to get the logits for the index in the sequence
            logits, _ = self(idx_cond)
            # pluck the logits at the final step and scale by desired temperature
            """
            temperature的是作用在logits上，通过temperature的大小来调整模型输出概率分布的"尖锐程度"。
                当temperature值较高（大于1）时，将使得概率分布变得更加均匀，模型的预测结果将更加多样化，但可能不太准确。
                换句话说，模型会有更大的概率去尝试预测不太可能的结果。

                相反，当temperature值较低（小于1）时，将使得概率分布变得更加尖锐，模型的预测结果将更加聚焦于最有可能的结果，但可能牺牲多样性。
                换句话说，模型会更倾向于预测最可能的结果。
            """
            logits = logits[:, -1, :] / temperature


            # optionally crop the logits to only the top k options
            """
                top_k在这段代码中的作用就是在生成预测结果后，只保留得分最高的前k个选项，以减少计算量，并且可能提高模型的生成质量。
            """
            if top_k is not None:
                # 取概率最大的topk个
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                # 取topk个概率中最小的值，如果小于此值就将概率置为负无穷，这样softmax后就趋近于0
                logits[logits < v[:, [-1]]] = -float('Inf')
            # apply softmax to convert logits to (normalized) probabilities
            probs = F.softmax(logits, dim=-1)
            # sample from the distribution
            # 函数作用是从模型输出的概率中采样，num_samples是采样次数
            idx_next = torch.multinomial(probs, num_samples=1)

            # append sampled index to the running sequence and continue
            # 在得到idx_next后，我们会将它和idx合在一起去预测新的下一个idx
            idx = torch.cat((idx, idx_next), dim=1)

        return idx

3.2 数据集的准备

我们运行下面的代码就可以下载Shakespeare数据集，并且将数据分为训练集和验证集。

python data/shakespeare_char/prepare.py

执行这个命令后，命令行会输出下面的内容：

length of dataset in characters:  1115394 # 统计所有数据有多少个字母
all the unique characters: 
!$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz # 统计有哪些字母和符号
vocab size: 65 # 构建vocabulary，即就是上面的这些符号
train has 1003854 tokens
val has 111540 tokens

运行完后，我们也会得到train.bin和val.bin。这两个.bin文件其实就是将对应字母/符号映射到vocabulary后的序号存储而形成的。train.bin包含了100w个token，val.bin包括10w个token。
因为我们训练的是character-level的GPT，所以我们的vocabulary就是26个英文字母的大小写和一些特殊符号组成，vocabulary的大小是65。运行完后，会将vocabulary信息：vocab_size、itos(id2token)、stoi(token2id)封装为dict，保存到meta.pkl文件中。
具体可以看data/shakespeare_char/prepare.py源码。

3.3 模型训练过程

其中如果pytorch版本＜2.0，在训练时需要指定compile为False
如果没有GPU，但是又想训练尝试下的话，需要设置device为cpu
train.py一共只有300多行代码，写的通俗易懂，支持多卡训练，混合精度训练，断点训练，梯度积累，动态学习率变化，wandb记录log，具体可以看源码。

这里主要介绍下DDP分布式训练的相关概念：

rank
- 进程号，在多进程上下文中，我们通常假定rank 0是第一个进程或者主进程
- 其它进程分别具有1，2，3不同rank号，这样总共具有4个进程
node
- 物理节点，可以是一个容器也可以是一台机器，节点内部可以有多个GPU
- nnodes指物理节点数量
- nproc_per_node指每个物理节点上面进程的数量
local_rank
- 指在一个node上进程的相对序号
- local_rank在node之间相互独立
- 这里需要注意的是，在torch.distributed.launch中，local_rank是隐式参数，即torch自动分配的。local_rank可以通过自动注入命令行参数来获得。
- torch1.10开始，官方建议使用环境变量的方式来获取local_rank。用终端命令torchrun来代替torch.distributed.launch，local_rank不再支持用命令行隐式传递的方式，完全使用环境变量配置各类参数。
world_size
- 全局进程总个数，即在一个分布式任务中rank的数量
group
- 进程组，一个分布式任务对应了一个进程组。
- 只有用户需要创立多个进程组时才会用到group来管理
- 默认情况下只有一个group

如下图所示，共有3个节点(node)，每个节点上有4个GPU，每台机器上起4个进程，每个进程占一块GPU，那么图中一共有12个rank(world_size)，nproc_per_node=4，nnodes=3，每个节点都一个对应的node_rank。

在这里插入图片描述

backend
- 通信后端，可选的包括：nccl（NVIDIA推出）、gloo（Facebook推出）、mpi（OpenMPI）。
- 一般建议GPU训练选择nccl，CPU训练选择gloo
master_addr与master_port
- 主节点的地址以及端口，供init_method 的tcp方式使用。
- 因为pytorch中网络通信建立是从机去连接主机，运行ddp只需要指定主节点的IP与端口，其它节点的IP不需要填写。
- 这个两个参数可以通过环境变量或者init_method传入

# 方式1：
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '12345'
dist.init_process_group(“nccl”, rank=rank, world_size=world_size)

# 方式2：
dist.init_process_group(“nccl”,init_method=“tcp://localhost:12345”
                        ,rank=rank,world_size=world_size)