论文
BluTrain:面向AI系统的C++/CUDA框架
BluTrain: A C++/CUDA Framework for AI Systems
摘要
大规模深度学习的进步更多是系统工程而非建模问题:模型在训练中的行为(吞吐、内存占用与结果的数值保真)较少取决于架构本身——更多取决于该架构如何在硬件上表达。为取得对这种硬件表达的绝对控制——同时抽象掉系统复杂性使建模无缝、并消除重复编排逻辑的需要——BluTrain从第一性原理被架构为标准C++与核心CUDA编程模型中的鲁棒、轻量、架构通用的训练框架。每一层均原生实现:带反向模式自动微分的类型化张量模块、线性代数库、缓存分配器、多模式分布式执行模块——以及基于MLIR的深度学习编译器。在8-GPU 6000 Ada系统上以FP32训练124M参数GPT-2基线的正式评估中——BluTrain在吞吐(平均407K token/s对PyTorch的395K token/s)与内存效率(占用最多降22%)上都超越行业标准基线——同时严格保持数值保真并收敛到略低的最终验证损失。每一层都显式开放原生调优——性能上限由框架自己抬升。
