LoRA를 사용하여 Qwen2-0.5B의 분산 미세 조정

이 Notebook은 서버리스 GPU 컴퓨팅에서 매개 변수 효율적인 기술을 사용하여 Qwen2-0.5B 큰 언어 모델을 효율적으로 미세 조정하는 방법을 보여 줍니다. 다음을 배우게 됩니다:

• 모델 품질을 유지하면서 학습 가능한 매개 변수를 최대 99% 줄이기 위해 LoRA(Low-Rank 적응) 를 적용합니다.
• 최적화된 Triton 커널을 사용하여 메모리 효율적인 학습에 Liger 커널 사용
• 감독된 미세 조정을 위해 TRL(변환기 보충 학습) 을 활용합니다.
• 거버넌스 및 배포를 위해 Unity 카탈로그에 미세 조정된 모델 등록

주요 개념:

• LoRA: 기본 모델을 동결하고 작은 어댑터 계층을 학습하여 메모리 요구 사항 및 학습 시간을 크게 줄이는 기술
• 리거 커널: 융합 작업을 통해 메모리 사용량을 최대 80% 줄이는 GPU 최적화 커널
• TRL: 강화 학습 및 감독된 미세 조정을 사용하여 언어 모델을 학습하기 위한 라이브러리
• Databricks 관리형 컴퓨팅: GPU 리소스를 자동으로 확장하는 서버리스 GPU 컴퓨팅

LoRA 및 전체 미세 조정 의사 결정 매트릭스

LoRA(Low-Rank Adaptation) 는 기본 모델을 고정하고 작은 어댑터 레이어만 학습하여 학습 가능한 매개 변수를 최대 99% 감소시킵니다. 이렇게 하면 학습이 더 빠르고 메모리 효율적입니다.

Liger 커널 혜택

리거 커널 은 여러 단계를 단일 커널로 융합하여 메모리 전송을 줄이고 효율성을 향상시키는 GPU 최적화 작업입니다. 기술 문서는 상당한 성능 향상을 보여주는 자세한 벤치마크를 제공합니다.

• 융합 작업: 작업(예: 선형 + 손실)을 결합하여 메모리 오버헤드를 최대 80%까지 줄일 수 있습니다.
• Triton 커널: 변환기 작업에 최적화된 사용자 지정 GPU 커널(RMSNorm, RoPE, SwiGLU, CrossEntropy)
• 메모리 효율성: GPU 메모리에 맞지 않는 더 큰 일괄 처리 크기 또는 모델을 허용합니다.
• 단일 GPU 최적화: A10/A100 단일 GPU 학습 시나리오에 특히 효과적

이 Notebook은 TRL 라이브러리 를 사용하여 학습 구성을 간소화하고 이러한 최적화를 자동으로 적용합니다.

서버리스 GPU 컴퓨팅에 연결

이 Notebook에는 서버리스 GPU 컴퓨팅이 필요합니다. 연결하려면 다음을 수행합니다.

1. 오른쪽 상단의 노트북 컴퓨팅 옵션을 클릭하고 서버리스 GPU를 선택합니다.
2. 오른쪽에서 환경 단추를 클릭합니다.
3. 액셀러레이터로 8xH100 선택
4. 기본 환경에서 AI v4 선택
5. 적용을 클릭합니다.

학습 함수는 분산 학습을 위해 8개의 H100 GPU를 자동으로 프로비전합니다.

필수 라이브러리 설치

다음 셀은 분산 미세 조정에 필요한 Python 패키지를 설치합니다.

핵심 학습 라이브러리:

• trl: 감독된 미세 조정과 RLHF를 위한 변환기 강화 학습 라이브러리
• peft: LoRA 구현을 제공하는 Parameter-Efficient Fine-Tuning 라이브러리
• liger-kernel: 효율적인 변환기 학습을 위한 메모리 최적화 GPU 커널

지원 라이브러리:

