@tom_doerr: 在单个 GPU 上从头训练十亿参数的大语言模型 https://github.com/FareedKhan-dev/train-llm-from-scratch…

使用单个GPU从头训练十亿参数级大型语言模型 https://github.com/FareedKhan-dev/train-llm-from-scratch… — # FareedKhan-dev/train-llm-from-scratch 来源：https://github.com/FareedKhan-dev/train-llm-from-scratch 主图 # 从头训练LLM Python 许可证 贡献 文档 我正在寻找人工智能领域的PhD职位。请查看我的简历 (https://drive.google.com/file/d/1Q_iklJ1RVGSb-Pdey8BHy3k8IF3UJv0z/view?usp=sharing) 或 GitHub (https://github.com/FareedKhan-dev) 我基于论文《注意力就是一切》(https://arxiv.org/abs/1706.03762)（Attention is All You Need），使用PyTorch从头实现了Transformer模型。你可以使用我的脚本，在单张GPU上训练你自己的十亿或百万参数级别的LLM。以下是训练好的1300万参数LLM的输出示例： 在 ***1978 年，公园被归还给了工厂板，公众共享于电子围栏的下方，该电子围栏跟随于车站城市。古代西方国家的运河被限制在城市地点。村庄直接与中国的城市相连，这些城市反抗美国预算，而在 Odambinais 则是不确定的，并建立在农村地区的财富之上。 ## 目录 - 训练数据信息 - 先决条件与训练时间 - 代码结构 - 使用方法 - 分步代码解释 - 导入库 - 准备训练数据 - Transformer 概述 - 多层感知机 (MLP) - 单头注意力 - 多头注意力 - Transformer 块 - 最终模型 - 批量处理 - 训练参数 - 训练模型 - 保存训练好的模型 - 训练损失 - 生成文本 - 下一步 ## 训练数据信息 训练数据来自于 Pile 数据集，这是一个多样化、开源、大规模的语言模型训练数据集。Pile 数据集包含 22 个不同的子数据集，包括书籍、文章、网站等文本。Pile 数据集总大小为 825 GB，以下是训练数据的一个样本： python 行号: 0 { "text": "睡眠质量对...癫痫的影响。", "meta": { "pile_set_name": "PubMed 摘要" } } 行号: 1 { "text": "LLMops 一个新的 GitHub 仓库 ...", "meta": { "pile_set_name": "Github" } } ## 先决条件与训练时间 请确保你具备面向对象编程（OOP）、神经网络（NN）和 PyTorch 的基本知识，以便理解代码。以下是一些入门资源： | 主题 | 视频链接 | |———————|———————————————————————| | 面向对象编程 | OOP 视频 (https://www.youtube.com/watch?v=Ej_02ICOIgs&pp=ygUKb29wIHB5dGhvbg%3D%3D) | | 神经网络 | 神经网络视频 (https://www.youtube.com/watch?v=Jy4wM2X21u0&pp=ygUbbmV1cmFsIG5ldHdvcmsgcHl0aG9uIHRvcmNo) | | PyTorch | PyTorch 视频 (https://www.youtube.com/watch?v=V_xro1bcAuA&pp=ygUbbmV1cmFsIG5ldHdvcmsgcHl0aG9uIHRvcmNo) | 你需要一张 GPU 来训练模型。Colab 或 Kaggle 的 T4 可以用于训练 1300 万以上参数的模型，但对于十亿参数的训练则会失败。请参考以下对比表： | GPU 名称 | 显存 | 数据大小 | 2B LLM 训练 | 13M LLM 训练 | 最大实用 LLM 规模（训练） | |–––––––––––––|––––|———–|———––|–––––––|–––––––––––––| | NVIDIA A100 | 40 GB | 大 | ✔ | ✔ | ~6B–8B | | NVIDIA V100 | 16 GB | 中 | ✘ | ✔ | ~2B | | AMD Radeon VII | 16 GB | 中 | ✘ | ✔ | ~1.5B–2B | | NVIDIA RTX 3090 | 24 GB | 大 | ✔ | ✔ | ~3.5B–4B | | Tesla P100 | 16 GB | 中 | ✘ | ✔ | ~1.5B–2B | | NVIDIA RTX 3080 | 10 GB | 中 | ✘ | ✔ | ~1.2B | | AMD RX 6900 XT | 16 GB | 大 | ✘ | ✔ | ~2B | | NVIDIA GTX 1080 Ti | 11 GB | 中 | ✘ | ✔ | ~1.2B | | Tesla T4 | 16 GB | 小 | ✘ | ✔ | ~1.5B–2B | | NVIDIA Quadro RTX 8000 | 48 GB | 大 | ✔ | ✔ | ~8B–10B | | NVIDIA RTX 4070 | 12 GB | 中 | ✘ | ✔ | ~1.5B | | NVIDIA RTX 4070 Ti | 12 GB | 中 | ✘ | ✔ | ~1.5B | | NVIDIA RTX 4080 | 16 GB | 中 | ✘ | ✔ | ~2B | | NVIDIA RTX 4090 | 24 GB | 大 | ✔ | ✔ | ~4B | | NVIDIA RTX 4060 Ti | 8 GB | 小 | ✘ | ✔ | ~1B | | NVIDIA RTX 4060 | 8 GB | 小 | ✘ | ✔ | ~1B | | NVIDIA RTX 4050 | 6 GB | 小 | ✘ | ✔ | ~0.75B | | NVIDIA RTX 3070 | 8 GB | 小 | ✘ | ✔ | ~1B | | NVIDIA RTX 3060 Ti | 8 GB | 小 | ✘ | ✔ | ~1B | | NVIDIA RTX 3060 | 12 GB | 中 | ✘ | ✔ | ~1.5B | | NVIDIA RTX 3050 | 8 GB | 小 | ✘ | ✔ | ~1B | | NVIDIA GTX 1660 Ti | 6 GB | 小 | ✘ | ✔ | ~0.75B | | AMD RX 7900 XTX | 24 GB | 大 | ✔ | ✔ | ~3.5B–4B | | AMD RX 7900 XT | 20 GB | 大 | ✔ | ✔ | ~3B | | AMD RX 7800 XT | 16 GB | 中 | ✘ | ✔ | ~2B | | AMD RX 7700 XT | 12 GB | 中 | ✘ | ✔ | ~1.5B | | AMD RX 7600 | 8 GB | 小 | ✘ | ✔ | ~1B | 13M LLM 训练指的是训练 1300 万以上参数的模型，2B LLM 训练指的是训练 20 亿以上参数的模型。数据大小分为小、中、大三类。小数据大小约为 1 GB，中数据大小约为 5 GB，大数据大小约为 10 GB。 ## 代码结构 代码库的组织结构如下： bash train-llm-from-scratch/ ├── src/ │ ├── models/ │ │ ├── mlp.py # 多层感知机模块的定义 │ │ ├── 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 模型、多层感知机、注意力机制和 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 亿参数的模型，专门针对我们的特定用例，尤其是在安全地运行私有数据时。我建议你首先使用我的 GitHub 仓库中的脚本训练一个 1300 万以上参数的模型。你将在一天内获得结果，而不是等待更长时间，或者如果你的本地 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 而不是三个。数据已经分为 train/val/test。下载完成后，确保将文件正确放置在对应的目录中。 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）压缩。让我们读取其中一个下载文件的样本，看看它的样子。 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 #### 输出 #### 行号: 0 { "text": "睡眠质量对...癫痫的影响。", "meta": { "pile_set_name": "PubMed 摘要" } } 行号: 1 { "text": "LLMops 一个新的 GitHub 仓库 ...", "meta": { "pile_set_name": "Github" } } 现在我们需要对数据集进行编码（tokenize）。我们的目标是让 LLM 至少能输出正确的单词。为此，我们需要使用一个已有的 tokenizer。我们将使用 OpenAI 的开源 tokenizer tiktoken。我们将使用 r50k_base tokenizer，它用于 ChatGPT (GPT-3) 模型，来对我们的数据集进行 tokenize。我们需要创建一个函数来避免重复，因为我们同时要对训练数据集和验证数据集进行 tokenize。 ``python def process_files(input_dir, output_file): “”“ 处理指定输入目录中的所有 .zst 文件，并将编码后的 token 保存至 HDF5 文件。 参数: input_dir (str): 包含输入 .zst 文件的目录