Abstract:In real applications, it is common to train deep neural networks (DNNs) on modest clusters. With the continuous increase of model size and batch size, the training of DNNs becomes challenging under restricted memory budget. The tensor partition and tensor rematerialization are two major memory optimization techniques to enable larger model size and batch size within the limited-memory constrain. However, the related algorithms failed to fully extract the memory reduction opportunity, because they ignored the invariable characteristics of dynamic computational graphs and the variation among the same size tensors at different memory locations. In this work, we propose MegTaiChi, a dynamic tensor-based memory management optimization module for the DNN training, which first achieves an efficient coordination of tensor partition and tensor rematerialization. The key feature of MegTaiChi is that it makes memory management decisions based on dynamic tensor access pattern tracked at runtime. This design is motivated by the observation that the access pattern to tensors is regular during training iterations. Based on the identified patterns, MegTaiChi exploits the total memory optimization space and achieves the heuristic, adaptive and fine-grained memory management. The experimental results show, MegTaiChi can reduce the memory footprint by up to 11% for ResNet-50 and 10.5% for GL-base compared with DTR. For the training of 6 representative DNNs, MegTaiChi outperforms MegEngine and Sublinear by 5x and 2.4x of the maximum batch sizes. Compared with FlexFlow, Gshard and ZeRo-3, MegTaiChi achieves 1.2x, 1.8x and 1.5x performance speedups respectively on average. For the million-scale face recognition application, MegTaiChi achieves 1.8x speedup compared with the optimal empirical parallelism strategy on 256 GPUs.

MAGIS: Memory Optimization Via Coordinated Graph Transformation and Scheduling for DNN

Memory-aware Scheduling for Complex Wired Networks with Iterative Graph Optimization

MegTaiChi: Dynamic Tensor-based Memory Management Optimization for DNN Training

STR: Hybrid Tensor Re-Generation to Break Memory Wall for DNN Training

Understanding GNN Computational Graph: A Coordinated Computation, IO, and Memory Perspective

G-NMP: Accelerating Graph Neural Networks with DIMM-based Near-Memory Processing

WELDER: Scheduling Deep Learning Memory Access Via Tile-graph

An Application-oblivious Memory Scheduling System for DNN Accelerators

ROAM: memory-efficient large DNN training via optimized operator ordering and memory layout

GATe: Streamlining Memory Access and Communication to Accelerate Graph Attention Network With Near-Memory Processing

Optimizing Off-Chip Memory Access for Deep Neural Network Accelerator

CMDS: Cross-layer Dataflow Optimization for DNN Accelerators Exploiting Multi-bank Memories

Memory System Optimization for Graph Processing: a Survey

TSPLIT: Fine-grained GPU Memory Management for Efficient DNN Training Via Tensor Splitting

Pinpointing the Memory Behaviors of DNN Training

Coop: Memory is not a Commodity

SmartMem: Layout Transformation Elimination and Adaptation for Efficient DNN Execution on Mobile

Stitching Weight-Shared Deep Neural Networks for Efficient Multitask Inference on GPU

Cocco: Hardware-Mapping Co-Exploration towards Memory Capacity-Communication Optimization

Memory and Computation Coordinated Mapping of DNNs Onto Complex Heterogeneous SoC.

FreshGNN: Reducing Memory Access via Stable Historical Embeddings for Graph Neural Network Training