论文

Puro-2B: 穷实验室在 RTX 5090 上以 5090 美元训练 Qwen2-1.5B

Puro-2B: 穷实验室在 RTX 5090 上以 5090 美元训练 Qwen2-1.5B

语言模型预训练几乎与高昂成本划等号,使得学术界和开源社区难以企及。尽管已有开源权重和训练配方等努力,但成本高效、硬件可及的开源预训练配方仍然缺失。即便在小规模下,训练 Llama-3.2-3B 也需超过 1.5M 美元,复现 SmolLM3-3B 则需超过 700K 美元。 本报告提出一种开源预训练配方,以降低这一门槛。利用该配方,我们在消费级 RTX 5090 GPU 上,以 FP8 精度从零训练了 Puro-2B 系列模型,最多达 1.4 万亿 tokens。该系列模型在 token 预算和配方变体上有所区别。我们的最佳模型训练成本低于 6.9K 美元,在我们评估协议下接近 Qwen2.5-1.5B 的性能。成本效率源于多种方法的结合,包括: - 硬件选择 - 低精度训练 - 超球优化 (hyperball optimization) - 课程模型平均 (curriculum model averaging) - 数据配方 除配方本身外,我们提供两项额外结果: 1. 基于 Puro-2B 系列,推导出 Puro 成本缩放定律,关联训练成本与平均模型性能;拟合定律表明约 4.4K 美元(低于 5,090 美元)即可达到 Qwen2-1.5B 的性能。 2. 作为端到端案例研究,我们考察了预训练数据课程如何影响后训练后的下游性能。此类受控研究得益于对完整预训练管道的访问,而非仅有模型权重。 我们以 Apache 2.0 协议在 https://huggingface.co/collections/thu-pacman/puro-2b 发布 Puro-2B 的完整训练配方,包括数据、代码和模型权重。

论文精读

TL;DR **Puro-2B** 在消费级 RTX 5090 GPU 上用 FP8 训练 2B 模型,总成本低于 $6.9K,性能接近 Qwen2.5-1.5B,并给出成本缩放定律与完整开源配方。

问题

问题背景:大语言模型预训练长期被高昂算力成本垄断,开源社区虽已获得开放权重与部分训练配方,但完整的低成本、消费级硬件可复现预训练流程依然缺失。当前领域关注如何将预训练门槛从数十万美元降至数千美元,同时保持可比的模型能力。

现有方法局限:已有开源配方如 Llama-3.2-3B 或 SmolLM3-3B 即便在小规模下也需要 $1.5M 或 $700K 以上,普通实验室无法承受。低精度训练(如 FP8)虽能加速但容易导致训练不稳定;优化器(如 MuonH)与学习率调度需要手动调参,缺乏系统指导;数据课程对下游性能的影响缺少受控实验,因为大多数团队只发布模型权重而非完整训练管线。此外,消费级 GPU 存在显存容量、多卡通信带宽与负载均衡等工程瓶颈,使得大规模 token 训练难以高效实施。

为什么难/重要:预训练是资源密集型试错,高成本阻碍了控制变量研究,限制了对优化器、数据混合、课程策略等关键选择的理解。降低门槛可以促进更多实验室验证假设、改进方法,并提高训练透明度,与开源精神一致。

行业类比:如同在边缘设备上部署个性化语言模型,需要在有限算力预算内完成训练,成本效率与稳定性同样关键。

核心洞察

  • 低成本预训练的关键在于多维度协同优化而非单一技术突破。论文将 RTX 5090 消费级 GPU、FP8 混合精度、MuonH 优化器有效学习率对齐、超球优化、课程模型平均等方法组合,实现低于 $7K 训练 2B 模型接近 Qwen2.5-1.5B 性能。相比已有开源预训练配方(Llama-3.2-3B 成本超 $150 万,SmolLM3-3B 超 $70 万),这种工程协同大幅降低门槛,为资源有限团队提供可复现路径。
  • 成本缩放定律将“花多少钱达到什么性能”从经验猜测变为可量化预测。论文从不同 token 预算和配方变体的 Puro-2B 模型集合中拟合出 Puro Cost Scaling Law,表明约 $4.4K 即可达到 Qwen2-1.5B 性能。这改变了预算决策方式,允许研究者按成本目标倒推训练配置,而不必依赖大规模模型外推的模糊估计,对预算规划具有直接工程价值。
  • 完整预训练管线的公开使得数据课程与后训练之间的因果关系可以被严谨研究。论文展示了预训练数据课程如何影响数学、指令遵循等下游能力,而以往仅有权重发布的模型无法进行此类反向分析。这一开放实践为社区提供了实验平台,推动数据配方设计从经验试错走向机制理解。

方法

方法流程

输入: Puro-2B 使用定制的大规模文本数据集,总量最高 1.4T tokens。数据经过采集、过滤与混合,形成多组件课程:每个组件内部按难度或主题排序,同时保持目标混合权重不变。

关键模块

  1. 硬件与精度: 选用消费级 RTX 5090 GPU,结合 FP8 混合精度训练 与通信感知并行策略,在有限算力下提升吞吐。
  2. 优化器: 采用 Hyperball Optimization(基于 MuonH),通过有效学习率缩放统一不同尺度的学习率调度,配合多幂律估计器设计开放端学习率曲线与终端衰减。
  3. 课程模型平均: 在课程排序与后期优化中联合设计,对后期连续性训练进行模型平均,提升最终性能。
  4. 后训练配方: 包含聚焦数学、扩展数学与广泛指令迁移三种设置,用于评测下游能力。

