FareedKhan-dev/train-llm-from-scratch 来源：https://github.com/FareedKhan-dev/train-llm-from-scratch 主图 # 从头训练 LLM Python 许可证 贡献 文档 我正在寻找人工智能领域的博士职位。GitHub (https://github.com/FareedKhan-dev) 我使用 PyTorch 从头实现了一个 Transformer 模型，基于论文《Attention is All You Need》(https://arxiv.org/abs/1706.03762)。你可以使用我的脚本，在单个 GPU 上训练自己的十亿或百万参数规模的 LLM。以下是训练好的 1300 万参数 LLM 的输出：

在 ***1978 年，公园被归还给了工厂板，公众共享到电子围栏的下方，该围栏跟随车站的城市而来。古代西方文明的运河被限制在城市地点。村庄直接连接着中国的城市，这些城市反抗美国预算，而在 Odambinais，情况不确定，财富建立在农村地区。

目录

• 训练数据信息
• 前提条件与训练时间
• 代码结构
• 使用方法
• 逐步代码解释
  • 导入库
  • 准备训练数据
  • Transformer 概述
  • 多层感知机 (MLP)
  • 单头注意力
  • 多头注意力
  • Transformer 块
  • 最终模型
  • 批处理
  • 训练参数
  • 训练模型
  • 保存训练好的模型
  • 训练损失
  • 生成文本
• 下一步计划

训练数据信息

训练数据来自 Pile 数据集，这是一个多样化、开源、大规模的语言模型训练数据集。Pile 数据集包含 22 个不同的子数据集，包括书籍、文章、网站等文本。Pile 数据集总大小为 825 GB。以下是训练数据的样本：

python Line: 0 { "text": "睡眠质量对……癫痫的影响。", "meta": { "pile_set_name": "PubMed Abstracts" } } Line: 1 { "text": "LLMops 一个新的 GitHub 仓库……", "meta": { "pile_set_name": "Github" } }

前提条件与训练时间

请确保你具备面向对象编程 (OOP)、神经网络 (NN) 和 PyTorch 的基本知识，才能理解代码。以下是一些帮助你入门的资源：

你需要一个 GPU 来训练模型。Colab 或 Kaggle 的 T4 可以训练 1300 万+ 参数的模型，但在训练十亿参数模型时会失败。请看对比表：

13M LLM 训练指训练 1300 万+ 参数的模型，2B LLM 训练指训练 20 亿+ 参数的模型。数据规模分为小、中、大三类：小规模约 1 GB，中规模约 5 GB，大规模约 10 GB。

代码结构

代码库的组织结构如下：

bash train-llm-from-scratch/ ├── src/ │ ├── models/ │ │ ├── mlp.py # 多层感知机 (MLP) 模块定义 │ │ ├── attention.py # 注意力机制定义（单头、多头） │ │ ├── transformer_block.py # 单个 Transformer 块定义 │ │ └── transformer.py # 主 Transformer 模型定义 ├── config/ │ └── config.py # 默认配置（模型参数、文件路径等） ├── data_loader/ │ └── data_loader.py # 创建数据加载器/迭代器的函数 ├── scripts/ │ ├── train_transformer.py # 训练 Transformer 模型的脚本 │ ├── data_download.py # 下载数据集的脚本 │ ├── data_preprocess.py # 预处理下载数据的脚本 │ └── generate_text.py # 使用训练好的模型生成文本的脚本 ├── data/ # 存储数据集的目录 │ ├── train/ # 训练数据 │ └── val/ # 验证数据 └── models/ # 保存训练好的模型的目录

scripts/ 目录包含下载数据集、预处理数据、训练模型和生成文本的脚本。src/models/ 目录包含 Transformer 模型、MLP、注意力机制和 Transformer 块的实现。config/ 目录包含默认参数的配置文件。data_loader/ 目录包含创建数据加载器的函数。

使用方法

克隆仓库并进入目录：

bash git clone https://github.com/FareedKhan-dev/train-llm-from-scratch.git cd train-llm-from-scratch

如果遇到导入问题，请将 Python 路径设置为项目根目录：

