Optimising the optimisation: performance analysis of memory management techniques for AI/ML training on non-NUMA hardware

Abstract

This research explores the potential of out-of-band optimisation techniques to enhance the efficiency of AI/ML training on non-NUMA hardware. By evaluating the impact of alternative memory allocators and transparent huge pages (THP) on training time across diverse AI/ML workloads, we aim to provide practical insights into memory management strategies. Our study focuses on the effectiveness of these techniques in improving training efficiency without modifying model architectures or algorithms. Through benchmarking different allocator configurations and THP settings, we aim to quantify their contributions to training performance, with the goal of informing future AI/ML development efforts. Our results shows that these low-level memory management strategies can squeeze out additional performance improvements, making them a valuable complement to traditional in-band optimisation methods.