간단한 소개: PEFT 또는 파라미터 효율적 미세 조정

PEFT는 Parameter Efficient Fine-tuning의 약자로, 사전 훈련된 언어 모델(PLMs)을 모든 모델 파라미터를 미세 조정하지 않고도 다양한 다운스트림 응용 프로그램에 효율적으로 적응할 수 있게 해주는 Hugging Face의 새로운 오픈 소스 라이브러리입니다. 현재 PEFT는 다음과 같은 기술들을 포함하고 있습니다:

LoRA: Low-Rank Adaptation of Large Language Models
Prefix Tuning: P-Tuning v2: Prompt Tuning Can Be Comparable to Fine-tuning Universally Across Scales and Tasks
P-Tuning: GPT Understanding, Too
Prompt Tuning: The Power of Scale for Parameter-Efficient Prompt Tuning

1. 개발환경 설정

!pip install "transformers==4.26.0" "datasets[s3]==2.9.0" sagemaker py7zr --upgrade --quiet

Sagemaker를 로컬 환경에서 사용하려면 Sagemaker에 필요한 권한을 가진 IAM 역할(IAM Role)에 접근할 수 있어야 합니다. 자세한 내용은 여기에서 확인할 수 있습니다.

import sagemaker
import boto3
sess = sagemaker.Session()
# sagemaker session bucket -> used for uploading data, models and logs
# sagemaker will automatically create this bucket if it not exists
sagemaker_session_bucket=None
if sagemaker_session_bucket is None and sess is not None:
    # set to default bucket if a bucket name is not given
    sagemaker_session_bucket = sess.default_bucket()
 
try:
    role = sagemaker.get_execution_role()
except ValueError:
    iam = boto3.client('iam')
    role = iam.get_role(RoleName='sagemaker_execution_role')['Role']['Arn']
 
sess = sagemaker.Session(default_bucket=sagemaker_session_bucket)
 
print(f"sagemaker role arn: {role}")
print(f"sagemaker bucket: {sess.default_bucket()}")
print(f"sagemaker session region: {sess.boto_region_name}")

2. 데이터셋 준비하기

우리는 약 16,000개의 메신저 형태의 대화와 요약본으로 구성된 Samsum 데이터셋을 사용할 것입니다. 이 대화들은 영어에 능통한 언어학자들이 만들어 작성한 것입니다.

{
  "id": "13818513",
  "summary": "Amanda baked cookies and will bring Jerry some tomorrow.",
  "dialogue": "Amanda: I baked cookies. Do you want some?\r\nJerry: Sure!\r\nAmanda: I'll bring you tomorrow :-)"
}

Samsum 데이터셋을 로드하려면 🤗 Datasets 라이브러리의 load_dataset() 메서드를 사용합니다.

from datasets import load_dataset
 
# Load dataset from the hub
dataset = load_dataset("samsum", split="train")
 
print(f"Train dataset size: {len(dataset)}")
# Train dataset size: 14732

훈련을 시작하기 전에 데이터를 전처리해야 합니다. 추상적 요약(Abstractive Summarization)은 텍스트 생성 작업입니다. 모델은 텍스트를 입력으로 받아 요약문을 출력으로 생성합니다. 데이터를 효율적으로 배치(batch)하기 위해 입력과 출력의 길이를 이해하는 것이 필요합니다.

prompt_template을 정의하여 모델의 성능을 높이기 위한 명령 프롬프트를 구성합니다. prompt_template은 “고정된” 시작과 끝 부분이 있으며, 문서는 그 중간에 위치합니다. 따라서 “고정된” 템플릿 부분과 문서가 모델의 최대 길이를 초과하지 않도록 해야 합니다. 우리는 훈련 전에 데이터셋을 전처리하고, 이를 디스크에 저장한 다음 S3에 업로드합니다. 이 단계를 로컬 머신이나 CPU에서 실행한 후 Hugging Face Hub에 업로드할 수 있습니다.

다음은 이를 수행하기 위한 단계별 절차 및 코드 입니다:

from random import randint
from itertools import chain
from functools import partial
 
# custom instruct prompt start
prompt_template = f"Summarize the chat dialogue:\n{{dialogue}}\n---\nSummary:\n{{summary}}{{eos_token}}"
 
# template dataset to add prompt to each sample
def template_dataset(sample):
    sample["text"] = prompt_template.format(dialogue=sample["dialogue"],
                                            summary=sample["summary"],
                                            eos_token=tokenizer.eos_token)
    return sample
 
 
