ForgeTrain：全自动 AI 生成训练框架，降低 10% 大模型预训练成本

原创发布日期：2026-05-29

一、ForgeTrain 是什么

ForgeTrain 是OpenBMB 开源的 v0.1.0 版本大模型预训练框架，是全球首个全程由 AI 智能体循环自主编写、零人工代码修改的生产级 LLM 预训练基础设施，核心面向 NVIDIA H100 系列显卡，专为 MiniCPM4‑0.5B/8B 模型预训练深度优化，采用 Apache‑2.0 开源协议，已完成真实模型权重训练验证，可直接用于工业级预训练任务。

该框架由 AI Agent 闭环自主完成代码编写、调试、优化与提交，彻底颠覆传统大模型训练框架依赖海量人工研发的模式，实现AI 制造 AI的工程落地，在算力利用率、训练速度上显著超越业界主流框架，为大模型预训练提供低成本、高效率、高可靠的全新解决方案。

二、功能特色

100% AI 自主研发，零人工介入
框架全程由 AI 智能体循环自主完成：读取参考代码→实现功能→启动训练任务→解析日志→定位故障→修复代码→通过校验→提交版本，无任何人工代码修改，是全球首款真正意义上全 AI 生成的生产级大模型训练框架。
算力利用率行业领先

在 64×H100、BF16 精度、纯数据并行场景下，MFU 达 44.13%，较 Megatron‑LM 基准（约 40%）提升约 10%
单算子 MFU 最高可达**90%**，自研 FlashAttention 算子性能比肩 FA4，超越 Transformer Engine 与 FA3

生产级稳定验证
已完成 MiniCPM4‑0.5B 全量预训练，产出真实可用模型权重，非演示项目；支持单机 8×H100、多机多卡分布式训练，满足工业级预训练稳定性与可靠性要求。
极致算子优化
AI 自主实现 5 个 CuTeDSL 自定义 GEMM 算子、从零编写 FlashAttention 内核，支持 CUDA Graph、Triton 融合核、通信‑计算重叠，兼顾训练速度与显存效率。
极简部署与使用
支持 AI 编程助手一键部署，无需手动逐行输入命令；提供一键导出 HuggingFace 格式权重，降低框架使用门槛。
双引擎适配不同规模模型

MiniCPM4‑0.5B 引擎：纯数据并行，64 卡 H100 达 44.13% MFU
MiniCPM4‑8B 引擎：张量并行 TP=2，单机 8 卡 H100 达 50.9% MFU

ForgeTrain：全自动 AI 生成训练框架，降低 10% 大模型预训练成本

三、技术细节

1. 核心架构：Agent Loop + Harness 脚手架

Harness（即将上线）：驱动 AI 智能体循环的核心脚手架，负责制定评测标准、控制智能体执行流程、校验代码质量
Training Engine：AI 智能体自主生成的训练引擎，分 0.5B/8B 两个版本，直接用于模型预训练
两者为严格的生产者‑产品关系，执行bash agent‑loop.sh即可启动全自动生成流程，无需人工干预。

2. AI 智能体闭环研发流程

读取参考训练框架（Megatron‑LM v0.15）代码与规范
自主实现训练核心模块、算子、分布式逻辑
启动训练任务，实时解析日志定位故障
自动修复代码漏洞、优化性能瓶颈
通过前向/反向对齐、长训练统计校验等关卡
自动提交合格代码，完成版本迭代

3. 两级关卡驱动收敛机制

Stage 1（M1‑M6）：实现比特级前向/反向对齐、多步训练、长训练统计校验（损失相对差异＜1%，MFU≥36%）
Stage 2：单算子 CUDA 内核优化，每个算子迭代 30 轮，择优选用，每次合并后通过分布式集成校验。

4. 高性能技术组件

CUDA Graph：支持 5 种捕获粒度，可自由组合 Bucket 梯度聚合、分片优化器、权重梯度重叠
Triton 融合核：自主实现 CE 前后向、SwiGLU、RMSNorm+残差、RoPE、融合 Adam+参数同步等融合算子
自定义算子：5 个 CuTeDSL GEMM 算子、自研 FlashAttention，性能比肩业界顶尖方案。

5. 环境依赖

GPU：NVIDIA H100 80GB（SM90，Hopper 架构）
GPU 数量：单机 8 卡用于完整预训练，单卡用于早期对齐
CUDA 版本：≥12.x
Python 版本：≥3.11
PyTorch 版本：≥2.4
适配模型：MiniCPM4‑0.5B（DP‑only）、MiniCPM4‑8B（TP=2）
精度：BF16

四、应用场景

大模型企业预训练
适合拥有 H100 集群的 AI 企业，用于 MiniCPM4 系列基座模型预训练，降低 10% 训练成本与耗时，提升研发效率。
科研机构模型研发
科研团队可快速复现预训练流程，聚焦模型结构创新，无需投入大量人力开发训练框架，缩短研发周期。
AI 框架技术验证
用于验证AI 自主生成工业级系统软件的可行性，探索 AI 辅助研发、全自动代码生成的前沿方向。
小参数模型高效训练
针对 0.5B‑8B 小参数模型，提供轻量化、高效率预训练方案，适配资源有限但需快速迭代模型的场景。
分布式训练教学与研究
作为分布式训练、算子优化、CUDA 编程的教学案例，直观展示高性能训练框架的核心原理。

五、使用方法

1. 环境准备

确保满足以下条件：

Python≥3.11、CUDA≥12.x、PyTorch≥2.4
NVIDIA H100 80GB 显卡，单机 8 卡用于完整训练

