[논문 리뷰] Memory-efficient DNN Training on Mobile Devices(In Gim, JeonGil Ko, ACM Mobisys 2022)

주요 요점
DNN(Deep Neural Network)를 Mobile에서 구현하는 경우 사생활 보호와 인프라 독립적인 장점.
그러나, Nontrivial Solution으로 취급되어 왔고, Memory Blowup 등의 문제점. Sage로 해결 제안

소개
Additional Training으로 정확도 개선과 최적화 효과 가능.
외부 데이터 교환 최소화.

1. 모바일 디바이스의 경우 서버에 비해 heterogeneous
2. 서버급 GPU와 모바일 GPU는 bandwidth 등에서 많은 차이
3. 모바일 메모리 자원은 한정되어 있음.

DNN Training - SAGE 제안

1. 유연한 자동 미분 계산
2. operator level optimization
3. hybrid approach for run-time memory management

Motivation

1. On-device DNN Training
   1. sensitive data에 대한 학습에 필요
   2. infrastructure dependency 최소화에 필요
   3. CPU의 낮은 성능, latency를 높이고 memory의 부족
   4. more simple model이 필요하다!
2. Memory Blowup 문제
   1. Chain-rule에 따른 gradient 계산에서 Automatic Differentiation -> output, original에 모두 의존하기 때문에 계산 복잡 -> final loss 계산 전까지 메모리에 머물러야 함. (가장 큰 사이즈의 활성화에 영향)
   2. lighter model을 사용할 경우 더 적은 용량으로 학습 가능한 장점
3. Memory Reduction Approach
   1. Gradient Accumulation
   2. Gradient Checkpointing
   3. Smaller Activation Size
   4. Number of Activation Reduce
   5. Memory Virtualization
4. 기술적 한계
   1. 같은 모델이라도 모바일 디바이스의 상황 모두 상이(버전에 따른 남은 메모리 용량 등)
   2. 프로세싱 파워 한계
   3. Scarce and Dynamic Memory

SAGE Design

1. Multiple Low-memory Training Method 조합
2. Pratical Latency Overhead
   Gradient Node를 별도로 분리해 놓았다.

3.1 Automatic Differentiation
AD -> 그래프를 분리하면 최적화 부분에서 저해될 위험이 있으므로 graph optimization을 통합하기 위한 프레임워크(DO->CO graph로 변환)

3.2 Graph level optimization
1. Sym operation
2. Asym operation
3. opaque operation
Just-in time compilation
OpenCL

3.3 Operator Level Optimizations

행렬 연산, Convolution에 많은 메모리 footprint가 소모

모바일 환경에 맞게 Single-Kernal Implementation구현

Copy Operation에 있어서 메모리 용량 많이 소모.

Memory Copy하지 않고 index translation을 통한 메모리 사용량 감소

3.4 Run-time Memory Management

1. Dynamic Gradient Checkpointing

on-line checkpointing 방법 사용
메모리 "가치"를 평가하고, 메모리상에 보존할지 결정 (Latency/Memory)

2. Dynamic Gradient Accumlation
   Mini-batch gradient 를 사용해 더함 (Micro-batch Size를 결정하는 과정 포함)

Hybrid- GPU최적화가 저해될 수 있다 (왜냐하면 low-computing power 때문에 micro-batch size가 너무 작아질 수 있기 때문)

메모리 관리 알고리즘

Algorithm 1 Sage’s runtime memory management
Require: 𝐺 ⊲ List of gradient nodes in the graph
Require: 𝑏𝑎𝑡𝑐ℎ ⊲ Mini-batch size
1: 𝑄 ← BinaryHeap(𝐺) ⊲ Sorted in topological order
2: 𝐷 ← Set ⊲ Prevents infinite loop in checkpointing
3: 𝑚 ← 𝑏𝑎𝑡𝑐ℎ
4: while 𝑏𝑎𝑡𝑐ℎ ≥ 0 do
5: while 𝑛𝑜𝑑𝑒 ← Pop(𝑄) do
6: if 𝑛𝑜𝑑𝑒 ∈ 𝑚𝑒𝑚𝑜𝑟𝑦 then continue
7: Insert(𝐷, Parents(𝑛𝑜𝑑𝑒))
8: if Parents(𝑛𝑜𝑑𝑒) ⊂ 𝑚𝑒𝑚𝑜𝑟𝑦 then
9: while not sufficient memory do
10: if 𝑚 ≥ 𝑀𝑇 𝑃 then ⊲ Accumulation
11: free all last batches in memory
12: 𝑚 ← 𝑚 − 1
13: else ⊲ Checkpointing
14: 𝐶 ← 𝑚𝑒𝑚𝑜𝑟𝑦 ∩ 𝐷
𝐶
15: free 𝑐 ∈ 𝐶 with the smallest Latency(𝑐)
Memory(𝑐)
16: end if
17: end while
18: Compute(𝑛𝑜𝑑𝑒)
19: Remove(𝐷, Parents(𝑛𝑜𝑑𝑒))
20: else
21: Push(𝑄, 𝑛𝑜𝑑𝑒)
22: Push(𝑄, Parents(𝑛𝑜𝑑𝑒))
23: end if
24: end while
25: 𝑏𝑎𝑡𝑐ℎ ← 𝑏𝑎𝑡𝑐ℎ − 𝑚
26: end while