ForgeTrain: Forging Production-Grade Training Frameworks via Harness-Driven AI Development

Abstract

Training large models still relies on general-purpose frameworks such as Megatron-LM, whose generality tax constrains scenario-specific optimization and adds runtime overhead through accumulated abstraction. AI code generation reduces the cost of building a framework, and makes it affordable to forge one per scenario. We propose Forge Engineering: building a dedicated implementation from scratch for each scenario and iteratively optimizing it toward peak performance under correctness and usability constraints. Dedicated implementations inherit no abstraction boundaries, so they can integrate optimizations across the stack and reach a higher performance ceiling. We instantiate this paradigm for training frameworks as ForgeTrain, which holds a trusted framework as a golden reference and relaxes equivalence monotonically from Bit-for-Bit to Surpass. Experiments across multiple model--hardware configurations show that ForgeTrain consistently produces correct training engines and improves MFU over established training frameworks by 4.7--33.2%. To our knowledge this is the first production-grade training framework forged end-to-end by AI to match or surpass its human reference.