QLoRA 실습 & Trainer vs SFTTrainer

QLoRA 실습 with MLLMs(InternVL)

Step 1. 필요 Library import:

import os

import torch
import torch.nn as nn
import bitsandbytes as bnb
import transformers

from peft import (
    LoraConfig,
    PeftConfig,
    PeftModel, 
    get_peft_model,
)
from transformers import (
    AutoConfig,
    AutoModel,
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    set_seed,
    pipeline,
    TrainingArguments,
)​

Step 2. 모델 불러온 후 prepare_model_for_kbit_training(model) 진행

devices = [0]#[0, 3]
max_memory = {i: '49140MiB' for i in devices}

model_name = 'OpenGVLab/InternVL2-8B'


model = AutoModelForCausalLM.from_pretrained(
    model_name, 
    cache_dir='/data/huggingface_models',
    trust_remote_code=True,
    device_map="auto",
    max_memory=max_memory,
    quantization_config=BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type='nf4'
        ),
)

# 모델 구조 출력
print(model)

# get_input_embeddings 메서드를 모델에 추가
def get_input_embeddings(self):
    if hasattr(self, 'embed_tokens'):
        return self.embed_tokens
    elif hasattr(self, 'language_model') and hasattr(self.language_model.model, 'tok_embeddings'):
        return self.language_model.model.tok_embeddings
    else:
        raise NotImplementedError("The model does not have an attribute 'embed_tokens' or 'language_model.model.tok_embeddings'.")

model.get_input_embeddings = get_input_embeddings.__get__(model, type(model))

# prepare_model_for_kbit_training 함수를 직접 구현
def prepare_model_for_kbit_training(model):
    for param in model.parameters():
        param.requires_grad = False  # 모든 파라미터의 기울기 계산을 비활성화

    if hasattr(model, 'model') and hasattr(model.model, 'tok_embeddings'):
        for param in model.model.tok_embeddings.parameters():
            param.requires_grad = True  # 임베딩 레이어만 기울기 계산 활성화
    elif hasattr(model, 'embed_tokens'):
        for param in model.embed_tokens.parameters():
            param.requires_grad = True  # 임베딩 레이어만 기울기 계산 활성화
    
    # 필요한 경우 다른 특정 레이어들도 기울기 계산을 활성화할 수 있음
    # 예시: 
    # if hasattr(model, 'some_other_layer'):
    #     for param in model.some_other_layer.parameters():
    #         param.requires_grad = True

    return model

model = prepare_model_for_kbit_training(model)​

Step 3. QLoRA를 붙일 layer 선택:

def find_all_linear_names(model, train_mode):
    assert train_mode in ['lora', 'qlora']
    cls = bnb.nn.Linear4bit if train_mode == 'qlora' else nn.Linear
    lora_module_names = set()
    for name, module in model.named_modules():
        if isinstance(module, cls):
            names = name.split('.')
            lora_module_names.add(names[0] if len(names) == 1 else names[-1])

    if 'lm_head' in lora_module_names:  # LLM의 Head부분에 속하는 애들 pass
        lora_module_names.remove('lm_head')
    
    return list(lora_module_names)


print(sorted(config.target_modules)) # ['1','output', 'w1', 'w2', 'w3', 'wo', 'wqkv']
config.target_modules.remove('1') # LLM의 Head부분에 속하는 애들 제거


config = LoraConfig(
    r=16,
    lora_alpha=16,
    target_modules=find_all_linear_names(model, 'qlora'),
    lora_dropout=0.05,
    bias="none",
    task_type="QUESTION_ANS" #CAUSAL_LM, FEATURE_EXTRACTION, QUESTION_ANS, SEQ_2_SEQ_LM, SEQ_CLS, TOKEN_CLS.
)

model = get_peft_model(model, config)

이후 trainer로 train진행.

QLoRA 붙인 결과:

 

 

 

 

 

 

 

trainer 종류? Trainer vs SFTTrainer

Trainer  v.s. SFTTrainer

∙ Trainer  v.s. SFTTrainer

 - 일반 목적의 훈련: 텍스트 분류, 질의응답, 요약 등의 지도 학습 작업에서 모델을 처음부터 훈련시키는 데 사용됩니다.
 - 높은 커스터마이징 가능성: hyperparameter, optimizer, scheduler, logging, metric 등을 미세 조정할 수 있는 다양한 구성 옵션을 제공합니다.
 - 복잡한 훈련 워크플로우 처리: 그래디언트 축적, 조기 종료, 체크포인트 저장, 분산 훈련 등의 기능을 지원합니다.
 - 더 많은 데이터 요구: 효과적인 훈련을 위해 일반적으로 더 큰 데이터셋이 필요합니다.

∙ SFTTrainer

 - 지도 학습 미세 조정 (SFT): 작은 데이터셋으로 PLMs Fine-Tuning에 최적화.
 - 간단한 인터페이스: 더 적은 configuration으로 간소화된 workflow를 제공.
 - 효율적인 메모리 사용: PEFT와 패킹 최적화와 같은 기술을 사용하여 훈련 중 메모리 소비를 줄입니다.
 - 빠른 훈련: 작은 데이터셋과 짧은 훈련 시간으로도 유사하거나 더 나은 정확도를 달성합니다.

∙ Trainer와 SFTTrainer 선택 기준:

 - Trainer 사용:
큰 데이터셋이 있고, 훈련 루프 또는 복잡한 훈련 워크플로우에 대한 광범위한 커스터마이징이 필요한 경우.
Data preprocessing, Datacollator는 사용자가 직접 설정해야 하며, 일반적인 데이터 전처리 방법을 사용

 - SFTTrainer 사용:
PLMS와 상대적으로 작은 데이터셋을 가지고 있으며, 효율적인 메모리 사용과 함께 더 간단하고 빠른 미세 조정 경험을 원할 경우.
PEFT를 기본적으로 지원, `peft_config`와 같은 설정을 통해 효율적인 파인 튜닝을 쉽게 설정할 수 있다.
Data preprocessing, Datacollator도 효율적인 FT를 위해 최적화되어 있음.
`dataset_text_field`와 같은 필드를 통해 텍스트 데이터를 쉽게 처리할 수 있음.

Feature	Trainer	SFTTrainer
목적	Gerneral Purpose training	Supervised Fine-Tuning of PLMs
커스텀 용도	Highly Customizable	Simpler interface with fewer options
Training workflow	Handles complex workflows	Streamlined workflow
필요 Data	Large Datsets	Smaller Datasets
Memory 사용량	Higher	Lower with PEFT & packing optimization
Training speed	Slower	Faster with smaller datasets