[RunPod] 런팟으로 LLM 파인튜닝

안녕하세요 jiogenes 입니다.

오늘은 저번 runpod 튜토리얼에 이어서 runpod 내에서 LLM을 학습해보겠습니다.

LLM을 학습하려면 큰 GPU 메모리가 필요한데요.

보통 7B(70억)파라미터 짜리 LLM을 fp32 데이터 타입으로 로드한다고 하면 70억(7B) * 4(Byte) ≈ 280억(28B) Byte ≈ 28GB 정도가 됩니다.

모델 훈련을 하고자 한다면 그레디언트와 optimizer의 파라미터로 인해 약 3배 정도 메모리가 더 필요하므로 훈련을 위해서 대략 80GB 이상의 메모리가 필요하다고 예상할 수 있습니다.

그런데 고작 7B모델만 훈련하려고 해도 현존하는 가장 큰 용량의 그래픽카드인 H100, A100(80GB) 조차도 1장으로는 부족합니다.

이 글에서 다 설명할 수 없지만 이러한 용량 부족을 극복하기 위해 다양한 방법이 존재합니다.

메모리 부족 해결 방법

Mixed precision

먼저 부동소수점의 정밀도를 낮춰서 fine-tuning을 하는 방법이 있습니다. 짧게 설명드리자면 기존 부동소수점 방식인 fp32는 32비트의 메모리를 통해 소수를 나타내는데 이를 fp16인 16비트만 사용해서 나타내는 방법입니다. 이로인해 소수점의 자릿수가 줄고 범위도 줄지만 pre-trained model에서 inference 혹은 fine-tuning을 하는 상황이라면 큰 손실없이 메모리 요구량이 절반 줄어드는 효과가 있습니다. 정확히는 절반이 아니라 precision을 낮춰도 되는 연산만 데이터 타입을 변경하기 때문에 mixed 라는 표현이 들어갑니다.

허깅페이스에서는 이러한 mixed precision 모델을 불러오는 간편한 API를 제공하고 있습니다.

# 모델을 불러올 때
model = AutoModel.from_pretrained("model_name", torch_dtype=torch.float16)

# Trainer를 사용할 때
training_args = TrainingArguments(fp16=True, **default_args)

trainer = Trainer(model=model, args=training_args, train_dataset=ds)
result = trainer.train()
print_summary(result)

더 자세한 설명을 보고싶다면 허깅페이스 문서를 읽어보시면 좋을것 같습니다.

PEFT

그리고 fine-tuning시 모든 파라미터를 전부 fine-tuning 하는게 아니라 일부 파라미터만 튜닝하는 PEFT(Parameter Efficient Fine-Tuning)방법이 있습니다.

PEFT 방법 중에서 가장 널리 사용되는 LoRA에 대해 짧게 설명하고 바로 실습해보도록 하죠.

LoRA는 위 그림처럼 모델의 기존 웨이트 옆에 r만큼 차원을 줄여주는 행렬과 원래 크기만큼 복구시키는 행렬을 학습하는 방법입니다.

적은 수의 파라미터를 학습하지만 hidden state에 더해지는 값으로 인해 fine-tuning의 효과는 상당합니다.

코드를 보자면 다음과 같습니다.

class Linear(nn.Linear, LoRALayer):
    # LoRA implemented in a dense layer
    def __init__(
        self,
        in_features: int,
        out_features: int,
        r: int = 0,
        lora_alpha: int = 1,
        lora_dropout: float = 0.,
        fan_in_fan_out: bool = False, # Set this to True if the layer to replace stores weight like (fan_in, fan_out)
        merge_weights: bool = True,
        **kwargs
    ):
        nn.Linear.__init__(self, in_features, out_features, **kwargs)
        LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout,
                           merge_weights=merge_weights)

        self.fan_in_fan_out = fan_in_fan_out
        # Actual trainable parameters
        if r > 0:
            self.lora_A = nn.Parameter(self.weight.new_zeros((r, in_features)))
            self.lora_B = nn.Parameter(self.weight.new_zeros((out_features, r)))
            self.scaling = self.lora_alpha / self.r
            # Freezing the pre-trained weight matrix
            self.weight.requires_grad = False
        self.reset_parameters()
        if fan_in_fan_out:
            self.weight.data = self.weight.data.transpose(0, 1)

    def reset_parameters(self):
        nn.Linear.reset_parameters(self)
        if hasattr(self, 'lora_A'):
            # initialize B the same way as the default for nn.Linear and A to zero
            # this is different than what is described in the paper but should not affect performance
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B)

    def train(self, mode: bool = True):
        def T(w):
            return w.transpose(0, 1) if self.fan_in_fan_out else w
        nn.Linear.train(self, mode)
        if mode:
            if self.merge_weights and self.merged:
                # Make sure that the weights are not merged
                if self.r > 0:
                    self.weight.data -= T(self.lora_B @ self.lora_A) * self.scaling
                self.merged = False
        else:
            if self.merge_weights and not self.merged:
                # Merge the weights and mark it
                if self.r > 0:
                    self.weight.data += T(self.lora_B @ self.lora_A) * self.scaling
                self.merged = True

    def forward(self, x: torch.Tensor):
        def T(w):
            return w.transpose(0, 1) if self.fan_in_fan_out else w
        if self.r > 0 and not self.merged:
            result = F.linear(x, T(self.weight), bias=self.bias)
            result += (self.lora_dropout(x) @ self.lora_A.transpose(0, 1) @ self.lora_B.transpose(0, 1)) * self.scaling
            return result
        else:
            return F.linear(x, T(self.weight), bias=self.bias)

