Memory Planning for Deep Neural Networks

TL;DR

This paper introduces MemoMalloc, a memory allocation technique for DNN inference that reduces latency caused by mutex contention, achieving up to 40% faster performance with moderate memory overhead.

Contribution

The paper presents MemoMalloc, a novel runtime and static analysis-based memory planning method for DNN inference that significantly improves latency.

Findings

MemoMalloc reduces DNN inference latency by up to 40%.

It outperforms existing general-purpose memory allocators.

The approach balances latency improvements with moderate memory increases.

Abstract

We study memory allocation patterns in DNNs during inference, in the context of large-scale systems. We observe that such memory allocation patterns, in the context of multi-threading, are subject to high latencies, due to \texttt{mutex} contention in the system memory allocator. Latencies incurred due to such \texttt{mutex} contention produce undesirable bottlenecks in user-facing services. Thus, we propose a "memorization" based technique, \texttt{MemoMalloc}, for optimizing overall latency, with only moderate increases in peak memory usage. Specifically, our technique consists of a runtime component, which captures all allocations and uniquely associates them with their high-level source operation, and a static analysis component, which constructs an efficient allocation "plan". We present an implementation of \texttt{MemoMalloc} in the PyTorch deep learning framework and evaluate memory consumption and execution performance on a wide range of DNN architectures. We find that \texttt{MemoMalloc} outperforms state-of-the-art general purpose memory allocators, with respect to DNN inference latency, by as much as 40\%.