Deep Learning 학습 시 메모리 절감 방법

학습을 위해 필요한 메모리 요구량은 추론에 비해 약 10배가 많다. 따라서 모델 사이즈를 증가시켜 모델 성능을 증가시키는 현재 AI 동향과 맞물려 학습 시 메모리 요구량을 줄이는 것이 매우 중요한 과제가 되었다.

(1) 추론에 필요한 메모리 요구량

• parameters (FP16): x2 → Model Memory
• activation (FP16): batch size와 parameters에 의존 → Activations Memory

(2) 학습에 필요한 메모리 요구량

• parameters (FP16): x2 → Model Memory
• gradients(FP16): x2 → Model Memory
• optimizer States(FP32): x16 → Optimizer Memory
• activation (FP16): batch size와 parameters에 의존 → Activations Memory

학습 시 메모리 사용량을 줄이기 위해 4가지 방법(Sparsity, Low Precision, Microbatching, Checkpointing)을 사용할 수 있다.

Sparsity, Low Precision, Micro-batching, Checkpointing별 Memory/Computation/Accuracy에 영향도 비교 table

Sparsity

Sparse tensor는 대부분의 값이 zero인 tensor로 zero인 tensor는 명시적으로 저장될 필요가 없어 메모리 사용량을 줄일 수 있다.

Technique

• 정해진 sparsity 패턴으로 네트워크를 초기화하고 약간의 iteration을 실행한 후, 네트워크에서 가장 작은 크기의 weights를 pruning(zero로 설정)한다.
• 동일한 수의 새로운 non-zero가 이전 non-zero 값을 대체하기 위해 네트워크 전체에 도입된다.
• Parameters는 global threshold를 기반으로 pruning 되는데 이 global threshold는 매번 고정된 parameters를 pruning하기 위해 매 rewiring 단계에서 그 크기가 조정된다.
• 위와 같은 prune & replace 동작을 rewiring이라고 불리운다.

Accuracy Tradeoff

• WideReseNet 모델에서 테스트한 결과, baseline(0% sparsity) 대비 20% non-zero(80% sparsity)까지 0.5% 정확도 저하만 발생하나 non-zero가 20%보다 낮으며 정확도는 급격하게 낮아진다.

Test accuracy vs. percentage of nonzero parameters, for WRN-28–2 on CIFAR-10

Computation

• Sparse 연산을 사용하면 zero 엔트리를 무시하여 FLOPS를 줄일 수 있다.
• 하지만 Sparsity는 불규칙한 메모리 접근 패턴때문에 실질적으로 speedup이 되지 않을 수 있다. 따라서 Sparsity는 Sparsity의 수준과 사용된 HW에 강하게 의존한다.
• Sparse Matrix Multiply 커널은 CPU와 GPU 모두에 존재하지만, Sparse Conv를 위한 HW 지원은 현재 제한적으로 제공된다. 하지만 최근 연구는 direct sparse conv를 사용하면 현재 세대의 CPU와 GPU HW에서도 조차 dense conv보다 추가적인 속도 향상을 얻을 수 있을 수 있음을 보여준다.
• 현 Sparsity 연구는 미래 또는 Special-purpose HW 구조를 통해 Sparsity 계산 성능이 추가적으로 향상될 수 있음을 제안한다.

Low Precision

일반적으로 뉴럴네트워크 학습에 FP32를 사용하였으나 최근 FP16이 학습에 충분하다라는 것이 제안(Nvidia의 Mixed precision training 논문. ‘17) 되면서 널리 사용되고 있다.

Lower-precision 연산을 사용하는 것은 메모리 사용량과 연산량을 둘다 줄이는 효과를 거둘 수 있다.

Technique

• 기본적인 FP16 학습은 모든 weights, optimizer buffer, input에 FP16을 사용한다.
• 추가적으로 gradients의 numerical underflow와 overflow를 회피하기 위해 동적 loss scaling을 사용하며 모든 batch normalization parameters는 FP32를 유지한다.

Accuracy Tradeoff

• FP16을 사용하였을 때 마지막 정확도에 의미있는 영향을 미치지 않는다.
• WideResNet 모델을 CIFAR-10으로 학습 시 FP32을 사용하였을 때, 테스트 정확도가 94.37 ± 0.0.7%, FP16을 사용하였을 때 테스트 정확도가 94.43 ± 0.17%로 차이가 거의 나지 않는다.