pytorch의 nn.Linear를 상속받아 Linear의 기존 행동을 그대로 유지할 수 있으며 추가적인 self.lora_A와 self.lora_B로 인해 이 Linear클래스를 사용하면 자동적으로 원래 파라미터는 고정되고 self.lora_A와 self.lora_B만 학습하게 됩니다. 특히 train메소드에서 model.train()과 model.eval()(=model.train(False))일 때 원래 웨이트에서 추가된 웨이트를 빼고 학습을 할 것인지 원래 웨이트에서 추가된 웨이트를 더하고 그냥 linear 연산만 할 것인지 구현해 놓은 부분을 보시면 이해가 빠르실 것 같습니다.

자세한 설명을 보고싶다면 LoRA 논문을 참고하시면 좋을것 같습니다.

LLM 파인튜닝

우선 파인튜닝을 하기 위해 이번 실습에서 사용할 모델과 데이터 및 라이브러리들은 다음과 같습니다.

• Llama2
• 한국어 챗봇 데이터(open korean instructions)
• bitsandbytes (mixed precision)
• peft (LoRA)

팟 생성

저는 3090 1개로 학습을 진행해 보겠습니다. LoRA와 Mixed precision 기술에 힙입어 3090 1개로도 Llama2 7B 모델의 파인튜닝이 가능해 졌기 때문입니다!

디스크는 데이터와 웨이트 파일을 저장할 수 있도록 넉넉하게 50GB로 준비해 줍니다.

가상환경 구성 및 라이브러리 설치

터미널 혹은 주피터 터미널을 들어가서 워크스페이스에 venv로 가상환경을 설치하고 필요한 라이브러리를 설치해봅시다.

root@4197544e2316:/workspace# python -m venv llama
root@4197544e2316:/workspace# source llama/bin/activate
(llama) root@4197544e2316:/workspace# pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu118
(llama) root@4197544e2316:/workspace# pip install trl transformers accelerate peft datasets bitsandbytes wandb ipykernel ipywidgets

runpod에는 이미 기본적으로 pytorch가 깔려있기 때문에 venv 없이 바로 pip install을 통해 설치해도 상관없지만 프로젝트별로 가상환경을 나눠 관리하는 습관을 가지면 좋습니다.

이제 venv 가상환경으로 주피터 노트북을 실행하기 위해 jupyter kernel에 가상환경을 추가하겠습니다.

(llama) root@4197544e2316:/workspace# pip install ipykernel ipywidgets
(llama) root@4197544e2316:/workspace# deactivate
root@4197544e2316:/workspace# python -m ipykernel install --user --name llama --display-name llama

그리고 새로고침 F5키를 눌러준 후 주피터 노트북을 실행하면 다음과 같이 가상환경 커널의 노트북을 만들 수 있습니다.

llama 노트북을 만들고 실습을 진행해 봅시다.

학습 코드

먼저 실습에 사용할 라이브러리를 설치하겠습니다.

!pip install trl transformers accelerate peft datasets bitsandbytes wandb

허깅페이스의 캐시 디렉토리를 변경해 줍시다. 디폴트 경로는 사용자 홈디렉토리 안에 생성되기 때문에 우리가 늘려놓은 Volume Disk에 저장되지 않기 때문입니다. Volume Disk는 /workspace의 용량을 조절합니다.

