SageMaker : QLora 기법으로 프라이빗 LLM 강화하기 (Fine-tuning with SageMaker HuggingFace DLC)
Written by Hyeonmin Kim
인공지능 모델에 데이터 셋을 활용, 이를 학습함으로써 특정 도메인을 강화하는 방법을 파인튜닝이라고 부릅니다.
오늘은 위 아키텍처를 토대로
SageMaker Jumpstrart 모델(Mistral)을 QLora 기법을 활용하여,
SageMaker HuggingFace DLC(Deep Learning Containers)를 사용,
nlpai-lab/databricks-dolly-15k-ko 데이터 셋을 사용하여 LLM을 파인튜닝하고,
최종적으로 한국어 사용이 불가능했던 Mistral-7B 모델을 한국어로 추론해 보는 작업을 해보겠습니다.
대략적인 플로우는 다음과 같습니다.
데이터 셋을 불러온 후 이를 Mistral 모델에 맞게 Training Dataset을 만들어준 후에 SageMaker Training을 사용하여 QLora 기법의 스크립트를 돌리게 됩니다. 이때 Sagemaker HuggingFace DLC를 사용하여 ML학습 환경을 컨테이너로 구성하고, 학습을 마치면 HuggingFace Inference Container를 사용하여 Endpoint로 배포하게 됩니다. 사용자는 이 Endpoint에 요청을 하기만 하면 됩니다.
환경 준비하기
먼저, ML 파이프라인을 구성하기 위한 환경이 필요한데, 이전 포스트를 참조하여 Sagemaker Canvas 환경을 구성하고 ML 파이프라인을 구축하기 위한 NoteBook을 생성합니다.
환경이 준비되었다면, 해당 환경에 필요한 파이썬 패키지와 HuggingFace Hub을 사용하기 위한 huggingface-cli를 설치합니다.
!pip install "transformers==4.34.0" "datasets[s3]==2.13.0" "sagemaker>=2.190.0" "gradio==3.50.2" "huggingface_hub[cli]" --upgrade --quiet설치가 완료되면, huggingface-cli에 로그인해야 하는데, 이때 필요한 토큰값은 huggingface.co 의 Profile - edit profile - Access Tokens에서 확인할 수 있습니다.
토큰값을 확인했으면, 로그인을 진행합니다.
!huggingface-cli login --token hf_xxxxxxxxxxxxxxxxxxx마지막으로 IamRole과 디폴트 버킷을 지정해 줍니다.
import sagemaker
import boto3
sess = sagemaker.Session()
sagemaker_session_bucket=None
if sagemaker_session_bucket is None and sess is not None:
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}")데이터 준비하기
먼저, 한국어 학습 셋을 준비해야하는데, Kullm(구름)모델을 개발한 고려대학교 연구소에서 제공하는 databricks-dolly 데이터 셋을 한국어로 번역한 nlpai-lab/databricks-dolly-15k-ko를 사용하도록 하겠습니다.
databricks-dolly 데이터 셋은 Databricks에서 생성한 오픈소스로, 브레인스토밍, 분류, 비공개 QA, 생성, 정보추출, 공개 QA 및 요약 등의 지침을 포함한 데이터 셋입니다. 데이터의 개수는 총 15,011개 입니다.
Datasets 라이브러리를 설치하였기 때문에 쉽게 데이터를 로드할 수 있습니다.
from datasets import load_dataset
from random import randrange
# Load dataset from the hub
dataset = load_dataset("nlpai-lab/databricks-dolly-15k-ko", split="train")
print(f"dataset size: {len(dataset)}")
print(dataset[randrange(len(dataset))])
mistral 모델이 데이터를 학습할 때, 특정 포맷을 지켜주어야 데이터를 인식할 수 있는데요, dolly 데이터 셋의 포멧을 이에 맞게 변경해 주도록 하겠습니다. mistral 모델은 ### Instruction, ### Context (옵션), ### Answer 로 데이터를 각기 구분하며, 이에 맞게 포맷을 변경해 주어야 합니다.
def format_dolly(sample):
instruction = f"### Instruction\n{sample['instruction']}"
context = f"### Context\n{sample['context']}" if len(sample["context"]) > 0 else None
response = f"### Answer\n{sample['response']}"
# join all the parts together
prompt = "\n\n".join([i for i in [instruction, context, response] if i is not None])
return promptprint(format_dolly(dataset[randrange(len(dataset))]))
이제 토크나이저를 생성하고 초기화를 진행해야 하는데, HuggingFace의 AutoTokenizer를 사용하여 Mistral 모델에 맞는 오토 토크나이저를 생성하도록 하겠습니다.
from transformers import AutoTokenizer
model_id = "mistralai/Mistral-7B-v0.1"
tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=True)이제, 위에서 생성한 토크나이저를 사용하여 데이터 전처리 파이프라인을 생성합니다. 데이터 셋에 템플릿을 적용하고 토크나이징한 이후 데이터 셋을 청크로 나눕니다.
from random import randint
import sys
sys.path.append("../scripts/utils")
from pack_dataset import pack_dataset
def template_dataset(sample):
sample["text"] = f"{format_dolly(sample)}{tokenizer.eos_token}"
return sample
dataset = dataset.map(template_dataset, remove_columns=list(dataset.features))
# print random sample
print(dataset[randint(0, len(dataset))]["text"])
# tokenize dataset
dataset = dataset.map(
lambda sample: tokenizer(sample["text"]), batched=True, remove_columns=list(dataset.features)
)
# chunk dataset
lm_dataset = pack_dataset(dataset, chunk_length=2048)
# Print total number of samples
print(f"Total number of samples: {len(lm_dataset)}")
사용된 util인 pack_dataset의 소스코드는 다음과 같습니다.
from itertools import chain
from functools import partial
remainder = {"input_ids": [], "attention_mask": [], "token_type_ids": []}
# empty list to save remainder from batches to use in next batch
def pack_dataset(dataset, chunk_length=2048):
print(f"Chunking dataset into chunks of {chunk_length} tokens.")
def chunk(sample, chunk_length=chunk_length):
# 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(
partial(chunk, chunk_length=chunk_length),
batched=True,
)
print(f"Total number of samples: {len(lm_dataset)}")
return lm_dataset훈련 데이터가 완성되었다면, 이를 S3에 저장하여 Sagemaker Training Job에서 사용할 수 있도록 하겠습니다.
training_input_path = f's3://{sess.default_bucket()}/processed/mistral/dolly-ko/train'
lm_dataset.save_to_disk(training_input_path)
print("uploaded data to:")
print(f"training dataset to: {training_input_path}")훈련 진행하기
데이터 셋도 준비되었다면, Sagemaker의 training job을 사용하여 훈련을 진행하도록 하겠습니다.
훈련에 사용된 스크립트는 QLoRA 스크립트를 사용하였습니다. QLoRA는 전체 16비트 미세 조정 작업 성능을 유지하면서 단일 48GB GPU에서 65B 매개변수 모델을 미세 조정할 수 있을 만큼 메모리 사용량을 줄이는 효율적인 미세 조정 접근 방식입니다. QLoRA는 고정된 4비트 양자화 사전 학습된 언어 모델을 통해 그라데이션을 LoRA(Low Rank Adapter)로 역전파하게 됩니다.
스크립트는 다음 URL에서 확인할 수 있습니다. https://github.com/artidoro/qlora/blob/main/qlora.py
다시 Sagemaker로 돌아와서, 훈련을 위한 파라미터를 정의하도록 하겠습니다.
from huggingface_hub import HfFolder
# hyperparameters, which are passed into the training job
hyperparameters ={
'model_id': model_id,
'dataset_path': '/opt/ml/input/data/training',
'num_train_epochs': 3,
'per_device_train_batch_size': 6,
'gradient_accumulation_steps': 2,
'gradient_checkpointing': True,
'bf16': True,
'tf32': True,
'learning_rate': 2e-4,
'max_grad_norm': 0.3,
'warmup_ratio': 0.03,
"lr_scheduler_type":"constant",
'save_strategy': "epoch",
"logging_steps": 10,
'merge_adapters': True,
'use_flash_attn': True,
'output_dir': '/opt/ml/checkpoints',
if HfFolder.get_token() is not None:
hyperparameters['hf_token'] = HfFolder.get_token()Estimator를 정의해야 하는데, 저희는 Sagemaker HuggingFace DLC를 사용할 것이기 때문에 HuggingFace Estimator를 사용하도록 하겠습니다. Sagemaker SDK 내에 이를 쉽게 사용할 수 있도록 통합되어 있기 때문에, 쉽게 정의할 수 있습니다.
from sagemaker.huggingface import HuggingFace
# define Training Job Name
job_name = f'huggingface-qlora-{hyperparameters["model_id"].replace("/","-").replace(".","-")}'
chekpoint_s3 = f's3://{sess.default_bucket()}/checkpoints'
# create the Estimator
huggingface_estimator = HuggingFace(
entry_point = 'run_qlora.py',
source_dir = '../scripts',
instance_type = 'ml.g5.4xlarge',
instance_count = 1,
checkpoint_s3_uri = chekpoint_s3,
max_run = 2*24*60*60,
base_job_name = job_name,
role = role,
volume_size = 300,
transformers_version = '4.28',
pytorch_version = '2.0',
py_version = 'py310',
hyperparameters = hyperparameters,
environment = { "HUGGINGFACE_HUB_CACHE": "/tmp/.cache" },
disable_output_compression = True
)이제 모든 훈련 준비가 완료되었으므로, 데이터를 정의하고 학습(fitting)을 진행해 보도록 하겠습니다.
data = {'training': training_input_path}
huggingface_estimator.fit(data, wait=True)
학습을 위한 딥러닝 컨테이너가 프로비저닝되고, 학습을 진행하게 됩니다.
학습이 진행되면, 진행 사항은 Sagemaker Training에서 확인할 수 있으며, CloudWatch를 통해 이를 모니터링할 수 있습니다.
학습은 총 29,172초(대략 486.2분, 8.1시간) 소요되었고, g5.4xlarge를 사용하였으므로, $13.16 의 비용이 발생하였습니다. Sagemaker Training은 AWS Batch와 비슷하게 동작하여, 작업을 수행하는 시간만큼만 비용이 발생하므로, 인스턴스를 프로비저닝하여 띄워놓는 불필요한 시간을 줄일 수 있어 일반적인 GPU 인스턴스에서 학습을 진행하는 방법보다 훨씬 저렴하다고 볼 수 있습니다.
모델 배포
모델 학습이 완료되면 이를 배포해야겠죠, 먼저 모델이 정상적으로 생성되었는지 확인해 보도록 하겠습니다. 모델은 디폴트 s3 경로에서 Training Job을 prefix로 한 output/model에서 확인할 수 있습니다.
정상적으로 모델이 저장된 것을 확인하였다면, 이제는 정말로 배포를 진행해 보도록 하겠습니다.
먼저, 배포하기 위한 이미지가 필요한데, 이번에도 huggingface에서 제공하는 inference 환경 이미지를 불러오도록 하겠습니다.
from sagemaker.huggingface import get_huggingface_llm_image_uri
# retrieve the llm image uri
llm_image = get_huggingface_llm_image_uri(
"huggingface",
version="1.1.0",
session=sess,
)
# print ecr image uri
print(f"llm image uri: {llm_image}")
이 이미지를 사용하여 LLM 모델을 정의하겠습니다. 배포에 사용되는 환경변수를 정의하고, 모델데이터에 이전에 확인한 모델의 S3 URI를 입력합니다.
import json
from sagemaker.huggingface import HuggingFaceModel
instance_type = "ml.g5.2xlarge"
number_of_gpu = 1
health_check_timeout = 300
config = {
'HF_MODEL_ID': "/opt/ml/model",
'SM_NUM_GPUS': json.dumps(number_of_gpu),
'MAX_INPUT_LENGTH': json.dumps(1024),
'MAX_TOTAL_TOKENS': json.dumps(2048),
}
llm_model = HuggingFaceModel(
role=role,
image_uri=llm_image,
model_data={'S3DataSource':{'S3Uri': model_s3_path,'S3DataType': 'S3Prefix','CompressionType': 'None'}},
env=config
)이제 모델을 정의하였으므로 배포만 하면 됩니다.
llm = llm_model.deploy(
initial_instance_count=1,
instance_type=instance_type,
container_startup_health_check_timeout=health_check_timeout,
)배포 과정에 있어서 10분 정도 소요되었습니다.
한국어 추론 테스트
이제, 한국어가 제대로 학습되었는지 확인해 보기 위해 Mistral Jumpstart 모델과, 제가 학습을 진행한 모델을 비교하여 프롬프트를 요청해 보도록 하겠습니다.
추론하는 코드는 다음과 같이 작성하였습니다.
import json
import boto3
newline, bold, unbold = "\n", "\033[1m", "\033[0m"
endpoint_name = "엔드포인트 이름"
def query_endpoint(payload):
client = boto3.client("runtime.sagemaker")
response = client.invoke_endpoint(
EndpointName=endpoint_name, InferenceComponentName='추론 컴포넌트 이름(없다면 생략)', ContentType="application/json", Body=json.dumps(payload).encode("utf-8")
)
model_predictions = json.loads(response["Body"].read())
generated_text = model_predictions[0]["generated_text"]
print(f"Input Text: {payload['inputs']}{newline}" f"Generated Text: {bold}{generated_text}{unbold}{newline}")일반적인 Jumpstart Base 모델
한국어 데이터 셋을 적용한 모델
모두 동일한 “대한민국 수도 서울에 대해 알려줘”라는 질문을 하였을 때, 일반적인 모델은 영어로 대답을 하는 것을 볼 수 있지만, 베이스 모델에서는 한국어를 지원하지 않기 때문에, 영어로 내놓은 답변이 최선이고, 아예 한국어를 인식하지 못해 이상한 말만 내뱉는 경우도 있습니다. 하지만 한국어 데이터 셋을 학습시킨 모델의 경우 한국어로 올바른 대답을 내놓고 있는 것을 확인할 수 있습니다.
오늘은 한국어 미지원 모델인 Mistral 모델을 한국어 데이터 셋을 사용하여 훈련을 진행, 한국어를 사용할 수 있도록 파인튜닝을 진행해 보았습니다. 학습에 사용된 데이터 셋이 15k로 많은 데이터가 아니기 때문에, 현재로써는 한국어가 가능하다 정도로 확인할 수 있었고, 이 문제는 더 많은 데이터 셋을 사용하게 되면 해결될 것으로 생각됩니다.
더 많은 데이터 셋을 학습하는 과정에 있어서는 비용이 꽤나 발생할 것으로 생각되므로 스팟 인스턴스를 사용한 학습을 진행하여 비용을 줄이려고 합니다. 이 부분도 업데이트 되면 추후 공유해드리도록 하겠습니다.
이상으로 포스트를 마칩니다. 감사합니다.
Comments