Hugging Face API로 배우는 GPU 친화적 모델 학습

SSunny 25.03.13

3,072 5 1

facebook twitter kakao link

5 1

facebook twitter kakao link

DEVOTEE 요약: 딥러닝 모델을 GPU에서 학습하는 과정은 환경 설정, 모델 준비, 학습 실행, 결과 활용의 단계로 나눌 수 있으며, GPU 메모리 효율성을 높이기 위해 여러 가지 기술들이 사용됩니다. 학습 실행 단계에서는 순전파, 역전파, 모델 업데이트를 통해 정확성을 개선하며, Gradient Accumulation과 Gradient Checkpoint 등의 메모리 최적화 기법이 사용됩니다. QLoRA와 같은 기술은 모델의 양자화를 통해 메모리 사용량을 획기적으로 줄임으로써 대규모 모델을 효율적으로 학습할 수 있는 방안을 제공합니다.

DEVOTEE 추천 블로그

모델 학습 과정

딥러닝을 GPU에서 수행하는 과정은 크게 1. 환경 설정, 2. 모델 준비, 3. 학습 실행, 그리고 4. 결과 활용의 단계로 나눌 수 있습니다.

환경 설정 단계에서는 NVIDIA GPU가 장착된 컴퓨터에 CUDA와 cuDNN을 설치하여 GPU 사용 환경을 구축합니다.
이는 GPU를 딥러닝에 활용할 수 있게 해주는 기반 작업입니다. 다음으로, TensorFlow나 PyTorch 같은 딥러닝 프레임워크를 설치합니다.
이 프레임워크들은 GPU를 효율적으로 사용할 수 있는 도구를 제공합니다.
모델 준비 단계에서는 사용할 딥러닝 모델을 설계하고, GPU 메모리에 맞게 데이터 처리 방식을 조정합니다.
학습 실행 단계에서는 모델과 데이터를 GPU 메모리로 이동시키고, GPU를 사용해 실제 학습을 수행합니다.
이 과정에서 GPU의 병렬 처리 능력을 활용하여 대량의 계산을 빠르게 처리합니다. 학습 중에는 GPU 사용률과 메모리 상태를 모니터링하여 최적의 성능을 유지합니다.
학습이 완료된 모델은 CPU로 이동시켜 저장하거나 추론에 사용합니다.
이러한 과정을 통해 GPU의 강력한 연산 능력을 활용하여 효율적으로 딥러닝 모델을 학습하고 활용할 수 있습니다.

모델 학습 실행 단계

학습 실행단계를 좀 더 살펴보겠습니다. 딥러닝 모델의 학습 과정은 세 단계로 이루어집니다:

순전파(Forward Pass): 순전파는 입력 데이터를 모델에 통과시켜 예측 결과를 생성하는 과정입니다.
이는 마치 학생이 시험을 치는 것과 비슷합니다. 모델에 데이터를 입력하고, 모델이 예측을 만들어냅니다.
그리고 이 예측이 얼마나 정확한지 손실(오차)을 계산합니다. 순전파 동안 계산된 중간 결과들은 역전파를 위해 임시로 저장됩니다.
이는 마치 시험을 볼 때 중간 계산 과정을 메모해두는 것과 비슷합니다. 이렇게 저장된 값들을 '순전파 상태 값'이라고 부릅니다.
이 값들은 역전파 과정에서 그래디언트를 계산하는 데 필요합니다.이 과정을 반복하면서 모델은 점점 더 정확한 예측을 할 수 있게 됩니다.
- 배치 정규화 (Batch Normalization): 각 배치의 평균과 분산을 사용하여 데이터를 정규화합니다. 이는 학습 속도를 높이고 안정성을 향상시킵니다.
- 드롭아웃 (Dropout): 특정 뉴런을 무작위로 비활성화하여 과적합을 방지합니다. 이는 모델의 일반화 능력을 향상시킵니다.
- GPU 메모리 사용: 입력 데이터 크기 + 모델 파라미터 크기 + 중간 결과 크기
역전파(Backward Pass): 역전파는 손실 함수를 기반으로 모델의 가중치를 업데이트하는 과정입니다.
이는 학생이 시험 결과를 받고 어디서 실수했는지 분석하는 과정과 유사합니다. 계산된 손실을 바탕으로, 모델의 각 부분이 이 오차에 얼마나 기여했는지 계산합니다.
이 과정에서 '그래디언트'라는 값이 생성되는데, 이는 모델의 각 부분을 어떻게 조정해야 할지를 나타냅니다.
- 그래디언트 클리핑 (Gradient Clipping): 그래디언트 값이 너무 커지는 것을 방지하여 학습이 불안정해지는 것을 막습니다. 이를 통해 학습 속도를 안정적으로 유지할 수 있습니다.
- 혼합 정밀도 학습 (Mixed Precision Training): FP16과 FP32를 혼합하여 사용함으로써 메모리 사용량을 줄이고 계산 속도를 높입니다. 이는 GPU의 연산 성능을 극대화하는 데 도움이 됩니다.
- GPU 메모리 사용: 순전파 때 사용한 메모리 + 그래디언트 저장 공간
모델 업데이트: 옵티마이저는 모델의 가중치를 업데이트하는 알고리즘입니다.
마지막으로, 학생이 분석을 바탕으로 공부 방법을 개선하는 것처럼, 옵티마이저라는 도구를 사용해 모델의 파라미터(가중치)를 업데이트합니다.
- Adam: 모멘텀과 적응형 학습률을 결합한 옵티마이저로, 빠른 수렴 속도를 자랑합니다. 메모리 사용량은 파라미터 수에 따라 다르며, 일반적으로 각 파라미터에 대해 12바이트의 메모리를 소모합니다.
- SGD (Stochastic Gradient Descent): 기본적인 경량 옵티마이저로, 메모리 사용량이 적고 단순하지만, 최적화 과정에서 더 많은 에포크가 필요할 수 있습니다.
- GPU 메모리 사용: 업데이트된 모델 파라미터 저장 공간