import os
cache_dir = '/workspace/cache'

if not os.path.exists(cache_dir):
    os.makedirs(cache_dir)

os.environ['HF_HOME'] = cache_dir

필요한 라이브러리를 임포트해줍니다.

from datasets import load_dataset

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TrainingArguments
from peft import LoraConfig, AutoPeftModelForCausalLM, prepare_model_for_kbit_training, get_peft_model
from trl import SFTTrainer, DataCollatorForCompletionOnlyLM

from huggingface_hub import notebook_login

import wandb

앞서 말씀드린것 처럼 llama2 모델과 open-korean-instruction 데이터셋의 리포지토리 이름을 설정합니다. 그리고 저장할 모델의 이름도 설정해 놓습니다.

model_name = 'meta-llama/Llama-2-7b-hf'
data_name = 'heegyu/open-korean-instructions'
fine_tuning_model_name = f'{model_name}-finetuned-open-korean-instructions'
device_map = 'auto'

LoRA의 하이퍼파라미터를 설정합니다. 알파값을 16으로 설정하여 스케일링을 해줍니다. 그리고 r은 64로 설정하여 입력 임베딩 사이즈를 64랭크까지 압축합니다.

peft_config = LoraConfig(
    lora_alpha=16,
    lora_dropout=0.1,
    r=64,
    bias='none',
    task_type='CAUSAL_LM'
)

Mixed precision도 설정해줍니다. 4bit 정밀도로 모델을 로드합니다. bnb_4bit_use_double_quant 파라미터를 통해 중복 양자화 할 수 있도록 설정해 줍니다. 양자화에 사용되는 스케일 펙터를 다시 양자화 함으로써 파라미터당 0.4bit 정도 더 정보를 압축할 수 있습니다. 4bit 정밀도의 데이터 타입은 nf4, 역양자화 할 때 사용할 데이터 타입은 float16으로 만들어 줍니다.

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type='nf4',
    bnb_4bit_compute_dtype='float16',
)

위와 같은 과정을 거치게 되면 파인튜닝에 150GB 넘게 필요한 7B 모델도 24GB 그래픽카드 한 장으로 파인튜닝 할 수 있게 됩니다! 심지어 성능 저하도 없는 상태로 말이죠.

이제 모델을 불러오겠습니다. 라마2 모델을 사용하려면 허깅페이스에 가입이 되어 있어야 합니다. 다음 셀을 실행해보면 아래와 같은 위젯이 나타납니다.

notebook_login()

아직 허깅페이스에 가입이 안되어있다면 허깅페이스 홈페이지에 들어가서 구글계정(혹은 다른 계정)으로 가입을 합니다. 가입한 후 페이지의 왼쪽 탭에서 Settings 라는 메뉴를 클릭해 들어갑니다.

Settings에 들어가면 다음과 같은 메뉴가 나타나는데요. 여기서 Access Tokens를 클릭합니다.

New token을 클릭해서 이름을 적고 Role을 write로 하여 새로운 토큰을 만들어 줍니다.

생성된 토큰을 복사해서 아까 만들어진 입력창에 복사해줍시다.

그러면 다음과 같이 위젯이 바뀌면서 로그인이 됩니다.

혹시 라마2 모델을 사용해본 경험이 없다면 허깅페이스 라마2 모델 홈페이지를 들어가서 모델 접근에 대한 동의를 얻어야 합니다.

위 그림 에서 아래 그림 처럼 바뀌면 됩니다.

다시 코드로 돌아와서 우리가 학습하는 모델이 어떻게 변화하는지 알아보기 위해 wandb에 로그인 합니다. wandb역시 홈페이지에서 가입이 필요합니다.

wandb.login()
wandb.init(project=fine_tuning_model_name.split('/')[-1])

가입 후 다음 링크에서 엑세스 토큰을 얻어 입력창에 붙여넣기 해 줍니다.

자 이제 귀찮은 작업들이 모두 완료됐습니다. 본격적인 학습 코드를 작성하겠습니다.

먼저, 데이터셋을 불러오고 데이터를 한번 찍어보겠습니다. 데이터셋이 많을수록 당연히 좋겠지만 학습 시간이 더 오래 걸리므로 데이터셋의 10%만 사용하겠습니다.

dataset = load_dataset(data_name, split='train[:10%]')
print(dataset[0]['text'])

