WHERE DID ALL THE MEMORY GO?

참고: ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
공부한거 나중에 보기 편하게 적는다.

어떤 특정 딥러닝 모델을 학습시킬 때 매우 많은 메모리가 필요하다 

그래서 요즘엔 32bit 실수 FP32 를 FP16으로의 mixed-precision training을 이용한 최적화가 많이 진행된다.

 

1.5B 딥러닝 모델은 파라미터(가중치 행렬)은 FP16으로 3gb의 메모리를 필요로 한다. 32gb의 single gpu에서 이 모델을 학습 시킬 수 없다. 어떤 부분에서 이러한 메모리 누수(라고 표현하는게 옳은지는 모르겠지만)가 발생하는 것일까?

(밑에 내용은 GPT로 번역시킨 것을 적당히 수정)

 

1. Model States: Optimizer States, Gradients and Parameters

모델 훈련 중에 대부분의 장치 메모리는 모델 상태에 의해 소모된다. Adam 옵티마이저는 두 가지 옵티마이저 상태를 저장해야 한다. 시간에 따른 평균 모멘텀, 기울기의 분산. 이를 통해 모델 업데이트를 계산하므로 모델을 훈련하려면 기울기와 가중치뿐만 아니라, 이 두 옵티마이저 상태도 저장할 만큼 충분한 메모리가 필요합니다.

 

cf) 옵티마이저들의 계산은 fp32이기 때문에 메모리 사용량이 더 높다. 옵티마이저는 훈련 중 매우 작은 값을 업데이트하거나, 변화량을 추적하는데 사용되므로, 작은 수치 차이도 훈련 안전성에 크게 영향을 미친다.

 

 

 

2. Residual Memory Consumption

 - 훈련 중에 활성화값(Activations value)은 상당한 양의 메모리를 소모한다.

예를 들어, 1.5B 파라미터를 가진 GPT-2 모델은 시퀀스 길이 1K, 배치 크기 32로 훈련할 때 약 60GB의 메모리를 사용한다. 활성화 체크포인팅(activation checkpointing) 또는 재계산(recomputation)은 활성화 메모리를 줄이기 위한 일반적인 방법이다. 이 방법은 활성화 총량의 제곱근 정도만큼 메모리 사용량을 줄여주지만, 대신 33%의 재계산 오버헤드가 발생한다. 

 

cf) 활성화값 = 각 layer 출력 이라 생각하면 편하다. 체크포인팅은 적당히 훈련 중 나온 값의 일부를 저장하고 나머지는 저장하지 않는 방식. 저장하지 않는 것은 backpropagation 도중에 다시 계산 한다. 예를 들면 Relu layer에서 input x만 들고 있고 Relu(x)는 저장하지 않고 있어도, backpropagation 할 때 Relu(x)로 그때 계산하면 된다.

 

- Temporary buffers 큰 모델에서 메모리 소모가 크다.

임시 버퍼는 여러 결과를 합칠 때 사용하는데, 병렬 계산에서 결과를 all-reduce해서 합치거나, 기울기나 모델의 값을 정규화 할 때 하나의 버퍼에 모은 후 계산하는데, 이런 값들은 연산에 따라 fp32로 해야한다. (optimizer와 마찬가지로 영향이 클 듯...)

 

- Memory Fragmentation도 문제가 될 수 있다. 한마디로 사용 가능한 메모리가 충분하더라도 연속된 메모리가 부족하면 메모리 요청이 실패할 수 있습니다. 이는 사실 시스템적인 관점이 필요할 것 같기는 하다.