[AI System] Efficient On-Device Training - Part 2

"Efficient On-Device Training - Part 2"

[LoRA(Low-Rank Adaptation)]

· 대형 Model을 Fine-Tuning(미세 조정)하려면 각 Layer의 큰 Weight 행렬 전체를 Update

→ But, LoRA는 Weight 전체를 직접 바꾸지 않고, 작은 두 행렬로 분해한 Low-Rank Update만 학습

△ 학습해야 할 Parameter 수↓

△ 메모리 사용량↓

△ 연산량↓

→ 대형 Foundation Model Fine-Tuning에 매우 적합

** LoRA 이외에 메모리, 연산, Delay를 줄이는 다양한 효율적인 방법 탐색

[Activation Checkpointing]

· Backpropagation 진행 시 각 Layer의 Gradient를 계산하려면, 해당 Layer의 Forward 결과가 필요

    = Forward Propagation 때 계산했던 Activation이 Backpropagation 때 다시 필요

※ Activation : 각 Layer를 통과한 후 나오는 중간 출력값 (∝ Layer 수, Batch Size, Feature Map 크기, Channel 수)

※ Activation Checkpointing : 모든 Activation을 저장 X → 일부 Activation만 Checkpoint로 정한 뒤 나머지는 필요할 때 재계산

▣ Without Checkpointing

· Without Checkpointing : Backpropagation 전에 모든 Intermediate Activation 저장

→ (이전 Layer의 Activation) x (현재  Layer의 Local Derivative) x (다음 Layer에서 계산되어 전달된 Loss Gradient)

Ex. Layer 1 출력 저장 → Layer 2 출력 저장 → … → 마지막까지 전부 저장

▣ With Checkpointing

· With Checkpointing : 몇 개의 중요한 지점만 저장하고, 그 사이의 값은 버림

→ Backward 때 필요해지면 저장된 Checkpoint부터 다시 Forward 계산해서 복원

Ex. a^[l－2], a^[l＋1]만 저장, a^[l－1], a^[l]은 Backward 시점에 다시 계산

△ 메모리 절약

▼ 같은 Forward 일부를 한 번 더 계산하므로 연산량↑ & 학습 시간 다소↑

** Memory-Computation Trade-Off

→ Edge Device 메모리는 보통 아주 제한적이므로 연산을 조금 더 하더라도 메모리를 줄일 수 있으면 실용적

▣ Checkpointing 적용

> Block Method

· 특정 수의 Transformer Layer를 한 Block 처럼 묶어서 재계산

Ex. 4개의 Layer 단위로 Checkpoint 설정 후 Block 경계만 저장 → Block 내부 Activation은 Backward 때 다시 계산

> Uniform Method

· 전체 Layer를 균등하게 나눠서 Checkpoint 배치 → 전체 모델을 여러 Chunk로 쪼개고, 각 Chunk 입력 Activation만 저장

** Checkpoint 너무 자주 설정 → 재계산량↓ But, 메모리 절약↓ / Checkpoint 너무 적게 설정 → 메모리 절약↑ But, 재계산량↑

∴ Checkpointing을 통해 Activation Memory를 줄일 수 있지만 Coumputation은 증가

    → But, Accuracy는 변하지 않고 정확히 같은 모델을 얻음

        (∵ 저장해둘 값을 저장하지 않고 나중에 똑같이 다시 계산 → 메모리와 연산 스케줄만 바뀌고 학습 결과 자체는 동일)

[Gradient Accumulation]

· Activation Memory는 Batch Size(동시에 처리하는 데이터 개수)에 비례해서 증가

※ Gradient Accumulation : 큰 Batch를 여러 개의 작은 Batch로 나눠서 처리하고, Gradient는 합쳐서 한 번만 Update

▣ Without Gradient Accumulation

· 전체 Data Set을 한 번에 다 학습하면 메모리 부족 & 비효율적 → Mini-Batch로 나눠서 학습

① Mini-Batch 1 Forward + Backward → Gradient 계산 & Update

② Mini-Batch 2 Forward + Backward → Gradient 계산 & Update (Mini-Batch 1로 Update된 Model 사용)

③ Mini-Batch 3 Forward + Backward → Gradient 계산 & Update (Mini-Batch 2로 Update된 Model 사용)

▣ With Gradient Accumulation

· Mini-Batch를 Micro-Batch로 나눠서 계산하고 한 번에 Update

① Micro-Batch 1 Forward + Backward → Gradient g₁ 계산 & Update X

② Micro-Batch 2 Forward + Backward → Gradient g₂ 계산 후 누적 & Update X

③ Micro-Batch 3 Forward + Backward → Gradient g₃ 계산 후 누적 & Update X

④ Micro-Batch 4 Forward + Backward → Gradient g₄ 계산 후 누적 & Update X