이 과정을 거치면서 총 GPU 메모리 사용량은 대략 ‘(입력 데이터 크기 + 모델 파라미터 크기 + 중간 결과 크기) × 2 (순전파, 역전파 사용 메모리 )+ 옵티마이저 상태 + 임시 변수 .. ‘ 입니다.

예를 들어, 100MB 크기의 모델을 배치 크기 32로 학습할 때:

입력 데이터: 32 × 입력 크기 (예: 1MB) = 32MB
모델 파라미터: 100MB
중간 결과: 대략 모델 크기와 비슷하다고 가정하면 100MB
총 사용량 ≈ (32 + 100 + 100) × 2 + 50 (추가 버퍼) ≈ 514MB

이 과정을 반복하면서 모델은 점점 더 정확한 예측을 할 수 있게 됩니다.

실제 모델 크기의 약 5배에 해당하는 GPU 메모리가 필요한 것은 일반적인 현상입니다.

허깅페이스에서 "bert-base-uncased"100mb 정도되는 모델을 로딩해서 학습을 수행하면 5배 정도의 결과가 나오는지 확인해 보겠습니다.

https://huggingface.co/google-bert/bert-base-uncased

from transformers import AutoModel
import torch

def calculate_gpu_memory(model_name, batch_size, input_size_mb):
    # 모델 로드
    model = AutoModel.from_pretrained(model_name)
    
    # 모델 크기 계산 (MB)
    model_size_mb = sum(p.numel() for p in model.parameters()) * 4 / 1024 / 1024  # float32 기준
    
    # 메모리 계산
    input_memory = batch_size * input_size_mb
    model_memory = model_size_mb
    intermediate_memory = model_size_mb  # 중간 결과를 모델 크기와 비슷하다고 가정
    total_memory = (input_memory + model_memory + intermediate_memory) * 2 + 50  # 추가 버퍼 포함
    
    return total_memory, model_size_mb

# 예시 사용
model_name = "bert-base-uncased"  # 허깅페이스 모델 이름
batch_size = 32
input_size_mb = 1  # 입력 데이터 크기 (MB)

required_memory, model_size = calculate_gpu_memory(model_name, batch_size, input_size_mb)

print(f"모델 크기: {model_size:.2f} MB")
print(f"필요한 GPU 메모리: {required_memory:.2f} MB")
print(f"모델 크기 대비 필요 메모리 비율: {required_memory / model_size:.2f}배")

"bert-base-uncased"를 사용하여 허깅페이스 모델을 로드하고 모델의 실제 크기를 계산합니다 (float32 기준).
입력 데이터, 모델 파라미터, 중간 결과에 대한 메모리를 계산합니다.
총 필요 메모리를 계산하고, 추가 버퍼를 포함합니다.

모델 크기: 433.85 MB
필요한 GPU 메모리: 1851.70 MB
모델 크기 대비 필요 메모리 비율: 4.27배

사용하려는 모델의 정보에 크기가 나와 있지 않은 경우 허깅페이스에서 제공하는 메모리 크기 계산기를 이용할 수 있습니다.

모델 이름을 넣고 검색하면 아래처럼 결과가 나옵니다. 우리가 예시로 사용한 모델 "bert-base-uncased" 은 코드를 통해 알아본 바 (433.85MB)와 유사한 크기(417.65MB)로 나오네요.

