Low-Rank Adaptation of Large Language Models (LoRA)

LoRA(Low-Rank Adaptation of Large Language Models)는 메모리를 적게 사용하면서 대규모 모델의 학습을 가속화하는 학습 방법입니다. 이는 rank-decomposition weight 행렬 쌍(업데이트 행렬이라고 함)을 추가하고 새로 추가된 가중치만 학습합니다. 여기에는 몇 가지 장점이 있습니다.

• 이전에 미리 학습된 가중치는 고정된 상태로 유지되므로 모델이 치명적인 망각 경향이 없습니다.
• Rank-decomposition 행렬은 원래 모델보다 파라메터 수가 훨씬 적으므로 학습된 LoRA 가중치를 쉽게 끼워넣을 수 있습니다.
• LoRA 매트릭스는 일반적으로 원본 모델의 어텐션 레이어에 추가됩니다. 🧨 Diffusers는 load_attn_procs() 메서드를 제공하여 LoRA 가중치를 모델의 어텐션 레이어로 불러옵니다. scale 매개변수를 통해 모델이 새로운 학습 이미지에 맞게 조정되는 범위를 제어할 수 있습니다.
• 메모리 효율성이 향상되어 Tesla T4, RTX 3080 또는 RTX 2080 Ti와 같은 소비자용 GPU에서 파인튜닝을 실행할 수 있습니다! T4와 같은 GPU는 무료이며 Kaggle 또는 Google Colab 노트북에서 쉽게 액세스할 수 있습니다.

cloneofsimo는 인기 있는 lora GitHub 리포지토리에서 Stable Diffusion을 위한 LoRA 학습을 최초로 시도했습니다. 🧨 Diffusers는 text-to-image 생성 및 DreamBooth을 지원합니다. 이 가이드는 두 가지를 모두 수행하는 방법을 보여줍니다.

모델을 저장하거나 커뮤니티와 공유하려면 Hugging Face 계정에 로그인하세요(아직 계정이 없는 경우 생성하세요):

huggingface-cli login

Text-to-image

수십억 개의 파라메터들이 있는 Stable Diffusion과 같은 모델을 파인튜닝하는 것은 느리고 어려울 수 있습니다. LoRA를 사용하면 diffusion 모델을 파인튜닝하는 것이 훨씬 쉽고 빠릅니다. 8비트 옵티마이저와 같은 트릭에 의존하지 않고도 11GB의 GPU RAM으로 하드웨어에서 실행할 수 있습니다.

학습

Pokémon BLIP 캡션 데이터셋으로 stable-diffusion-v1-5를 파인튜닝해 나만의 포켓몬을 생성해 보겠습니다.

시작하려면 MODEL_NAME 및 DATASET_NAME 환경 변수가 설정되어 있는지 확인하십시오. OUTPUT_DIR 및 HUB_MODEL_ID 변수는 선택 사항이며 허브에서 모델을 저장할 위치를 지정합니다.

export MODEL_NAME="runwayml/stable-diffusion-v1-5"
export OUTPUT_DIR="/sddata/finetune/lora/pokemon"
export HUB_MODEL_ID="pokemon-lora"
export DATASET_NAME="lambdalabs/pokemon-blip-captions"

학습을 시작하기 전에 알아야 할 몇 가지 플래그가 있습니다.

• --push_to_hub를 명시하면 학습된 LoRA 임베딩을 허브에 저장합니다.
• --report_to=wandb는 학습 결과를 가중치 및 편향 대시보드에 보고하고 기록합니다(예를 들어, 이 보고서를 참조하세요).
• --learning_rate=1e-04, 일반적으로 LoRA에서 사용하는 것보다 더 높은 학습률을 사용할 수 있습니다.

이제 학습을 시작할 준비가 되었습니다 (전체 학습 스크립트는 여기에서 찾을 수 있습니다).

accelerate launch train_dreambooth_lora.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --instance_prompt="a photo of sks dog" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --checkpointing_steps=100 \
  --learning_rate=1e-4 \
  --report_to="wandb" \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=500 \
  --validation_prompt="A photo of sks dog in a bucket" \
  --validation_epochs=50 \
  --seed="0" \
  --push_to_hub

추론

이제 StableDiffusionPipeline에서 기본 모델을 불러와 추론을 위해 모델을 사용할 수 있습니다:

>>> import torch
>>> from diffusers import StableDiffusionPipeline

>>> model_base = "runwayml/stable-diffusion-v1-5"

>>> pipe = StableDiffusionPipeline.from_pretrained(model_base, torch_dtype=torch.float16)

기본 모델의 가중치 위에 파인튜닝된 DreamBooth 모델에서 LoRA 가중치를 불러온 다음, 더 빠른 추론을 위해 파이프라인을 GPU로 이동합니다. LoRA 가중치를 프리징된 사전 훈련된 모델 가중치와 병합할 때, 선택적으로 ‘scale’ 매개변수로 어느 정도의 가중치를 병합할 지 조절할 수 있습니다:

>>> pipe.unet.load_attn_procs(model_path)
>>> pipe.to("cuda")
# LoRA 파인튜닝된 모델의 가중치 절반과 기본 모델의 가중치 절반 사용

>>> image = pipe(
...     "A picture of a sks dog in a bucket.",
...     num_inference_steps=25,
...     guidance_scale=7.5,
...     cross_attention_kwargs={"scale": 0.5},
... ).images[0]
# 완전히 파인튜닝된 LoRA 모델의 가중치 사용

>>> image = pipe("A picture of a sks dog in a bucket.", num_inference_steps=25, guidance_scale=7.5).images[0]
>>> image.save("bucket-dog.png")