• hf_transfer: Rust 기반 전송을 활용한 허깅 페이스 허브의 가속 다운로드
• mlflow>=3.0: 실험 추적 및 모델 레지스트리 통합

%restart_python 명령은 새로 설치된 패키지가 제대로 로드되도록 Python 인터프리터를 다시 시작합니다.

%pip install --upgrade peft==0.17.1
%pip install --upgrade hf_transfer==0.1.9
%pip install --upgrade transformers==4.56.1
%pip install trl==0.18.1
%pip install liger-kernel
%pip install mlflow==3.7.0
%restart_python

구성 설정

Unity 카탈로그 통합

다음 셀은 미세 조정된 모델을 저장하고 등록할 위치를 구성합니다.

• 카탈로그 스키마: Unity 카탈로그 네임스페이스 내에서 모델 구성(기본값: main.default)
• 모델 이름: 거버넌스 및 배포를 위해 Unity 카탈로그에 등록된 모델 이름
• 볼륨: 학습 중에 모델 검사점을 저장하기 위한 Unity 카탈로그 볼륨

이러한 위젯을 사용하면 코드를 편집하지 않고도 스토리지 위치를 사용자 지정할 수 있습니다. 간편한 {catalog}.{schema}.{model_name} 액세스 및 버전 제어를 위해 모델이 등록됩니다.

하이퍼 매개 변수 학습

또한 셀은 주요 학습 매개 변수를 정의합니다.

• 모델 및 데이터 세트: Capybara 대화형 데이터 세트가 있는 Qwen2-0.5B
• 일괄 처리 크기(8): 학습 단계당 GPU당 예제 수
• 그라데이션 누적(4): 유효 일괄 처리 크기 32를 위해 4개 일괄 처리보다 그라데이션을 누적합니다.
• 학습률(1e-4): LoRA 학습에 대해 자동으로 10배 더 높게 조정된 보수적 속도
• Epoch (1): 과잉 맞춤을 방지하기 위해 데이터 세트를 통과하는 단일 패스
• 로깅 및 검사점: 100단계마다 진행률을 저장하고 25단계마다 메트릭을 기록합니다.

dbutils.widgets.text("uc_catalog", "main")
dbutils.widgets.text("uc_schema", "default")
dbutils.widgets.text("uc_model_name", "qwen2_liger_lora_assistant")
dbutils.widgets.text("uc_volume", "checkpoints")

UC_CATALOG = dbutils.widgets.get("uc_catalog")
UC_SCHEMA = dbutils.widgets.get("uc_schema")
UC_MODEL_NAME = dbutils.widgets.get("uc_model_name")
UC_VOLUME = dbutils.widgets.get("uc_volume")

print(f"UC_CATALOG: {UC_CATALOG}")
print(f"UC_SCHEMA: {UC_SCHEMA}")
print(f"UC_MODEL_NAME: {UC_MODEL_NAME}")
print(f"UC_VOLUME: {UC_VOLUME}")

# MLflow and Unity Catalog configuration

# Model selection - Choose based on your compute constraints
MODEL_NAME = "Qwen/Qwen2-0.5B"
DATASET_NAME = "trl-lib/Capybara"
OUTPUT_DIR = f"/Volumes/{UC_CATALOG}/{UC_SCHEMA}/{UC_VOLUME}/qwen2-0.5b-lora"

# Training hyperparameters
BATCH_SIZE = 8
GRADIENT_ACCUMULATION_STEPS = 4
LEARNING_RATE = 1e-4
NUM_EPOCHS = 1
EVAL_STEPS = 100
LOGGING_STEPS = 25
SAVE_STEPS = 100

LoRA 구성

다음 셀은 모델을 미세 조정하는 방법을 제어하는 LoRA(Low-Rank 적응) 매개 변수를 구성합니다. LoRA는 기본 모델 가중치를 동결하고 작은 어댑터 매트릭스만 학습하여 메모리 요구 사항을 크게 줄입니다.

