清华团队开源 Puro-2B:RTX 5090 预训练 20 亿参数模型,成本约 4400 美元
小团队也能预训练!普罗2B:RTX5090训练成本不到5090美元
清华大学陈文光、吕凯风团队推出全流程开源的 Puro-2B(普罗-2B)预训练方案,在 RTX 5090 集群上从头训练约 20 亿参数模型,最多用 1.4 万亿 token,约 4400 美元预训练成本下 15 项任务平均指标超过 Qwen2-1.5B,6900 美元版本接近 Qwen2.5-1.5B。
新智元报道 训不起2.8T的K3,玩够了124M的nanoGPT,想试试 十亿参数、万亿token的预训练 ,门槛有多高? 这类模型被广泛使用,例如经典的Qwen系列(Qwen2-1.5B,Qwen2.5-1.5B,Qwen3-1.7B… … )。 不过Qwen系列尽管开源了权重,却未完整公开预训练配方。 而OLMo、SmolLM等进一步开源了数据和训练细节,但复现依然昂贵:据估计,即使只有十亿级参数,复现OLMoE-1B/7B和SmolLM3-3B所需要的算力成本分别约为 20万美元和71.9万美元 。
对于普通实验室,知道「怎么训」还不够,还得真正「训得起」。 我们需要一套 低成本、可复现的「穷人版」预训练方案 。 Puro-2B 全流程开源的「穷人版」预训练方案 围绕这个问题,清华大学陈文光教授与吕凯风教授团队联合推出全流程开源的 Puro-2B(普罗-2B), 开源了训练数据、训练框架、训练配置和模型权重。 论文Figure 3 | Puro-2B「穷人版」训练流程 他们在消费级显卡RTX 5090的计算集群上,从头训练了约20亿参数的模型,最多使用了1.4万亿token。
在覆盖15项任务的综合性评测中,基于这套方案, 以约4400美元的预训练计算成本,模型在15项任务平均指标上超过Qwen2-1.5B 。 6900美元成本的最佳版本,平均性能接近Qwen2.5-1.5B 。 相比其他开源方案,Puro-2B显著降低了训练成本。
论文链接:https://arxiv.org/pdf/2608.27370 开源材料合集:https://huggingface.co/collections/thu-pacman/puro-2b 模型与中间检查点 :https://huggingface.co/collections/thu-pacman/puro-2b 训练数据 与配比 :https://huggingface.co/datasets/thu-pacman/Puro-2B 训练代码 :https://github.com/thu-pacman/Puro-Megatron 数据处理代码 :h ttps://github.com/thu-pacman/Kaiyuan-Spark 论文Figure1|模型性能与复现成本 如果预算再少一些,能训练模型到什么水平?
他们进一步拟合了 成本—性能经验曲线(Puro Cost Scaling Law) ,用来估计不同预算下的模型能训出的水平。 他们从同一个第一阶段训练的检查点出发,调整第二阶段均匀数据配方的训练预算,来拟合平均性能得分与总成本的关系。 这条针对Puro训练方案的经验曲线,让研究者能够在正式训练之前,估计不同预算成本所能达到的能力水平,并结合需求调整训练方案,而不需要复现最长的模型训练过程。
论文Figure 2b|Puro Cost Scaling Law 这套低成本方案,来自 从硬件选择、低精度训练、优化器改进、高质量数据利用效率、数据选择这五个部分的协同设计。 单位成本,更高算力 相比于其他开源的训练方案,Puro-2B团队从更基础、更底层、更贴近实际的问题出发: 同样的预算能买到多少算力? 按论文采用的GPU租赁价格计算,他们发现,RTX 5090的BF16峰值算力/美元约为H200的 2.77倍。
这意味着单位成本,RTX5090能比广泛使用的数据中心卡获得更多算力, 同时RTX5090作为消费级显卡还兼顾了硬件的可获取性。 论文Table2简化|不同显卡的计算价格比值的统计与对比 但RTX5090只有32GB显存,也没有NVLink,不能直接照搬数据中心GPU的训练配置。 Puro-2B团队通过驱动和系统配置调整,在节点内启用PCIe P2P通信,将GPU间单向带宽从31.5 GB/s 提升到 56GB/s ,通信延迟从14.3µ s降低到 0.4µs ;
节点间启用GPU Direct RDMA,让24卡AllReduce总线带宽从8.87 GB/s 提升到 19.93GB/s 。 Puro的训练策略也针对十亿级参数规模的计算负载与RTX5090的硬件特点,进行优化:他们采用数据并行与流水线并行组合,避免通信频繁的张量并行; 通过重新分配流水线各阶段的计算量,调整微批次大小,并按显存占用情况均衡放置优化器状态,平衡不同节点的显存负载。 通过上述优化, 综合MFU能超过70% ,消费卡的性价比通过这些工程适配与系统优化转化为训练成本的优势。
有了更划算的算力,下一步是提高计算精度效率。 FP8发挥RTX 5090的低精度效率优势 RTX 5090的Blackwell架构还为FP8训练提供了基础。 他们沿用DeepSeek-V3的细粒度分块量化思路,将主要线性层的矩阵乘法放在FP8路径上,同时让数值敏感操作和训练状态保留BF16或FP32精度。 FP8会带来一定精度损失。 为判断是否划算,他们通过不同模型规模的对照实验(Scaling ladder),来验证了FP8训练带来的净效率收益。
通过五个模型规模的实验对照,他们发现,相比BF16,FP8的验证损失高出约 0.0031—0.0039 ,但差距较小且跨规模稳定。 与此同时,FP8的训练吞吐效率得到 1.36倍 提升; 把精度损失折算为达到相同损失所需的额外计算后,据估计仍有约 1.34倍净效率收益 。 论文Figure 10b|FP8的损失差距与质量校正后的效率 从普通学习率到有效学习率 确定了计算精度后,如何设计优化器和学习率会影响每步更新的效果。 优化器决定了每一步的更新方向,学习率决定了每一步的更新步长。
他们采用了新兴的MuonH优化器:在Muon基础上将权重的更新量做归一化处理,并在权重更新后,将权重投影回固定半径,从而同时控制更新范数和权重范数。 这样一来,衡量更新量相对权重尺度的 有效学习率(Effective LR, ELR), 就能被显式预设,而不随权重尺度在训练过程中的变化而间接漂移。 论文Figure 4|MuonH原理示意图以及与Muon的对照 他们的关键观察来自一组对照实验:用普通Muon不做固定半径的权重投影,只通过动态调整学习率,使训练过程的ELR轨迹与MuonH对齐,则损失曲线便几乎重合。
这说明,实际中MuonH的优势不能仅归因于范数约束 ,有效学习率ELR的变化过程才是理解超球优化有效性的重要线索 。 综合论文Figure 5与Figure 15|ELR对齐实验的解释与对照 进一步比较发现,在weight decay对权重范数约束不足时,普通Muon优化器的ELR在前期下降更快,损失也更早降低; MuonH则保留更平缓的衰减,在后期表现更好。 这让他们把注意力从「用哪个优化器」,转向「如何设计更合理的有效学习率调度(ELR schedule)」。
随后,他们针对训练的衰减阶段占比对峰值学习率和训练数据量进行消融。 结果显示,较高的峰值学习率和更长的训练周期,通常更需要较长的衰减过程。 因此,他们在训练第一阶段采用不依赖预设终点的 power schedule ,支持持续训练; 第二阶段采用 长线性衰减 ,为最终的损失收敛留出空间。 基于Scaling Ladder的比较进一步显示,MuonH方案相比调参后的Muon基线,训练效率约为 1.19倍 。
论文Figure 6|学习率设计 论文Figure 10a|MuonH的训练效率消融 把好数据放到训练后段 还要让模型学得进去 除优化器之外,数据效率还受数据排布影响。 高质量数据在训练数据集中占比不高但是作用关键,放到训练后段有助于提高利用效率。 许多开源数据集会同时提供样本的质量分数。 Puro团队利用这些信号,在各数据源内部把更高质量的样本排到训练后段,再按训练进度对齐、混合不同来源的数据块。
这样既让训练过程中数据集逐步走向更高质量的区间,又保持局部训练区间的数据集配比近似稳定,避免质量变化同时造成领域数据比例的突变。 但数据排好了,并不等于模型能充分学到。 Puro团队此前的研究发现: 当训练后期 学习率 过低时,最后出现的优质数据难以推动参数更新。 为此,他们采用了 课程模型平均(Curriculum Model Averaging,CMA) 在训练后期保持恒定学习率,使模型继续学习优质数据,再对末端的检查点做权重平均,来缓解较高学习率带来的训练噪声。
论文Figure 7|课程数据构造与模型平均 进一步的消融也显示,直接在学习率衰减的训练轨迹上做模型平均,并不会稳定获益。 真正需要配合的是数据顺序、后期学习率与模型平均策略。 Figure 13b|对数据顺序、后期学习率与模型平均策略的消融实验 按Puro Cost Scaling Law这一经验曲线估算换算,相比于随机数据排布(Uniform, UD),数据课程加学习率衰减会相当于把成本效率提高到约1.65倍,完整CMA方案则达到约 2.40倍 。 数据评测 数据顺序解决后,还要决定选哪些数据。
他们建立了「 数据源—数据切片—能力评测 」的分层比较流程。 对于大规模且有评分的数据集,分别抽取不同分数区间的数据切片; 其他符合条件的数据源采用随机切片。 每个候选切片都从同一个检查点出发,在相同训练配置下继续训练,再用固定评测得到数学、代码、中文和通用能力的多维表现。 这既能比较不同数据源,也能发现同一数据源内部的质量差异。 小规模替代实验(代理实验)还显示,相对于数学能力,代码能力对通用能力有更多削弱。 据此,他们在第二阶段适当提高了数学数据比重,同时结合领域覆盖和质量筛选确定最终配比。
这些小规模的实验为「选哪些数据、保留哪些分数区间、给不同来源分配多少比例」提供了一定参考,相关数据集评估结果也公开在报告中。 论文Figure 9|代理评测流程; Figure 8|两阶段数据配比 从节省算力到覆盖训练流程 回看整套方案,各环节解决的是不同层次的效率问题: 硬件降低每单位计算的价格,FP8提高计算峰值以缩短执行时间,MuonH提高 损失 值收敛速度,CMA提高 训练数据 的利用效率,代理评测帮助选择更值得训练的数据。 这些改进共同组成了Puro的低成本训练流程。
论文Figure 2a|各环节效率收益汇总 那么,这套训练配方可以怎样用于研究或实践呢? 例如,研究者可以控制预训练过程,再检验后训练结果。 一个例子是: 研究者常关心,预训练阶段的方法改进,经过后训练之后还是否有效。 如果只有开源权重,研究者很难控制模型此前的预训练过程; 而Puro-2B让修改预训练的方法、继续做微调、检验最终效果成为一条可运行且成本可承担的实验流程。 Puro团队对CMA与Uniform两个预训练权重采用相同的微调流程。
在两组数学能力的微调实验中,基于CMA初始化微调的GSM8K准确率分别高出 1.77和2.02个百分点 ; 在Tulu-3混合领域微调中,15项基准平均分提高 1.17个百分点 ,IFEval提高 1.36个百分点 。 三组对照均呈现总体提升,说明CMA的优势整体能够延续到后训练之后。 由此而言,Puro-2B提供的不仅是一个分数更高的模型,还是一套预训练的基础流程,用于检验预训练方法能否带来最终收益。
论文Figure 11|CMA与Uniform数据预训练的权重做微调后的对照 Puro-2B普罗模型, 名称就似乎和普罗大众有所关联 ,那么未来,究竟是模型厂商在scaling规模的路上一路狂飙,还是模型训练逐步成为普罗大众可及的一部分呢? 普罗的尝试,也许正在使更多团队参与预训练研究成为可能。 参考资料: https://arxiv.org/pdf/2608.27370 编辑:LRST 秒追ASI ⭐ 点赞、转发、在看一键三连 ⭐ 点亮星标,锁定新智元极速推送!
来源:北京智源研究院 · hub.baai.ac.cn