输出

训练得到一系列 Puro-2B 模型(不同 token 预算和配方变体),最佳模型在评测协议下接近 Qwen2.5-1.5B 性能,总计算成本低于 $6.9K。同时推导出 Puro Cost Scaling Law,表明约 $4.4K 即可达到 Qwen2-1.5B 水平。

与同类开放预训练配方(如 Llama-3.2-3B 成本超 $1.5M、SmolLM3-3B 超 $700K)相比,本方法通过低精度训练、超球优化、课程模型平均与数据课程的系统组合,将 2B 模型预训练成本降至消费级硬件可复现的范围。

实验

实验设计

作者训练了一组 Puro-2B 模型,从零开始在最多 1.4T tokens 上以 FP8 精度 训练,使用消费级 RTX 5090 GPU。模型集合在 token 预算和配方变体上有所差异,采用了 硬件选择、低精度训练、hyperball 优化、curriculum model averaging 以及数据配方。此外,还考察了预训练数据课程如何影响后训练的下游性能。

关键发现

最优模型训练成本低于 $6.9K,性能接近 Qwen2.5-1.5B。通过 Puro Cost Scaling Law 发现,约 $4.4K 即可达到 Qwen2-1.5B 的性能水平。成本效率源于硬件性价比、低精度加速、MuonH 优化器带来的计算等效偏移,以及块级 FP8 在加速的同时保持质量。

与基线对比

与 Llama-3.2-3B(成本超过 $1.5M)和 SmolLM3-3B(成本超过 $700K)相比,Puro-2B 的成本降低了一个数量级以上,同时性能接近同级别的 Qwen2.5-1.5B。这证明在小规模上通过精心设计的配方可以达到极具竞争力的性能,为资源受限的研究者提供了可行的开源路径。

行业影响

落地场景

Puro-2B 的低成本预训练方案适用于中小型团队、学术机构或企业部门构建垂直领域语言模型,例如电商客服对话、内容平台评论审校、教育辅导解题、金融文本分析等。其 2B 参数规模在消费级 GPU 上即可完成训练,适合对数据隐私和定制化要求高的场景。

  • 电商客服机器人:用商品评论和客服对话语料预训练一个 2B 模型,成本仅数千美元,可替代通用大模型 API,降低延迟和费用。
  • 教育解题助手:针对数学推理数据定制预训练,再结合小型指令微调,得到轻量解题模型,部署在边缘设备或本地服务器。

商业价值

核心价值在降本:论文证明用不到 $7K 即可训出接近 Qwen2.5-1.5B 性能的模型,而传统同规模训练成本高达数十万美元。这让预算有限的企业能自主迭代模型,无需依赖高价云 API,同时数据不出域,满足合规要求。

体验提升体现在推理侧:2B 模型推理延迟和吞吐显著优于 7B+ 模型,适合实时交互场景。此外,通过成本缩放定律,团队可以在训练前估算所需预算,合理规划资源。

跟现有产品/工作流的接口

Puro-2B 的代码基于 Megatron-LM,与主流 NVIDIA 训练生态兼容;FP8 混合精度和超球优化可直接嵌入 PyTorch 训练脚本。数据配方和课程模型平均策略可作为独立模块接入企业已有的数据管线。

  • 模型权重以 Apache 2.0 发布在 HuggingFace,可直接用 transformers 加载微调或蒸馏。
  • 训练脚本开源在 GitHub,便于 CI/CD 集成到内部 MLOps 流程。
  • 成本估算协议可帮助技术决策者在采购硬件或云资源时做预算对标。

局限

  • - **评估协议与数据污染**:论文在有限 benchmark 集合上评测,主要覆盖数学、代码、推理和知识,缺乏更广泛的语言理解、多语言、安全性和对话能力评估。作者在 Limitations 中承认存在数据污染风险和 processed-data 的不确定性,因此性能对比可能存在高估。与 Qwen2.5-1.5B 的对比仅在特定协议下进行,无法全面反映模型综合能力。
  • - **硬件依赖与可复现性**:整个配方高度依赖 **RTX 5090** 的 FP8 支持和特定通信拓扑,其他消费级 GPU(如 4090、A100 等)上 FP8 效率差异显著,迁移需要重新调优并行策略和微批量大小。论文仅验证 2B 参数规模,未展示扩展到更大模型时成本效率是否保持;且 RTX 5090 供货不稳定,$5090 的硬件成本假设在实际采购中可能波动。
  • - **对比同类工作的边际优势**:与 **SmolLM3-3B**、**TinyLlama** 等已有低成本配方相比,Puro-2B 在成本上确实更低,但其模型性能仅接近 **Qwen2.5-1.5B**,而后者本身参数量更小、训练资源消耗也较低。算法层面的创新(hyperball optimization、curriculum model averaging)多为工程技巧的组合,缺乏全新方法贡献,因此对研究社区的指导意义主要在于工程优化而非理论突破。
论文Kairong Luo2026-08-27原文

相关内容