Nova用编译器技术把Transformer模型全图融合,提升大模型训练速度。
Nova: An End-to-End MLIR Compiler for Deep Learning

- 从计算结构直接生成细粒度内核,实现端到端自动编译。
- 在GPT-2上达421K tokens/秒,比原生快3.7%,比torch.compile快0.2%。
- 适合需要高效部署大模型的工程团队和研究者使用。
大规模深度学习模型的性能高度依赖于高层数学操作对底层硬件的有效映射。尽管高层张量框架提供灵活抽象,但其执行模型缺乏全局视图,难以最大化硬件利用率,常需依赖封闭的手写内核库处理复杂操作如Attention。为此,我们提出Nova的下一代版本,一个自动化端到端即时编译器,通过直接从计算结构合成细粒度内核,实现对硬件映射的绝对控制。本工作将Nova编译流程扩展至原生支持完整Transformer架构。通过捕获即时执行并统一前向与反向传播为单一值语义方言,Nova实现了激进的全图优化。不再依赖固定预编译库调用,而是广泛进行跨操作融合,将复杂的因果注意力子图、逐元素操作及内存密集型归一化直接合并为单一融合内核,显著减少全局内存往返次数。在Ada 6000 GPU上训练完整GPT-2架构的评估中,Nova展现更优端到端吞吐率,平均达421K tokens/秒,优于自身原生执行的406K tokens/秒,以及torch.compile的405K tokens/秒。通过编译器原生融合大幅降低内存瓶颈开销,Nova实现了现代硬件上高效的大语言模型全编译,且严格保持数值一致性。
原文摘要 · Abstract (English)
The performance of deep learning models at scale relies heavily on how effectively high-level mathematical operations are mapped to underlying physical hardware. While high-level tensor frameworks provide flexible abstractions, their execution models inherently lack the whole-graph visibility required to maximize hardware utilization, often forcing a reliance on opaque, hand-written kernel libraries for complex operations like Attention. To bridge this gap, we present the next iteration of Nova, an automated end-to-end JIT compiler that achieves absolute control over hardware mapping by synthesizing fine-grained kernels directly from the computation's structure. In this work, we extend Nova's compilation pipeline to natively support full Transformer architectures. By capturing eager executions and unifying forward and backward passes into a single value-semantic dialect, Nova unlocks aggressive whole-graph optimizations. Rather than relying on rigid, pre-compiled library calls, Nova focuses on extensive cross-operator fusions, collapsing complex causal attention sub-graphs, element-wise operations, and memory-bound normalizations directly into single fused kernels to drastically reduce global memory roundtrips. In our evaluations training a full GPT-2 architecture on Ada 6000 GPUs, Nova demonstrates superior end-to-end throughput, averaging 421K tokens/second compared to 406K for our own eager execution and 405K for torch.compile. By drastically reducing memory-bound overheads through compiler-native fusion, Nova enables efficient full LLM compilation on modern hardware while strictly maintaining numerical parity.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。