매개 변수 선택

• 순위(r=8): 성능과 매개 변수의 적절한 균형을 제공합니다.
• 알파 (32): 배율 요인, 일반적으로 2-4 배 등급
• 드롭아웃(0.1): 과잉 맞춤을 방지하기 위한 정규화

Qwen2의 대상 모듈

이 예제에서는 모든 키 변환 계층을 대상으로 합니다.

• 주의: q_proj, k_proj, v_projo_proj
• MLP: gate_proj, , up_projdown_proj

LORA_R = 8
LORA_ALPHA = 32
LORA_DROPOUT = 0.1
LORA_TARGET_MODULES = [
    "q_proj", "k_proj", "v_proj", "o_proj",
    "gate_proj", "up_proj", "down_proj"
]

학습 함수 정의

다음 셀은 여러 GPU에서 실행되는 분산 학습 함수를 만듭니다. 다음과 같은 작업을 수행합니다.

분산 학습 설정

데코레이터는 @distributed 서버리스 GPU 컴퓨팅을 구성합니다.

• 8개의 GPU: 더 빠른 학습을 위해 8개의 H100 GPU에 학습 배포
• 자동 오케스트레이션: GPU 프로비저닝, 데이터 배포 및 동기화 처리

학습 워크플로

함수는 다음 단계를 실행합니다.

1. 데이터 세트 로드: Capybara 대화형 데이터 세트 다운로드 및 준비
2. 모델 초기화: 채팅 서식을 사용하여 Qwen2-0.5B 및 토큰화기 로드
3. LoRA 적용: 어댑터 레이어를 연결하여 학습 가능한 매개 변수를 최대 99% 줄입니다.
4. 학습 구성: 일괄 처리 크기, 학습 속도 및 Liger 커널 최적화 설정
5. 모델 학습: 자동 체크포인트 저장 및 로깅을 사용하여 학습 루프를 실행합니다.
6. 아티팩트 저장: LoRA 어댑터 및 토케나이저를 Unity 카탈로그 볼륨에 저장
7. MLflow 실행 ID 반환: 모델 등록에 대한 실행 ID를 제공합니다.

주요 최적화가 활성화됨

• 리거 커널: 결합된 GPU 연산으로 메모리 사용량을 최대 80%까지 줄일 수 있습니다
• 혼합 정밀도(FP16): 메모리 사용량이 적은 더 빠른 계산
• Gradient checkpointing: 보다 큰 배치를 처리하기 위해 메모리와 연산을 교환합니다.
• 그라데이션 누적: 안정적인 학습을 위해 더 큰 일괄 처리 크기를 시뮬레이션합니다.

from serverless_gpu import distributed
from serverless_gpu import runtime as rt

