PaperlayerView exact revision on arXiv

Source-linked AI summary

TrainBridge: Agentic NPU Training Optimization Guided by GPU Compilation Evidence

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

arXiv:2610.29171v1cs.CLcs.AI

TL;DR

NPU training programs retain compilation-related inefficiencies, and profiling-guided tuning can require costly trial and error. TrainBridge transfers optimization intents from mature GPU compilation to NPU rewrites and refines them through profiling. It reports average speedups of 1.7× across four workloads and nearly 2.0× across 27 Modded-NanoGPT records.

  • Problem

    NPU compilation optimizations are only partially applied, leaving fragmented computation, while profiling-guided tuning requires costly iterative analysis and validation.

    Source: §1 · p. 1

  • Method

    TrainBridge extracts optimization intents from eager-versus-compiled GPU execution to guide NPU rewrites, then iteratively profiles the updated program to address residual bottlenecks.

    Source: §1 · pp. 1–2

  • Results

    TrainBridge achieves an average 1.7× speedup across four workloads and nearly 2.0× across 27 Modded-NanoGPT records, relative to initial NPU implementations.

    Source: §1 · p. 2, §7 · p. 10

  • Takeaways & Limitations

    The results show that GPU compilation evidence can guide NPU-specific program optimization and help address compilation-related inefficiencies in NPU training.

    Source: §7 · p. 10

    TrainBridge requires a source GPU implementation with useful framework-level compilation evidence, limiting it to workloads supported by a mature source stack.

    Source: §6 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

NPU training programs retain compilation-related inefficiencies, while profiling-guided tuning can require costly iteration. TrainBridge uses optimization evidence from compiled GPU execution to guide NPU rewrites and profiling, achieving speedups across additional workloads and Modded-NanoGPT records.

Motivation

  • NPU training programs retain fragmented computation, redundant operations, and materialized intermediates because framework-level compilation optimizations are only partially applied.Profiling-guided tuning can diagnose bottlenecks, but repeated analysis, modification, execution, and validation makes optimization time-consuming under limited budgets.Source: §1 · p. 1

Evaluation scope

  • The evaluation spans LLM pretraining, image classification, and vision-language training, and includes an NPU Modded-NanoGPT speedrun benchmark adapted from historical GPU implementations.The speedrun targets end-to-end pretraining efficiency under a fixed validation target.Source: §1 · p. 2

Approach

  • TrainBridge extracts optimization intents from eager-versus-compiled GPU execution, guides NPU rewrites, then iteratively profiles the rewritten program to address remaining bottlenecks.The approach uses GPU compilation decisions as task-specific guidance for NPU tuning.Source: §1 · pp. 1–2

Results

  • 1.7× average speedup across four workloads and nearly 2.0× across 27 Modded-NanoGPT records, relative to initial NPU implementations.Across the four additional workloads, the average NPU-to-GPU training-time ratio falls from 2.8× to 1.5×.Source: §1 · p. 2

2 NPU Training: Compilation Gaps and Optimization Challenges

The section establishes the migration setting and shows that NPU training retains optimization opportunities that GPU compilation has already addressed. These gaps make agent tuning difficult, while GPU compilation evidence offers concrete patterns for NPU-specific rewrites.

2.1 Training Workloads and Migration Setting

  • The migration study covers five workloads across language-model pretraining, image classification, and vision-language training, evaluated on one to eight devices.For Modded-NanoGPT, migrated programs preserve the training schedule, hyperparameters, and model dynamics and must reach validation loss below 3.28 at the same training step.Source: §2.1 · p. 2

2.2 Compilation Gaps in NPU Training

  • Across four workloads besides Modded-NanoGPT, the initial NPU-to-GPU end-to-end training-time ratio averages 2.8×; Record 50 has a 4.40× ratio.The reported throughput comparison suggests hardware capability alone cannot fully explain the end-to-end gaps, motivating investigation of software-level optimization.Source: §2.2 · p. 3

2.3 Challenges in Agentic Performance Tuning

  • A generic coding agent produced valid changes, but many did not improve end-to-end training time; profiling identifies expensive regions without specifying their acceleration potential or best optimization.Repeated analysis, modification, execution, and validation can be inefficient under a limited tuning budget.Source: §2.3 · p. 3

