TrainBridge speeds up training on NPUs using GPU optimization clues
TrainBridge: Agentic NPU Training Optimization Guided by GPU Compilation Evidence
Computation and LanguageArtificial Intelligence
Summary
Training machine learning models on special chips called NPUs often runs slower because some software optimizations are missing. The authors found a way to make NPU training faster by learning from how GPUs, a different kind of chip, run the same tasks very efficiently. Their method, TrainBridge, copies the GPU’s smart tricks for optimizing programs and adjusts them for NPUs, boosting speed by about 1.7 to 2 times.
What this means in practice
- •For machine learning engineers: Optimize training workflows on Ascend NPUs by transferring GPU compilation insights to improve speed and reduce overhead.
- •For hardware acceleration developers: Enhance NPU compiler design by integrating profiling guided by GPU reference executions to achieve better operator fusion and program efficiency.
Authors
Yupeng Chen, Kunxiang Zhao, Xingjian Wang, Yue Zhao, Yijing Wu, Ye Liu, Kai Li, Congliang Chen, Ruoyu Sun
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.