>>>
<usr> 유언장이 있는 것이 좋다는 말을 들었습니다. 유언장이란 무엇입니까?
<bot> 유언장은 귀하가 사망한 후 귀하의 재산이 어떻게 분배되어야 하는지를 지정하는 법적 문서입니다. 또한 귀하가 가질 수 있는 자녀나 기타 부양가족을 누가 돌봐야 하는지 명시할 수 있습니다. 유언장에 적용되는 법률이 주마다 다르기 때문에 귀하의 유언장이 유효하고 최신인지 확인하는 것이 중요합니다.

<usr>토큰 과 <bot>토큰 으로 사용자 프롬프트와 챗봇 대답이 나뉘어 지는것을 확인할 수 있습니다.

모델을 불러오겠습니다. 라마2 모델을 불러오고 peft와 mixed precision을 적용한 모델로 바꿔줍니다.

base_model = AutoModelForCausalLM.from_pretrained(model_name,
                                             quantization_config=bnb_config,
                                             use_cache=False,
                                             device_map=device_map)
base_model.config.pretraining_tp = 1
base_model.gradient_checkpointing_enable()
base_model = prepare_model_for_kbit_training(base_model)
peft_model = get_peft_model(base_model, peft_config)

토크나이저도 불러옵니다.

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = 'right'

학습 하이퍼파라미터를 설정해 줍시다. 여기서 특히 per_device_train_batch_size는 총 배치 사이즈를 나타내고 gradient_accumulation_steps는 원하는 배치 사이즈를 한번에 넣을 메모리가 부족할 때 나눠서 넣고 그레디언트를 누적해서 마지막에 최적화 할 수 있도록 해줍니다.

training_args = TrainingArguments(
    output_dir=fine_tuning_model_name,
    num_train_epochs=3,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=2,
    gradient_checkpointing=True,
    optim='paged_adamw_32bit',
    logging_steps=5,
    save_strategy='epoch',
    learning_rate=2e-4,
    weight_decay=0.001,
    max_grad_norm=0.3,
    warmup_ratio=0.03,
    group_by_length=False,
    lr_scheduler_type='cosine',
    disable_tqdm=True,
    report_to='wandb',
    seed=42
)

학습은 기본 트레이너로도 잘 동작하지만 SFTTrainer를 사용하면 더 효과적으로 학습할 수 있습니다. 자세한 내용은 여기를 참조해 주세요.

trainer = SFTTrainer(
    model=peft_model,
    train_dataset=dataset,
    dataset_text_field='text',
    max_seq_length=min(tokenizer.model_max_length, 2048),
    tokenizer=tokenizer,
    packing=True,
    args=training_args
)

이제 학습을 해봅시다.

trainer.train()

모델 저장 및 평가

학습 중에는 이렇게 완디비에서 학습률과 로스, 하드웨어 사용량 등등을 확인할 수 있습니다.

학습이 끝나면 완디비를 종료하고 모델을 저장해 줍니다. 이 모델은 LoRA의 업데이트된 가중치만 저장합니다.

wandb.finish()
trainer.save_model()

학습된 모델을 다시 불러와 merge_and_unload메소드를 통해 LoRA로 업데이트 된 웨이트벡터를 원래 웨이트에 더해줍니다. 더해진 최종 모델을 저장해 줍니다.

trained_model = AutoPeftModelForCausalLM.from_pretrained(
    training_args.output_dir,
    low_cpu_mem_usage=True,
    return_dict=True,
    torch_dtype=torch.float16,
    device_map=device_map
)

lora_merged_model = trained_model.merge_and_unload()
lora_merged_model.save_pretrained('merged', safe_serialization=True)
tokenizer.save_pretrained('merged')

저장된 모델을 허깅페이스에 업데이트 합니다.

lora_merged_model.push_to_hub(training_args.output_dir)
tokenizer.push_to_hub(training_args.output_dir)

허깅페이스 허브에 업데이트하면 다음과 같이 허깅페이스 허브에 자신의 모델이 올라와 있는것을 확인할 수 있습니다. 이 모델은 다른 허깅페이스 모델 처럼 언제 어디서든 API로 불러와 사용할 수 있습니다!

마지막으로 모델의 성능을 평가해 봅니다.

prompt = '<usr> 누가 "공산당 선언" 이라는 책을 썼습니까?n<bot>'
input_ids = tokenizer(prompt, return_tensors='pt', truncation=True).input_ids.cuda()

