적은 자원으로도 모델의 재학습이 가능하다? Unsloth로 효율적인 CPT 구현하기

머신러닝에서 모델을 업데이트하거나 새로운 도메인의 지식을 추가하기 위해 파인튜닝(fine-tuning) 기법을 많이 사용합니다. 그러나 기존의 파인튜닝 기법들은 많은 자원이 필요하거나 한계가 뚜렷했는데요, 특히 LoRA와 같은 PEFT(Parametric Efficient Fine-Tuning) 기법은 모델의 일부만 학습하기 때문에 자원의 효율성은 좋지만 한계도 명확합니다. 이러한 문제를 해결하기 위해 최근 Unsloth라는 솔루션이 등장했는데, 적은 자원으로도 모델의 성능을 효율적으로 개선할 수 있어 주목받고 있습니다.

이번 글에서는 기존 파인튜닝 방식의 한계를 짚어보고, Unsloth가 어떻게 Continued Pre-Training(CPT)을 효율적으로 수행하는지, 실제 예시를 통해 살펴보겠습니다.

 

반응형

---

기존 파인튜닝의 한계

기존의 파인튜닝 기법들은 모델에 새로운 지식을 주입하기보다는, 모델의 출력을 원하는 방향으로 조정하는 정도의 역할을 해왔습니다. LoRA와 같은 PEFT 기법은 원래 모델 파라미터의 수보다 훨씬 적은 수의 파라미터만을 학습합니다. 이는 모델이 학습할 수 있는 정보의 양을 제한하며, 충분한 데이터를 가지고 있음에도 학습 효과가 제한적일 수밖에 없는 한계로 이어집니다.

특히, 모델에 새로운 지식을 주입하기 위해서는 모델 전체의 파라미터를 재학습하는 것이 필요합니다. 이를 **Continued Pre-Training(CPT)**이라고 합니다. 하지만 CPT는 대량의 데이터셋을 활용하기 때문에 GPU 메모리 요구량이 상당히 높습니다. 예를 들어, 9B 파라미터 모델을 Full 파인튜닝하려면 A100 80G GPU 두 장이 필요할 정도입니다. 더 큰 모델의 경우, 학습 비용은 더욱 증가할 수밖에 없죠. 이러한 이유로 인해 많은 경우 LoRA와 같은 파인튜닝 기법을 통해 타협점을 찾곤 했습니다.

---

Unsloth의 제안: 효율적인 CPT 방식

Unsloth는 기존 LoRA 방식에서 한 걸음 더 나아가 훨씬 적은 자원으로도 CPT를 효율적으로 수행할 수 있는 방안을 제시합니다. Unsloth는 다음과 같은 두 가지 주요 개선점을 강조합니다.

1. 모델 구조 전체 학습: 모델을 대규모의 새로운 도메인 또는 언어로 학습하려면, embed_tokens와 lm_head 레이어도 함께 학습해야 성능이 크게 개선된다고 주장합니다. 이로 인해 모델이 새로운 토큰과 표현을 더 잘 이해하게 됩니다.
2. 파라미터의 한계 극복: 기존 LoRA 방식에서는 전체 파라미터 중 1% 미만만 학습할 수 있었습니다. Unsloth는 이를 극복하기 위해 rsLoRA 기법을 도입하여 학습 가능한 파라미터의 비율을 최대 20%까지 끌어올렸습니다. 이를 통해 대규모 데이터셋에서도 안정적으로 학습이 가능해졌습니다.

---

728x90

Unsloth를 이용한 CPT 간단 예시

이제 Unsloth를 사용하여 모델에 새로운 지식을 주입하는 과정을 간단히 살펴보겠습니다. 아래 예시는 unsloth/Qwen2-7B 모델을 사용하여 학습을 진행하는 예제입니다.

모델 로드

from unsloth import FastLanguageModel
import torch

