BladeDISC++: Memory Optimizations Based On Symbolic Shape

TL;DR

BladeDISC++ introduces memory optimization techniques for dynamic shape graphs in deep learning, leveraging symbolic shapes and a combined compile-time and runtime strategy to reduce memory usage without precise shape information.

Contribution

It presents novel op scheduling and rematerialization methods based on symbolic shapes, addressing memory optimization challenges in dynamic shape compilers.

Findings

Reduces memory consumption in dynamic shape graphs

Achieves memory efficiency comparable to precise shape optimizations

Enhances adoption of dynamic shape compilers in large models

Abstract

Recent deep learning workloads exhibit dynamic characteristics, leading to the rising adoption of dynamic shape compilers. These compilers can generate efficient kernels for dynamic shape graphs characterized by a fixed graph topology and uncertain tensor shapes. However, memory optimization, although particularly crucial in this large model era, remains relatively underexplored for dynamic shape graphs. The fundamental challenge lies in the lack of precise tensor shapes which are essential in conventional methods such as operation scheduling(op scheduling) and rematerialization. To address this challenge, we propose op scheduling and rematerialization approaches based on symbolic shapes and developed BladeDISC++. Besides, since rematerialization decisions cannot be made solely at compile time when tensor shapes are unknown, BladeDISC++ employs a compilation-runtime combined strategy to optimally address shape dynamics. Evaluations indicate that BladeDISC++ effectively reduces memory usage for dynamic shape graphs, achieving memory consumption comparable to optimizations using precise shapes, thereby promoting the broader adoption of dynamic shape compilers.