⑤ 한 번에 Update → g = (g₁ + g₂ + g₃ + g₄) / 4 = (1/B)∑∇L_i로 큰 Batch와 완전히 동일한 Gradient

** 단순히 Batch Size를 B/4로 줄이면 메모리는 해결되지만 모델 성능이 달라지므로, 누적한 후 한 번에 Update가 핵심

    → 유효 Batch Size는 B로 유지

** Activation Memory 감소 이유

· 원래 방식 : Batch = 100 → Activation = 100개 저장

· Gradient Accumulation 방식 : Batch = 25 x 4 → Activation = 25개만 저장 (Iteration 수 증가에 의한 약간의 시간 증가)

∴ Iteration 수 증가로 인해 시간이 약간 증가할 수 있지만, Activation Memory가 크게 감소

[Training Only Certain Layers Or Parameters]

▣ Training Full Layer

· 일반적인 Fine-Tuning → 모든 Weight Update

△ 높은 표현력

△ 새로운 Feature 학습 가능

▼ 많은 양의 계산 필요 (모든 Layer의 Backpropagation)

▼ 큰 메모리 요구 (모든 Activation 저장)

▼ 많은 데이터 필요 (Overfitting 방지를 위한 데이터 필요)

→ 성능은 좋지만 비용이 너무 큼

· 메모리 병목 문제는 Parameter 크기 영향보다 Activation Memory의 영향이 매우 큼
· Parameter의 크기를 줄인 MobileNets(빨간색 그래프)를 봐도 Activation Memory는 거의 동일한 것을 확인 가능

▣ Training Only the Last Layer

· 앞쪽 Layer : Edges, Textures, Shapes 같은 일반 Feature 추출

· 마지막 Layer : 현재 Task에 맞는 Class 분류

    → 새로운 Task에서는 앞부분 Feature Extractor는 그대로 두고, 마지막 Classifier Head만 새 Task에 맞게 변경

△ 필요한 데이터의 양↓

△ 계산량과 필요한 메모리↓

▼ 새로운 Feature를 잘 학습하지 못함

→ 기존 Task와 새 Task가 비슷하면 학습이 잘 되지만, Task가 다르면 성능 저하↑

** Activation Memory 감소 이유

· 전체 Model 학습 : 모든 Layer Activation 저장

· 마지막 Layer만 학습 : 마지막 Layer 입력만 저장

∵ 마지막 Layer만 Update 시, 앞쪽 Layer들의 Activation을 Backward용으로 저장할 필요 X (마지막 Layer의 입력만 필요)

▣ Training Only the Last Layer의 성능

<Full Layer Training / Last Layer Training>

▶ 왼쪽 그래프 : 정확도

· Last Layer Training이 Full Layer Training보다 훨씬 낮음 → 마지막 Layer만으로는 충분한 학습 불가

▶ 오른쪽 그래프 : 학습 Parameter 수

· Last Layer Training이 Full Layer Training보다 훨씬 적음 → 계산량 & 필요한 메모리↓

▣ Training Batch Normalization Layers

· W(Convolution Weight) 또는 B(Bias)는 고정하고 Batch Normalization의 β, γ만 학습

※ Batch Normalization : Activation 정규화 → Scale/Shift → 학습 가능한 Parameter {β, γ}

△ Full Model 학습과 비교했을 때 Parameter의 수↓

▼ 여전히 제한된 성능

▣ Training Batch Normalization Layers + Last Layer의 성능

<Full Layer Training / Last Layer Training / BN+Last Layer Training>

▶ 왼쪽 그래프 : 정확도

· BN+Last Layer Training이 Last Layer Training보다 훨씬 높음

▶ 오른쪽 그래프 : 학습 Parameter 수

· Last Layer Training과 비슷하게 BN+Last Layer Training이 Full Layer Training보다 훨씬 적음 → 계산량 & 필요한 메모리↓

· Last Layer Training : 마지막 Decision Boundary만 변경

· BN+Last Layer Training : Feature Distribution도 조금 변경 → 표현력 증가

<Full Layer Training / Last Layer Training / BN+Last Layer Training>

▶ 왼쪽 그래프 : 정확도

· BN+Last Layer Training이 Full Layer Training보다 여전히 낮음

▶ 오른쪽 그래프 : Memory 비용

· BN+Last Layer Training의 Parameter 수가 Full Layer Training 대비 12배 줄었는데, Memory 비용은 1.8배 정도만 감소

∵ 학습 메모리의 큰 병목은 Parameter보다 Activation에 의함 (Backpropagation 경로가 깊으면 Activation 저장↑)

    → BN만 학습하더라도 Gradient 계산을 위해 여전히 많은 Activation을 저장해야 하는 상황 발생

