BluTrain:基于C++/CUDA的AI训练框架

BluTrain: A C++/CUDA Framework for AI Systems

精选理由

这个新框架用C++从头写,训练GPT-2比PyTorch快3%且省内存22%,适合追求极致性能的开发者。

AI 摘要

BluTrain是一个用标准C++和CUDA实现的AI训练框架。在8-GPU 6000 Ada系统上训练124M参数GPT-2模型(FP32),其吞吐量达407K tokens/s,比PyTorch的395K tokens/s高约3%。同时内存占用减少22%,且严格保持数值精度。框架包含原生实现的张量模块、反向模式自动微分、线性代数库、缓存分配器、分布式执行和MLIR编译器。

原文 · arXiv cs.AI

BluTrain: A C++/CUDA Framework for AI Systems

Progress in deep learning is, at scale, more a matter of systems engineering than of modelling: the behaviour of a model in training (its throughput, its memory footprint, and the numerical fidelity of the result) is determined less by the architecture itself than by how that architecture is expressed on the hardware. To achieve absolute control over this hardware expression while abstracting away systems complexity to make modelling seamless and eliminating the need for repetitive orchestration logic, BluTrain was architected from first principles as a robust, lightweight, and architecture-general training framework in standard C++ and the core CUDA programming model. Every layer is implemented natively: a typed tensor module with reverse-mode autograd, a linear-algebra library, a caching allocator, a multi-mode distributed-execution module, and an MLIR-based deep-learning compiler. In formal evaluations training a 124M-parameter GPT-2 baseline in FP32 on an 8-GPU 6000 Ada system, BluTrain outperforms industry-standard baselines in both throughput (sustaining an average of 407K tokens/s versus PyTorch's 395K tokens/s) and memory efficiency (achieving up to a 22% footprint reduction), while strictly preserving numerical fidelity and converging to a marginally lower final validation loss. With every layer explicitly open to native tuning, the performance ceiling is the framework's own to raise.