一块RTX 5090训练25小时:从零构建2亿参数GPT全记录

从零手搓一个GPT有多难
在大模型动辄千亿参数、训练成本数百万美元的今天,一位开发者选择了一条截然不同的道路:几乎从零开始,用一块租来的RTX 5090显卡,训练出一个2亿(200M)参数的GPT解码器模型。GPT系列模型采用的是Transformer架构中的解码器(Decoder-only)部分——原始Transformer由Vaswani等人在2017年的论文《Attention Is All You Need》中提出,包含编码器和解码器两部分。GPT的创新在于只使用解码器部分,通过因果掩码(causal mask)确保每个位置只能看到它之前的token,从而实现自回归式的文本生成。这种架构的简洁性使其在规模扩展时表现出色,成为当前主流LLM的首选架构。这个项目在Reddit上引发关注,不仅因为它的完成度,更因为它系统性地复现了现代大语言模型的核心技术栈。
据这位开发者介绍,该模型在fineweb-edu数据集的40亿(4B)tokens上完成训练,最终达到约3.3的平均损失(mean loss)和28的困惑度(perplexity)。这里需要解释token的概念:token是语言模型处理文本的基本单位,现代LLM通常不按单词分割文本,而是使用BPE(Byte Pair Encoding)或类似的子词分词算法,将文本切分为介于字符和单词之间的片段。例如,"unhappiness"可能被切分为"un"、"happiness"两个token。典型的词表大小在32k到100k之间,4B tokens的训练数据大约相当于30亿英文单词或约2000万页标准文档的文本量。
fineweb-edu是由HuggingFace团队发布的高质量教育类文本数据集,从CommonCrawl中经过多层过滤和质量评分筛选而来,总量超过1.3万亿tokens,专门面向训练语言模型设计,文本质量显著高于未经筛选的互联网爬取数据。CommonCrawl是一个非营利组织,自2008年起持续爬取互联网网页并免费公开数据,其数据集包含数十亿网页、数万亿tokens的原始文本,是几乎所有大型语言模型训练数据的上游来源。然而原始CommonCrawl数据质量参差不齐,包含大量广告、垃圾邮件、重复内容和低质量文本,因此需要经过多轮过滤(语言识别、质量评分、去重、有害内容过滤等)才能用于训练高质量模型。fineweb-edu在此基础上进一步筛选出教育价值高的内容。
关于这些训练指标,平均损失3.3意味着模型在交叉熵损失函数上的表现——交叉熵损失衡量的是模型预测的概率分布与真实分布之间的距离,数值越低说明模型预测下一个token的准确度越高。在实践中,经过充分训练的大型语言模型在通用文本上通常能达到2.0-2.5的损失值,而小模型或训练不充分的模型则在3.0以上。困惑度是损失的指数形式(e^loss),28的困惑度意味着模型在预测下一个词时,平均面临约28个等概率候选词的不确定性。作为参考,GPT-3在类似评测上的困惑度约为20以下,而随机猜测词表中任意词的困惑度则高达数万。虽然这些指标与商业级模型相去甚远,但对于一个「几乎全部手写」的教学级实现来说,已经相当扎实。
值得注意的是,本项目2亿参数配合40亿tokens的训练配置,与DeepMind在2022年发表的Chinchilla研究所提出的最优比例高度吻合。Chinchilla研究指出,给定固定计算预算,模型参数量和训练token数应大致保持1:20的比例才能达到最优性能——2亿参数对应的最优训练数据量恰好约为40亿tokens,说明作者在规划训练时参考了这一重要的Scaling Laws研究成果。

