시간, 메모리 효율적으로 LLM 학습하기 (1) (Gradient Accumulation, Gradient Checkpointing, Mixed Precision Training ... )

 

 

시간, 메모리 효율적으로 LLM 학습하기 (1) 
(Gradient Accumulation, Gradient Checkpointing, Mixed Precision Training ... )

 

 

 

 

 

 

모델을 학습시키다 보면 OOM 문제를 맞닥뜨리게도 되고, 또 학습하는 시간 때문에도 골머리를 앓게 된다...! 

 

본인에게 가능한 환경에서 최대한의 퍼포먼스를 낼 수 있는 방법을 찾아보자! 

 

본 포스트에서는 Single GPU 환경에서 할 수 있는 방안에 대해서 살펴본다

 

 

https://huggingface.co/docs/transformers/en/perf_train_gpu_one

Methods and tools for efficient training on a single GPU

 

 

 

 

 

1. Gradient Accumulation 

- 시간 👎 메모리 👍

- Batch size를 작게 하여 LLM을 돌리게 되면, 모델의 성능이 좋지 않게 되는데, 이를 극복하기 위한 방법

- Backward Pass시에 Gradient Update 매번 하는 것(1 step 마다)이 아니라 몇 번 쌓아뒀다가 하는 방법 

- 연산 시간이 늘어나는 대신 메모리 사용량이 줄어든다  (Gradient checkpointing은 메모리 효율성을 증가시킬 수 있으나, 학습 속도를 약 20%가량 저하시킨다.)

- accelerate를 사용하면 아래와 같이 간단하게 !

 

from accelerate import Accelerator
accelerator = Accelerator(gradient_accumulation_steps=2)
model, optimizer, training_dataloader, scheduler = accelerator.prepare(
    model, optimizer, training_dataloader, scheduler
)
for batch in training_dataloader:
    with accelerator.accumulate(model):
        inputs, targets = batch
        outputs = model(inputs)
        loss = loss_function(outputs, targets)
        accelerator.backward(loss)
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()

 

 

https://huggingface.co/docs/accelerate/en/usage_guides/gradient_accumulation

Performing gradient accumulation with Accelerate

 

 

 

 

2. Gradient Checkpointing

- 시간 👎 메모리 👍

- Gradient Accumulation을 통해서 step을 모아서 가중치를 업데이트 한다고 해도, 역전파 과정에서 메모리가 터져버리는 현상이 발생 할 수 있음 

- 매번 Backward Pass에서 필요한 Gradient를 다 저장하는 것이 아니라 일부만 저장, 필요한 것은 그때그때 계산

 

 

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

trainer = Trainer(model=model, args=training_args, train_dataset=ds)
result = trainer.train()
print_summary(result)

 

 

 

 

 

3. Mixed Precision Training

 

- 시간 👍 메모리 👍

- 대부분 우리가 모델을 학습할 때 사용하게 되는 Data type은 32bit (FP32) - weight와 input 모두 ! 

- 하지만 메모리와 학습 시간 모두 효율적으로 하기 위해서 FP32와 FP16을 적절히 섞어서 학습 수행

- 그럼 어떤 건 FP32으로 표현하고 어떤 건 FP16으로 표현하는가? 

• activations
• activation gradients
• weights
• weight gradients

- 위에서 weights와 weight gradients에는 FP16의 범위 안에서 잘 표현이 되나 activation gradients의 경우에는 FP16으로 표현하게 된다면 FP16 으로 표현할 수 있는 범위를 넘어서면서 강제적으로 0이 되어버리는 현상이 발생하였다(underflow).

 

- PyTorch에서 mixed precision learning을 위해 제공하는 package인 AMP(Automatic Mixed Precision)을 사용하여 Mixed Precision Training할 수 있다 

import torch
 
scaler = torch.cuda.amp.GradScaler() # Training시에 생성
 
for data, label in data_iter:
   optimizer.zero_grad()
   with torch.cuda.amp.autocast(): # Mixed precision으로 operation들을 casting 
      outputs = model(data)
 
   scaler.scale(loss).backward() # Loss를 scaling한 후에 backward진행
   scaler.step(optimizer) # 원래 scale에 맞추어 gradient를 unscale하고 optimizer를 통한 gradient update
   scaler.update() # 다음 iteration을 위해 scale update

 

 

 

 

4. 배치 사이즈 줄이기 

 

 

 

위의 방법으로도 효과가 없을 경우, 더 좋은 사양의 GPU를 사용하거나 multi-GPU 환경을 사용할 것을 고려해야 한다.

 

여러 개의 GPU를 사용하는 상황에도 물론, 위 사항들은 적용 가능하다.

 

다음 포스트에서는 multi GPU 환경에서 취할 수 있는 방안들을 살펴본다. 

 

 

 

Reference

 

 

https://littlefoxdiary.tistory.com/126

[Huggingface] Single GPU에서 효율적인 모델 학습을 하는 방법

https://jihan819.tistory.com/entry/AI-Mixed-Precision-Training-%EC%9D%B4%EB%9E%80

[AI] Mixed Precision Training 이란?

https://velog.io/@twinjuy/Auto-Mixed-Precision%EC%9D%B4%EB%9E%80

Auto Mixed Precision이란?