ForgeTrain is an LLM pretraining framework entirely authored by an autonomous AI agent loop, offering a novel approach to framework development. It targets researchers and engineers seeking highly optimized LLM training infrastructure, providing a production-validated, AI-generated solution that claims superior performance (MFU) over established baselines like Megatron-LM.

How It Works

The core innovation is a "Self-Diagnosing Agent Loop" that autonomously iterates through reference implementation, job launching, log parsing, root-cause analysis, and patching, with zero human intervention. This loop produces the training framework, including custom GEMM and FlashAttention kernels written in CuTeDSL and Triton. The process employs a two-stage gate-driven convergence: Stage 1 focuses on alignment and distributed training stability, while Stage 2 optimizes per-operator CUDA kernels and integrates them. This AI-driven optimization aims for high hardware utilization (MFU) and performance.

Quick Start & Requirements

• Primary Install: pip install -e . within the exports/train_engine_0.5B/ directory.
• Prerequisites: NVIDIA H100 80GB (SM90), CUDA 12.x, PyTorch >= 2.4, Python >= 3.11. Full pretraining requires 8x H100; early stages can run on a single GPU.
• Setup: Precompiling operators (scripts/precompile_ops.py) is recommended for the first run, taking only a few seconds after initial build. Full pretraining commands are provided for single-node (8x H100) and multi-node setups.
• Links: exports/train_engine_0.5B/README.md, exports/train_engine_8b/README.md for detailed CLI documentation.

Highlighted Details

• Achieves 44.13% MFU on H100 for MiniCPM4-0.5B (BF16, DP-only), approximately 10% above the Megatron-LM baseline.
• The entire framework and custom GEMM/Attention kernels (up to 90% per-op MFU) are 100% AI-authored with zero manual edits.
• Self-built FlashAttention implementation is on par with FA4 and outperforms Transformer Engine/FA3.
• Production-validated: MiniCPM4-0.5B was fully pretrained, yielding real model weights.

Maintenance & Community

No specific details on contributors, sponsorships, or community channels (Discord/Slack) are provided in the README. The project is actively developed, with v0.1.0 released in May 2026.

Licensing & Compatibility

Licensed under the Apache License 2.0. Vendored code snapshots retain their upstream copyright headers, with specific notices in NOTICE.md files within the respective quack/ subdirectories. Compatible with commercial use under Apache 2.0 terms.

Limitations & Caveats

The "Harness" scaffolding, which drives the autonomous agent loop, is explicitly marked as "coming soon." The current release (v0.1.0) focuses on the training engine for MiniCPM4-0.5B (DP-only) and MiniCPM4-8B (TP=2), with broader parallelism strategies and model families supported by Megatron-LM not yet covered. The agent-friendly deploy instructions suggest it's best suited for AI agents.