2.4 GPU Compilation as Optimization Evidence

  • On eight NVIDIA H100 SXM GPUs, torch.compile reduces Record 50 training time from around 338 seconds to 128 seconds.Table 1 reports GPU and NPU training times, with ratios referenced to the compiled GPU program.Source: Table 1 · p. 3, §2.4 · p. 3
  • For Record 50 KeyShift, overlapping-slice assignment remains as ViewCopy operations on NPU, while compiled GPU execution removes the inefficiency through functionalization and kernel fusion.TrainBridge extracts the intent to avoid overlapping view updates and proposes an NPU-specific expression; Table 2 reports local KeyShift latency before and after rewriting for BF16 input shape [1, 24576, 6, 128].Source: §2.4 · pp. 3–4, Table 2 · p. 4

3 TrainBridge

TrainBridge uses eager-to-compiled GPU changes as evidence for NPU-specific rewrites, then profiles the rewritten NPU program to address residual bottlenecks. The formulation accounts for the limits of GPU evidence, while Stage 1 aligns related computations and Stage 2 targets NPU-specific overheads.

3.1 Problem Formulation

  • The difference between eager and compiled GPU computation regions provides evidence of optimization intent that can guide NPU rewrites when corresponding inefficiencies remain after migration.The goal is to improve end-to-end NPU training time while preserving correctness under a limited tuning budget.Source: §3.1 · p. 4
  • GPU evidence cannot address every migrated-NPU inefficiency: backend-specific transformations may not help, and GPU evidence does not expose host-scheduling, communication, or memory-management overheads.Source: §3.1 · pp. 4–5

3.2 Stage 1: NPU Optimization Guided by GPU Compilation Evidence

  • TrainBridge relates eager-to-compiled GPU changes to corresponding NPU regions, then proposes validated NPU-specific rewrites that preserve computation semantics.The agent examines operator, dtype, layout, intermediate, and kernel-boundary differences to check whether an inefficiency remains.Source: §3.2 · p. 5

3.3 Stage 2: Profiling-Guided Residual Optimization

  • Stage 2 combines operator statistics, host and device timelines, communication events, memory activity, and source inspection to identify residual NPU bottlenecks.It revalidates candidate changes and repeats profiling until performance no longer improves or the tuning budget is exhausted.Source: §3.3 · pp. 5–6

4 Experimental Setup

TrainBridge is evaluated on Ascend 910C migrations with shared initial NPU implementations, fixed tuning budgets, and validation criteria. Across Modded-NanoGPT records and four additional workloads, it improves validated NPU training time; stage comparisons examine GPU evidence and NPU profiling.

4.1 Tasks and Migration Protocol

  • The evaluation migrates five workloads to Ascend 910C NPUs; 27 randomly sampled historical Modded-NanoGPT records are treated as separate tuning tasks.A valid speedrun optimization preserves the original training schedule and reaches validation loss below 3.28 at the same training step.Source: §4.1 · p. 6

4.2 Hardware and Software

  • Experiments use up to eight H100 SXM GPUs or eight Ascend 910C NPUs, and report end-to-end training time excluding tuning time.Modded-NanoGPT uses eight devices of each type in its respective experiments; the main agent harness is Codex with GPT-5.6-SOL at medium reasoning effort.Source: §4.2 · p. 6

4.3 Research Questions

  • The research questions test whether TrainBridge finds faster validated NPU programs under a fixed budget and whether each optimization stage contributes.Methods start from the same initial NPU implementation; controlled Record 50 comparisons use a 12-hour budget, with three runs for RQ1 and one per method for RQ2.Source: §4.3 · p. 6

4.4 Overall Agentic Tuning Effectiveness

  • Across three independent runs on Record 50, TrainBridge’s average best validated training time was 283.2 seconds versus 419.2 seconds for the generic agent, a 32.4% reduction.Figure 2 reports results under a 12-hour budget using the same Codex harness, with standard-deviation error bars.Source: §4.4 · p. 7, Figure 2 · p. 7

4.5 Contribution of the Two Stages

  • Stage 1 reduces Record 50 training time from 563.0 to 299.7 seconds, and Stage 2 further lowers it to 285 seconds; profiling-only Stage 2 reaches around 334 seconds.The comparison indicates GPU compilation evidence adds guidance beyond profiling, while profiling further improves the Stage 1 program.Source: §4.5 · pp. 7–8, Figure 3 · p. 7

4.6 Results across Modded-NanoGPT Records

  • 1.95× average speedup across 27 Modded-NanoGPT records, reducing the average NPU/GPU training-time ratio from 4.0× to 2.0× versus initial NPU implementations.All optimized programs satisfy the validation-loss requirement; Figure 4 displays ten sampled records.Source: §4.6 · p. 8, Figure 4 · p. 8