# 설정
max_seq_length = 2048 
dtype = None  # None은 자동 감지. Float16은 Tesla T4, V100, Bfloat16은 Ampere+ 용
load_in_4bit = True  # 메모리 사용량을 줄이기 위한 4bit 양자화 설정

# 모델 로드
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="unsloth/Qwen2-7B",  # 사용할 모델 이름
    max_seq_length=max_seq_length,
    dtype=dtype,
    load_in_4bit=load_in_4bit,
)

FastLanguageModel 클래스를 통해 모델을 로드합니다. 또한 max_seq_length를 지정하고, 메모리 절약을 위해 4bit 양자화를 활성화합니다.

모델 파인튜닝 설정

# PEFT 설정
model = FastLanguageModel.get_peft_model(
    model,
    r=128,  # 학습할 공간을 충분히 확보하기 위한 r 값 설정
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", 
    "up_proj", "down_proj", "embed_tokens", "lm_head"],  # 학습 대상 모듈
    lora_alpha=32,
    lora_dropout=0,
    bias="none",
    use_gradient_checkpointing="unsloth",  # 매우 긴 문맥 지원
    random_state=3407,
    use_rslora=True,  # 안정적인 학습을 위한 rsLoRA 활성화
    loftq_config=None,
)

# 학습 가능한 파라미터 확인
model.print_trainable_parameters()

위 코드는 embed_tokens와 lm_head 레이어도 학습 대상으로 지정하여 모델이 새로운 도메인의 지식을 더 잘 이해하도록 합니다. rsLoRA를 활성화하여 안정적인 학습을 보장합니다.

데이터셋 처리

# 위키피디아 프롬프트 템플릿
wikipedia_prompt = """위키피디아 기사
### 제목: {}

### 기사:
{}"""

EOS_TOKEN = tokenizer.eos_token  # EOS 토큰 추가

def formatting_prompts_func(examples):
    titles = examples["title"]
    texts = examples["text"]
    outputs = []
    for title, text in zip(titles, texts):
        # EOS 토큰 추가
        text = wikipedia_prompt.format(title, text) + EOS_TOKEN
        outputs.append(text)
    return {"text": outputs}

프롬프트 포맷팅 함수에서는 위키피디아 데이터를 적절히 포맷팅하여 모델 학습에 사용합니다. EOS_TOKEN을 추가하여 생성이 끝날 수 있도록 합니다.

모델 학습

from transformers import TrainingArguments
from unsloth import is_bfloat16_supported, UnslothTrainer, UnslothTrainingArguments

# 학습 설정
trainer = UnslothTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,  # 사용할 데이터셋
    dataset_text_field="text",
    max_seq_length=max_seq_length,
    dataset_num_proc=2,

    args=UnslothTrainingArguments(
        per_device_train_batch_size=2,
        gradient_accumulation_steps=8,
        warmup_ratio=0.1,
        num_train_epochs=1,
        learning_rate=5e-5,
        embedding_learning_rate=1e-5,  # 임베딩 레이어에 대해 더 낮은 학습률
        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()

UnslothTrainer 클래스를 사용하여 모델을 학습합니다. 임베딩 레이어에 대해 더 낮은 학습률을 적용하고, 8bit AdamW 옵티마이저를 사용하여 메모리 효율성을 높입니다.

---

Unsloth는 "가난한 자의 파인튜닝 솔루션"이라는 수식어에 걸맞게 적은 자원으로도 효율적인 파인튜닝을 가능하게 해줍니다. 특히 기존 LoRA 방식의 한계를 극복하고, 대규모 데이터셋을 활용하여 모델을 더 깊이 있게 학습할 수 있도록 개선된 점이 매우 인상적이었습니다.

여러분도 적은 자원으로 대규모 모델을 효과적으로 재학습하고 싶다면, Unsloth를 한 번 시도해 보세요! 최신 기술을 활용해 성능을 극대화할 수 있는 좋은 도구가 될 것입니다.