AWS 기술 블로그

QLoRA 기법으로 Falcon-40B 및 기타 대규모 모델(LLM)을 Amazon SageMaker Studio 노트북의 대화형 환경에서 파인튜닝하기

이 글은 AWS Machine Learning Blog에 게시된 Interactively fine-tune Falcon-40B and other LLMs on Amazon SageMaker Studio notebooks using QLoRA by Sean Morgan, Philipp Schmid, and Lauren Mullennex를 한국어로 번역 및 편집하였습니다.

대규모 언어 모델(LLM; Large Language Models)을 파인튜닝(Fine-tuning) 하면 오픈 소스 파운데이션 모델(Foundation model)을 개선하여 도메인별 작업에서 더욱 향상된 성능을 끌어낼 수 있습니다. 이 게시물에서는 최신 오픈 소스 모델을 파인튜닝 하기 위해 Amazon SageMaker 노트북을 사용할 때의 이점에 관해 설명합니다. 단일 노트북 인스턴스에서 대규모 모델의 대화형 파인튜닝 작업을 하기 위하여 허깅페이스(Hugging Face)의 효율적인 파라미터 파인튜닝(PEFT; Parameter-efficient Fine-tuning) 라이브러리와 bitsandbytes 패키지를 통한 양자화 기법을 활용합니다. 본 게시물에서는 NVIDIA A10G GPU 4장이 탑재된 단일 ml.g5.12xlarge 인스턴스로 Falcon-40B 모델을 파인튜닝 하는 방법을 보여드리지만, 같은 전략으로 p4d/p4de 노트북 인스턴스에서 훨씬 더 큰 모델을 파인튜닝 할 수 있습니다.

보통 이러한 대규모 모델의 전체 정밀도 표현의 모델 파라미터는 한 장의 GPU나 여러 장의 GPU의 메모리에 모두 들어가지 않습니다. 이러한 크기의 모델에 대한 파인튜닝하고 추론을 수행하는 대화형 노트북 환경을 위해 저희는 Quantized LLMs with Low-Rank Adapters (QLoRA) 라는 신규 기법을 사용합니다. QLoRA는 LLM의 메모리 사용량을 줄이면서도 견고한 성능을 유지하는 효율적인 파인튜닝 방식입니다. 허깅페이스와 앞서 언급된 논문의 저자들은 기본 개념과 트랜스포머 및 PEFT 라이브러리와의 연동 내용을 다루는 자세한 블로그 게시물을 게시했습니다.

SageMaker 노트북으로 LLM 파인튜닝 하기

SageMaker에는 데이터를 탐색하고 기계 학습(ML) 모델을 구축하기 위한 완전 관리형 노트북을 가동하는 두 가지 옵션을 제공합니다. 첫 번째 옵션은 빠른 시작을 제공하는 협업 노트북으로, ML을 위한 완전 통합 개발 환경(IDE; Integrated Development Environment)인 Amazon SageMaker Studio 내에서 액세스합니다. SageMaker Studio에서 노트북을 빠르게 시작하고, 작업을 중단하지 않고 기본 컴퓨팅 리소스에 접속하거나 중단할 수 있으며, 필요에 따라 컴퓨팅 리소스를 변경할 수 있습니다. (예: ml.t3.medium에서 ml.g4dn.xlarge로 변경) 또한, 노트북에서 실시간으로 공동 편집 및 공동 작업도 가능합니다. 노트북 생성 외에도 SageMaker Studio의 단일 창에서 모델 빌드, 훈련, 디버그, 추적, 배포 및 모니터링까지의 일련의 ML 개발 단계를 수행할 수 있습니다. 두 번째 옵션은 클라우드에서 노트북을 실행하는 완전 관리형 단일 ML 컴퓨팅 인스턴스인 SageMaker 노트북 인스턴스로, 노트북 구성을 더 잘 제어할 수 있습니다.

