100行Numpy代码实现GPT2推理：从Transformer原理到KV缓存优化

GPT2推理Numpy实现Transformer架构

于 2026-07-07 15:20:38 修改

·本内容遵循CC 4.0 BY-SA版权协议

这次我们来看一个很有意思的开源项目——不用任何深度学习框架，仅用100行Numpy代码实现GPT2推理引擎。这个项目对于想深入理解大语言模型底层原理的开发者来说，是个很好的学习材料。

项目重点不是追求高性能，而是通过最基础的Numpy操作还原GPT2的核心推理过程。如果你对Transformer架构、自注意力机制、KV缓存等概念感兴趣，但觉得直接读Huggingface源码太复杂，这个精简实现能帮你快速抓住关键点。

本文会带你完成这个Numpy版GPT2的完整部署和测试流程。我们会从环境准备开始，逐步讲解模型加载、前向推理、文本生成等关键步骤，最后还会讨论如何在这个基础上添加KV缓存优化。整个代码只有100行左右，但包含了GPT2推理的核心逻辑。

1. 核心能力速览

能力项	说明
项目类型	教育性质的技术演示项目
核心功能	GPT2模型推理、文本生成
依赖框架	仅需Numpy，无需PyTorch/TensorFlow
模型支持	GPT2-small (124M参数)
硬件要求	CPU即可运行，无GPU要求
代码规模	约100行核心代码
适合场景	学习Transformer原理、理解KV缓存、教学演示

2. 适用场景与使用边界

这个项目主要适合以下几类开发者：

适合场景：

想从零理解GPT2推理过程的初学者
需要教学演示材料的讲师或技术作者
希望了解大模型底层计算逻辑的工程师
在资源受限环境中验证基础推理逻辑

使用边界：

不支持模型训练，仅用于推理演示
性能远低于优化框架，不适合生产环境
仅实现基础生成，不支持批量推理
模型规模有限，只包含GPT2-small

重要提醒： 虽然这是教育项目，但涉及的语言模型技术同样需要注意合规使用。生成内容时请遵守相关法律法规，避免产生不当内容。

3. 环境准备与前置条件

准备环境很简单，只需要基础的Python环境：

3.1 系统要求

操作系统：Windows/Linux/macOS均可
Python版本：3.7及以上
内存：至少2GB可用内存（用于加载模型权重）

3.2 必要依赖

只需要安装Numpy库：

BASH

pip install numpy

3.3 模型文件准备

项目需要GPT2的模型权重文件，可以从Huggingface转换获取：

PYTHON

# 需要先安装transformers库来获取权重

pip install transformers torch

4. 模型权重获取与转换

由于项目只使用Numpy，我们需要先将PyTorch格式的权重转换为Numpy格式：

4.1 权重转换脚本

PYTHON

import numpy as np

import torch

from transformers import GPT2LMHeadModel, GPT2Tokenizer

# 加载原始GPT2模型和tokenizer

model = GPT2LMHeadModel.from_pretrained('gpt2')

tokenizer = GPT2Tokenizer.from_pretrained('gpt2')

# 创建权重字典存储numpy数组

numpy_weights = {}

# 转换所有权重

for name, param in model.named_parameters():

numpy_weights[name] = param.detach().numpy()

# 保存为npz文件

np.savez('gpt2_weights.npz', **numpy_weights)

# 测试保存结果

weights = np.load('gpt2_weights.npz')

print("权重文件包含的数组:", list(weights.keys())[:5]) # 打印前5个键名

4.2 关键权重结构说明

GPT2的权重主要包含以下几类：

wte.weight: 词嵌入矩阵
wpe.weight: 位置编码矩阵
h.{层号}.attn.c_attn.weight: 注意力层的QKV投影矩阵
h.{层号}.attn.c_proj.weight: 注意力输出投影矩阵
h.{层号}.mlp.c_fc.weight: MLP层第一层权重
h.{层号}.mlp.c_proj.weight: MLP层第二层权重

5. 核心推理引擎实现

下面是精简版的GPT2推理引擎核心代码：

5.1 基础架构定义

PYTHON

import numpy as np

class NumpyGPT2:

def __init__(self, weights_path):

self.weights = np.load(weights_path)

self.n_layer = 12 # GPT2-small有12层

self.n_head = 12 # 12个注意力头

self.d_model = 768 # 隐藏层维度

def layer_norm(self, x, g, b, eps=1e-5):

# 层归一化实现