print(f"-------------------------\n")
print(f"Prompt:\n{prompt}\n")
print(f"-------------------------\n")

print(f"Base Model Response :\n")
output_base = base_model.generate(input_ids=input_ids, max_new_tokens=500, do_sample=True, top_p=0.9,temperature=0.5)
print(f"{tokenizer.batch_decode(output_base.detach().cpu().numpy(), skip_special_tokens=True)[0][len(prompt):]}")
print(f"-------------------------\n")

print(f"Trained Model Response :\n")
trained_model = lora_merged_model.generate(input_ids=input_ids, max_new_tokens=500, do_sample=True, top_p=0.9,temperature=0.5)
print(f"{tokenizer.batch_decode(trained_model.detach().cpu().numpy(), skip_special_tokens=True)[0][len(prompt):]}")
print(f"-------------------------\n")

print(f"LORA Model Response :\n")
output_trained_lora = lora_merged_model.generate(input_ids=input_ids, max_new_tokens=500, do_sample=True, top_p=0.9,temperature=0.5)
print(f"{tokenizer.batch_decode(output_trained_lora.detach().cpu().numpy(), skip_special_tokens=True)[0][len(prompt):]}")
print(f"-------------------------\n")

>>>
-------------------------

Prompt:
<usr> 누가 "공산당 선언" 이라는 책을 썼습니까?
<bot>

-------------------------

Base Model Response :

공산당 선언은 러시아 공산당이 발행한 책입니다. 공산당은 1917년 러시아 혁명 이후 러시아 제국을 몰아내고 러시아 공화국을 세웠습니다. 공산당은 또한 소비에트 연방의 창설자였습니다.
-------------------------

Trained Model Response :

공산당 선언은 1917년 11월 7일 볼셰비키가 작성한 공식 선언문이다. 공산당 선언은 볼셰비키가 러시아 제국에서 독재 정권을 장악했다는 것을 선언한 첫 번째 문서였다.
-------------------------

LORA Model Response :

공산당 선언은 1918년 러시아 공산당이 작성한 책입니다. 공산주의자들은 작업자 클럽과 소비자 클럽을 통해 러시아 사회를 통제하고 싶어했습니다. 공산당은 러시아 혁명에서 승리한 뒤 실질적으로 집권하기 위해 실행되었습니다.
-------------------------

런팟을 이용해서 LLM을 학습하는 실습을 진행해 보았습니다. 저는 약 20시간 정도가 걸렸고 시간당 0.29달러로 약 6달러가 안되는 돈으로 Llama2 모델을 파인튜닝 할 수 있었습니다. 결과는 엄청난 할루시네이션의 향연입니다. 하지만 학습에 사용한 데이터가 적고 한번도 학습한 적 없는 데이터니 이정도면 감지덕지라고 해야할까요.

더 많은 데이터와 데이터 전처리 그리고 다양한 학습 방법들을 적용한다면 이것보다 더 좋은 결과가 나올것입니다. 좋은 LLM은 단지 데이터만 많이 넣고 돌리는 수준에서 만들어 지는게 아니라 수많은 삽질(데이터 수집 및 전처리)에 삽질(아키텍쳐 수정)을 거듭해서 나온 결과물일 것입니다.

GPU 메모리가 부족해서 LLM을 건드려볼 시도조차 못해본 사람도 많겠지만 우리는 런팟을 이용해서 (돈이 조금 들긴 하겠지만) 삽질을 해볼 수 있는 기회를 얻은 것입니다. 이 기회를 잘 활용해서 LLM연구가 더 활발히 진행됐으면 좋겠습니다.

읽어주셔서 감사합니다 🤗

참고자료

• https://abvijaykumar.medium.com/fine-tuning-llm-parameter-efficient-fine-tuning-peft-lora-qlora-part-2-d8e23877ac6f
• https://huggingface.co/datasets/heegyu/open-korean-instructions
• https://huggingface.co/meta-llama/Llama-2-7b-hf
• https://github.com/huggingface/peft
• https://huggingface.co/docs/diffusers/main/en/training/lora
• https://github.com/microsoft/LoRA
• https://arxiv.org/abs/2106.09685
• https://huggingface.co/docs/transformers/main/en/perf_train_gpu_one
• https://huggingface.co/docs/trl/v0.7.4/en/sft_trainer#trl.SFTTrainer
• https://huggingface.co/blog/4bit-transformers-bitsandbytes