단일 GPU 효율적으로 사용하기
GPU 자원은 한정되어 있기 때문에, 언제든 GPU 메모리를 효율적으로 사용하는 방법을 탐구해야 한다. 이번에는 GPU 메모리 사용 효율화를 위한 방법인 그레이디언트 누적 / 그레이디언트 체크포인팅에 대하여 알아본다.
- 그레이디언트 누적 (Gradient Accumulation) - 큰 batch size를 사용할 수 있도록 보조하는 방법론이다. batch size를 크게 잡을수록 모델 학습이 더 안정적으로 이루어지지만, 이전 글에서 확인했듯이 batch size를 늘릴수록 GPU 자원을 더 많이 소모한다. Gradient Accumulation은 작은 배치를 여러 번 처리한 후, 각 배치에서 계산된 Gradient를 누적하여 하나의 큰 배치처럼 처리한다.
- 그레이디언트 체크포인팅 (Gradient Checkpointing) - 학습 중 저장된 중간 계산값(활성화 값)을 메모리에서 의도적으로 삭제한 후 다시 계산하는 방법이다. 일반적으로 DL 모델의 층이 깊어질수록 활성화 값을 계산하는 데 많은 메모리가 소요되기 때문에, 일부 활성화 값을 저장하지 않고 스킵하여 GPU 메모리 사용량을 줄인다.
그레이디언트 누적
사실 그레이디언트 누적은 이미 전 게시글에서 한 번 사용했었다.
def train_model(model, dataset, training_args):
if training_args.gradient_checkpointing:
model.grdient_checkpointing_enable()
train_dataloader = DataLoader(dataset, batch_size=training_args.per_device_train_batch_size)
optimizer = AdamW(model.parameters())
model.train()
gpu_utilization_printed = False
for step, batch in enumerate(train_dataloader, start=1):
batch = {k: v.to(model.device) for k, v in batch.items()}
outputs = model(**batch)
loss = outputs.loss
loss = loss / training_args.gradient_accumulation_steps
loss.backward()
if step % training_args.gradient_accumulation_steps == 0:
optimizer.step()
gradients_memory = estimate_memory_of_gradients(model)
optimizer_memory = estimate_memory_of_optimizer(optimizer)
if not gpu_utilization_printed:
print_gpu_utilization()
gpu_utilization_printed = True
optimizer.zero_grad()
loss = loss / training_args.gradient_accumulation_steps 부분에서 분모를 4로 설정한다면, 역전파를 수행할 때 loss를 4로 나누어 수행하게 된다. 즉, gradient_accumulation_steps=4 라면, 4개의 배치에 대해 각각 손실을 계산한 후에 그 값들의 평균을 사용한다. 이 방식을 이용하면 실제로는 4개의 배치를 처리한 것처럼 학습이 수행된다.
if step % training_args.gradient_accumulation_steps == 0: 위 조건을 통해 설정한 steps 만큼의 배치가 지났을 때만 optimizer.step()을 호출하여 모델 파라미터를 업데이트 한다. 즉, 기존 배치마다 파라미터 업데이트를 수행하는 것이 아니라, steps 만큼 누적된 그레이디언트를 기준으로 한 번에 업데이트를 수행한다.
단점은?
- 계산 시간 증가 - 파라미터를 업데이트 하는 시점까지 Gradient를 계속 누적시킨 후 기다려야하기 때문에, 계산 시간이 더 오래 소요된다는 단점이 있다.
- 또 다른 하이퍼파라미터 추가 - 그레이디언트 누적의 효과를 최적화하려면 gradient_accumulation_steps 를 적절히 조절해야 한다. step값이 너무 적으면 누적의 효과가 적고, 너무 크면 학습 시간이 너무 오래걸린다. 따라서 이 또한 튜닝이 필요한 부분이므로, 전체적으로 학습 구조가 복잡해질 수 있다.
그레이디언트 체크포인팅
그레이디언트 체크포인팅을 사용하지 않은 DL 모델 학습 과정은 위와 같다. 오른쪽 방향은 순전파이고, 왼쪽 방향은 역전파이다. (보라색이 값을 메모리에 저장하고 있는 상태이다.) 위 그림에서 알 수 있는 것은, n번째 층에서 역전파를 수행하기 위해 n-1번째 층까지의 순전파 메모리를 저장하고 있어야 한다는 점이다.
위 방식을 개선하기 위해 도입된 방식이다. 역전파를 수행할 때 모든 순전파 메모리를 기억하고 있지 않아도 된다. 즉, GPU에는 최소한의 메모리만 저장하고 있다가 나머지는 필요할 때 다시 계산하여 가져오는 것이다. 그림에서 보면 알 수 있듯이, 순전파 메모리를 모든 층에서 저장하고 있지 않는다. 층이 깊어질수록 얕은 층의 메모리는 지워버린다.
하지만 메모리가 필요해질 때마다 순전파를 다시 수행한다는 것 자체가 굉장한 단점이다. 기존 방식을 개선하려 도입되었지만 치명적 단점이 있기에 이대로 사용하긴 어려워보인다.
이에 따라 그레이디언트 체크포인팅이 도입되었다. 위 그림에서 윗 줄의 세 번째 층이 체크포인트가 되는데, 체크포인트 층에 해당하는 값은 날려버리지 않고 항상 저장하고 있는다. 이에 따라 네 번째 층의 순전파 데이터가 필요해지는 시점이 오면, 맨 처음 층부터 계산을 다시 하는게 아니라 체크포인트 층부터 순전파 계산을 이어나간다.
사용 방식 또한 간단하다. (전 게시글의 메모리 사용량 측정 코드에서 이미 활용했다!)
분산 학습
위에서는 단일 GPU를 효율적으로 사용하는 방법에 대해 알아보았다면, 이번에는 여러 개의 GPU를 이용하는 방식이다. 모델의 크기가 하나의 GPU에 올릴 수 있을 정도로 작은 경우에, 여러 GPU에 각각 해당 모델을 올리고 학습 데이터를 병렬로 처리하여 학습 속도를 높이는 방식을 데이터 병렬화(Data Parallelism)라고 한다.
- 데이터 병렬화(Data Parallelism) : 여러 모델에 동일한 모델을 올리고, 데이터를 병렬로 처리하는 방법. 개별 GPU에 모델을 전부 올리는 것이 가능해야 한다.
- 모델 병렬화(Model Parallelism) : 모델 자체가 너무 커서 하나의 GPU에 올리기 힘든 경우에 여러 개의 GPU에 모델을 나눠 올릴 수 있다.
- 파이프라인 병렬화(Pipeline Parallelism) : 모델의 layer를 나눠서 GPU에 올리는 방식
- 텐서 병렬화(Tensor Parallelism) : 하나의 층도 나눠서 GPU에 올리는 방식
텐서 병렬화는 아래와 같이 열 병렬화(Column-Wise Parallel), 행 병렬화(Row-Wise Parallel)로 구분할 수 있다.
열 병렬화, 행 병렬화는 이름부터 직관적이다. Hidden layer 행렬을 행(열)로 나누어 각 GPU에서 따로 처리하면 행(열) 병렬화가 되는 것이다. 처리한 후에는 Aggregation sum을 통하여 단일 GPU에서 수행한 것과 동일한 결과를 얻을 수 있다.
LoRA (Low-Rank Adaptation)
일반적으로 LLM 모델의 크기는 굉장히 크고, 지속적으로 커지고 있기 때문에 단일 GPU를 이용하여 전체 미세 조정(full fine-tuning)을 수행하기는 어렵다. 따라서 일부 파라미터만 학습하는 PEFT(Parameter Efficient Fine-Tuning)에 대한 연구가 활발히 이루어지고 있다. 그 중에서 LoRA(Low-Rank Adaptation)과 QLoRA(Quantized LoRA)에 대해 알아본다.
LoRA 특징
- 저랭크 행렬 분해 (Low-Rank Matrix Decomposition) ⇢ LoRA는 파라미터 matrix를 Low Rank matrix로 분해한다. 이에 따라 학습에 필요한 파라미터 수는 대폭 줄이면서도, 기존 모델의 성능은 유지한다.
- 효율적인 fine tuning ⇢ 기존 모델은 그대로 두고, 추가적인 Low Rank matrix를 도입하여 학습하는 것이므로 모델이 새로운 task에 더 잘 대응할 수 있도록 할 수 있다. 또한, 기존 모델을 그대로 유지하기 때문에 모델의 재사용성이 좋다.
위 도식에서 파란색 matrix는 LLM 모델의 기존 parameter matrix를 의미한다. LoRA에서는 이 거대한 파라미터를 전부 학습하는 게 아닌, 그보다 Rank 수가 훨씬 적은 A, B matrix를 도입하여 학습한다. A, B matrix를 곱한 결과는 기존 parameter matrix의 크기 (d, k)와 동일하므로, 이 둘을 더할 수 있게 된다.
기존 parameter인 W matrix는 그대로 유지하되, 사용자가 원하는 새로운 task에 대하여 해당 모델이 더 잘 작동할 수 있도록 A, B matrix를 도입한 것이다. 이때 A, B matrix의 크기는 W에 비하여 굉장히 작게 설정되기 때문에, 연산량을 대폭 줄일 수 있다.
A, B matrix의 크기는 각각 (d, r), (r, k)로 이루어지는데, d와 k는 W matrix의 행과 열의 크기이므로 반드시 고정된다. 이때 r의 크기를 사용자가 직접 설정할 수 있다.
- r이 너무 큰 경우 → 기존 파라미터 W를 학습하는 게 오히려 더 낫다.
- r이 너무 작은 경우 → 연산량은 대폭 줄지만, 그만큼 모델이 학습할 수 있는 capacity가 낮아져 표현력이 낮아질 우려가 있다.
따라서 RoLA를 이용할 때는 모델의 성능을 크게 해치지 않으면서도, 학습 비용을 효율적으로 줄일 수 있는 적절한 r값을 찾아내는 것이 중요하다.
또한, Layer마다 개별적으로 LoRA를 적용하는 것이 가능하다. 예를 들어 self attention 연산의 query, key, value에 LoRA를 적용할 수도 있고, feed forward에만 적용할 수도 있다. 미세 조정 과정에서 LoRA를 어느 층에 적용할지도 적절히 선택해야 한다.
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model
def load_model_and_tokenizer(model_id, peft=None):
tokenizer = AutoTokenizer.from_pretrained(model_id)
if peft is None:
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto",
device_map={"":0})
elif peft == "lora":
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto",
device_map={"":0})
lora_config = LoraConfig(
r=8,
lora_alpha=32,
target_modules=["query_key_value"],
lora_dropout = 0.05,
bias="none",
task_type="CAUSAL_LM"
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
print_gpu_utilization()
return model, tokenizer
import gc, torch
from transformers import TrainingArguments, Trainer
import numpy as np
from datasets import Dataset
def make_dummy_dataset():
seq_len, dataset_size = 256, 64
dummy_data = {
"input_ids" : np.random.randint(100, 30000, (dataset_size, seq_len)),
"labels" : np.random.randint(100, 30000, (dataset_size, seq_len))
}
dataset = Dataset.from_dict(dummy_data)
dataset.set_format("pt")
return dataset
def cleanup():
if 'model' in globals():
del globals()['model']
if 'dataset' in globals():
del globals()['dataset']
gc.collect()
torch.cuda.empty_cache()
def print_gpu_utilization():
if torch.cuda.is_available():
used_memory = torch.cuda.memory_allocated() / 1024**3
print(f"GPU 메모리 사용량 : {used_memory:.3f} GB")
else:
print("런타임 유형을 GPU로 변경하세요")
def gpu_memory_experiment(batch_size,
gradient_accumulation_steps=1,
gradient_checkpointing=False,
model_id="EleutherAI/polyglot-ko-1.3b",
peft=None):
print(f"배치 크기: {batch_size}")
model, tokenizer = load_model_and_tokenizer(model_id, peft=peft)
if gradient_checkpointing == True or peft == 'qlora':
model.config.use_cache = False
dataset = make_dummy_dataset()
training_args = TrainingArguments(
per_device_train_batch_size = batch_size,
gradient_accumulation_steps=gradient_accumulation_steps,
gradient_checkpointing=gradient_checkpointing,
output_dir="./result",
num_train_epochs=1
)
try:
train_model(model, dataset, training_args)
except RuntimeError as e:
if "CUDA out of memory" in str(e):
print(e)
else:
raise e
finally:
del model, dataset
gc.collect()
torch.cuda.empty_cache()
print_gpu_utilization()
from transformers import AdamW
from torch.utils.data import DataLoader
def estimate_memory_of_gradients(model):
total_memory = 0
for param in model.parameters():
if param.grad is not None:
total_memory += param.grad.nelement() * param.grad.element_size()
return total_memory
def estimate_memory_of_optimizer(optimizer):
total_memory = 0
for state in optimizer.state.values():
for k, v in state.items():
if torch.is_tensor(v):
total_memory += v.nelement() * v.element_size()
return total_memory
def train_model(model, dataset, training_args):
if training_args.gradient_checkpointing:
model.grdient_checkpointing_enable()
train_dataloader = DataLoader(dataset, batch_size=training_args.per_device_train_batch_size)
optimizer = AdamW(model.parameters())
model.train()
gpu_utilization_printed = False
for step, batch in enumerate(train_dataloader, start=1):
batch = {k: v.to(model.device) for k, v in batch.items()}
outputs = model(**batch)
loss = outputs.loss
loss = loss / training_args.gradient_accumulation_steps
loss.backward()
if step % training_args.gradient_accumulation_steps == 0:
optimizer.step()
gradients_memory = estimate_memory_of_gradients(model)
optimizer_memory = estimate_memory_of_optimizer(optimizer)
if not gpu_utilization_printed:
print_gpu_utilization()
gpu_utilization_printed = True
optimizer.zero_grad()
print(f"옵티마이저 상태의 메모리 사용량: {optimizer_memory / (1024 ** 3):.3f} GB")
print(f"그레이디언트 메모리 사용량 : {gradients_memory / (1024 ** 3):.3f} GB")
cleanup()
print_gpu_utilization()
gpu_memory_experiment(batch_size=16, peft='lora')
torch.cuda.empty_cache()
위 코드는 이전 게시글에서 GPU 사용량을 측정했던 "EleutherAI/polyglot-ko-1.3b" 모델에 LoRA 기법을 적용한 후, GPU 사용량을 측정한다. 결과는 다음과 같다.
GPU 메모리 사용량 : 2.581 GB
GPU 메모리 사용량 : 4.345 GB
옵티마이저 상태의 메모리 사용량: 0.012 GB
그레이디언트 메모리 사용량 : 0.006 GB
LoRA 적용 이후, GPU 사용량이 크게 줄어든 것을 확인할 수 있다.
QLoRA (Quantized LoRA)
QLoRA는 기존 LoRA에 양자화 기법을 추가한 방식이다. LoRA는 16bits로 모델을 저장하는 반면, QLoRA에서는 4bits 형식으로 모델을 저장한다.
QLoRA에서는 NF4(Normal Float 4-bit) 비트 표현을 이용하여 양자화를 수행한다. QLoRA 논문 팀이 LLaMa 모델을 확인했을 때 92.5% 정도의 모델 파라미터가 정규 분포를 따랐다고 한다. 따라서 입력 데이터가 정규분포라고 가정했을 때 모델의 성능을 유지하면서도 빠른 양자화가 가능해지기에, NF4 양자화를 이용해 메모리 사용량을 더욱 줄일 수 있다.
정규분포와 양자화의 관계
- 기존 16bits 양자화에 비하여, 4bits 양자화는 숫자의 정밀도가 상당히 떨어진다. 하지만 데이터가 정규분포를 따른다고 가정하면, 분포가 상대적으로 균등하기 때문에 양자화 비트 수를 줄이더라도 정보 손실을 최소화 할 수 있다. 즉, 적은 비트 수로 양자화된 값들이 거의 모든 값들을 고르게 커버할 수 있게 된다.
- 또한, 데이터가 정규분포를 따를 경우 값들이 대체로 평균 근처에 집중되어 있기 때문에 양자화 후에도 중요한 정보가 잘 보존되는 경향이 있다. (만약 데이터가 정규분포를 따르지 않는다면 적은 비트 수로 양자화 한 이후 정보 손실이 클 가능성이 높다.)
페이지 옵티마이저 (Paged Optimizer)
QLoRA 논문에서는 Gradient checkpointing 과정에서 발생할 수 있는 OOM(Out Of Memory)를 방지하기 위해 Paged Optimizer를 활용한다.
Gradient Checkpointing 과정은 위와 같았다. 애초에 이 기법이 메모리 사용량을 줄이기 위해 등장한 것인데, 모델의 크기는 갈수록 커져만 가므로.. Gradient checkpointing을 사용하더라도 OOM이 일어날 수 있다.
Paged Optimizer는 엔비디아의 통합 메모리를 통해 GPU가 CPU 메모리인 RAM을 공유하는 기법을 의미한다. 여기서 통합 메모리는 컴퓨터의 가상 메모리 시스템과 비슷한 개념을 GPU 메모리 관리에 적용한 개념이다. 가상 메모리를 통해 RAM과 디스크를 활용하여 더 많은 GPU를 사용하는 것처럼 동작한다.
'딥러닝' 카테고리의 다른 글
| [LLM] 3. 좋은 데이터 셋, GPU (0) | 2024.11.03 |
|---|---|
| [LLM] 2. 허깅페이스 트랜스포머 모델 학습하기 (0) | 2024.11.01 |
| [LLM] 1. 임베딩, 어텐션, 트랜스포머 모델들 (1) | 2024.11.01 |
| 14. LSTM (Long Short Term Memory) 개념 + GRU (0) | 2023.11.15 |
| 13. RNN 코드실습 (0) | 2023.10.31 |