mean = np.mean(x, axis=-1, keepdims=True)

variance = np.var(x, axis=-1, keepdims=True)

x = (x - mean) / np.sqrt(variance + eps)

return g * x + b

def softmax(self, x):

# 稳定的softmax实现

exp_x = np.exp(x - np.max(x, axis=-1, keepdims=True))

return exp_x / np.sum(exp_x, axis=-1, keepdims=True)

def gelu(self, x):

# GELU激活函数近似

return 0.5 * x * (1 + np.tanh(np.sqrt(2 / np.pi) * (x + 0.044715 * np.power(x, 3))))

5.2 自注意力机制实现

PYTHON

def self_attention(self, x, layer_idx):

# 计算QKV投影

c_attn_weight = self.weights[f'h.{layer_idx}.attn.c_attn.weight']

c_attn_bias = self.weights[f'h.{layer_idx}.attn.c_attn.bias']

qkv = np.dot(x, c_attn_weight.T) + c_attn_bias

# 分割Q、K、V

d_head = self.d_model // self.n_head

qkv = qkv.reshape(x.shape[0], x.shape[1], 3, self.n_head, d_head)

q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]

# 计算注意力分数

attn_scores = np.matmul(q, k.swapaxes(-1, -2)) / np.sqrt(d_head)

attn_probs = self.softmax(attn_scores)

# 应用注意力权重到V

attn_output = np.matmul(attn_probs, v)

attn_output = attn_output.reshape(x.shape[0], x.shape[1], self.d_model)

# 输出投影

c_proj_weight = self.weights[f'h.{layer_idx}.attn.c_proj.weight']

c_proj_bias = self.weights[f'h.{layer_idx}.attn.c_proj.bias']

return np.dot(attn_output, c_proj_weight.T) + c_proj_bias

5.3 前向传播完整流程

PYTHON

def forward(self, input_ids, past_key_values=None):

# 词嵌入

wte = self.weights['wte.weight']

token_embeddings = wte[input_ids]

# 位置编码

wpe = self.weights['wpe.weight']

position_ids = np.arange(len(input_ids))

position_embeddings = wpe[position_ids]

# 组合嵌入

x = token_embeddings + position_embeddings

# 逐层处理

for layer_idx in range(self.n_layer):

# 层归一化1

ln1_gamma = self.weights[f'h.{layer_idx}.ln_1.weight']

ln1_beta = self.weights[f'h.{layer_idx}.ln_1.bias']

x_ln1 = self.layer_norm(x, ln1_gamma, ln1_beta)

# 自注意力

attn_output = self.self_attention(x_ln1, layer_idx)

x = x + attn_output # 残差连接

# 层归一化2

ln2_gamma = self.weights[f'h.{layer_idx}.ln_2.weight']

ln2_beta = self.weights[f'h.{layer_idx}.ln_2.bias']

x_ln2 = self.layer_norm(x, ln2_gamma, ln2_beta)

# MLP层

c_fc_weight = self.weights[f'h.{layer_idx}.mlp.c_fc.weight']

c_fc_bias = self.weights[f'h.{layer_idx}.mlp.c_fc.bias']

mlp_hidden = np.dot(x_ln2, c_fc_weight.T) + c_fc_bias

mlp_hidden = self.gelu(mlp_hidden)

c_proj_weight = self.weights[f'h.{layer_idx}.mlp.c_proj.weight']

c_proj_bias = self.weights[f'h.{layer_idx}.mlp.c_proj.bias']

mlp_output = np.dot(mlp_hidden, c_proj_weight.T) + c_proj_bias

x = x + mlp_output # 残差连接

# 最终层归一化

ln_f_gamma = self.weights['ln_f.weight']

ln_f_beta = self.weights['ln_f.bias']

x = self.layer_norm(x, ln_f_gamma, ln_f_beta)

# 输出投影

output = np.dot(x, wte.T)

return output

6. 文本生成功能实现

有了前向传播，接下来实现文本生成逻辑：

6.1 基础生成函数

PYTHON

def generate(self, prompt, max_length=50, temperature=1.0):

# 简单的tokenization（实际使用时需要完整的tokenizer）

# 这里假设输入已经是token ID列表

if isinstance(prompt, str):

# 简易空格分词，实际项目应该使用完整tokenizer

words = prompt.split()

input_ids = [hash(word) % 1000 for word in words] # 模拟token化

else:

input_ids = prompt

generated = input_ids.copy()

