Show HN: 使用图形着色器实现的 GPT-2
该项目基于 WebGL2,在浏览器中实现了 GPT-2 small 模型,并可视化了转换块和注意力矩阵。它利用 GPU 加速推理,无需 WASM 即可在浏览器中使用 `js-tiktoken` 进行 BPE 分词。项目包含 Python 脚本用于下载预训练权重,并使用 Vite 打包前端代码。 开发者提供了详细的安装和运行步骤,以及项目结构说明。 该项目遵循 MIT 许可证。
nathan-barry / **gpt2-webgl ** Public
基于浏览器的、使用 WebGL2 实现的 GPT-2,带有转换块和注意力矩阵可视化。
GPT-2 WebGL 推理演示
基于浏览器的 WebGL2 实现的 GPT-2。
🚀 特性
- 完整的 GPT-2 small (117M) 前向传递,通过 WebGL2 着色器在 GPU 中进行。
- 使用
js-tiktoken
在浏览器中进行 BPE 分词 (无需 WASM 获取)。 - 简单的 Python 脚本来下载预训练权重。
📋 前置条件
- Node.js ≥ 16.x 和 npm
- Python ≥ 3.8
- 支持 WebGL2 的现代浏览器 (Chrome, Firefox, Safari, Edge)。
🐍 下载 GPT-2 权重
我们依赖 HuggingFace 的 transformers
来拉取官方 GPT-2 权重,并生成原始的 Float32Array
blobs:
-
安装 Python 依赖:
pip install torch numpy transformers
-
运行下载器:
python download_weights.py
这将获取:
wte.bin
(token embeddings)wpe.bin
(positional embeddings)c_attn_q_w_0.bin
…c_attn_q_w_11.bin
c_attn_k_w_0.bin
… etc.lm_head_w.bin
,lm_head_b.bin
- 以及生成的
manifest.json
,映射名称 → URL。
⚙️ 使用 Vite 设置前端
我们使用 Vite 来打包 TS,提供 ESM 模块 & 处理 js-tiktoken
:
-
安装 JS 依赖:
npm install
-
启动本地开发服务器:
npm run dev
-
在浏览器中打开
http://localhost:5173
src/
下的任何更改都将触发 HMR 和实时重新加载。
📁 项目结构
.
├── public/ # 静态资源,在 `/` 下提供
│ └── weights/ # GPT-2 权重二进制文件 + manifest.json
├── src/
│ ├── gpt2_webgl.ts # WebGL2 推理 + 着色器 + 分词器
│ └── main.ts # 引导程序:加载 manifest, 设置 UI
├── download_weights.py # Python 脚本,用于获取 & 转储权重
├── index.html # (由 Vite 复制) 入口 HTML
├── vite.config.ts # Vite 配置
├── package.json
└── tsconfig.json
📄 许可证
MIT