이 게시물에서는 허깅페이스 트랜스포머의 텐서보드(TensorBoard) 지원과 함께 SageMaker Studio의 관리형 텐서보드 실험 추적을 활용하고자 하기에 SageMaker Studio 노트북을 사용합니다. 그러나 모든 예제 코드는 conda_pytorch_p310 커널을 사용하는 SageMaker 노트북 인스턴스에서도 작동합니다. 한 가지 주목할 만한 점은 SageMaker Studio의 Amazon EFS(Amazon Elastic File System) 볼륨을 사용하면 미리 정해진 Amazon EBS(Amazon Elastic Block Store) 볼륨 크기를 프로비저닝할 필요가 없다는 것인데, 이는 LLM의 모델 용량이 큰 경우에 유용합니다. (예컨대 GPT-3는 1,750억 개의 파라미터로 모델 가중치 크기가 약 800기가 이상 필요하며, Floating point 16을 적용한다 해도 약 400기가 정도의 모델 용량을 점유합니다)

다만, SageMaker Studio 노트북은 컴퓨팅 리소스를 변경하거나 추가할 때, 인스턴스 프로비저닝 및 컨테이너 런칭까지 소요되는 콜드 스타트가 존재합니다. 이 경우 대규모 GPU 인스턴스가 지원되는 SageMaker 노트북을 런칭하면 콜드 스타트 없이 신속한 프로토타이핑 및 디버깅이 가능합니다. 하지만 노트북 인스턴스 사용이 끝나면 종료해야 추가 비용이 발생하지 않습니다. 파인튜닝을 위해 Amazon SageMaker JumpStart 및 SageMaker 허깅페이스 컨테이너와 같은 다른 옵션을 사용할 수 있으며, 아래 게시물들을 참조하여 여러분과 여러분의 팀에게 가장 적합한 옵션을 선택하는 것이 좋습니다.

사전 요구 사항

SageMaker Studio를 처음 사용하는 경우, 먼저 SageMaker 도메인을 생성해야 합니다. 또한 선택 사항이지만, 이 튜토리얼에서는 실험 추적을 위한 관리형 텐서보드 인스턴스를 사용합니다.

또한, 해당 SageMaker Studio KernelGateway 앱에 대한 서비스 할당량(quota) 증가를 요청해야 할 수도 있습니다. Falcon-40B 모델의 파인튜닝을 위해 ml.g5.12xlarge 인스턴스를 사용합니다.

서비스 할당량 증가를 요청하려면 AWS 서비스 할당량 콘솔에서 AWS 서비스, Amazon SageMaker로 이동한 다음 ml.g5.12xlarge 인스턴스에서 실행 중인 Studio KernelGateway 앱(Studio KernelGateway Apps running on ml.g5.12xlarge instance)을 선택합니다.

시작하기

이 게시물의 코드 샘플은 GitHub 저장소에 있습니다. 먼저, 패키지를 설치하기 위해 최신 Python 3.10 환경을 갖출 수 있도록 SageMaker Studio에서 Data Science 3.0 이미지와 Python 3 커널을 선택합니다.

파이토치(PyTorch)와 필수 허깅페이스 라이브러리 및 bitsandbytes 라이브러리를 설치합니다:

%pip install -q -U torch==2.0.1 bitsandbytes==0.39.1
%pip install -q -U datasets py7zr einops tensorboardX
%pip install -q -U git+https://github.com/huggingface/transformers.git@850cf4af0ce281d2c3e7ebfc12e0bc24a9c40714
%pip install -q -U git+https://github.com/huggingface/peft.git@e2b8e3260d3eeb736edf21a2424e89fe3ecf429d
%pip install -q -U git+https://github.com/huggingface/accelerate.git@b76409ba05e6fa7dfc59d50eee1734672126fdba

그런 다음 파이토치의 종속 요소로 함께 설치된 CUDA를 사용하여 LD_LIBRARY_PATH 환경 변수를 설정합니다. 이는 bitsandbytes 라이브러리가 올바른 CUDA 공유 라이브러리(CUDA shared object binary)를 찾아 로드하는데 필요한 단계입니다.

# Add installed cuda runtime to path for bitsandbytes
import os
import nvidia

cuda_install_dir = '/'.join(nvidia.__file__.split('/')[:-1]) + '/cuda_runtime/lib/'
os.environ['LD_LIBRARY_PATH'] = cuda_install_dir

사전 훈련된 파운데이션 모델 로드