for _ in range(max_length):

# 前向传播

logits = self.forward(np.array([generated]))

# 获取最后一个token的logits

next_token_logits = logits[0, -1, :]

# 应用温度调节

if temperature != 1.0:

next_token_logits = next_token_logits / temperature

# softmax计算概率

probs = self.softmax(next_token_logits)

# 基于概率采样

next_token = np.random.choice(len(probs), p=probs)

generated.append(int(next_token))

# 简单停止条件（遇到特定token停止）

if next_token == 0: # 假设0是结束token

break

return generated

6.2 使用示例

PYTHON

# 初始化模型

model = NumpyGPT2('gpt2_weights.npz')

# 生成文本

prompt = "人工智能是"

output_ids = model.generate(prompt, max_length=20, temperature=0.8)

print("生成的token IDs:", output_ids)

7. KV缓存优化实现

为了提高生成效率，我们可以实现KV缓存机制：

7.1 带KV缓存的注意力实现

PYTHON

def self_attention_with_cache(self, x, layer_idx, past_key_values=None):

c_attn_weight = self.weights[f'h.{layer_idx}.attn.c_attn.weight']

c_attn_bias = self.weights[f'h.{layer_idx}.attn.c_attn.bias']

qkv = np.dot(x, c_attn_weight.T) + c_attn_bias

d_head = self.d_model // self.n_head

qkv = qkv.reshape(x.shape[0], x.shape[1], 3, self.n_head, d_head)

q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]

if past_key_values is not None:

# 合并过去的KV和当前的KV

past_k, past_v = past_key_values

k = np.concatenate([past_k, k], axis=1)

v = np.concatenate([past_v, v], axis=1)

# 保存当前KV用于下一次生成

new_past_key_values = (k, v)

# 计算注意力（只计算最后一个token的注意力）

attn_scores = np.matmul(q[:, -1:], k.swapaxes(-1, -2)) / np.sqrt(d_head)

attn_probs = self.softmax(attn_scores)

attn_output = np.matmul(attn_probs, v)

attn_output = attn_output.reshape(x.shape[0], 1, self.d_model)

c_proj_weight = self.weights[f'h.{layer_idx}.attn.c_proj.weight']

c_proj_bias = self.weights[f'h.{layer_idx}.attn.c_proj.bias']

output = np.dot(attn_output, c_proj_weight.T) + c_proj_bias

return output, new_past_key_values

7.2 带缓存的生成函数

PYTHON

def generate_with_cache(self, prompt, max_length=50, temperature=1.0):

input_ids = prompt if isinstance(prompt, list) else [hash(prompt) % 1000]

generated = input_ids.copy()

# 初始化KV缓存

past_key_values = [None] * self.n_layer

for step in range(max_length):

# 如果是第一步，处理整个序列；否则只处理最后一个token

if step == 0:

current_input = np.array([generated])

else:

current_input = np.array([[generated[-1]]])

x = self._forward_with_cache(current_input, past_key_values, step == 0)

# 后续逻辑与基础生成函数相同...

logits = np.dot(x, self.weights['wte.weight'].T)

next_token_logits = logits[0, -1, :]

if temperature != 1.0:

next_token_logits = next_token_logits / temperature

probs = self.softmax(next_token_logits)

next_token = np.random.choice(len(probs), p=probs)

generated.append(int(next_token))

if next_token == 0:

break

return generated

def _forward_with_cache(self, input_ids, past_key_values, is_first_step):

# 简化的带缓存前向传播

wte = self.weights['wte.weight']

wpe = self.weights['wpe.weight']

token_embeddings = wte[input_ids]

position_ids = np.arange(past_key_values[0][0].shape[1] if not is_first_step else 0,

past_key_values[0][0].shape[1] + input_ids.shape[1])

position_embeddings = wpe[position_ids]

x = token_embeddings + position_embeddings

new_past_key_values = []

for layer_idx in range(self.n_layer):

# ... 层归一化等操作

if is_first_step:

attn_output = self.self_attention(x, layer_idx)

else:

attn_output, new_past = self.self_attention_with_cache(

x, layer_idx, past_key_values[layer_idx])

new_past_key_values.append(new_past)

# ... 残差连接和MLP处理

return x

8. 性能测试与效果验证

8.1 基础功能测试

PYTHON

def test_basic_functionality():

"""测试模型基础功能"""

model = NumpyGPT2('gpt2_weights.npz')

# 测试前向传播