2. 一键部署（推荐 AI 编程助手执行）

将以下命令粘贴至 Cursor/Claude Code 等编程助手，自动完成部署与测试：

Following this project's exports/train_engine_0.5B/README.md,
run a 5-step minimal pretraining demo on the current node:

1. Check the environment (Python ≥ 3.11, CUDA ≥ 12.x, H100, PyTorch ≥ 2.4)
  and install anything missing;
2. Install the repo: pip install -e . and HF deps: pip install datasets transformers;
3. Import smoke test:
  PYTHONPATH=src python -c "from training_engine_tensor import config; print('OK')"
4. Run 5 steps on HF GSM8K:
  torchrun --standalone --nproc-per-node=1 \
   -m training_engine_tensor pretrain \
   --num-steps 5 --global-batch-size 1 --micro-batch-size 1 \
   --seq-length 4096 \
   --hf-dataset openai/gsm8k --hf-dataset-config main \
   --hf-text-template "Question: {question}\nAnswer: {answer}" \
   --tokenizer-path openbmb/MiniCPM4-0.5B \
   --save-dir ./checkpoints/demo
5. Print the final loss, step time, and MFU.

3. 手动部署流程

克隆仓库

git clone https://github.com/OpenBMB/ForgeTrain.git
cd ForgeTrain/exports/train_engine_0.5B

安装依赖

pip install -e .
pip install datasets transformers

验证安装

PYTHONPATH=src python -c "from training_engine_tensor import config; print('OK')"

预编译算子（首次运行）

PYTHONPATH=src CUSTOM_GEMM=1 OP_ATTENTION=v1 python scripts/precompile_ops.py

单机 8×H100 启动预训练

torchrun --standalone --nproc-per-node=8 \
  -m training_engine_tensor pretrain \
  --num-steps 200 \
  --global-batch-size 1280 --micro-batch-size 10 \
  --seq-length 4096 \
  --hf-dataset openai/gsm8k \
  --hf-dataset-config main \
  --hf-text-template "Question: {question}\nAnswer: {answer}" \
  --tokenizer-path openbmb/MiniCPM4-0.5B \
  --save-dir ./checkpoints/run1

运行后输出格式：[STEP 200] loss=X.XXX | step_time=XXXms | mfu=44.XX%，MFU 稳定在 44% 左右即为正常。

六、竞品对比

选取业界主流大模型预训练框架 Megatron‑LM、DeepSpeed 与 ForgeTrain 对比，核心差异如下：

对比维度	ForgeTrain	Megatron‑LM	DeepSpeed
代码生成方式	100% AI 智能体自主生成，零人工修改	人工研发为主	人工研发为主
H100 MFU	44.13%（MiniCPM4‑0.5B，DP‑only）	约 40%	约 38%
性能优势	较 Megatron‑LM 提升约 10%	业界基准	通用性强，性能中等
自定义算子	自主实现 CuTeDSL GEMM、FlashAttention，比肩 FA4	依赖 TE/FA 上游算子	支持自定义，优化程度一般
模型适配	专注 MiniCPM4‑0.5B/8B	支持多模型家族	支持全尺寸模型
部署难度	极简，AI 一键部署	复杂，需人工配置	中等，需手动调参
权重导出	单脚本一键导出 HuggingFace	手动操作	需额外配置
开源协议	Apache‑2.0	Apache‑2.0	Apache‑2.0

七、常见问题解答

Q：运行时提示“CUDA 版本不兼容”怎么办？

A：ForgeTrain 要求 CUDA≥12.x，需升级 CUDA 版本，并确保 PyTorch 与 CUDA 版本匹配，重新安装依赖后重试。

Q：非 H100 显卡可以运行吗？

A：当前 v0.1.0 仅支持 NVIDIA H100（SM90）显卡，暂不兼容其他型号显卡，后续版本将逐步扩展适配范围。

Q：MFU 低于 40% 是什么原因？

A：可能是算子未预编译、batch size 设置不合理、显卡驱动版本过低，建议执行预编译脚本、调整 batch size、更新显卡驱动后重新运行。

Q：如何导出 HuggingFace 格式权重？

A：框架提供专用导出脚本，直接运行对应目录下的导出脚本，即可将训练 checkpoint 转换为 HuggingFace 格式，无需手动转换。

Q：单卡 H100 可以完整训练吗？

A：单卡仅支持早期对齐与简易测试，完整预训练需单机 8×H100 或多机多卡集群，确保算力与显存满足需求。

Q：训练过程中断后如何恢复？

A：框架支持断点续训，启动命令中添加续训参数，指定已保存的 checkpoint 目录，即可从中断步骤继续训练。

八、相关链接

ForgeTrain GitHub 仓库：https://github.com/OpenBMB/ForgeTrain
MiniCPM4 模型官方地址：https://huggingface.co/openbmb/MiniCPM4-0.5B
OpenBMB 官方社区：https://openbmb.github.io

九、总结

ForgeTrain 作为 OpenBMB 推出的全球首款AI 全自主研发生产级大模型预训练框架，以零人工代码介入、领先的算力利用率、极简的部署流程、生产级的稳定性，重新定义大模型训练框架研发模式，在 H100 集群上实现超越 Megatron‑LM 的训练性能，已完成 MiniCPM4 系列模型真实预训练验证，兼顾高效、低成本与易用性，为大模型工业级预训练、AI 自主研发探索提供关键支撑，是小参数模型高效训练与 AI 生成系统软件的标杆级开源项目。