Purpose
ForgeTrain is a training framework for large language model (LLM) pretraining. Its uniqueness lies in the fact that the entire framework was autonomously written by an AI Agent loop, without any human modification. The project also provides the Harness scaffold that produced this framework (to be released soon), for reproducing this Agent-driven development process on other training stacks.
Main capabilities
- 100% written by an Agent loop: The entire framework was generated by an AI Agent running in an automatic loop mode, with zero human editing.
- Self-diagnosing Agent loop: The loop process includes reading reference code → implementation → launching jobs → parsing logs → root cause analysis → patching → passing gates → committing, all autonomously.
- High performance: Achieves 44.13% MFU on H100, about 10% higher than the Megatron-LM baseline (about 40%), validated under 64 H100s, BF16, pure data parallel configuration.
- Production-validated: Complete pretraining of MiniCPM4-0.5B, producing real model weights (not a demo).
- GEMM + Attention kernels written by the Agent loop: Single-operator MFU reaches up to 90%; the from-scratch FlashAttention implementation surpasses Transformer Engine / FA3 and is on par with FA4.
- Out of the box: Supports CUDA Graph, Triton fused kernels, and communication-computation overlap.
- 8B version: Runs MiniCPM4-8B on a single 8×H100 host, tensor parallelism=2, 50.9% MFU, about 8% higher than the Megatron-LM baseline.
Usage
The project contains two usable training framework export directories:
exports/train_engine_0.5B/: 0.5B training framework, ready to use.exports/train_engine_8b/: 8B training framework.
Environment requirements: Python ≥ 3.11, CUDA 12.x, PyTorch ≥ 2.4, NVIDIA H100 80GB (SM90). Full pretraining requires 8×H100; the early alignment stage can run on a single GPU.
Quick deployment is designed for AI Agents, allowing prompts to be pasted into Cursor / Claude Code / Codex / Cline, where the Agent automatically reads the README, installs dependencies, runs smoke tests, and reports MFU. The minimal pretraining demo contains 5 steps. Dataset argument --hf-dataset accepts a HuggingFace Hub name (such as openai/gsm8k) or a local Parquet / Arrow / JSON directory.
On 8×H100 with BF16, the expected MFU is about 44%, and with GBS=1280 / MBS=10 / seq=4096 the step time is about XXXms.
For full documentation (including multi-node training, checkpoint recovery, and HuggingFace export), see the README in each export directory.
Use Cases
- Teams that need to perform LLM pretraining on H100 clusters.
- Developers who want to study or reproduce the development paradigm of "AI Agents autonomously writing training frameworks."
- Researchers using the MiniCPM4 series model architecture.
Limitations
ForgeTrain v1 is limited in scope to MiniCPM4-0.5B (data parallelism only) and MiniCPM4-8B (tensor parallelism=2) × BF16; Megatron-LM supports a broader range of model families and parallelism strategies. The comparison is based on Megatron-LM v0.15 on the same hardware (H100, SM90). The Harness scaffolding and the accompanying documentation for the 8B framework have not yet been released (coming soon).
Version and License
Current version: v0.1.0 (released 2026-05). The technical report paper "ForgeTrain: Forging Production-Grade Training Frameworks via Harness-Driven AI Development" was released in 2026-08. Licensed under the Apache License 2.0. The project is built on top of open-source work such as CUTLASS / CuTeDSL, FlashAttention-4, TransformerEngine, Megatron-LM, Cursor, and MiniCPM4.