https://huggingface.co/docs/accelerate/usage_guides/model_size_estimator

역전파 단계의 최적화 기술

딥러닝 모델을 학습할 때, GPU 메모리 부족 문제는 흔히 발생합니다. GPU 메모리 사용량을 이해하고 관리하는 것은 효율적인 딥러닝 모델 학습에 매우 중요합니다.

1. 메모리 최적화 기술1) Gradient Accumulation (그래디언트 누적)

그래디언트 누적은 작은 배치 크기로 여러번 그래디언트를 계산하고 누적하여 큰 배치 효과를 얻습니다.

이를 통해 메모리 사용량을 줄일 수 있습니다. 마치 여러 번 나눠서 장을 보는 것과 비슷합니다.

한 번에 큰 배치의 데이터를 처리하기 힘들 때, 작은 배치로 나누어 여러 번 처리하고 그 결과를 모아서 한 번에 모델을 업데이트합니다.

이렇게 하면 큰 배치를 사용한 효과를 얻으면서도 메모리 사용량을 줄일 수 있습니다. 아래는 허깅페이스의 Trainer 클래스를 사용하여 구현한 예시입니다.

from transformers import Trainer, TrainingArguments, AutoModelForSequenceClassification, AutoTokenizer
import datasets

# 데이터셋 로드
dataset = datasets.load_dataset("glue", "mrpc")

# 모델과 토크나이저 초기화
model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased")
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")

# 데이터 전처리
def preprocess_function(examples):
    return tokenizer(examples['sentence1'], examples['sentence2'], truncation=True)

tokenized_datasets = dataset.map(preprocess_function, batched=True)

# TrainingArguments 설정
training_args = TrainingArguments(
    output_dir="./results",
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,  # 그래디언트 누적 스텝 수 설정
    num_train_epochs=3,
    logging_dir='./logs',
)

# Trainer 초기화
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets["train"],
)

# 모델 학습
trainer.train()

**TrainingArguments**에서 **gradient_accumulation_steps**를 설정하여 여러 미니배치에 대해 그래디언트를 누적하고, 지정된 스텝 수에 도달했을 때만 가중치를 업데이트합니다.

2. 메모리 최적화 기술2) Gradient Checkpoint (그래디언트 체크포인팅)

그래디언트 체크포인트는 순전파 중 일부 중간 활성화를 저장하고 나머지는 다시 계산하여 메모리 사용량을 줄이는 기법입니다.

이는 긴 여행 중 모든 경유지의 사진을 찍는 대신 주요 지점의 사진만 찍고, 필요할 때 그 지점부터 다시 여행을 재현하는 것과 유사합니다.

이 방법은 메모리 사용량을 크게 줄이지만, 일부 계산을 다시 해야 하므로 학습 시간이 약간 늘어날 수 있습니다.

from transformers import Trainer, TrainingArguments, AutoModelForSequenceClassification, AutoTokenizer
import datasets

# 데이터셋 로드
dataset = datasets.load_dataset("glue", "mrpc")

# 모델과 토크나이저 초기화
model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased")
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")

# 데이터 전처리
def preprocess_function(examples):
    return tokenizer(examples['sentence1'], examples['sentence2'], truncation=True)

tokenized_datasets = dataset.map(preprocess_function, batched=True)

# TrainingArguments 설정
training_args = TrainingArguments(
    output_dir="./results",
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,
    num_train_epochs=3,
    logging_dir='./logs',
)

# Gradient Checkpointing 활성화
model.gradient_checkpointing_enable()  # 체크포인팅 활성화

# Trainer 초기화
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets["train"],
)

# 모델 학습
trainer.train()

모델의 gradient_checkpointing_enable() 메서드를 호출하여 활성화합니다.

Gradient Accumulation은 배치 크기와 관련된 메모리 문제를 해결하고, Gradient Checkpoint는 모델의 깊이나 시퀀스 길이와 관련된 메모리 문제를 해결하는 데 도움이 됩니다.

이를 통해 제한된 GPU 메모리로도 더 큰 모델이나 더 긴 시퀀스를 다룰 수 있게 됩니다. 이는 마치 작은 주방에서도 효율적인 요리 방법을 통해 큰 파티를 준비할 수 있는 것과 비슷합니다.

결과적으로, 이러한 기법들은 연구자들이 더 복잡하고 강력한 딥러닝 모델을 개발하는 데 큰 도움을 줍니다.

3. 파라미터 조정 기술1) LoRA (Low-Rank Adaptation)

