从零开始200行python代码实现LLM！

Datawhale

编者荐语：

Datawhale干货推荐。

以下文章来源于阿里云开发者，作者思潜

[

阿里云开发者 .

阿里巴巴官方技术号，关于阿里的技术创新均呈现于此。

](#)

前言

大语言模型（LLM）很火，讨论的文章铺天盖地，但对于没有机器学习背景的人来说，看多了只是粗浅了解了一堆概念，疑惑只增不减。

本文尝试从零开始，用python实现一个极简但完整的大语言模型，在过程中把各种概念“具象化”，让大家亲眼看到、亲手写出self-attention机制、transformer模型，亲自感受下训练、推理中会遇到的一些问题。

本文适用范围及目标：

‒✅只需会写基本的python代码；

‒✅尝试实现完整的语言模型（但由于层数、dataset限制，只会写诗词）；

‒❌不解释数学、机器学习原理性的知识，只做到“能用”为止；

‒❌不依赖抽象层次高的框架，用到的部分也会做解释；

声明：文章绝大部分内容来自ak大神的nanoGPT[1]。

相关代码都在Github仓库：

simpx/buildyourownllm [2]上，建议先clone下来，并通过 `pip install torch` 安装唯一的依赖后，在仓库目录下运行各个代码体验过程。

动手写代码最容易把抽象的概念具象化，非常建议使用vscode + ipynb的组合调试文中的代码，鉴于篇幅，不额外介绍工具。

本文先介绍“从零基础到Bigram模型”，下一篇文章再介绍“从Bigram模型到LLM”。

先用传统方式实现一个“诗词生成器”

让我们忘记机器学习，用传统思路来实现一个“诗词生成器”。

观察一下我们的数据集 ci.txt ，里面包含了宋和南唐的词，我们的目标是实现一个生成类似诗词的工具。

$ head -n 8 ci.txt

词是由一堆字组成的，那么一个简单的想法，我们可以通过计算每个字后面出现各个字的概率。

然后根据这些概率，不断的递归生成“下一个字”，生成的字多了，截断一部分，就是一首词了。

具体思路为：

准备词汇表：将ci.txt 出现的所有字去重，得到我们的词汇表，长度为vocab_size；
统计频率：准备一个vocab_size * vocab_size 的字典，统计每个词后出现别的词的频率；
计算概率，生成新“字”：根据频率计算概率，并随机采样，生成下一个字；

完整的代码如下（带注释版的见simplemodel\_with\_comments.py[3]）：

simplemodel.py

import random

直接通过python simplemodel.py 即可运行，去掉random.seed(42) 可以看到不同的输出结果。

在我的mac电脑上耗时2秒，效果如下：

$ python simplemodel.py

这像是一首名为“春江月”、作者为“张先生疑被。”的词，但其实我们只是实现了一个“下一个词预测器”。

在代码的眼里，只不过“春”字后面大概率是“江”，而“江”字后面大概率是“月”而已，它不知道什么是词，甚至不知道什么是一首词的开头、结尾。

这个字符序列层面的“意义”，实际上是由读者赋予的。

词汇表 - tokenizer

我们的“词汇表”，相当于LLM里的tokenizer，只不过我们直接使用ci.txt 里出现过的所有字符当做词汇表用。我们的词汇表只有6418个词汇，而真正的LLM有更大的vocab\_size，以及更高效的编码，一些常用词组直接对应1个token，比如下面是qwen2.5的tokenizer。

>>> from transformers import AutoTokenizer

qwen2.5使用了一个大小为151643的词汇表，其中常见的词汇“阿里巴巴”、“人工智能”都只对应1个token，而在我们的词汇表里，1个字符永远对应1个token，编码效率较低。

模型、训练、推理

我们刚刚实现的“模型”，实际是就是自然语言N-gram模型中的“Bigram模型”。这是一种基于统计的语言模型，用于预测一个词出现的概率，在这个模型中，假设句子中的每个字只依赖于其前面的一个字。具体的实现就是一个词频字典transition，而所谓的“训练”过程就是遍历所有数据，统计“下一个词”出现的频率。但我们的“推理”过程还是非常像真正的LLM的，步骤如下：

1.我们从transition 中获取下一个token的logits（logits是机器学习中常用的术语，表示最后一层的原始输出值），我们可以把logits[i]简单理解为“下一个token\_id是i的得分”，因此logits肯定是长度为vocab\_size的字典；

2.获得“得分字典”后，使用[logit / total for logit in logits] 做归一化处理，这是为了下一步更好的做随机采样。在这里我们使用最简单的线性归一，不考虑total为0的情况；

3.根据归一后的“得分字典”，使用random.choices 随机获取一个token id并返回；

4.循环反复，直到获得足够多的token。

进行重构，更加有“机器学习风格”

接下来我们把Bigram模型的实现变得更加“机器学习风格”，以便帮助我们理解后面真实的pytorch代码，有pytorch背景的同学可以直接跳过本节。

完整的代码码如下（带注释版的见simplebigrammodel\_with\_comments.py[4]）：

simplebigrammodel.py

import random

虽然有100多行代码，但实际上功能和上一个50行代码几乎是一样的，稍微运行、调试一下就能明白。

直接通过python simplebigrammodel.py 即可运行，这一次会生成2个字符序列：

$ python simplebigrammodel.py

解释一下这100多行代码的实现：

机器学习风格的一些约定

我们用Tokenizer 类封装了词汇表，以便它能像qwen的词汇表一样被使用。

同时，我们实现了一个BigramLanguageModel 类，这模仿pytorch里的nn.Module 写法，即：

1.参数在__init__ 中初始化；

2.推理在forward 函数中实现，并通过__call__ 允许对象被直接调用；

3.序列生成在generate 函数中实现；

最后，我们修改了数据加载的机制，如下：

def get_batch(tokens, batch_size, block_size):

每次调用get_batch 的时候，会随机返回两份数据，其中y 数组中的每一个token，都是x 数组内对应位置的token的下一个token。采用这样的写法，是为了方便后续操作。

批处理in，批处理out

这一个版本最难懂的地方，是数据都以多维数组的方式呈现，连推理结果返回的都是2个。

实际上，我们这里的“多维数组”，就是机器学习中的“张量”（Tensor），是为了最终方便GPU处理而准备的。

张量（Tensor）是数学和物理学中用于表示多维数据的对象，广泛应用于机器学习、深度学习和计算机视觉等领域。在深度学习框架（如 TensorFlow 和 PyTorch）中，张量是数据的基本结构。

而我们代码中低效的for循环，未来在GPU中都会被高效的并行计算。

我们先以传统思维来仔细看一下forward 函数的实现，以进一步理解“张量”和“批处理”。

    def forward(self, idx: List[List[int]]) -> List[List[List[float]]]:

forward 函数的入参是一个大小为B * T的二维数组，按照机器学习领域的说法，就是一个形状为(B, T)的“张量”，表示输入了“B”批次的数据，每个批次包含“T”个token。

这里B、T、C都是机器学习里的常用变量名，B（Batch Size）代表批量大小、T（Time Steps or Sequence Length）对于序列数据来说代表序列的长度、C（Channels）在图像处理中代表通道数，在语言模型中可以表示特征维度。

返回值logits 是一个形状为(B, T, C)的张量（C等于vocab\_size），它表示了“每个批次”的序列中，“每个token”的下一个token的频率。这么说起来很绕，其实只要想象成：“所有B*T个数的token，都有一张独立的表，表中记录了下一个出现的token是X的频率”。

logits 的大小为B * T * C，由于我们是Bigram模型，每个token的概率只和它上一个token有关，所以实际上我们只需要计算批次中最后一个token的logit就可以了，但为了和以后的模型统一，依旧保留了这些冗余计算。

好消息，我们现在已经有了一个能跑的玩具“模型”，它能根据概率预测下一个词，但却缺乏了真正的训练过程。

坏消息，在实现真正的机器学习之前，我们还是绕不开pytorch。不过幸运的是，我们只需要做到“知其然”即可。

5分钟简明pytorch教程

PyTorch 是一个开源的深度学习库，提供一系列非常方便的基础数据结构和函数，简化我们的操作。

下面是一个使用pytorch实现线性回归的简单例子：

pytorch\_5min.py

import torch

通过python pytorch_5min.py 即可运行：

$ python pytorch_5min.py

这个例子中，最特别的是有真正的“训练”过程，“训练”究竟是什么？我们经常听到的“反向传播”、“梯度下降”、“学习率”又是什么？

鉴于这只是5分钟教程，我们只要记住后面我们所有的机器学习代码都是这样的结构即可。

tensor操作

这一部分详见代码，看完代码后才发现，大学时候的《线性代数》课程是多么重要。

这里最值得注意的是“矩阵相乘”，即“点积”、matmul操作，简写为“@”符号，是后面self-attention机制的核心。

矩阵乘还经常用作张量形状的变换，如形状为(B, T, embd)的张量和形状为(embd, C)的张量相乘，结果为(B, T, C)的张量 —— 这一点也经常被用到。

此外，tensor.to(device) 可以把tensor数据移动到指定的设备，如GPU。

模型、神经网络的layer

我们的模型内部只有一个简单的线性层nn.Linear(1, 1) ，它输入输出都是一维张量。(1,1)的线性层实际上内部就是一个线性方程，对于输入任何数字x，它会输出x * w + b，实际上神经网络中的“layer”就是内含了一系列参数、可被训练的单元。通过输出nn.Linear 可以更清晰的看出实现。

>>> layer = nn.Linear(1, 1)

手动计算一下就能发现，实际上layer的输出值，就是输入x * weight + bias的结果。

其中grad_fn 是pytorch用来反向传播的关键，pytorch记住了这个tensor是怎么计算出来的，在后面的反向传播中被使用，对pytorch用户不可见。

反向传播和梯度下降

5分钟的教程只需要我们先硬记住一点，机器学习的“训练”就是这样一个过程：

1.先“前向传播”，计算出输出（如Linear层输出结果）。

2.再“反向传播”。

a.通过“损失函数”计算出模型的输出和真实数据之间的“损失值”loss（如例子中的MSELoss损失函数）；

b.计算“梯度”，利用损失函数对输出层的梯度进行计算，接着向前传播（反向传播）计算前一层的梯度，直到输入层（这一步pytorch能自动处理，不需要我们关心。可以简单理解为，“梯度”就是损失函数对各个参数的导数。核心目的就是为了计算出“如何调整w和b的值来减少损失”）；

c.更新参数，“梯度”是一个向量，把“梯度”乘上我们的“学习率”再加上原来的参数，就是我们新的参数了。如果学习率大，那么每次更新的多，学习率小，每次更新的就少。“梯度下降”，就是我们通过迭代更新参数，以寻找到损失函数最小的过程；

这中间最复杂的求导、算梯度、更新每一层参数的操作，pytorch都自动完成了（前面看到的grad_fn 就是用于这个过程），我们只需要知道在这个结构下，选择不同的优化器算法、损失函数实现、模型结构即可，剩下的交给pytorch。

而“推理”，就只有“前向传播”，计算出输出即可。

实现一个真正的Bigram模型

5分钟“精通”完pytorch，接下来我们来实现真正的pytorch版Bigram模型。

首先，我们把前面的simplebigrammodel.py ，用pytorch的tensor数据结构改造成一个新版本，代码见simplebigrammodel_torch.py [5]，这里不再展开。通过这份代码，能在熟悉算法的基础上，进一步深刻理解tensor。

然后，我们基于它进一步实现Bigram模型，后续我们的代码都将基于这个为基础，逐渐改出完整的gpt。

完整代码如下，也可以看babygpt\_v1.py[6]。

babygpt\_v1.py

import torch

在我的mac上通过 python babygpt_v1.py 运行，大概60k t/s的训练速度，而在4090上这个速度可以达到180k t/s。

$ python babygpt_v1.py

这份代码也没有难点，实际上就是前面pytorch实现的线性回归模型和我们自己土方法实现的bigram模型的结合体，尤其是训练部分，基本上和前面线性回归是一样的，差别主要在模型上。

模型

Embedding层

这次我们的模型由一个nn.Embedding(vocab_size, n_embd) 层和一个nn.Linear(n_embd, vocab_size) 层组成。

nn.Embedding(vocab\_size, n\_embd) 可以简单理解成一个映射表，只不过它的key取值为0 ~ vocab\_size-1，而它的value是一个n\_embd维的参数。简单的理解为，通过embedding操作（嵌入操作），我们把一个离散的token，映射为了一个密集的向量。

实际上Embedding的实现真的就是一个lookup-table，如下所示：

>>> layer = nn.Embedding(10, 3)

Embedding内部就是保存了一个(vocab_size, n_embd)的张量，“对tensor X执行嵌入操作”和“在weight中取key为X的值”效果是一样的。

Embedding通常作为各种模型的第一层，因为我们要把离散的“token”，映射为一些连续的“数值”，才可以继续后续的操作。两个token id之间是没有关系的，但两个Embedding的向量可以有距离、关联度等关系。

由于我们只实现了一个Bigram模型，下一个词只和上一个词有关，而Embedding内部恰好能表示一种A到B的映射关系，所以这里我们的模型主体就是Embedding本身，我们训练的直接就是Embedding内的参数。

lm\_head层

lm\_head(Language Model Head)是我们的输出层，几乎所有模型最后一层都是这么一个Linear 层，它的用途是把我们中间各种layer算出来的结果，最终映射到vocab_size 维的向量里去。因为我们最终要算的，就是vocab_size 个词里，每个词出现的概率。

语言模型的常见流程如下示意图，模型间主要的差异都在中间层上，LLM也不例外：

损失函数、归一函数和采样

在forward 实现中，我们使用交叉熵函数作为损失函数，且为了满足交叉熵函数对于参数的要求，我们把(B, T, C)的张量，变形为(B * T, C)，不需要理解交叉熵函数计算方式，只需知道它得出了两个tensor的差值即可。

我们使用softmax 代替前面的线性归一函数做归一化，也省去了考虑total 值为0的情况，并且用torch.multinomial 代替random.choices 作为采样函数。

训练

训练部分代码和5分钟pytorch教程中的没太多差别，我们用AdamW 优化器替换了SGD 优化器，具体原因这里不展开解释，只要知道这就是不一样的调整参数的算法即可。

并且我们每处理一些数据，就尝试输出当前模型，在训练数据和校验数据上的损失值。以便我们观察模型是否过拟合了训练数据。

如果数据足够多、耗时足够久的话，我们在这里可以用torch.save 方法把参数保存下来，也就是checkpoint。

回顾和Next

令人兴奋，目前为止，我们用131行python代码，实现了一个语言模型，居然能生成看起来像是词的东西，It just works。

这个模型目前参数量为 Embedding层:6148 (vocab_size) * 32 (n_embd) + Linear层6148 * 32 + 6148 = 399620 ，消耗399620 * 4字节 = 1.52MB 空间，即一个0.0004B的参数，而qwen2.5最小的也是0.5B。

我们亲眼看到了模型的参数、layer、学习率、正向传播、反向传播、梯度等一堆概念。

如果对于模型流程和结构没太理解，可以问AI实现各种简单的demo，会发现结构大差不大；如果对于中间各种变量转换没太理解，强烈建议在调试中通过.shape 观察各种tensor的形状变化、通过.weight 观察各个layer的参数变量，来体会其中的细节。

下一篇文章，我们会基于babygpt_v1.py 开始实现“自注意力机制”，进而实现完整的GPT，Happy Hacking。