Falcon-40B 모델을 4장의 A10G GPU가 탑재된 ml.g5.12xlarge 인스턴스에서 파인튜닝 하기 위해서는 QLoRA 기법으로 4비트 정밀도로 양자화를 수행해야 합니다. bitsandbytes와 허깅페이스 Acclerate의 파이프라인 병렬화를 같이 사용하면 4비트 정밀도로 양자화한 모델을 각 A10G GPU의 메모리에 로드할 수 있습니다. 앞서 언급한 허깅페이스 게시물에서 설명한 대로, 모델 가중치는 4비트 NormalFloat로 저장되지만, 필요에 따라 forward 및 backward pass에서 bfloat16으로 역양자화(dequantize)합니다. 그 덕분에 다양한 실험에서 QLoRA 튜닝은 16비트 파인튜닝의 성능과 동등한 수준을 보이는 것으로 나타났습니다.

model_id = "tiiuae/falcon-40b"
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

사전 훈련된 가중치를 로드할 때 device_map=”auto"로 지정하여 허깅페이스 Accelerate가 모델의 각 레이어에 배치할 GPU를 자동으로 결정하도록 합니다. 이 프로세스를 모델 병렬화(model parallelism)라고 합니다.

# Falcon requires you to allow remote code execution. This is because the model uses a new architecture that is not part of transformers yet.
# The code is provided by the model authors in the repo.
model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True, quantization_config=bnb_config, device_map="auto")

허깅페이스의 PEFT 라이브러리를 사용하면 대부분의 원래 모델 가중치를 고정하고 훨씬 작은 추가 파라미터 세트를 훈련하여 모델 레이어를 교체하거나 확장할 수 있습니다. 따라서 훈련에 필요한 컴퓨팅 비용이 현저히 줄어듭니다. 파인튜닝 하려는 Falcon 모델의 트랜스포머 블록 모듈을 LoRA 설정에서 target_modules로 설정합니다:

from peft import LoraConfig, get_peft_model

config = LoraConfig(
    r=8,
    lora_alpha=32,
    target_modules=[
        "query_key_value",
        "dense",
        "dense_h_to_4h",
        "dense_4h_to_h",
    ],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)

model = get_peft_model(model, config)
print_trainable_parameters(model)
# Output: trainable params: 55541760 || all params: 20974518272|| trainable%: 0.2648058910327664

모델 파라미터의 0.26%만 파인튜닝하고 있으므로 짧은 시간 내에 훈련할 수 있습니다.

데이터 세트 로드

파인튜닝에 samsum 데이터 세트를 사용합니다. Samsum은 16,000개의 레이블링 된 요약문이 포함된 메신저 스타일의 대화문으로 이루어져 있습니다. 다음은 데이터 세트의 예시입니다.

{
    "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 :-)"
}

실무에서는 모델을 튜닝하고자 하는 작업에 대한 특정 정보가 담긴 데이터 세트를 사용하기를 원할 겁니다. 이러한 데이터 세트를 구축하는 프로세스는 High-quality human feedback for your generative AI applications from Amazon SageMaker Ground Truth Plus 게시물에 설명한 대로 Amazon SageMaker Ground Truth Plus를 사용하여 가속할 수 있습니다.

모델 파인튜닝

모델을 파인튜닝 하기 전에 사용할 하이퍼파라미터를 정의하고 모델을 훈련합니다. 또한 logging_dir 매개변수를 정의하고 허깅페이스 트랜스포머에 report_to="tensorboard"를 요청하여 지표(metric)를 텐서보드에 로깅 할 수도 있습니다:

bucket = ”<YOUR-S3-BUCKET>”
log_bucket = f"s3://{bucket}/falcon-40b-qlora-finetune"

import transformers

# We set num_train_epochs=1 simply to run a demonstration

trainer = transformers.Trainer(
    model=model,
    train_dataset=lm_train_dataset,
    eval_dataset=lm_test_dataset,
    args=transformers.TrainingArguments(
        per_device_train_batch_size=8,
        per_device_eval_batch_size=8,
        logging_dir=log_bucket,
        logging_steps=2,
        num_train_epochs=1,
        learning_rate=2e-4,
        bf16=True,
        save_strategy = "no",
        output_dir="outputs",
        report_to="tensorboard",
    ),
    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
)

파인튜닝 모니터링

앞선 설정을 통해 실시간으로 파인튜닝을 모니터링할 수 있습니다. GPU 사용량을 실시간으로 모니터링하려면 커널의 컨테이너에서 직접 nvidia-smi를 실행하면 됩니다. 이미지 컨테이너에서 실행되는 터미널을 시작하려면 노트북 상단의 터미널 아이콘을 클릭하세요.