모델 파인튜닝 기법 중 하나로 LoRA는 모델에 일부 파라미터를 추가(재구성)하고 그 부분만 학습하는 방식입니다. 이는 마치 거대한 책의 일부분만을 수정하는 것과 비슷합니다.

모델의 전체 파라미터를 조정하는 대신, 작은 행렬을 추가하여 모델의 특정 부분만 학습 시킵니다. 이로 인해 학습에 필요한 메모리와 계산량을 크게 줄일 수 있습니다.

from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, DataCollatorForLanguageModeling
from datasets import load_dataset
from peft import get_peft_model, LoraConfig, TaskType

# 모델과 토크나이저 로드
model_name = "gpt2"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# LoRA 설정
peft_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    r=8,
    lora_alpha=32,
    lora_dropout=0.1,
    bias="none",
)

# LoRA 모델 생성
model = get_peft_model(model, peft_config)

# 데이터셋 로드 및 전처리
dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")
def tokenize_function(examples):
    return tokenizer(examples["text"], truncation=True, max_length=512)

tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=["text"])

# 트레이닝 인자 설정
training_args = TrainingArguments(
    output_dir="./lora_gpt2",
    per_device_train_batch_size=4,
    num_train_epochs=1,
    logging_steps=100,
    save_steps=1000,
    learning_rate=3e-4,
)

# 트레이너 초기화 및 학습
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False),
)

trainer.train()

LoraConfig에서 LoRA의 설정을 정의합니다. **r=8**은 LoRA의 랭크를 설정합니다. 이는 얼마나 많은 새로운 파라미터를 추가할지 결정합니다. **lora_alpha**와 **lora_dropout**은 LoRA의 학습 방식을 조절합니다.

4. 파라미터 조정 기술2) QLoRA (Quantized Low-Rank Adaptation)

QLoRA는 LoRA를 더욱 발전시킨 기술입니다. 이는 모델의 가중치를 4비트로 양자화하여 메모리 사용량을 획기적으로 줄입니다.

이는 마치 고해상도 사진을 낮은 해상도로 압축하는 것과 유사하지만, 중요한 정보는 유지합니다.

QLoRA는 다음과 같은 핵심 기술을 사용합니다:

4비트 NormalFloat: 모델 가중치를 효율적으로 저장하는 새로운 데이터 형식
이중 양자화: 양자화 과정에서 사용되는 상수도 추가로 양자화하여 메모리 사용을 더욱 줄임
페이지드 옵티마이저: 메모리 사용을 최적화하여 대규모 모델 학습을 가능하게 함

QLoRA의 주요 장점은 일반적인 GPU로도 대규모 언어 모델을 학습할 수 있게 해준다는 것입니다.

이는 마치 작은 노트북으로도 슈퍼컴퓨터급의 작업을 할 수 있게 되는 것과 비슷합니다.

이러한 기술들은 AI 연구의 민주화를 촉진하며, 더 많은 연구자와 개발자들이 대규모 언어 모델을 다룰 수 있게 합니다.

결과적으로, 다양한 분야에서 AI 기술의 혁신과 응용이 가능해지며, 의료, 금융, 교육 등 여러 산업 분야에서 고성능 AI 모델을 활용할 수 있는 길을 열어줍니다.

from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, DataCollatorForLanguageModeling
from datasets import load_dataset
from peft import get_peft_model, LoraConfig, TaskType, prepare_model_for_kbit_training
import torch

# 모델과 토크나이저 로드 (4비트 양자화 적용)
model_name = "facebook/opt-350m"
model = AutoModelForCausalLM.from_pretrained(model_name, load_in_4bit=True, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 모델을 4비트 훈련을 위해 준비
model = prepare_model_for_kbit_training(model)

# LoRA 설정
peft_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    r=8,
    lora_alpha=32,
    lora_dropout=0.1,
    bias="none",
)

# LoRA 모델 생성
model = get_peft_model(model, peft_config)

# 데이터셋 로드 및 전처리
dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")
def tokenize_function(examples):
    return tokenizer(examples["text"], truncation=True, max_length=512)

tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=["text"])

# 트레이닝 인자 설정
training_args = TrainingArguments(
    output_dir="./qlora_opt",
    per_device_train_batch_size=4,
    num_train_epochs=1,
    logging_steps=100,
    save_steps=1000,
    learning_rate=3e-4,
)

# 트레이너 초기화 및 학습
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False),
)

trainer.train()

model = AutoModelForCausalLM.from_pretrained(model_name, load_in_4bit=True, device_map="auto") 에서 **load_in_4bit=True**를 사용하여 모델을 4비트로 양자화하여 모델의 크기를 대폭 줄여 메모리 사용량을 감소시킵니다.