Computation

• FP16 연산을 사용하는 것은 FLOPS의 개수를 줄이는 것이 아니라 floating-point를 지원하는 HW의 각 floating-point 연산의 cost를 줄인다.
• 정확한 Speed-up 수치는 사용되는 HW에 의존한다.

Micro-batching

네트워크의 Activations는 mini-batch size에 직접적으로 비례한다. 메모리 사용량을 줄이기 위해 mini-batch size를 줄일 수 있으나, 이는 학습 최적화 속성을 변경시키므로 잠재적으로 상당한 hyper-parameters 튜닝을 필요로 한다. 이에 대한 대안으로 mini-batch를 더 작은 그룹으로 나누는 micro-batching가 제안되었다.

Test accuracy vs. microbatch size, for WRN-28–2 on CIFAR-10

Technique

• Micro-batching에서 micro-batch는 forward & backward를 독립적으로 실행되며 수학적으로 본래 mini-batch의 결과와 일치한다.
• 하지만 batch normalization 레이어를 포함하는 모델은 mini-batch 이내의 예제가 더 이상 독립적이지 않다. 그러므로 microbatch normalization 통계는 전체 minibatch의 통계와 달라 모델의 최종 정확도에 영향을 미칠 수 있다.

Accuracy Tradeoff

• 다른 모든 hyper-parameters를 고정한 후 micro-batch size만 변화하였을 때, baseline인 micro-batch size=100를 1/10로 줄여 micro-batch size=10으로 학습하여도 loss의 차이가 없음을 관찰할 수 있다.
• 비록 더 작은 micro-batch가 더 높은 가변 normalization 통계를 갖지만 micro-batch size가 매우 작아질 때까지 정확도에 나쁜 영향을 미치지 않는다.

Computation

• Micro-batching은 총 FLOPS의 양에 영향을 미치지 않지만, parallelism을 줄이는 대신 메모리 사용량을 줄인다.
• GPU는 일반적으로 병렬 연산에 매우 효율적으로 최적화되어 있어서 가능한한 많은 batch size를 사용하는 것이 좋다. (IPU와 같이 small batch size에서도 빠르고 효율적으로 동작하도록 디자인된 가속기도 있다.)
• 하지만 어떠한 모델과 데이터셋은 최신 GPU 메모리에 large batch를 넣을수 없기 때문에 이와 같이 GPU 메모리 용량이 제한적인 조건에서는 micro-batching이 필요하다.

Checkpointing

Gradient checkpointing (또는 Activation checkpointing)은 뉴럴 네트워크의 activations 중 일부만을 저장하여 메모리를 절약하는 알고리즘이다. 버려진 activations은 backward propagation 중에 다시 재계산을 통해 복원해 낼 수 있다.

Checkpointing은 memory-accuracy tradeoff가 없으나, 데이터셋, 모델 구조, mini-batch size에 의존하여 단일 레이어의 activations 조차 대용량의 메모리를 차지할 수 있다.

Technique

• checkpointing은 forward pass 중에서 뉴럴네트워크의 m_th node의 입력 activations을 저장하는 것이다.
• 예를들어 N_1, . . . , N_mn의 m×n nodes가 존재한다면 N_m, N_2m, . . . , N_m(n−1)만을 저장한다. 즉 mn개의 전체 activations 대신 n+m개의 activations만 저장한다.
• checkpointing은 activiations를 위한 메모리 공간을 절감하는 대신 backward propagation 중에 forward propagation 만큼의 computation양이 추가된다.

Computation Tradeoff

• checkpoint 방식에 따라 총 연산량과 activation 저장 공간은 trade-off 관계에 있다. 즉, 연산량을 증가시켜 activation용 저장 공간을 절감시키는 것이 가능하다.
• checkpoint-residual-2*는 총 FLOPS가 30% 증가하는 반면 activations을 위한 storage 공간을 5.8x까지 줄일 수 있다.

FLOPs ratio to baseline for various degrees of checkpointing, for training WideResNet on CIFAR-10. (A) denotes checkpoint-none, (B) denotes checkpoint-no-bn, © denotes checkpoint-residual-1*, and (D) denotes checkpointresidual-2*.

레퍼런스

[1] Low-Memory Neural Network Training: A Technical Report

Low-Memory Neural Network Training: A Technical Report