``bash export PYTHONPATH=“${PYTHONPATH}:/path/to/train-llm-from-scratch”

或者如果你已经在 train-llm-from-scratch 目录中

export PYTHONPATH=“$PYTHONPATH:.” ``

安装所需依赖：

bash pip install -r requirements.txt

你可以在 src/models/transformer.py 中修改 Transformer 架构，在 config/config.py 中修改训练配置。

下载训练数据：

bash python scripts/data_download.py

脚本支持以下参数：

• --train_max：下载的最大训练文件数，默认 1（最大 30）。每个文件约 11 GB。
• --train_dir：训练数据存储目录，默认 data/train。
• --val_dir：验证数据存储目录，默认 data/val。

预处理下载的数据：

bash python scripts/data_preprocess.py

脚本支持以下参数：

• --train_dir：训练数据文件所在目录（默认 data/train）。
• --val_dir：验证数据文件所在目录（默认 data/val）。
• --out_train_file：处理后训练数据的 HDF5 文件路径（默认 data/train/pile_train.h5）。
• --out_val_file：处理后验证数据的 HDF5 文件路径（默认 data/val/pile_dev.h5）。
• --tokenizer_name：用于处理数据的 tokenizer 名称（默认 r50k_base）。
• --max_data：每个数据集（训练和验证）处理的最大 JSON 对象（行）数，默认 1000。

现在数据已预处理，你可以通过修改 config/config.py 中的配置来训练 1300 万参数的 LLM：

``python

定义词汇表大小和 Transformer 配置（30 亿）

VOCAB_SIZE = 50304 # 词汇表中的唯一 token 数 CONTEXT_LENGTH = 128 # 模型的最大序列长度 N_EMBED = 128 # 嵌入空间的维度 N_HEAD = 8 # 每个 Transformer 块的注意力头数 N_BLOCKS = 1 # 模型中的 Transformer 块数 ``

训练模型：

bash python scripts/train_transformer.py

这将开始训练模型，并将训练好的模型保存在 models/ 默认目录（或配置文件中指定的目录）。

使用训练好的模型生成文本：

bash python scripts/generate_text.py --model_path models/your_model.pth --input_text hi

脚本支持以下参数：

• --model_path：训练好的模型路径。
• --input_text：用于生成新文本的初始提示。
• --max_new_tokens：生成的最大 token 数（默认 100）。

它将基于输入提示使用训练好的模型生成文本。

逐步代码解释

