[Fine-Tuning] LLM 파인튜닝 솔루션 - Unsloth

LLM을 클라우드 환경이나 Google Colab에서 파인튜닝하는 과정에서 종종 라이브러리 간 의존성 충돌로 인해 실행 오류가 발생하고, 높은 메모리 사용량과 긴 학습 시간이 문제가 되곤 합니다. 이러한 문제를 해결할 수 있는 보다 효율적인 방법을 찾던 중, 단일 GPU 환경에서도 최적의 성능을 제공하는 "Unsloth"를 접하게 되어 소개해 보겠습니다.

1. Unsloth란 무엇인가?

Unsloth는 LLM(대형 언어 모델) 파인튜닝을 보다 효율적으로 수행할 수 있도록 설계된 혁신적인 도구입니다. Michael과 Daniel Han 형제가 개발한 이 프로젝트는 적은 자원으로도 강력한 성능을 발휘할 수 있도록 최적화되어 있으며, 학습 속도 향상과 메모리 사용량 절감을 주요 목표로 하고 있습니다.

Unsloth는 NVIDIA의 GTX 1070과 같은 저사양 GPU부터 최신 H100까지 폭넓은 하드웨어를 지원하며, 허깅페이스(Hugging Face) 생태계와 완벽히 호환됩니다. 이를 통해 연구자와 개발자들은 기존 환경에서 쉽게 활용할 수 있습니다.

1.1 Unsloth의 주요 특징

• 고속 학습: 기존 기법보다 빠른 학습 속도를 제공하며, 대부분의 GPU 환경에서 뛰어난 성능을 발휘합니다.
• 메모리 최적화: 대형 모델을 단일 GPU에서도 학습할 수 있도록 메모리 사용량을 최소화합니다.
• 다양한 호환성: 허깅페이스의 SFTTrainer, DPOTrainer, PPOTrainer 등과 연동되어 익숙한 환경에서 활용할 수 있습니다.

2. Unsloth의 최적화 기술

Unsloth는 다양한 최적화 기술을 활용하여 학습 속도를 향상시키고 메모리 사용량을 줄일 수 있도록 설계되었습니다. 대표적인 기술들을 살펴보겠습니다.

2.1 Intelligent Weight Upcasting (지능형 가중치 업캐스팅)

기존 QLoRA(양자화된 Low-Rank Adaptation) 기법에서는 모델 안정성을 위해 일부 계층을 FP32로 변환합니다. Unsloth는 특정 모델(Mistral, LLaMA 등)에 최적화된 방식으로 업캐스팅을 수행하여 메모리 및 연산 효율성을 극대화합니다.

2.2 Manual Autograd (수동 그래디언트 계산)

Pytorch의 자동 미분 시스템(AutoGrad)은 LoRA 기반 파인튜닝 과정에서 비효율적일 수 있습니다. Unsloth는 자동 미분을 수동으로 최적화하여 연산 비용을 줄이고 속도를 높였습니다.

2.3 Triton 커널 최적화

Unsloth는 OpenAI에서 개발한 Triton을 활용하여 GPU 연산 성능을 극대화하였습니다. 이를 통해 연산 코드의 핵심 부분을 최적화하여 높은 학습 성능을 제공합니다.

2.4 xFormers 프레임워크 활용

Unsloth는 xFormers 프레임워크와 Flash-Attention을 활용하여 메모리 사용량을 절감하고 연산 속도를 가속화합니다.

3. Unsloth를 활용한 LLM 파인튜닝

Unsloth를 사용하면 복잡한 최적화 작업 없이 간편하게 LLM을 파인튜닝할 수 있습니다. 아래는 Gemma-2-9b 모델을 Google Colab에서 파인튜닝하는 예제입니다.

3.1 Unsloth 설치

!pip install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"
!pip install --no-deps "xformers<0.0.27" "trl<0.9.0" peft accelerate bitsandbytes

# %%capture
# !pip install unsloth
# # Also get the latest nightly Unsloth!
# !pip uninstall unsloth -y && pip install --upgrade --no-cache-dir --no-deps git+https://github.com/unslothai/unsloth.git

# # Install Flash Attention 2 for softcapping support
# import torch
# if torch.cuda.get_device_capability()[0] >= 8:
# !pip install --no-deps packaging ninja einops "flash-attn>=2.6.3"

3.2 FastLanguageModel을 이용한 모델 로딩

Unsloth 라이브러리에서 FastLanguageModel 클래스를 import 합니다. FastLanguageModel 클래스는 Unsloth에서 가장 중요한 클래스 입니다. 허깅페이스 Transformers의 다양한 라이브러리들을 기본으로 하여 추가적인 최적화 및 패치 작업을 진행하게 됩니다.

from unsloth import FastLanguageModel
import torch

max_seq_length = 2048
dtype = None
load_in_4bit = True

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="unsloth/gemma-2-9b",
    max_seq_length=max_seq_length,
    dtype=dtype,
    load_in_4bit=load_in_4bit,
)

3.3 PEFT 설정

FastLanguageModel 클래스의 get_peft_model 메서드는 주어진 모델과 설정을 기반으로 PEFT 객체를 반환합니다.

from trl import SFTTrainer
from transformers import TrainingArguments
from unsloth import is_bfloat16_supported

