Methods and tools for efficient training on a single GPU

<aside>

Gradient Accumulation

The gradient accumulation method aims to calculate gradients in smaller increments instead of computing them for the entire batch at once.

→ 한 번에 전체 배치에 대한 가중치를 계산하는 대신, 더 작은 단위로 가중치를 계산하는 것을 목표로 함

This approach involves iteratively calculating gradients in smaller batches by performing forward and backward passes through the model and accumulating the gradients during the process.

→ 더 작은 배치에서 가중치를 반복적으로 계산하고, 이를 축적하는 방식

예를 들어, batch_size = 64로 학습하고 싶었으나, gpu 메모리가 부족해서 16으로 나누어 4번 계산하면서 가중치를 축적

즉, 매번 가중치를 따로 계산하는 게 아니라, 계산된 가중치를 계속 더해서 저장해두는 방식

Once a sufficient number of gradients have been accumulated, the model’s optimization step is executed.

→ 충분한 개수의 가중치가 축적되면, 모델의 최적화 단계 실행(가중치를 업데이트 한다는 말)

→ 그럼 여러 번 가중치를 누적하다가 충분하다 싶으면, 그때 가중치를 업데이트 하는 것인데

By employing gradient accumulation, it becomes possible to increase the effective batch size beyond the limitations imposed by the GPU’s memory capacity.

→ 이렇게 하면, 실제 gpu가 처리할 수 있는 크기보다 더 큰 배치 크기로 학습하는 효과

However, it is important to note that the additional forward and backward passes introduced by gradient accumulation can slow down the training process.

→ 하지만 추가적인 순전파, 역전파 연산 초래되기 때문에 학습 속도는 느려짐

<aside>

🎯 요약

  1. 배치 크기가 너무 크면 GPU 메모리가 부족해 학습이 어렵다
  2. 그래서 한 번에 전체 배치를 학습하는 대신, 작은 단위로 나누어 가중치를 계산하고 누적
  3. 일정 횟수만큼 누적한 후, 그제야 가중치를 업데이트 한다
  4. 덕분에 GPU 메모리 한계를 넘어서 더 큰 배치 크기로 학습하는 효과를 얻을 수 있다!
  5. 하지만 여러 번 forward + backward 연산을 수행해야 해서 속도가 느려짐 ㅠ </aside>

</aside>

<aside>

You can enable gradient accumulation by adding the gradient_accumulation_steps argument to TrainingArguments:

training_args = TrainingArguments(per_device_train_batch_size=1, gradient_accumulation_steps=4, **default_args)

In the above example, your effective batch size becomes 4.

Alternatively, use 🤗 Accelerate to gain full control over the training loop. Find the 🤗 Accelerate example further down in this guide.

While it is advised to max out GPU usage as much as possible, a high number of gradient accumulation steps can result in a more pronounced training slowdown.

→ gpu 사용량을 최대한 활용하는 것이 좋지만(지난 스터디 내용), 가중치 축적 단계가 많아질 수록 학습 속도 느려짐

Consider the following example. Let’s say, the per_device_train_batch_size=4 without gradient accumulation hits the GPU’s limit.

If you would like to train with batches of size 64, do not set the per_device_train_batch_size to 1 and gradient_accumulation_steps to 64.

Instead, keep per_device_train_batch_size=4 and set gradient_accumulation_steps=16. This results in the same effective batch size while making better use of the available GPU resources.

For additional information, please refer to batch size and gradient accumulation benchmarks for RTX-3090 and A100.

그렇다면 효과는 완전 동일할까..? → 동일하지 않을 수도 있는 이유

<aside>

  1. Floating Point 연산 오차
  2. 정규화(BatchNorm, Dropout) 차이

그래도 일반적으로 거의 동일한 효과를 얻을 수 있으니까 이렇게 적었다고 함 ㅋ

</aside>

<aside>

Gradient Checkpointing

Some large models may still face memory issues even when the batch size is set to 1 and gradient accumulation is used.

→ 일부 대형 모델에서 배치 크기 1로 하고, 가중치 축적을 사용해도 여전히 메모리 문제 발생할 수 있음

This is because there are other components that also require memory storage.

→ 모델이 크면 가중치 외에도 추가적인 메모리 사용 요소 발생

Saving all activations from the forward pass in order to compute the gradients during the backward pass can result in significant memory overhead.

→ 역전파에서 가중치 계산을 위해, 순전파 동안 모든 활성화 값을 저장하면 상당한 오버헤드 발생

The alternative approach of discarding the activations and recalculating them when needed during the backward pass, would introduce a considerable computational overhead and slow down the training process.

→ 활성화 값을 버리고 필요할 때 다시 계산하는 대안이 있지만, 학습 속도 엄청 느려짐

Gradient checkpointing offers a compromise between these two approaches and saves strategically selected activations throughout the computational graph so only a fraction of the activations need to be re-computed for the gradients.

→ 계산 그래프에서 전략적으로 선택된 활성화 값만 저장하여, 역전파 시 일부 활성화 값은 저장하고, 나머지는 필요할 때 다시 계산하는 방식

메모리 절약 + 속도 저하 최소

For an in-depth explanation of gradient checkpointing, refer to this great article.

To enable gradient checkpointing in the Trainer, pass the corresponding a flag to TrainingArguments:

training_args = TrainingArguments(
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,
    gradient_checkpointing=True, # 이거!
    **default_args
)

Alternatively, use 🤗 Accelerate - find the 🤗 Accelerate example further in this guide.

Accelerate를 사용하면 학습 루프를 완전히 제어 가능

While gradient checkpointing may improve memory efficiency, it slows training by approximately 20%.

→ Gradient Checkpointing은 메모리 효율을 개선하지만, 학습 속도를 약 20% 정도 느리게 함

<aside>

🎯 요약

  1. Gradient Checkpointing을 사용하면 일부 활성화 값은 저장되지 않고 다시 계산
  2. 그 때문에, 연산량이 증가하여 학습 속도가 평균적으로 20% 느려짐
  3. 하지만
  4. 한마디로, 속도가 조금 느려지더라도, 모델 크기를 키울 수 있는 장점이 있음! </aside>

</aside>

image.png

image.png