项目简介
基于《Attention is All You Need》论文，从零构建Transformer模型的教学教程。涵盖数据准备、模型架构、在Pile数据集上训练以及文本生成。适合理解大语言模型内部原理的实践代码。
README 预览
\n\n\n\n\n# Train LLM From Scratch\n \n [](#step-by-step-code-explanation)\n\n**I am Looking for a PhD position in AI**. [GitHub](https://github.com/FareedKhan-dev)\n\n\n\nI implemented a transformer model from scratch using PyTorch, based on the paper [Attention is All You Need](https://arxiv.org/abs/1706.03762). You can use my scripts to train your own **billion** or **million** parameter LLM using a single GPU.\n\nBelow is the output of the trained 13 million parameter LLM:\n\n```\nIn ***1978, The park was returned to the factory-plate that \nthe public share to the lower of the electronic fence that \nfollow from the Station's cities. The Canal of ancient Western \nnations were confined to the city spot. The villages were directly \nlinked to cities in China that revolt that the US budget and in\nOdambinais is uncertain and fortune established in rural areas.\n```\n\n## Table of Contents\n- [Training Data Info](#training-data-info)\n- [Prerequisites and Training Time](#prerequisites-and-training-time)\n- [Code Structure](#code-structure)\n- [Usage](#usage)\n- [Step by Step Code Explanation](#step-by-step-code-explanation)\n - [Importing Libraries](#importing-libraries)\n - [Preparing the Training Data](#preparing-the-training-data)\n - [Transformer Overview](#transformer-overview)\n - [Multi Layer Perceptron (MLP)](#multi-layer-perceptron-mlp)\n - [Single Head Attention](#single-head-attention)\n - [Multi Head Attention](#multi-head-attention)\n - [Transformer Block](#transformer-block)\n - [The Final Model](#the-final-model)\n - [Batch Processing](#batch-processing)\n - [Training Parameters](#training-parameters)\n - [Training the Model](#training-the-model)\n - [Saving the Trained Model](#saving-the-trained-model)\n - [Training Loss](#training-loss)\n - [Generating Text](#generating-text)\n- [What’s Next](#whats-next)\n\n## Training Data Info\n\nTraining data is from the Pile dataset, which is a diverse, open-source, and large-scale dataset fo
常见问题 (5)
故障排除如何修复在加载使用 torch.compile 和 DDP 保存的 PyTorch 检查点时，'_orig_mod.' 前缀导致的键缺失问题？
反复剥离每个 state-dict 键上的 DDP 和 torch.compile 包装器前缀，因为它们可能以任意顺序嵌套。将原来的单个 module. 检查替换为：
def _strip_ddp_prefix(state_dict: dict) -> dict:
out = {}
for k, v in state_dict.items():
while k.startswith("module.") or k.startswith("_orig_mod."):
k = k.removeprefix("module.").removeprefix("_orig_mod.")
out[k] = v
return out在所有检查点加载路径中使用同一个函数，包括推理和评估：
from src.post_training.utils import _strip_ddp_prefix
state = ck["model_state_dict"] if "model_state_dict" in ck else ck
state = _strip_ddp_prefix(state)
model.load_state_dict(
{k: v for k, v in state.items() if k in model.state_dict()},
strict=False,
)对于新的检查点，保存未包装的模型，使键从一开始就保持干净。加载器还应将缺失的预期参数视为错误，以防止意外地使用随机初始化的权重运行。修复后，compiled/DDP 检查点加载时应达到 missing: 0 个键。
故障排除如何在仅配备 CPU 的笔记本电脑上从零开始训练一个小型 LLM？
使用仓库中对 CPU 友好的 laptop/tiny 预设。首先准备 TinyStories 数据集并训练 BPE 分词器，然后训练 369K 参数的模型，最后生成文本：
python scripts/prepare_tiny_data.py
python scripts/train_transformer.py --preset tiny --arch modern
python scripts/generate_text.py --model_path models/tiny.pttiny 预设专为仅使用 CPU 的机器设计，在典型的笔记本 CPU 上大约需要 5 分钟。此外还提供更大的 student 和 small 预设。运行 python scripts/benchmark.py 可查看在你的硬件上的训练耗时，引导式学习路径请参阅 docs/student。
故障排除如何修复在 A100 40GB 上训练默认的 3B 参数 Transformer 时出现的 CUDA out of memory 问题？
出现 OOM 是意料之中的，因为默认配置使用了约 3B 参数（N_EMBED=2048、N_BLOCKS=64），优化器状态加上激活值会超过 40GB。请更新到包含 PR #29 的版本，并启用可选的内存优化来运行训练：
python scripts/train_transformer.py --amp --amp-dtype bf16 --grad-checkpointing --grad-accum 8在开始训练之前估算显存（VRAM）使用量：
python scripts/train_transformer.py --report-memory--amp 启用混合精度，--grad-checkpointing 通过重新计算 transformer 块来减少激活内存占用，--grad-accum 8 则使用更小的 micro-batch 来模拟更大的 batch。这些更改使 3B 配置能够在 A100 40GB 上运行，还可以降低 RTX 3090/4090/5090 级别 GPU 的显存占用。如果报告了显存碎片化问题，也可以尝试：
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python scripts/train_transformer.py --amp --amp-dtype bf16 --grad-checkpointing --grad-accum 8故障排除如何从最新的检查点恢复中断的 LLM 训练？
检查点恢复功能已在 main 分支的 scripts/train_transformer.py 中提供。训练过程中，使用 --checkpoint-every N 启用周期性检查点保存；使用 --keep-last K 仅保留最近的检查点。训练中断后，使用 --resume latest 重新启动以加载最新的检查点，或使用 --resume /path/to/checkpoint.pt 加载指定的检查点。恢复的检查点包含模型权重、优化器状态、LR 调度器状态、损失历史记录以及最后完成的步数，因此训练会从下一步继续。示例：python scripts/train_transformer.py --checkpoint-every 1000 --keep-last 3 --resume latest。检查点通过临时文件配合 os.replace 以原子方式写入，以降低进程在保存过程中崩溃导致文件损坏的风险。
故障排除如何修复 README.md 中返回 404 的失效 Google Drive 链接？
404 表示该 Google Drive 文件已被删除,或其共享权限设置为私密。请从 README.md 中移除该失效链接,或在 Google Drive 中启用链接共享后,将其替换为有效的共享链接。在所报告的案例中,修复方法是移除 https://drive.google.com/file/d/1Q_iklJ1RVGSb-Pdey8BHy3k8IF3UJv0z/view 并保留 GitHub 引用。提交更改,例如:git add README.md && git commit -m "fix: remove broken resume Google Drive link"。