4.7 Results on Additional Training Workloads

  • 2.22×, 1.68×, 1.25×, and 1.84× speedups over initial NPU implementations for CIFAR-10 AirBench, Qwen3, ViT, and nanoVLM, respectively.The four workloads span image classification, language-model training, and vision-language training; Figure 5 compares training time relative to the compiled GPU baseline.Source: §4.7 · p. 9, Figure 5 · p. 9

5 Related Works

Prior work applies agents to kernel and systems optimization, using execution feedback, profiling, search structures, and reusable experience. TrainBridge adds same-workload GPU compilation evidence as guidance for NPU tuning, followed by target-side profiling.

Agentic performance optimization

  • Agentic optimization has addressed GPU kernels, NPU kernel synthesis, broader ML systems, and end-to-end LLM deployment on AMD NPUs.Examples include Ascend C kernel optimization in large-scale training and GPU training-engine construction.Source: §5 · p. 9

Guiding the optimization search

  • Training-free search guidance includes hardware profiles, candidate search trees, and retrieved expert knowledge or optimization experience.Training-based guidance is also represented by CUDA Agent, which uses agentic reinforcement learning with execution-based rewards [Dai et al., 2026].Source: §5 · p. 9
  • TrainBridge uses eager-to-compiled differences from the same GPU workload to guide where and what to tune on the NPU, then uses target-side profiling to expose further bottlenecks.Source: §5 · p. 9

6 Limitations and Future Work

The current evaluation is limited to single-node training, while multi-node settings introduce network, communication, and parallelization factors. TrainBridge also requires useful framework-level compilation evidence from a source GPU implementation.

Scope

  • The current tuning setting focuses on single-node training; multi-node performance factors include network topology, cross-node communication, and parallelization strategy.Evaluation in multi-node settings with frameworks such as MindSpeed-LLM and Megatron-LM is left to future work.Source: §6 · p. 9

Source-side evidence

  • TrainBridge requires a source GPU implementation that provides useful framework-level compilation evidence, limiting it to workloads supported by a mature source stack.The authors propose a retrieval database of source-side optimization evidence and validated NPU implementations as a possible extension.Source: §6 · p. 10

7 Conclusion

TrainBridge addresses compilation-related inefficiencies in NPU training by using mature GPU compilation evidence to guide NPU-specific optimization. Across the evaluated workloads, it improves NPU training speed.

Conclusion

  • TrainBridge extracts optimization intents from eager-versus-compiled GPU differences to construct NPU-specific rewrites, then profiles the updated program for residual bottlenecks.Source: §7 · p. 10
  • 2.0× speedup across Modded-NanoGPT speedrun records and 1.7× speedup across four additional workloads.Source: §7 · p. 10
  • The results show GPU compilation evidence can guide NPU-specific program optimization and help address compilation-related inefficiencies in NPU training.Source: §7 · p. 10

A.1 Additional Training Workloads

The additional workload evaluation covers AirBench, Qwen3, ViT, and nanoVLM, with task-specific hardware, training configurations, and validation or timing criteria.

A.1 Additional Training Workloads

  • For Qwen3 and nanoVLM, final validation loss may differ from the initial NPU port by at most 0.01; ViT uses the same threshold for final training loss.Source: §A.1 · p. 12
  • AirBench uses airbench96_fast, with GPU timing remeasured on one H100 and the NPU port evaluated on one Ascend 910C.The optimized NPU program must achieve mean accuracy of at least 96%.Source: §A.1 · p. 12
  • Qwen3 training uses a randomly initialized 0.6B model on SlimPajama-6B for 2,000 steps, comparing final-step validation loss.The sequence length is 4,096, global batch size 64, and each run uses eight H100 GPUs or eight Ascend 910C NPUs.Source: §A.1 · p. 12
  • ViT-Base is trained on CIFAR-10 for 200 epochs, with end-to-end training time compared across four H100 GPUs and four Ascend 910C NPUs.Source: §A.1 · p. 12
  • nanoVLM uses its default training configuration on eight H100 GPUs or eight Ascend 910C NPUs, measuring 200 optimizer steps after 10 warmup steps.Source: §A.1 · p. 12

B Additional Evaluation with Claude Opus 4.8

With Claude Opus 4.8 and Claude Code, the evaluation compares tuning with and without TrainBridge from the same initial NPU implementation. In a single 8-hour run, TrainBridge finds a faster validated implementation.

B Additional Evaluation with Claude Opus 4.8

C Per-Record Modded-NanoGPT Results

Table 3 reports per-record Modded-NanoGPT training times and NPU speedups for records with complete measurements.

C Per-Record Modded-NanoGPT Results

Loading 2610.29171v1…