Kan training scales efficiently on multi gpu high performance computers
Scalability Analysis of Distributed Kolmogorov-Arnold Network Training on High-Performance Computing Systems
Distributed, Parallel, and Cluster ComputingMachine LearningPerformance
Summary
Kolmogorov-Arnold Networks (KANs) are a special type of neural network that use learnable functions instead of fixed math rules to make decisions. The authors studied how well KANs train when the work is split across many GPUs in a supercomputer. They found that KAN training gets faster almost as expected when adding more GPUs, with some small delays due to communication between machines. Their experiments help show how to best set up hardware and model size for training KANs efficiently.
What this means in practice
- •For deep learning engineers: Optimize the use of multi-GPU HPC systems for training interpretable neural networks like KANs by following deployment guidelines on GPU setup and model sizing.
- •For high performance computing teams: Improve scheduling and resource allocation strategies by understanding KAN training communication patterns and scalability behavior on multi-node GPU clusters.
Tested on simulated data.
Authors
Guangneng Chen, David Garcia Selfa, Pablo Quesada Barriuso
Abstract
Kolmogorov-Arnold Networks (KANs) replace the fixed activation functions and linear weights of Multi-Layer Perceptrons (MLPs) with learnable univariate functions on network edges, offering improved interpretability and, in some settings, competitive parameter efficiency. While the approximation properties of KANs have received considerable attention, their behavior under distributed, multi-GPU training has not been systematically characterized. This paper presents an empirical scalability study of data-parallel KAN training on multi-node, multi-GPU high-performance computing (HPC) infrastructure, evaluated along four dimensions: strong scaling, weak scaling, communication overhead, and model-size scaling. Experiments were conducted on the FinisTerrae III supercomputer using up to 8 NVIDIA A100 GPUs across 4 nodes with PyTorch Distributed Data Parallel (DDP). KAN training reaches 74.7% parallel efficiency at 8 GPUs with a 5.97x speedup, consistent with conventional deep learning workloads. Weak scaling shows an initial single-to-multi-GPU throughput drop followed by strong stability. Communication overhead follows a non-monotonic pattern (1.3%-6.1%), driven primarily by All-Reduce algorithm selection and inter-node latency rather than KAN's edge-wise gradient structure. The parameter-to-memory ratio improves with model size even as training time scales unfavorably. These results indicate that operator-level and data-parallel optimizations for KAN are complementary. We provide deployment guidelines for GPU topology and model-size selection, and discuss the limitations of a synthetic-regression evaluation.