대규모 언어 모델을 Fine tuning할 때 HuggingFace에서는 AutoTrain과 SFTTrainer 두 가지를 사용 할 수 있다. AutoTrain은 CLI 기반으로 손쉽게 학습을 시작할 수 있는 도구이고, SFTTrainer는 파이썬 코드 기반으로 세밀한 제어가 가능한 도구이다.
AutoTrain은 학습 환경을 자동으로 구성해주기 때문에 초보자나 빠르게 결과를 확인하고 싶은 경우에 적합하다. 몇 가지 옵션만 주면 모델과 데이터 로딩, 학습 파라미터 세팅, 저장 경로까지 모두 알아서 설정해주는 것이 장점이다. 하지만 CLI 환경이라서 옵션의 유연성이 제한되고, 체크포인트 재개 기능이 옵티마이저 상태까지 복원되지 않는 제약이 있는 것이 단점이다.
SFTTrainer는 HuggingFace의 trl 라이브러리에서 제공하는 학습기반 클래스이다.
Trainer를 확장하여 Supervised Fine-Tuning에 맞게 최적화되어 있으며, resume_from_checkpoint 옵션을 통해 학습을 완전히 재개할 수 있는 것이 장점이다.
Data preprocessing, Training loop, Logging, Callback 등을 세밀하게 커스터마이징할 수 있기 때문에 연구 목적이나 장기 학습에서 강력한 도구이다. 다만 AutoTrain에 비해 설정을 직접 챙겨야 하므로 알아야 할 것이 많다.
AutoTrain의 CLI 실행 예시는 다음과 같다.
autotrain llm \
--train \
--model {base model name} \
--project-name {fine tuning model name} \
--data-path data/ \
--text-column text \
--lr 2e-4 \
--batch-size 1 \
--epochs 1 \
--block-size 512 \
--warmup-ratio 0.1 \
--lora-r 16 \
--lora-alpha 32 \
--lora-dropout 0.05 \
--weight-decay 0.01 \
--gradient-accumulation 8 \
--mixed-precision fp16 \
--peft \
--quantization int4 \
--trainer sft
위 명령어는 SFTTrainer 코드로 다음과 같이 변환할 수 있다.
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments
from trl import SFTTrainer
from peft import LoraConfig, get_peft_model
from datasets import load_dataset
base_model = f"{base_model_name}"
finetuned_model = f"{fine_tuning_model_name}"
dataset = load_dataset("csv", data_files="data/train.csv")
train_dataset = dataset["train"]
tokenizer = AutoTokenizer.from_pretrained(base_model, use_fast=True)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(
base_model,
load_in_4bit=True,
device_map="auto"
)
training_args = TrainingArguments(
output_dir=finetuned_model,
num_train_epochs=1,
per_device_train_batch_size=1,
gradient_accumulation_steps=8,
learning_rate=2e-4,
weight_decay=0.01,
warmup_ratio=0.1,
logging_dir=f"{finetuned_model}/logs",
logging_steps=50,
save_strategy="steps",
save_steps=200,
save_total_limit=2,
fp16=True,
report_to="none",
)
peft_config = LoraConfig(
r=16,
lora_alpha=32,
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
target_modules=["q_proj", "k_proj", "v_proj"]
)
model = get_peft_model(model, peft_config)
trainer = SFTTrainer(
model=model,
tokenizer=tokenizer,
train_dataset=train_dataset,
max_seq_length=512,
packing=False,
args=training_args,
)
trainer.train()
trainer.save_model(finetuned_model)
tokenizer.save_pretrained(finetuned_model)
이와 같이 AutoTrain CLI는 빠른 실험에 적합하고, SFTTrainer는 완전한 재개와 세밀한 제어가 필요한 상황에 적합하다. 상황에 따라 두 가지 도구를 병행하여 사용하는 것이 가장 효율적이다.