# apply prompt template per sample
dataset = dataset.map(template_dataset, remove_columns=list(dataset.features))
 
print(dataset[randint(0, len(dataset))]["text"])
 
# empty list to save remainder from batches to use in next batch
remainder = {"input_ids": [], "attention_mask": []}
 
 
def chunk(sample, chunk_length=2048):
    # define global remainder variable to save remainder from batches to use in next batch
    global remainder
    # Concatenate all texts and add remainder from previous batch
    concatenated_examples = {k: list(chain(*sample[k])) for k in sample.keys()}
    concatenated_examples = {k: remainder[k] + concatenated_examples[k] for k in concatenated_examples.keys()}
    # get total number of tokens for batch
    batch_total_length = len(concatenated_examples[list(sample.keys())[0]])
 
    # get max number of chunks for batch
    if batch_total_length >= chunk_length:
        batch_chunk_length = (batch_total_length // chunk_length) * chunk_length
 
    # Split by chunks of max_len.
    result = {
        k: [t[i : i + chunk_length] for i in range(0, batch_chunk_length, chunk_length)]
        for k, t in concatenated_examples.items()
    }
    # add remainder to global variable for next batch
    remainder = {k: concatenated_examples[k][batch_chunk_length:] for k in concatenated_examples.keys()}
    # prepare labels
    result["labels"] = result["input_ids"].copy()
    return result
 
 
# tokenize and chunk dataset
lm_dataset = dataset.map(
    lambda sample: tokenizer(sample["text"]), batched=True, remove_columns=list(dataset.features)
).map(
    partial(chunk, chunk_length=2048),
    batched=True,
)
 
# Print total number of samples
print(f"Total number of samples: {len(lm_dataset)}")

데이터셋을 전처리한 후, 새로운 FileSystem 통합 기능을 사용하여 데이터셋을 S3에 업로드할 것입니다. 우리는 sess.default_bucket()을 사용하여 S3 버킷을 지정합니다. 다른 S3 버킷에 저장하려면 이 부분을 조정하면 됩니다. 이 S3 경로는 나중에 훈련 스크립트에서 사용됩니다.

다음은 S3에 데이터셋을 업로드하는 단계별 절차입니다:

# save train_dataset to s3
training_input_path = f's3://{sess.default_bucket()}/processed/samsum-sagemaker/train'
lm_dataset.save_to_disk(training_input_path)
 
print("uploaded data to:")
print(f"training dataset to: {training_input_path}")

3. Amazon SageMaker에 LoRA와 bnb int-8로 BLOOM Fine-tuning하기

LoRA(Low-Rank Adaptation) 기법 외에도, bitsandbytes의 LLM.int8()을 사용하여 고정된 LLM을 int8로 양자화할 것입니다. 이를 통해 BLOOMZ 모델의 메모리 사용량을 약 4배 줄일 수 있습니다.

우리는 모델 훈련을 위해 PEFT(Parameterized Efficient Fine-Tuning)를 사용하는 run_clm.py 스크립트를 준비했습니다. 이 스크립트의 작동 방식에 관심이 있다면, LoRA와 Hugging Face 블로그에서 제공하는 ‘Efficient Large Language Model training with LoRA‘를 참조하세요.

SageMaker 훈련 작업을 생성하기 위해서는 HuggingFace Estimator가 필요합니다. Estimator는 Amazon SageMaker의 훈련 및 배포 작업을 관리합니다. Estimator는 필요한 인프라를 관리하며, SageMaker는 필요한 모든 EC2 인스턴스를 시작하고 관리하며, 올바른 HuggingFace 컨테이너를 제공하고, 제공된 스크립트를 업로드하고, S3 버킷에서 데이터를 컨테이너의 /opt/ml/input/data 경로로 다운로드합니다. 그런 다음 훈련 작업을 시작합니다.

다음은 SageMaker 훈련 작업을 설정하는 단계별 절차입니다:

import time
# define Training Job Name
job_name = f'huggingface-peft-{time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())}'
 
from sagemaker.huggingface import HuggingFace
 
# hyperparameters, which are passed into the training job
hyperparameters ={
  'model_id': model_id,                                # pre-trained model
  'dataset_path': '/opt/ml/input/data/training', # path where sagemaker will save training dataset
  'epochs': 3,                                         # number of training epochs
  'per_device_train_batch_size': 1,                    # batch size for training
  'lr': 2e-4,                                          # learning rate used during training
}
 
# create the Estimator
huggingface_estimator = HuggingFace(
    entry_point          = 'run_clm.py',      # train script
    source_dir           = 'scripts',         # directory which includes all the files needed for training
    instance_type        = 'ml.g5.2xlarge', # instances type used for the training job
    instance_count       = 1,                 # the number of instances used for training
    base_job_name        = job_name,          # the name of the training job
    role                 = role,              # Iam role used in training job to access AWS ressources, e.g. S3
    volume_size          = 300,               # the size of the EBS volume in GB
    transformers_version = '4.26',            # the transformers version used in the training job
    pytorch_version      = '1.13',            # the pytorch_version version used in the training job
    py_version           = 'py39',            # the python version used in the training job
    hyperparameters      =  hyperparameters
)

네, 이제 .fit() 메서드를 사용하여 S3 경로를 훈련 스크립트에 전달하고 훈련 작업을 시작할 수 있습니다. 아래는 전체 프로세스를 요약한 코드입니다:

# define a data input dictonary with our uploaded s3 uris
data = {'training': training_input_path}
 
# starting the train job with our uploaded datasets as input
huggingface_estimator.fit(data, wait=True)

훈련 시간은 20,632초로 약 5.7시간이 소요되었습니다. 사용한 ml.g5.2xlarge 인스턴스의 비용은 시간당 $1.515입니다. 따라서 BLOOMZ 7B 모델을 훈련하는 총 비용은 $8.63입니다. 비용 절감을 위해 스팟 인스턴스를 사용할 수 있지만, 스팟 인스턴스는 대기 시간이나 재시작으로 인해 훈련 시간이 길어질 수 있습니다.

4. 모델을 Amazon SageMaker Endpoint에 배포하기

PEFT를 사용하여 훈련할 때, 보통 어댑터 가중치(adapter weights)가 생성됩니다. 이를 편리하게 배포하기 위해 베이스 모델과 어댑터를 병합하고 언로드하는 merge_and_unload() 메서드를 추가했습니다. 이제 transformers 라이브러리의 파이프라인 기능을 사용할 수 있습니다.

우리는 이제 HuggingFace Estimator 객체에서 deploy()를 사용하여 모델을 배포할 수 있습니다. 원하는 인스턴스 수와 인스턴스 타입을 전달하면 됩니다.

from sagemaker.huggingface import HuggingFaceModel
 
# create Hugging Face Model Class
huggingface_model = HuggingFaceModel(
   model_data=huggingface_estimator.model_data,
   #model_data="s3://hf-sagemaker-inference/model.tar.gz",  # Change to your model path
   role=role,
   transformers_version="4.26",
   pytorch_version="1.13",
   py_version="py39",
   model_server_workers=1
)
# deploy model to SageMaker Inference
predictor = huggingface_model.deploy(
   initial_instance_count=1,
   instance_type= "ml.g5.4xlarge"
)

SageMaker 는 배포 프로세스를 시작할 때 SageMaker 엔드포인트 구성과 SageMaker 엔드포인트를 생성합니다. 엔드포인트 구성은 모델과 인스턴스 타입을 정의합니다.
테스트 데이터셋에서 예제를 사용하여 테스트해 보겠습니다.

from random import randint
from datasets import load_dataset
 
# Load dataset from the hub
test_dataset = load_dataset("samsum", split="test")
 
# select a random test sample
sample = test_dataset[randint(0,len(test_dataset))]
 
# format sample
prompt_template = f"Summarize the chat dialogue:\n{{dialogue}}\n---\nSummary:\n"
 
fomatted_sample = {
  "inputs": prompt_template.format(dialogue=sample["dialogue"]),
  "parameters": {
    "do_sample": True,
    "top_p": 0.9,
    "temperature": 0.1,
    "max_new_tokens": 100,
  }
}
 
# predict
res = predictor.predict(fomatted_sample)
 
print(res[0]["generated_text"].split("Summary:")[-1])
# Kirsten and Alex are going bowling this Friday at 7 pm. They will meet up and then go together.

테스트 결과와 비교해 보겠습니다.

print(sample["summary"])
# Kirsten reminds Alex that the youth group meets this Friday at 7 pm and go bowling.

마지막으로, 엔드포인트를 다시 삭제합니다.

predictor.delete_model()
predictor.delete_endpoint()

Reference

Train and Deploy BLOOM with Amazon SageMaker and PEFT

Learn how to fine-tune BLOOMZ 7B with Amazon SageMaker on a Single GPU using LoRA Hugging Face Transformers.