从零使用纯Python实现OpenAI的GPT-OSS大型语言模型

GPT-OSS从零开始，使用纯Python实现OpenAI的大型语言模型，提供一个教育性强、深入解析大语言模型内部工作原理的开源实践，完整实现了GPT-OSS的Transformer架构，用C++实现的自定义BFloat16数据类型以确保数值精度，集成了专家混合（Mixture of Experts, MoE）架构，采用带有NTK感知缩放的旋转位置嵌入（RoPE），引入了支持注意力汇和滑动窗口的分组查询注意力（Grouped Query Attention），提供功能性的SwiGLU、RMSNorm、Softmax和线性层。GPT-OSS能够与PyTorch参考实现保持一致，通过逐token验证测试来保障其实现质量。

• 代码清晰，附带详细注释，适合学习研究

• 数值精度与 PyTorch 参考实现完全匹配

• 提供逐 token 验证机制，确保与参考模型结果一致

• 模块化设计，代码结构清晰，方便修改和扩展

• 核心功能不依赖 PyTorch，安装灵活轻便

基础参数配置

• 模型规模：200 亿参数解码器架构

• transformer 块数量：36 层

• 隐藏层维度：2880

• 注意力机制：分组查询注意力 + 滑动窗口

• 前馈网络：32 个专家的混合专家系统（MoE），采用 top-4 路由策略

核心技术组件

• C++ 自定义 BFloat16 实现，保障数值计算精度

• 混合专家系统（MoE），提升模型并行计算效率

• 支持 NTK 感知缩放的旋转位置编码（RoPE）

• 带注意力槽和滑动窗口的分组查询注意力

• 功能性 SwiGLU 激活函数、RMSNorm 归一化、Softmax 函数及线性层

模型运行流程

输入 Token
↓
嵌入层（词汇表大小 → 隐藏层大小）
↓
┌─────────────────────────────────────┐
│ Transformer 块（×36） │
│ ┌───────────────────────────────┐ │
│ │ 注意力机制 │ │
│ │ • RMSNorm 归一化 │ │
│ │ • QKV 投影 │ │
│ │ • RoPE 位置编码 │ │
│ │ • 缩放点积运算 │ │
│ │ • 输出投影 │ │
│ │ • 残差连接 │ │
│ └───────────────────────────────┘ │
│ ┌───────────────────────────────┐ │
│ │ 混合专家系统 │ │
│ │ • RMSNorm 归一化 │ │
│ │ • 专家路由（top-4） │ │
│ │ • 专家计算 │ │
│ │ • 加权组合 │ │
│ │ • 残差连接 │ │
│ └───────────────────────────────┘ │
└─────────────────────────────────────┘
↓
最终 RMSNorm 归一化
↓
解嵌入层（隐藏层大小 → 词汇表大小）
↓
输出 Logits

GPT-OSS 运行环境要求

仅支持 Ubuntu 22.04 或 Ubuntu 24.04 系统

GPT-OSS 安装

1、克隆仓库

git clone https://github.com/projektjoe/gptoss.git
cd gptoss

2、安装系统依赖

sudo apt update
sudo apt install -y \
python3-dev \
libopenblas-dev \
build-essential \
libdnnl-dev \
cmake

3、配置 Python 环境并安装

方式 A：基础安装（默认无 PyTorch 支持）

1、安装 UV 包管理器

curl -LsSf https://astral.sh/uv/install.sh | sh

2、重启终端后创建并激活虚拟环境

uv venv .venv
source .venv/bin/activate

3、安装项目（自动构建 C++ 扩展）

uv pip install -e .

方式 B：带 PyTorch 支持的安装（可选）

如需通过 PyTorch 层匹配 OpenAI 官方实现的数值精度，可按以下步骤操作：

1、完成方式 A 中的所有步骤

2、安装 PyTorch 及相关依赖

# 方法 1：先安装 PyTorch，再禁用构建隔离
uv pip install torch
uv pip install scikit_build_core
uv pip install --no-build-isolation -e .

# 方法 2：通过环境变量禁用构建隔离
export UV_NO_BUILD_ISOLATION=1
uv pip install -e ".[torch]"

4、下载模型权重

从 Hugging Face Hub 下载 GPTOSS-20B 模型权重，放置在项目根目录：

hf download openai/gpt-oss-20b --include "original/*" --local-dir gpt-oss-20b/

5、运行项目

# 启动主程序
python main.py

# 运行测试（验证与官方实现的一致性）
python test/test.py

测试说明

测试脚本支持两种验证模式，可通过修改 VERIFY_LAYER_BY_LAYER 参数控制：

• VERIFY_LAYER_BY_LAYER = True：将官方实现的输出作为下一层输入，逐层隔离测试

• VERIFY_LAYER_BY_LAYER = False：测试整个模型，错误会向后传播

测试会验证以下模块的数值一致性：

• 嵌入层查找 • RMSNorm 计算 • QKV 投影 • RoPE 应用 • 注意力机制 • MoE 路由与专家计算 • 最终 Logits 输出

测试输出示例

[OK] block[0].attn.norm token 0 passed.
[OK] qkv layer 0 token 0 passed.
[OK] rope q layer 0 token 0 passed.
[OK] rope k layer 0 token 0 passed.
[OK] att layer 0 token 0 passed.
[OK] linear & residual layer 0 token 0 passed.
[OK] gate layer 0 token 0 passed.
[OK] moe layer 0 token 0 passed.