test_input = [1, 2, 3, 4] # 示例token IDs

output = model.forward(np.array([test_input]))

print("输入形状:", np.array([test_input]).shape)

print("输出形状:", output.shape)

print("输出范围: [{:.3f}, {:.3f}]".format(np.min(output), np.max(output)))

# 测试生成功能

generated = model.generate([1, 2, 3], max_length=10)

print("生成结果:", generated)

if __name__ == "__main__":

test_basic_functionality()

8.2 内存占用分析

PYTHON

import psutil

import os

def analyze_memory_usage():

"""分析内存使用情况"""

process = psutil.Process(os.getpid())

memory_before = process.memory_info().rss / 1024 / 1024 # MB

model = NumpyGPT2('gpt2_weights.npz')

memory_after_load = process.memory_info().rss / 1024 / 1024

# 执行一次推理

model.forward(np.array([[1, 2, 3, 4]]))

memory_after_inference = process.memory_info().rss / 1024 / 1024

print(f"加载前内存: {memory_before:.1f} MB")

print(f"加载模型后: {memory_after_load:.1f} MB")

print(f"推理后内存: {memory_after_inference:.1f} MB")

print(f"模型占用: {memory_after_load - memory_before:.1f} MB")

9. 常见问题与排查方法

问题现象	可能原因	排查方式	解决方案
导入错误：找不到权重文件	文件路径错误或文件未生成	检查文件路径和生成脚本	重新运行权重转换脚本
形状不匹配错误	权重形状与预期不符	打印各层权重形状	检查GPT2模型版本是否匹配
内存不足	模型太大或序列过长	监控内存使用	减小序列长度或使用内存映射
生成结果无意义	Tokenization不匹配	检查输入token化方式	实现完整tokenizer或使用匹配的ID
数值不稳定	浮点数精度问题	检查中间结果范围	添加数值稳定性处理

9.1 权重文件验证

PYTHON

def validate_weights(weights_path):

"""验证权重文件完整性"""

weights = np.load(weights_path)

required_keys = [

'wte.weight', 'wpe.weight', 'ln_f.weight', 'ln_f.bias'

]

# 添加各层需要的权重键

for i in range(12): # GPT2-small有12层

required_keys.extend([

f'h.{i}.ln_1.weight', f'h.{i}.ln_1.bias',

f'h.{i}.attn.c_attn.weight', f'h.{i}.attn.c_attn.bias',

f'h.{i}.attn.c_proj.weight', f'h.{i}.attn.c_proj.bias',

f'h.{i}.ln_2.weight', f'h.{i}.ln_2.bias',

f'h.{i}.mlp.c_fc.weight', f'h.{i}.mlp.c_fc.bias',

f'h.{i}.mlp.c_proj.weight', f'h.{i}.mlp.c_proj.bias'

])

missing_keys = [key for key in required_keys if key not in weights]

if missing_keys:

print("缺少的权重键:", missing_keys)

return False

print("权重文件验证通过")

return True

10. 扩展与优化建议

在这个基础版本之上，可以考虑以下几个优化方向：

10.1 性能优化

PYTHON

# 1. 使用内存映射减少内存占用

weights = np.load('gpt2_weights.npz', mmap_mode='r')

# 2. 批量推理支持

def batch_forward(self, batch_input_ids):

"""支持批量推理"""

batch_size = len(batch_input_ids)

max_len = max(len(ids) for ids in batch_input_ids)

# 填充批次

padded_batch = np.zeros((batch_size, max_len), dtype=np.int32)

for i, ids in enumerate(batch_input_ids):

padded_batch[i, :len(ids)] = ids

return self.forward(padded_batch)

10.2 功能扩展

完整的tokenizer集成：集成Huggingface tokenizer支持完整文本处理
采样策略：实现top-k、top-p等高级采样方法
停止条件：添加更智能的生成停止条件
长度惩罚：实现生成长度惩罚机制

10.3 工程化改进

模型序列化：添加模型保存和加载功能
配置化：通过配置文件管理模型参数
性能监控：添加推理时间和内存监控
单元测试：建立完整的测试覆盖

这个Numpy实现的GPT2虽然性能有限，但作为学习工具非常有价值。通过这100行代码，你可以清晰看到Transformer架构的每个关键组件如何工作，为后续理解更复杂的大模型打下坚实基础。

建议从基础版本开始，逐步添加KV缓存、批量处理等优化功能，这样可以深入理解每个优化技术背后的原理。