@distributed(gpus=8, gpu_type="H100")
def run_train(use_lora=True):
    import logging
    from datasets import load_dataset
    from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
    from peft import LoraConfig, TaskType, get_peft_model
    from trl import (
        SFTConfig,
        SFTTrainer,
        setup_chat_format
    )
    import json
    import os
    import mlflow

    dataset = load_dataset(DATASET_NAME)
    logging.info(f"✓ Dataset loaded: {dataset}")

    if "test" not in dataset:
        logging.info("Creating validation split from training data...")
        dataset = dataset["train"].train_test_split(test_size=0.1, seed=42)
        logging.info("✓ Data split: 90% train, 10% validation")

    # model and tokenizer initialization
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        trust_remote_code=True,
    )

    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_NAME,
        trust_remote_code=True,
        use_fast=True
    )

    # Chat template formatting for conversational fine-tuning
    if tokenizer.chat_template is None:
        logging.info("Adding chat template for proper conversation formatting...")
        model, tokenizer = setup_chat_format(model, tokenizer, format="chatml")
        logging.info("✓ ChatML format applied for structured conversations")

    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        logging.info("✓ Padding token set to EOS token")

    logging.info("✓ Model and tokenizer loaded successfully")

    # PEFT
    peft_config = None
    if use_lora:
        try:
            logging.info("Configuring LoRA for parameter-efficient fine-tuning...")

            peft_config = LoraConfig(
                task_type=TaskType.CAUSAL_LM,
                inference_mode=False,
                r=LORA_R,
                lora_alpha=LORA_ALPHA,
                lora_dropout=LORA_DROPOUT,
                target_modules=LORA_TARGET_MODULES,
                bias="none",
                use_rslora=False,
                modules_to_save=None,
            )

            logging.info(f"LoRA configuration: rank={LORA_R}, alpha={LORA_ALPHA}, dropout={LORA_DROPOUT}")
            logging.info(f"Target modules: {', '.join(LORA_TARGET_MODULES)}")

            original_params = model.num_parameters()
            model = get_peft_model(model, peft_config)

            trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
            total_params = sum(p.numel() for p in model.parameters())
            efficiency_ratio = 100 * trainable_params / total_params

            logging.info(f"✓ LoRA applied successfully:")
            logging.info(f"  • Original parameters: {original_params:,}")
            logging.info(f"  • Trainable parameters: {trainable_params:,}")
            logging.info(f"  • Training efficiency: {efficiency_ratio:.2f}% of parameters")
            logging.info(f"  • Memory savings: ~{100-efficiency_ratio:.1f}% reduction in gradient memory")

        except Exception as e:
            logging.info(f"✗ LoRA configuration failed: {e}")
            logging.info("Falling back to full fine-tuning...")
            peft_config = None
    else:
        logging.info("Full fine-tuning mode selected")
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        logging.info(f"Trainable parameters: {trainable_params:,} (100% of model)")

    # Learning rate adjustment for LoRA
    adjusted_lr = LEARNING_RATE * 10 if use_lora else LEARNING_RATE
    logging.info(f"Learning rate: {adjusted_lr} ({'LoRA-adjusted' if use_lora else 'standard'})")

    training_args_dict = {
        "output_dir": OUTPUT_DIR,
        "per_device_train_batch_size": BATCH_SIZE,
        "per_device_eval_batch_size": BATCH_SIZE,
        "gradient_accumulation_steps": GRADIENT_ACCUMULATION_STEPS,
        "learning_rate": adjusted_lr,
        "num_train_epochs": NUM_EPOCHS,
        "eval_steps": EVAL_STEPS,
        "logging_steps": LOGGING_STEPS,
        "save_steps": SAVE_STEPS,
        "save_total_limit": 2,
        "report_to": "mlflow",
        "run_name": f"{MODEL_NAME}_fine-tuning",
        "warmup_steps": 50,
        "weight_decay": 0.01,
        "metric_for_best_model": "eval_loss",
        "greater_is_better": False,
        "dataloader_pin_memory": False,
        "remove_unused_columns": False,
        "use_liger_kernel": True,  # Enable Liger kernel optimizations
        "fp16": True,  # Mixed precision training
        "gradient_checkpointing": True,
        "gradient_checkpointing_kwargs": {"use_reentrant": False}, # Required for LORA with DDP
    }

    logging.info("✓ Liger kernel optimizations enabled")

    training_args = SFTConfig(**training_args_dict)

    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset["train"],
        eval_dataset=dataset["test"],
        processing_class=tokenizer,
        peft_config=peft_config,
    )

    logging.info("\n" + "="*50)
    logging.info("STARTING TRAINING")
    logging.info("="*50)

    logging.info("🚀 Training with Liger kernels for memory-efficient single GPU training")
    if use_lora:
        logging.info("🎯 Using LoRA for parameter-efficient fine-tuning")

    trainer.train()
    logging.info("\n✓ Training completed successfully!")
    if rt.get_global_rank() == 0:
        logging.info("\nSaving trained model...")

        logging.info("Saving LoRA adapter weights...")
        trainer.save_model(training_args.output_dir)
        logging.info("✓ LoRA adapters saved - use with base model for inference")
        tokenizer.save_pretrained(training_args.output_dir)
        logging.info("✓ Tokenizer saved with model")
        logging.info(f"\n🎉 All artifacts saved to: {training_args.output_dir}")

    mlflow_run_id = None
    if mlflow.last_active_run() is not None:
        mlflow_run_id = mlflow.last_active_run().info.run_id

    return mlflow_run_id

