arXiv:2606.24780cs.AIcs.LG2026-06被引 1

用原生C++/CUDA打造高效训练框架,性能超越PyTorch

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

论文配图:BluTrain: A C++/CUDA Framework for AI Systems
图 1 · 摘自论文原文
  • 全栈自研:从张量运算到编译器均用原生代码实现
  • 8卡环境下吞吐达40.7万词/秒,内存占用减少22%
  • 适合追求极致性能与可控性的系统级深度学习开发者

深度学习的规模化进展更多依赖系统工程而非模型设计:模型训练时的吞吐、内存开销及数值精度主要取决于架构在硬件上的表达方式。为实现对硬件表达的绝对控制,同时抽象系统复杂性以保障建模流畅,并消除重复编排逻辑,BluTrain从零开始构建为一个基于标准C++和核心CUDA编程模型的轻量、通用训练框架。所有组件均原生实现:带反向自动微分的类型化张量模块、线性代数库、缓存分配器、多模式分布式执行模块,以及基于MLIR的深度学习编译器。在8卡6000 Ada系统上以FP32训练1.24亿参数GPT-2基线模型的正式评估中,BluTrain在吞吐上达到平均40.7万词/秒(优于PyTorch的39.5万词/秒),内存开销降低最高达22%,同时严格保持数值精度,并收敛至略低的最终验证损失。每个层级均可直接原生调优,性能上限由框架自身决定。

原文摘要 · Abstract (English)

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.

AI框架CUDA系统优化训练加速

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。