∴ Update하는 Parameter 수뿐만 아니라 Activation 저장 자체를 줄여야 메모리 절약 가능

▣ Training Only Bias

· 각 Layer의 Bias만 학습하고 Weight는 고정

<Weight Gradient>

· a^[l－1] 같은 Activation필요

<Bias Gradient>

· Bias만 Update할 때는 전체 Activation을 모두 저장할 필요 X

> In ReLU (입력이 0 초과 : 1 / 입력이 0 이하 : 0)

· Backward에 필요한 정보는 Activation 값 전체가 아니라, 해당 값이 양수인지 음수인지 정도의 정보 → Binary Mask로 저장 가능

· 보통 Activation은 Float32로 저장 : 32-bit / ReLU Mask 저장 : 1-bit

∴ Update Only Bias = Weight Gradient 계산에 필요한 Activation 필요 X, ReLU Derivative 계산에 필요한 Sign 정보만 필요

▣ Training Full Model vs Only Bias

<Fine-Tune Full>

·  Intermediate Feature Map을 메모리에 저장 = Block 내부의 중간 출력들을 전부 저장

∵ 각 Convolution Weight에 대한 Gradient 계산 & Backward 진행에 필요

<Fine-Tune Bias Only>

· ReLU 사용 시 Backward에 필요한 것은 부호 Mask → 중간 Feature Map 전부 저장할 필요 X

▣ Training Bias + Last Layer의 성능

<Full Layer Training / BN+Last Layer Training / Bias+Last Layer Training>

▶ 왼쪽 그래프 : Memory 비용

· Bias+Last Layer Training이 Full Layer Training보다 12배 감소

▶ 오른쪽 그래프 : 정확도

· Bias+Last Layer Training이 Full Layer Training보다 16.3% 낮음

△ Weight Update X → 계산량↓

△ Activation 저장↓ → 메모리 비용↓

** 정확도 감소 이유

· Bias는 단순히 Shift 역할 → Feature 자체 변경 X & Decision Boundary도 제한적으로만 변화 = 표현력 부족

∴ Training Only Bias는 효율성이 좋지만, 성능이 부족

▣ Add Lite Residual Learning Module

<Add Lite Residual Learning Module>

· Training Only Bias의 부족한 표현력을 보완하기 위해 Lite Residual Learning Module 추가

· Downsampling → Activation Size↓ → 메모리↓

· Group Convolution → Channel을 나눠서 계산 → 연산량↓ & 메모리↓

· 최소한의 추가 Parameter만 학습에 필요

· 최소한의 Activation 저장만 필요

∴ 큰 Feature Map을 직접 다루지 않고, 작은 Feature Map에서 보정만 수행

▣ Add Lite Residual Learning Module

<Full Layer Training / BN+Last Layer Training / Bias+Last Layer Training / Add Lite Residual Learning Module>

▶ 왼쪽 그래프 : Memory 비용

· Add Lite Residual Learning Module이 Bias+Last Layer Training보다 5MB만 증가

▶ 오른쪽 그래프 : 정확도

· Add Lite Residual Learning Module이 Bias+Last Layer Training보다 11.6% 증가

▣ TinyTL(Training Only Bias + Add Lite Residual Learning Module)

<Full Layer Training / Last Layer Training / BN+Last Layer Training / Add Lite Residual Learning Module>

▶ 왼쪽 그래프 : 정확도

· TinyTL이 Full Layer Training과 거의 비슷

▶ 오른쪽 그래프 : Memory 비용

· TinyTL이 Full Layer Training보다 6배 낮음

∴ TinyTL = Activation을 거의 저장하지 않으면서도 정확도를 유지하는 구조

▣ Adapter

<Trained / Frozen>

· 기존 Pre-Trained Model은 그대로 두고, Task마다 작은 Module을 끼워 넣어서 학습하는 방식

    = Task가 바뀔 때 전체 Model을 새로 저장하지 않고 작은 Adapter만 추가하여 학습

→ 정확도는 비슷하게 유지하면서, 학습해야 하는 Parameter 수를 크게 줄임

· Inference : 학습이 끝난 후, 학습된 Adapter가 Test 중 예측을 하는데 사용됨

** Layer Norm은 학습 가능한 Parameter {β, γ} 존재하여, Backward 진행 시 {β, γ}도 Update 됨

▶ Full Fine-Tuning : Task가 여러 개 있을 때 Task마다 Model 전체를 따로 저장

▶ Adapter : Pre-Trained Model 1개만 공통으로 저장 & Task별로는 작은 Adapter만 따로 저장

△ Trained Parameter 수를 감소 가능

△ Full Fine-Tuning과 비슷한 수준의 성능

▣ Adapter vs LoRA

▶ 공통점

· 전체 Model을 다 학습 X → Task별로 작은 Parameter만 추가

