动态性感知的深度学习内存分配器  

Dynamics-aware Memory Allocator for Deep Learning

作　　者：黄奕桐 张昱[1] HUANG Yitong;ZHANG Yu(School of Computer Science and Technology,University of Science and Technology of China,Hefei 230026,China)

机构地区：[1]中国科学技术大学计算机科学与技术学院,合肥230026

出　　处：《小型微型计算机系统》2024年第9期2049-2054,共6页Journal of Chinese Computer Systems

基　　金：国家自然科学基金项目(62272434)资助.

摘　　要：深度学习应用程序通常被认为在训练小批量步之间具有规律的内存分配模式.这种规律性在许多情况下是正确的,但由于未知的张量大小以及各种并行执行模式下不确定的执行顺序等因素,内存分配请求在运行时可能会变得动态和不规律.在这些因素作用下,一些基于内存分配的规律性所进行的优化在运行时存在许多不匹配的情况.为了解决这个问题,本文提出了一个动态性感知的深度学习内存分配器.这个分配器在运行时收集内存分配信息并生成内存分配计划,然后通过模拟分配的方式动态调整生成的内存分配计划.实验结果表明,与TensorFlow相比,本文所设计的分配器平均可以达到1.24的加速比,并减少56.34%的CPU内存使用量以及21.13%的GPU显存使用量.Deep learning applications are generally considered to have regular memory allocation patterns across the training mini-batches.Though this observation holds true in many cases,the memory allocation requests can actually become dynamic and irregular due to several factors in the execution,such as unknown tensor sizes and uncertain execution ordering in various parallel execution modes.These factors make the memory optimization methods that assume regular memory allocations ill-suited.To address the issue,we propose a dynamics-aware memory allocator for deep learning.This allocator collects memory allocation information at run time and generates a memory allocation plan,and then dynamically adjusts the generated memory allocation plan by simulating allocation.The experimental results show that our allocator can achieve an average speedup of 1.24,and reduce CPU memory usage by 56.34%and GPU memory usage by 21.13%.

关 键 词：深度学习 动态性 内存分配 模型训练 

分 类 号：TP183[自动化与计算机技术—控制理论与控制工程]