TrainBridge: Optimizing AI Model Training on NPUs

Paper: TrainBridge: Agentic NPU Training Optimization Guided by GPU Compilation Evidence

Authors: Yupeng Chen, Kunxiang Zhao, Xingjian Wang, Yue Zhao, Yijing Wu, Ye Liu, Kai Li, Congliang Chen, Ruoyu Sun

Executive Summary

This research focuses on enhancing the efficiency of training artificial intelligence models on Ascend Neural Processing Units (NPUs). Current training programs for NPUs fail to fully utilize advanced compilation techniques, leading to incomplete computations and unnecessary delays. The study introduces TrainBridge, a novel approach that leverages optimization strategies developed for Graphics Processing Units (GPUs) to improve NPU performance.

TrainBridge works by transferring the highly efficient execution patterns from GPU training to NPU tasks. It begins by analyzing the optimized processes used in GPU training to extract key optimization strategies. These insights guide the reconfiguration of NPU programs, which are then fine-tuned through a process of iterative profiling. This method aims to minimize inefficiencies that typically arise during NPU training, which often involves costly trial-and-error adjustments.

The experiments conducted in this study demonstrate significant improvements in training speed. Specifically, TrainBridge achieved an average speedup of 1.7 times across four different workloads, which included tasks such as large language model pretraining, image classification, and vision-language training. Moreover, when looking at 27 benchmark records related to Modded-NanoGPT, TrainBridge delivered an impressive average speedup of nearly 2.0 times compared to highly optimized GPU implementations.

These findings are important for companies looking to enhance their AI training processes. By adopting TrainBridge, organizations may see reduced training times and improved resource utilization, ultimately leading to faster model development and deployment. This research highlights the potential for cross-platform optimization strategies to enhance the performance of NPUs, which are becoming increasingly vital in the field of AI.

While the results stem from controlled experiments rather than real-world applications, they suggest that the implementation of TrainBridge could represent a practical solution for enterprises aiming to leverage NPUs more effectively. This approach could be particularly beneficial for sectors that rely heavily on AI, such as technology, healthcare, and finance, where training efficiency can significantly impact operational costs and project timelines.

Academic Abstract

We study how to improve training efficiency on Ascend NPUs. In the NPU training programs we study, framework-level compilation optimizations such as operator fusion are only partially applied, leaving fragmented computation and avoidable execution overheads. Profiling-guided tuning can address these inefficiencies, but it often requires costly trial and error. To make training on NPUs more efficient, we propose TrainBridge, an agentic approach for transferring GPU optimization guidance to NPU training. TrainBridge uses the highly optimized execution produced by mature GPU compilation stacks as a task-specific reference for NPU tuning. It extracts optimization intents from the GPU training execution to guide NPU program rewrites and further refines the rewritten program through iterative profiling. Our GPU-to-NPU migration experiments show that TrainBridge achieves an average 1.7x speedup across four workloads covering LLM pretraining, image classification, and vision-language training, and an average speedup of nearly 2.0x across 27 Modded-NanoGPT speedrun records with highly optimized GPU implementations.