▶ 차이점

> Adapter

· Module을 추가하여 기존 Model에 존재하지 않는 새로운 Parameter 도입 → 새로운 Parameter 미세 조정

· Forward 경로에 새 Layer가 추가되기 때문에 Inference 시 Latency 증가

> LoRA

· 보정을 위해 기존 Parameter에 더해지는 새로운 Parameter(Low Rank A, B 행렬) 도입 → 기존 Parameter 미세 조정

· 기존 연산에 합쳐질 수 있어 Inference 시 추가 Latency X

** 특정 Parameter 또는 Layer만 학습해도 성능이 좋으며, 효율성을 단순히 Trainable Parameter 개수만 적으면 된다고 생각 X

[Sparsified Backpropagation]

▣ Sparse Backpropagation

· Backpropagation 진행 시, 중요한 몇 개의 Gradient 성분만 남기고 나머지는 0으로 만들어 일부 Parameter만 Update

> Original Backpropagation Process

· 각 Layer의 Gradient 계산은 다음 Layer의 Gradient 결과에 의존 (Output → Input)
→ 각 Layer가 다음 Layer의 신호를 대기

> Sparse Backpropagation Process

· Gradient 값들 중에서 크기(절댓값)가 가장 큰 상위 k개의 값만 남기고, 나머지는 모두 0으로 변경

    (Ex. top_k([1, 2, 3, -4]) & k = 2 → [0, 0, 3, -4])

→ 곱해지는 행렬에 0이 매우 많아져 전체적인 연산량 대폭 감소

▣ Key Idea

<Forward Propagation>

<Original Backpropagation>

· 모든 Gradient 계산 → 모든 Neuron Update

<Sparse Backpropagation>

· 일부 Gradient만 계산 → 일부 Neuron만 Update → 계산량↓

▣ Mini-Batch Setting

<Forward Propagation>

<Original Backpropagation>

<Sparse Backpropagation>

· 여러 Sample이 있어도 각 Gradient에서 top_k만 선택

** 연산은 줄었지만, Activation Memory는 크게 줄지 않음 (∵ 이미 계산이 완료된 Gradient에서 top 선택)

→ 계산 최적화 O / 메모리 최적화 X

· 정확도 = 거의 동일 / 속도 = 훨씬 빠름

· Overfitting 감소 효과 존재 (∵ 작은 Gradient 제거 → Noise 감소)

∴ Sparse Backpropagation은 빠르면서 성능 유지 가능

[Local-Loss-Based Backpropagation]

▣ Backward Locking Problem

· Delay at Layer l = (Forward Propagation Delay) + (Backpropagation Delay from Last Layer to Layer l)

→ Layer는 다음 Layer에서 계산되는 Gradient를 대기 (Ex. Layer3은 Layer4 계산이 완료될 때까지 대기)

▼ 병렬 처리 불가

▼ Latency 증가

▼ 느린 학습

▣ Local-Loss-Based Training

· Model 신경망 중간 중간에 여러 개의 Local Loss를 삽입 → 각 Layer가 독립적으로 학습 가능

· 중간 Loss를 계산하기 위해 작고 간단한 CNN 또는 MLP 형태의 Auxiliary Networks(보조 네트워크) 연결

→ 각 보조 네트워크의 출력 차원은 최종 Model의 출력 차원과 동일하게 맞춰야 함

→ Cross Entropy Loss 등을 사용해 국소적으로 오차를 계산

> Loss1 : Layer1만 Update

> Loss2 : Layer2, Layer3 Update

> Loss3 : Layer4, Layer5 Update

> 병렬 Update를 통한 속도 향상

· 각 Local Loss를 기반으로 여러 Layer 그룹을 동시에(병렬) Update 가능

· 인접한 두 Local Loss 사이의 Layer들 간에는 여전히 대기 현상 존재

· 전체 Model 관점에서는 Backward Locking이 획기적으로 완화되어 훈련 속도 크게 향상

> Activation Memory 절약

· 특정 Layer의 Update 완료 시, Backward에 필요했던 해당 Layer의 Activation 값을 메모리에서 삭제 가능 → 메모리 절약

△ 병렬화를 통한 빠른 속도

△ 메모리 절약

▼ Local Loss 개수를 늘릴수록 성능 저하 (∵ Global Objective가 깨짐) ∴ 적절한 개수의 Local Loss가 중요

[Summary for "Efficient On-Device Training"]

· LoRA(Low-Rank Adaptation)

· Activation Checkpointing / Gradient Accumulation

· Training only certain layers or parameters

· Sparsified Backpropagation

· Local-Loss-Based Backpropagation

→ 함께 사용하여 효율성 극대화

[Reference]

· Efficient On-Device Training Part 2 (AI System : Week 4) - Dongjun Han