분산 학습 실행

이 셀은 8개의 H100 GPU에서 학습 함수를 실행합니다. 메서드 distributed()을(를) 처리합니다.

• 서버리스 GPU 컴퓨팅 리소스 프로비전
• 여러 GPU에 학습 워크로드 배포
• 모델 등록을 위한 MLflow 실행 ID 수집

학습은 일반적으로 데이터 세트 크기 및 컴퓨팅 가용성에 따라 15-30분이 걸립니다.

mlflow_run_id = run_train.distributed(use_lora=True)[0]
print(mlflow_run_id)

MLflow 및 Unity 카탈로그 등록

모델 등록 전략

• MLflow 추적: 로그 모델 아티팩트 및 메타데이터
• Unity 카탈로그: 거버넌스 및 배포를 위한 모델 등록
• 모델 버전 관리: 모델 수명 주기 관리를 위한 자동 버전 관리
• 메타데이터: 재현성을 위한 전체 모델 정보

print("\nRegistering model with MLflow and Unity Catalog...")

from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
import mlflow
from mlflow import transformers as mlflow_transformers

try:
    # Load the trained model for registration
    print("Loading LoRA model for registration...")
    # For LoRA models, we need both base model and adapter
    base_model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        trust_remote_code=True
    )
    # Load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    adapter_dir = OUTPUT_DIR
    peft_model = PeftModel.from_pretrained(base_model, adapter_dir)
    # Merge LoRA into base and drop PEFT wrappers
    merged_model = peft_model.merge_and_unload()

    components = {
        "model": merged_model,
        "tokenizer": tokenizer,
    }

    # Create Unity Catalog model name
    full_model_name = f"{UC_CATALOG}.{UC_SCHEMA}.{UC_MODEL_NAME}"

    print(f"Registering model as: {full_model_name}")

    # Start MLflow run and log model
    task = "llm/v1/chat"
    with mlflow.start_run(run_id=mlflow_run_id):
        model_info = mlflow.transformers.log_model(
            transformers_model=components,
            artifact_path="model",
            task=task,
            registered_model_name=full_model_name,
            metadata={
                "task": task,
                "pretrained_model_name": MODEL_NAME,
                "databricks_model_family": "QwenForCausalLM",
            },
        )

    print(f"✓ Model successfully registered in Unity Catalog: {full_model_name}")
    print(f"✓ MLflow model URI: {model_info.model_uri}")

    # Print deployment information
    print(f"\n📦 Model Registration Complete!")
    print(f"Unity Catalog Path: {full_model_name}")
    print(f"Model Type: {model_type}")
    print(f"Optimization: Liger Kernels + LoRA")

except Exception as e:
    print(f"✗ Model registration failed: {e}")
    print("Model is still saved locally and can be registered manually")
    print(f"Local model path: {OUTPUT_DIR}")

다음 단계

이제 모델을 미세 조정하고 등록했으므로 다음을 수행할 수 있습니다.

• 모델 배포: 모델 서비스를 사용하여 모델 제공
• 분산 학습에 대한 자세한 정보: 다중 GPU 분산 학습
• 서버리스 GPU 컴퓨팅 사용 최적화: 서버리스 GPU 컴퓨팅 모범 사례
• 문제 해결: 서버리스 GPU 컴퓨팅의 문제 해결

예제 노트

LoRA를 사용하여 Qwen2-0.5B의 분산 미세 조정

노트북 받기