실습 3 — SFTTrainer 로 LoRA 학습하기

# 앞에서 생성에 사용한 모델을 정리하고 학습할 모델을 새로 불러옵니다.
# SFTConfig = SFT 학습 설정표, SFTTrainer = supervised fine-tuning trainer(모범 답안 학습 담당)
from trl import SFTConfig, SFTTrainer  # SFT 학습 설정과 Trainer를 불러옵니다.

# generate: 6차시 셀 11의 생성 파이프라인. 모델을 붙잡고 있으므로 함께 지운다
del generate, model  # 기존 모델과 생성 파이프라인에 대한 참조를 삭제합니다.
gc.collect()  # 더 이상 참조하지 않는 파이썬 객체를 정리합니다.
# empty_cache(): PyTorch 가 잡아 두었던 빈 GPU 메모리를 돌려준다
torch.cuda.empty_cache()  # 새 모델을 불러오기 전에 사용하지 않는 GPU 캐시를 비웁니다.
# load_base_model(): 6차시 셀 6의 함수. Qwen 을 BF16 으로 불러와 GPU 에 올린다
model = load_base_model()  # 학습할 새 기반 모델을 불러옵니다.

# LoRA 학습에 사용할 배치 크기, 학습률, 평가 시점 등을 정합니다.
sft_config = SFTConfig(  # 첫 LoRA 학습의 실행 설정을 만듭니다.
    # output_dir: 학습 기록을 둘 폴더(save_strategy="no" 라 체크포인트는 남지 않는다)
    output_dir="outputs/persona-lora",  # 학습 결과를 저장할 경로입니다.
    # per_device_train_batch_size: GPU 하나가 한 번에 보는 학습 예시 수
    per_device_train_batch_size=1,  # 장치 하나에서 한 단계에 예시 1개를 학습합니다.
    # per_device_eval_batch_size: GPU 하나가 한 번에 보는 검증 예시 수
    per_device_eval_batch_size=1,  # 장치 하나에서 한 번에 예시 1개를 검증합니다.
    # gradient_accumulation_steps: 그래디언트를 8번 모아 한 번 고친다 → 실질 배치 1 × 8 = 8
    gradient_accumulation_steps=8,  # 8단계의 그래디언트를 모은 뒤 가중치를 한 번 업데이트합니다.
    # learning_rate: 한 걸음의 보폭. 2e-4 = 0.0002
    learning_rate=2e-4,  # LoRA 가중치의 학습률입니다.
    # lr_scheduler_type: 학습률을 코사인 곡선 모양으로 0 까지 줄인다
    lr_scheduler_type="cosine",  # 학습률을 줄이는 방식입니다.
    # warmup_ratio: 전체 99스텝의 5%(올림해서 5스텝) 동안 학습률을 0 에서 서서히 올린다
    warmup_ratio=0.05,  # 학습률을 서서히 올릴 구간입니다.
    # max_length: 이보다 긴 예시는 뒤를 자른다
    max_length=256,  # 템플릿 적용 후 최대 토큰 수입니다. 데이터 길이가 대부분 100토큰 안팎이라 여유 있게 잡았습니다.
    # completion_only_loss: 질문 토큰은 채점에서 뺀다(7차시의 labels -100 과 같은 일)
    completion_only_loss=True,  # 질문은 제외하고 답변 토큰에 대해서만 손실을 계산합니다.
    # eval_strategy: epoch 한 바퀴마다 검증 데이터 32개로 검증 손실을 잰다
    eval_strategy="epoch",  # 학습 데이터를 한 번 다 볼 때마다 검증 데이터를 평가합니다.
    # save_strategy: 중간 저장 안 함(어댑터는 실습 4에서 따로 저장)
    save_strategy="no",  # 중간 체크포인트는 저장하지 않습니다.
    # bf16 / fp16: 6차시에서 고른 MODEL_DTYPE 에 맞춘다. RTX 3080 은 BF16 → bf16=True
    bf16=MODEL_DTYPE == torch.bfloat16,  # MODEL_DTYPE이 BF16이면 BF16으로 계산합니다.
    fp16=MODEL_DTYPE == torch.float16,  # MODEL_DTYPE이 FP16이면 FP16으로 계산합니다.
)
# num_train_epochs 는 적지 않았으므로 기본값 3 을 쓴다

trainer = SFTTrainer(  # 모델·데이터·설정을 하나의 학습 객체로 묶습니다.
    # model: 1단계에서 새로 불러온 기반 모델
    model=model,  # LoRA를 붙일 기반 모델입니다.
    # args: 2단계의 설정표
    args=sft_config,  # 학습 설정입니다.
    # train_dataset: 7차시의 학습 데이터 263개(prompt·completion 꼴)
    train_dataset=persona_ds["train"],  # 학습 데이터입니다.
    # eval_dataset: 검증 데이터 32개
    eval_dataset=persona_ds["valid"],  # 검증 데이터입니다.
    # processing_class: 대본 양식 입히기와 토큰화를 맡을 토크나이저
    processing_class=tokenizer,  # 템플릿 적용과 토큰화를 수행합니다.
    # peft_config: 실습 1의 lora_config. SFTTrainer 가 get_peft_model 을 대신 부른다
    peft_config=lora_config,  # SFTTrainer가 LoRA를 붙입니다.
)
# 결과: trainable params: 8,798,208 || all params: 502,830,976 || trainable%: 1.7497  (실습 1과 같다)
trainer.model.print_trainable_parameters()  # LoRA 파라미터만 학습 대상인지 확인합니다.

# train(): 학습 반복을 돈다. epoch 마다 검증 손실을 재고 표로 찍는다
# 결과: [99/99 03:11, Epoch 3/3] 과 epoch 별 손실 표(학습 손실 3.09 → 2.42 → 1.87)
train_result = trainer.train()  # LoRA 학습을 실행합니다.