你可能没注意到,作者刻意限制了对PyTorch的依赖——PyTorch主要只用于自动微分(autograd)和GPU计算,其余的架构组件、优化器、采样逻辑几乎都是自行实现。自动微分(Automatic Differentiation)是深度学习框架的核心技术,区别于数值微分(有限差分法,精度有限且计算开销大)和符号微分(符号计算系统,表达式可能膨胀到不可处理)。PyTorch使用动态计算图(define-by-run),即在前向计算过程中即时构建计算图,记录每个运算操作及其输入输出关系。反向传播时,系统沿着这个图从输出反向遍历到输入,利用链式法则自动计算每个参数的梯度。这使得开发者只需定义前向计算逻辑,无需手动推导可能极其复杂的梯度公式。
在这个项目中,作者将PyTorch的角色严格限定为「计算基础设施」——即利用其自动微分和CUDA张量运算能力,但不使用torch.nn中预封装的Layer、Optimizer等高层抽象。这相当于在搭建房屋时只借用电动工具,但砖瓦结构全部自己设计砌筑。这种做法虽然增加了工作量,但对于深入理解Transformer的每一个齿轮如何咬合,具有无可替代的价值。
完整复现现代LLM技术栈
这个项目最大的看点,在于它并非停留在2017年原始Transformer论文的水平,而是集成了近年来大量被验证有效的工程改进。这些技术恰恰是当前主流开源模型(如LLaMA系列)的标配。
注意力机制与旋转位置编码RoPE
模型采用了**多头自注意力(Multi-head Self-Attention)作为核心,并使用旋转位置编码(RoPE, Rotary Positional Embeddings)**替代传统的绝对位置编码。RoPE由苏剑林在2021年提出,其核心思想是将位置信息编码为向量空间中的旋转操作。具体而言,对于位置m处的query向量和位置n处的key向量,RoPE通过将向量的每两个相邻维度视为一个二维平面,按照与位置成正比的角度进行旋转。这样,两个位置的内积自然地只依赖于相对距离(m-n),从而实现了相对位置编码。与传统的正弦位置编码相比,RoPE天然支持外推到训练时未见过的更长序列长度(通过NTK-aware缩放等技术),且不增加额外参数。Meta的LLaMA、Google的PaLM以及几乎所有2023年后的主流开源模型都采用了RoPE,使其成为现代LLM的事实标准。
此外,作者还实现了KV Cache(键值缓存)。这是推理加速的关键机制——在自回归生成过程中,模型每生成一个新token都需要对整个已生成序列计算注意力。如果不做缓存,生成第N个token时需要重新计算前N-1个token的key和value向量,导致计算量随序列长度呈二次方增长。KV Cache的策略是将每一层注意力机制中已计算过的key和value向量缓存起来,生成新token时只需计算当前token的query、key、value,并将新的key和value追加到缓存中。这将推理的计算复杂度从O(N²)降至每步O(N),使得长文本生成的速度提升数十倍。然而KV Cache也带来了显存压力——对于长序列和大模型,缓存的key-value矩阵可能占据数十GB显存,这也催生了GQA(分组查询注意力)等压缩KV Cache的技术。能主动实现KV Cache,说明作者的视野已经超越了「跑通训练」,延伸到了推理效率层面。
RMSNorm归一化与SwiGLU激活函数
在架构设计上,模型采用了**Pre-Norm(前置归一化)**结构,即将归一化层放在残差连接之前。原始Transformer采用Post-Norm结构,即在残差连接之后进行归一化,但研究发现Post-Norm在训练深层网络时容易出现梯度消失或训练不稳定的问题。Pre-Norm将归一化前置到子层输入之前,使得残差路径上的梯度可以直接回传而不经过归一化层的缩放,大幅改善了深层网络的训练稳定性。
归一化方式选用了RMSNorm而非传统LayerNorm。RMSNorm(Root Mean Square Normalization)是LayerNorm的简化版本:标准LayerNorm先减去均值再除以标准差,而RMSNorm省略了减均值步骤,直接用向量的均方根值进行归一化。实验表明这种简化几乎不影响模型性能,却减少了约10-15%的归一化层计算量,在大规模训练中累积效果显著,同样是LLaMA等模型的选择。
激活函数方面,作者使用了SwiGLU,这是由Google的Noam Shazeer在2020年提出的门控线性单元变体。其公式为SwiGLU(x) = Swish(xW₁) ⊗ (xW₂),其中⊗表示逐元素相乘,Swish(x) = x·sigmoid(βx)。这种门控机制允许网络学习性地「开关」信息通路——W₁分支产生的门控信号决定了W₂分支的信息有多少能通过。相比传统ReLU或GELU,SwiGLU在前馈网络的相同参数预算下通常带来0.5-1%的困惑度提升。需要注意的是,SwiGLU的前馈层因为有两个投影矩阵,其隐藏层维度通常设为传统FFN的2/3以保持总参数量一致。PaLM、LLaMA、Mistral等模型均采用SwiGLU。
此外还应用了权重绑定(Weight Tying),让输入嵌入层与输出投影层共享参数。输入嵌入矩阵将token ID映射为向量表示(维度为vocab_size × d_model),而输出投影层将隐藏状态映射回词表空间(维度为d_model × vocab_size),两者恰好互为转置关系。共享这些参数不仅节省了大量参数(对于50k词表和768维隐藏层,可节省约3800万参数),还在语义上创建了一种对称性:一个词的输入表示和输出预测使用相同的向量空间,这被认为有助于模型学习更一致的词义表征。GPT-2、T5等经典模型均使用了这一技术。
训练工程的细节把控
除了模型架构,训练环节的工程实现同样体现了作者的功底。
AdamW优化器与余弦学习率调度
作者自行实现了AdamW优化器,并配合**余弦学习率衰减(Cosine LR Decay)加预热(Warmup)**的调度策略。AdamW是Adam优化器的权重衰减修正版本,由Loshchilov和Hutter于2019年提出。标准Adam中的L2正则化与自适应学习率机制存在耦合问题——参数的梯度越大,其有效权重衰减反而越小,违反了正则化的初衷。AdamW将权重衰减从梯度更新中解耦出来,直接在参数上施加衰减,修复了这一问题。从更技术的角度看,Adam为每个参数维护一阶矩(梯度均值)和二阶矩(梯度方差)的指数移动平均,从而为不同参数自适应地调整学习率——梯度波动大的参数获得更小的有效学习率,梯度稳定的参数则获得更大的步长,这使得Adam比简单的SGD在处理稀疏梯度和非平稳目标时表现更好。
这套组合是训练Transformer的经典配方:预热阶段让学习率从小逐步爬升,避免训练初期模型在随机初始化的参数空间中以大步伐更新导致的不稳定(预热通常占总训练步数的1-5%);余弦衰减则让学习率按照余弦函数从峰值平滑下降至接近零的最小值,其曲线形状意味着训练中期保持较高学习率维持探索能力,后期快速衰减促进收敛到更好的解。Chinchilla、GPT-4等模型的训练均采用这套标准配方。
同时,模型使用了**梯度累积(Gradient Accumulation)**技术。其核心思想是:在不增加显存占用的情况下,等效增大batch size。假设目标batch size为1024个样本,但显存只能容纳128个样本,那么可以连续进行8次前向-反向传播(每次128样本),将8次计算得到的梯度相加,再执行一次参数更新。数学上,这与一次性处理1024个样本计算梯度完全等价(因为梯度关于样本是线性可加的)。大batch训练对语言模型至关重要——更大的batch提供更稳定的梯度估计,减少训练噪声,通常能带来更好的最终性能和更快的收敛。研究表明,语言模型训练的临界batch size(在此之上继续增大batch的收益递减)通常在数十万到数百万tokens的范围。在单卡显存受限的情况下,梯度累积允许通过多个小batch累加梯度来模拟大batch训练的效果,这对于在消费级/单卡环境下训练较大模型至关重要。
单卡训练的算力成本与效率分析
从算力数据看,作者在vast.ai平台租用了一块RTX 5090进行训练。RTX 5090是NVIDIA在2025年初发布的旗舰消费级显卡,基于Blackwell架构,配备32GB GDDR7显存,FP16算力约为100+ TFLOPS,相比上一代RTX 4090性能提升约50-70%。NVIDIA的GPU架构经历了从Maxwell、Pascal、Volta、Turing、Ampere、Ada Lovelace到Blackwell的演进,每一代都在计算核心数量、显存带宽、Tensor Core设计等方面有显著提升。Tensor Core是专为矩阵乘法设计的硬件单元,对深度学习训练至关重要——一个Tensor Core可以在单个时钟周期内完成一个4×4矩阵的乘加运算,相比通用CUDA Core在矩阵运算上快数十倍。
vast.ai是一个去中心化的GPU租赁市场,将全球闲置GPU资源汇集起来供用户按需租用,价格通常远低于AWS或Azure等主流云平台。其商业模式类似GPU算力的Airbnb——个人和小型数据中心可以将闲置GPU上架出租,用户则可以按需选择不同地区、不同型号的GPU,按小时计费。
每个更新步骤(step)耗时约2秒,吞吐量达到每步8.1万tokens,总训练时长约25小时。RTX 5090在vast.ai上的租用价格大约在每小时1-2美元左右,意味着本项目25小时的训练总成本约为25-50美元——这个价格训练出一个功能完整的2亿参数语言模型,在几年前是完全不可能的。作为对比,GPT-3的训练成本估计超过400万美元,而即使是较小的GPT-2(15亿参数)在2019年的训练也需要数万美元的算力。
这组数据颇具参考价值:它意味着以相对低廉的租用成本,个人开发者完全有能力从零训练出一个功能完整的小型语言模型。这在几年前几乎是不可想象的,也反映了硬件性能提升与开源生态成熟带来的门槛下降。
生成效果与模型局限性
在文本生成方面,作者实现了Top-k采样和**重复惩罚(Repetition Penalty)**来控制输出质量。Top-k采样是一种在文本生成时平衡多样性与质量的解码策略:在每一步生成中,模型输出词表上的概率分布后,Top-k只保留概率最高的k个候选词,将其余词的概率置零并重新归一化,再从这k个词中按概率采样。k值越小生成越确定性(k=1即为贪心解码),k值越大则越具创造性但也越不可控。在实际应用中,Top-k常与Top-p(nucleus sampling)配合使用——Top-p不是固定候选词数量,而是选择累积概率达到阈值p的最小词集合,能更好地适应不同置信度的预测场景。温度(temperature)参数则通过缩放logit来整体调节分布的尖锐程度:温度<1使分布更集中,温度>1使分布更平坦。
重复惩罚则通过降低已生成token再次被选中的概率来避免模型陷入重复循环——这是自回归语言模型的常见退化模式,尤其在小模型中更为严重。常见的实现方式是将已出现token的logit除以一个大于1的惩罚系数(如1.2),使模型倾向于生成新内容。重复退化的根本原因在于自回归模型的训练目标(教师强制)与推理方式(自由生成)之间的不匹配——训练时模型总是基于真实前文进行预测,但推理时基于自己之前可能有误的输出,错误会逐步累积放大。
从公开的样本输出来看,模型能够生成语法通顺、主题连贯的英文段落:
在被问及美国在AI领域的领导地位时,模型生成了一段关于「谷歌宣布4亿美元投资人工智能、成立名为Google Insights的新公司」的内容。
这段输出结构完整、行文流畅,具备基本的话题连贯性。但需要清醒认识到,这类内容存在明显的「一本正经胡说八道」现象——例如「Google Insights」这样的信息很可能是模型的幻觉编造。模型幻觉(Hallucination)是当前所有语言模型面临的核心挑战之一,其根源在于语言模型本质上是在学习文本的统计分布模式而非建立事实知识库。模型生成的每个token都是基于概率采样,它无法区分「看起来合理的内容」和「事实正确的内容」。对于小模型而言,由于参数容量有限,其能够记忆和表征的知识量远少于大模型,因此幻觉问题更为严重。
这正是2亿参数级别小模型的典型局限:语言形式的掌握远超过事实知识的可靠性。困惑度28的指标也印证了这一点——模型学会了如何组织语言,但离精准可靠还有很大距离。
对开发者的启示与后续计划
这个项目的意义并不在于训练出了多强的模型,而在于它提供了一条清晰、可复现的学习路径。对于希望深入理解大语言模型的开发者而言,与其停留在调用API,不如亲手实现RoPE、KV Cache、SwiGLU这些现代组件——这个过程带来的理解深度是阅读论文难以替代的。
作者透露了后续计划:一是添加自定义CUDA kernel以进一步提速,二是尝试在100亿(10B)tokens上训练一个5亿(500M)参数的更大模型。自定义CUDA kernel意味着绕过PyTorch的通用计算调度,针对特定运算(如融合注意力计算、自定义激活函数)编写专用的GPU代码,通常能带来2-5倍的特定运算加速。CUDA(Compute Unified Device Architecture)是NVIDIA推出的并行计算平台和编程模型,允许开发者直接控制GPU上数千个计算核心的行为。一个典型的优化案例是FlashAttention——通过重新安排注意力计算的IO模式,减少对GPU高带宽存储器(HBM)的访问次数,在不改变数学结果的情况下将注意力计算加速2-4倍并大幅降低显存占用。类似地,kernel融合(将多个连续小运算合并为一个大kernel)可以减少核心启动开销和中间结果的显存读写,是高性能深度学习框架的常用优化手段。
这种从小规模验证到逐步扩大的迭代思路,也正是工业界研发大模型的缩影。许多前沿实验室在训练大模型之前,都会先在小规模(如数亿参数)上验证架构设计和超参数选择,确认loss曲线的行为符合预期后,再依据Scaling Laws预测大规模训练的最终性能,最后才投入巨额资源进行全量训练。
对于囊中并不宽裕的个人研究者,这个案例传递出一个积极信号:借助按需租用的云GPU和成熟的开源数据集(如fineweb-edu),从零构建一个「麻雀虽小五脏俱全」的GPT,已经成为一件可及的事情。真正稀缺的,从来不是算力,而是把每个技术细节吃透并亲手实现的耐心与好奇心。
核心要点
核心要点
相关推荐

Sutura:Linux下STL/3MF模型修复开源工具详解
Sutura是一款专为Linux用户打造的开源3D模型修复工具,支持STL和3MF格式,基于PyMeshLab和manifold3d库,提供CLI、GUI和文件管理器右键菜单三种使用方式,填补Linux 3D打印工作流中的模型修复空白。

Human Behavior:AI智能体如何闭环处理产品分析问题
Human Behavior是一款AI驱动的产品分析工具,通过采集、理解、行动、闭环四步链路,让AI智能体自动识别用户体验问题并提交修复代码,彻底改变传统仪表盘模式。

Anthropic删除Claude Code 80%提示词的启示:上下文工程新规则
Anthropic将Claude Code系统提示词删除80%后性能不降反升。本文解析上下文工程六条新规则,包括精简禁令、按需加载Skills、善用参照物等实操建议,帮助你优化AI Agent的上下文管理策略。