여기에서 Linux의 watch 명령을 사용하여 0.5초마다 nvidia-smi를 반복적으로 실행할 수 있습니다.

watch -n 0.5 nvidia-smi

앞의 애니메이션에서 모델 가중치가 4장의 GPU에 분산되어 있고 레이어가 차례대로 처리되면서 계산이 분산되고 있는 것을 확인해 보세요.

훈련 지표를 모니터링하기 위해 지정된 Amazon S3(Amazon Simple Storage Service) 버킷에 기록하는 텐서보드 로그를 활용합니다. AWS SageMaker 콘솔에서 SageMaker Studio 도메인 사용자의 텐서보드를 실행할 수 있습니다.

텐서보드 로드 후, 허깅페이스 트랜스포머가 로그를 쓰도록 한 S3 버킷을 지정하여 훈련 및 검증 지표를 확인할 수 있습니다.

모델 평가

모델 훈련이 완료되면 아래 코드 예시로 체계적인 평가를 실행하거나 간단한 응답 결과를 생성합니다.

tokens_for_summary = 30
output_tokens = input_ids.shape[1] + tokens_for_summary

outputs = model.generate(inputs=input_ids, do_sample=True, max_length=output_tokens)
gen_text = tokenizer.batch_decode(outputs)[0]
print(gen_text)
# Sample output:
# Summarize the chat dialogue:
# Richie: Pogba
# Clay: Pogboom
# Richie: what a s strike yoh!
# Clay: was off the seat the moment he chopped the ball back to his right foot
# Richie: me too dude
# Clay: hope his form lasts
# Richie: This season he's more mature
# Clay: Yeah, Jose has his trust in him
# Richie: everyone does
# Clay: yeah, he really deserved to score after his first 60 minutes
# Richie: reward
# Clay: yeah man
# Richie: cool then
# Clay: cool
# ---
# Summary:
# Richie and Clay have discussed the goal scored by Paul Pogba. His form this season has improved and both of them hope this will last long

모델의 성능에 만족하면 모델을 저장합니다.

trainer.save_model("path_to_save")

이 모델을 전용 SageMaker 엔드포인트에서 호스팅하도록 선택할 수도 있습니다. 자세한 내용은 Falcon-40B 모델을 대규모 모델 추론 딥러닝 컨테이너(DLC)로 Amazon SageMaker에 배포하기 게시물을 참조하기를 바랍니다.

리소스 정리

리소스를 정리하려면 다음 단계를 완료하세요.

  1. 추가 비용이 발생하지 않도록 SageMaker Studio 인스턴스를 종료합니다.
  2. 텐서보드 애플리케이션을 종료합니다.
  3. 아래 커맨드로 허깅페이스 캐시 디렉터리를 삭제하여 EFS 디렉터리를 정리합니다:
rm -R ~/.cache/huggingface/hub

결론

SageMaker 노트북을 사용하면 대화형 환경에서 빠르고 효율적인 방식으로 LLM을 파인튜닝 할 수 있습니다. 이 포스팅에서는 bitsandbtyes와 허깅페이스 PEFT를 사용하여 SageMaker Studio 노트북에서 QLoRA를 사용하여 Falcon-40B 모델을 파인튜닝 하는 방법을 보여드렸습니다.

또한 SageMaker JumpStart, Amazon Titan 모델 및 Amazon Bedrock을 살펴보고 Amazon의 생성형 AI 기능에 대해 자세히 알아보길 권장합니다.

Daekeun Kim

Daekeun Kim

김대근 AI/ML 전문 솔루션즈 아키텍트는 다년간 스타트업, 제조 및 금융 업계를 거치며 컴퓨터 비전 엔지니어로서 다수의 1저자 특허를 등록하고 제품 양산에 기여했으며, 데이터 과학자로서 다양한 PoC와 현업 프로젝트를 수행했습니다. 현재는 고객들이 AWS 인프라 상에서 AI/ML 서비스를 더욱 효율적으로 사용할 수 있도록 기술적인 도움을 드리면서 AI/ML 생태계 확장에 기여하고 있습니다. 머신러닝을 공부하기 시작했을 때 접한 톰 미첼(Tom M. Mitchell)의 명언, “머신러닝으로 문제를 해결하려면 해결하고자 하는 문제를 명확히 정의해야 한다”라는 말을 상기하며 항상 초심을 잃지 않으려 합니다.