model = FastLanguageModel.get_peft_model(
    model,
    r=16,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    lora_alpha=16,
    lora_dropout=0,
    bias="none",
    use_gradient_checkpointing="unsloth",
    random_state=3407,
    use_rslora=False,
    loftq_config=None,
)

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    dataset_text_field="text",
    max_seq_length=max_seq_length,
    dataset_num_proc=2,
    packing=False,
    args=TrainingArguments(
        per_device_train_batch_size=2,
        gradient_accumulation_steps=4,
        warmup_steps=5,
        max_steps=60,
        learning_rate=2e-4,
        fp16=not is_bfloat16_supported(),
        bf16=is_bfloat16_supported(),
        logging_steps=1,
        optim="adamw_8bit",
        weight_decay=0.01,
        lr_scheduler_type="linear",
        seed=3407,
        output_dir="outputs",
    ),
)

trainer.train()

3.4 Data Prep

Alpaca dataset으로 샘플 데이터셋을 구성해줍니다.

alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

### Instruction:
{}


### Input:
{}


### Response:
{}"""



EOS_TOKEN = tokenizer.eos_token # Must add EOS_TOKEN
def formatting_prompts_func(examples):
instructions = examples["instruction"]
inputs = examples["input"]
outputs = examples["output"]
texts = []
for instruction, input, output in zip(instructions, inputs, outputs):
# Must add EOS_TOKEN, otherwise your generation will go on forever!
text = alpaca_prompt.format(instruction, input, output) + EOS_TOKEN
texts.append(text)
return { "text" : texts, }
pass

from datasets import load_dataset
dataset = load_dataset("yahma/alpaca-cleaned", split = "train")
dataset = dataset.map(formatting_prompts_func, batched = True,)

3.5 Train the model

Unsloth 모델은 SFTTrainer을 그대로 사용할 수 있으므로 별도의 랩핑된 학습 라이브러리가 필요하지 않습니다. SFTTrainer를 통해 기존에 사용하던 방식대로 모델, 토크나이저, 데이터셋 및 학습과 관련된 파라미터들을 설정하면 됩니다.

from trl import SFTTrainer
from transformers import TrainingArguments
from unsloth import is_bfloat16_supported

trainer = SFTTrainer(
model = model,
tokenizer = tokenizer,
train_dataset = dataset,
dataset_text_field = "text",
max_seq_length = max_seq_length,
dataset_num_proc = 2,
packing = False, # Can make training 5x faster for short sequences.
args = TrainingArguments(
per_device_train_batch_size = 2,
gradient_accumulation_steps = 4,
warmup_steps = 5,
max_steps = 60,
learning_rate = 2e-4,
fp16 = not is_bfloat16_supported(),
bf16 = is_bfloat16_supported(),
logging_steps = 1,
optim = "adamw_8bit",
weight_decay = 0.01,
lr_scheduler_type = "linear",
seed = 3407,
output_dir = "outputs",
report_to = "none", # Use this for WandB etc
),

)

trainer_stats = trainer.train()

 

QLoRA를 통해 학습될 파라미터 수의 비율을 확인할 수 있습니다.

model.print_trainable_parameters()

# trainable params: 54,018,048 || all params: 10,213,228,032 || trainable%: 0.5289

Show current memory stats

현재 사용하고 있는 GPU의 타입과 메모리에 대한 정보를 얻을 수 있습니다.

#@title Show current memory stats
gpu_stats = torch.cuda.get_device_properties(0)
start_gpu_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)
max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)
print(f"GPU = {gpu_stats.name}. Max memory = {max_memory} GB.")
print(f"{start_gpu_memory} GB of memory reserved.")

Show final memory and time stats

학습 시간, 사용 메모리 등 파인 튜닝 과정에 대한 전반적인 정보들을 알 수 있습니다.

#@title Show final memory and time stats
used_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)
used_memory_for_lora = round(used_memory - start_gpu_memory, 3)
used_percentage = round(used_memory /max_memory*100, 3)
lora_percentage = round(used_memory_for_lora/max_memory*100, 3)
print(f"{trainer_stats.metrics['train_runtime']} seconds used for training.")
print(f"{round(trainer_stats.metrics['train_runtime']/60, 2)} minutes used for training.")
print(f"Peak reserved memory = {used_memory} GB.")
print(f"Peak reserved memory for training = {used_memory_for_lora} GB.")
print(f"Peak reserved memory % of max memory = {used_percentage} %.")
print(f"Peak reserved memory for training % of max memory = {lora_percentage} %.")

4. 결론

Unsloth는 LLM 파인튜닝을 위한 최적화된 솔루션으로, 적은 자원으로도 강력한 성능을 제공합니다. 다양한 최적화 기술을 통해 속도와 메모리 효율을 극대화하며, 허깅페이스 생태계와의 높은 호환성 덕분에 손쉽게 적용할 수 있습니다.

LLM 파인튜닝을 보다 빠르고 효율적으로 수행하고 싶다면, Unsloth를 적극 활용해보는 것을 추천합니다.

 

참고자료

https://devocean.sk.com/blog/techBoardDetail.do?ID=166285&boardType=techBlog

https://digitalbourgeois.tistory.com/433

https://colab.research.google.com/drive/1vIrqH5uYDQwsJ4-OO3DErvuv4pBgVwk4?usp=sharing