本节面向希望详细了解代码的读者。我将从导入库到训练模型和生成文本，逐步解释代码。之前我在 Medium 上写过一篇关于使用 Tiny Shakespeare 数据集创建 230 万+ 参数 (https://levelup.gitconnected.com/building-a-million-parameter-llm-from-scratch-using-python-f612398f06c2) LLM 的文章，但输出没有意义。以下是输出样本：

``bash

230 万参数 LLM 输出

ZELBETH: Sey solmenter! tis tonguerered if Vurint as steolated have loven OID the queend refore Are been, good plmp: Proforne, wiftes swleen, was no blunderesd a a quain beath! Tybell is my gateer stalk smend as be matious dazest ``

我当时想，如果我把 Transformer 架构做得更小、更简单，同时让训练数据更多样化，结果会怎样？那么，一个普通人在近乎报废的 GPU 上，能训练出多大参数规模的模型，使其能够使用正确的语法并生成有意义的文本？我发现 1300 万+ 参数的模型已经足够在语法和标点方面开始有意义，这是一个积极的信号。这意味着我们可以使用一个非常特定的数据集，进一步微调之前训练的模型，以完成具体任务。最终我们可能会得到一个不到 10 亿参数，甚至约 5 亿参数的模型，完美适用于特定用例，尤其是在私有数据上安全运行。我建议你首先训练一个 1300 万+ 参数的模型，使用我 GitHub 仓库中的脚本。你将在一天内获得结果，而不是等待更长时间；或者如果你的本地 GPU 不够强大，无法训练十亿参数模型。

导入库

让我们导入本文中使用的必要库：

``python

PyTorch 用于深度学习和张量操作

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

数值计算和数组处理

import numpy as np

处理 HDF5 文件

import h5py

操作系统和文件管理

import os

命令行参数解析

import argparse

HTTP 请求和交互

import requests

循环进度条

from tqdm import tqdm

JSON 处理

import json

Zstandard 压缩库

import zstandard as zstd

大型语言模型的 Tokenization 库

import tiktoken

数学运算（用于高级数学函数）

import math ``

准备训练数据

我们的训练数据集需要多样化，包含来自不同领域的信息，而 Pile 是正确的选择。尽管它的大小为 825 GB，但我们只使用其中一小部分，即 5%–10%。让我们先下载数据集，看看它的结构。我将下载 HuggingFace (https://huggingface.co/datasets/monology/pile-uncopyrighted) 上的版本。

``python

下载验证数据集

!wget https://huggingface.co/datasets/monology/pile-uncopyrighted/resolve/main/val.jsonl.zst

下载训练数据集第一部分

!wget https://huggingface.co/datasets/monology/pile-uncopyrighted/resolve/main/train/00.jsonl.zst

下载训练数据集第二部分

!wget https://huggingface.co/datasets/monology/pile-uncopyrighted/resolve/main/train/01.jsonl.zst

下载训练数据集第三部分

!wget https://huggingface.co/datasets/monology/pile-uncopyrighted/resolve/main/train/02.jsonl.zst ``

下载需要一些时间，你也可以只限制训练数据集为一个文件 00.jsonl.zst 而不是三个。数据集已经划分为训练/验证/测试。下载完成后，确保将文件正确放入各自目录。

``python import os import shutil import glob

定义目录结构

train_dir = “data/train” val_dir = “data/val”

创建目录（如果不存在）

os.makedirs(train_dir, exist_ok=True) os.makedirs(val_dir, exist_ok=True)

移动所有训练文件（例如 00.jsonl.zst, 01.jsonl.zst, …）

train_files = glob.glob(“*.jsonl.zst”) for file in train_files: if file.startswith(“val”): # 移动验证文件 dest = os.path.join(val_dir, file) else: # 移动训练文件 dest = os.path.join(train_dir, file) shutil.move(file, dest) ``

我们的数据集是 .jsonl.zst 格式，这是一种常用于存储大型数据集的压缩文件格式。它结合了 JSON Lines (.jsonl，每行代表一个有效的 JSON 对象) 和 Zstandard (.zst) 压缩。让我们读取一个下载文件的样本，看看它的样子。

python in_file = "data/val/val.jsonl.zst" # 验证文件路径 with zstd.open(in_file, 'r') as in_f: for i, line in tqdm(enumerate(in_f)): # 读取前 5 行 data = json.loads(line) print(f"Line {i}: {data}") # 打印原始数据以检查 if i == 2: break

上述代码的输出如下：

``python

输出

Line: 0 { “text”: “睡眠质量对……癫痫的影响。”, “meta”: { “pile_set_name”: “PubMed Abstracts” } } Line: 1 { “text”: “LLMops 一个新的 GitHub 仓库……”, “meta”: { “pile_set_name”: “Github” } } ``

现在我们需要编码（tokenize）我们的数据集。我们的目标是让 LLM 至少能输出正确的单词。为此，我们需要使用一个已有的 tokenizer。我们将使用 OpenAI 的开源 tokenizer tiktoken，具体是用于 ChatGPT (GPT-3) 模型的 r50k_base tokenizer。我们需要创建一个函数来避免重复，因为我们将同时 tokenize 训练和验证数据集。

python def process_files(input_dir, output_file): """ 处理指定输入目录中的所有 .zst 文件，并将编码后的 tokens 保存到 HDF5 文件中。 参数: input_dir (str): 包含输入 .zst 文件的目录。 output_file (str): 输出 HDF5 文件的路径。 """ with h5py.File(output_file, 'w') as out_f: # 在 HDF5 文件中创建一个可扩展的数据集，命名为 'tokens' dataset = out_f.create_dataset('tokens', (0,), maxshape=(None,), dtype='i')