레이 튠을 이용한 하이퍼파라미터 탐색

Important

이 기능은 공개 미리보기 단계에 있습니다.

이 예시는 Ray Tune 을 사용하여 4개의 1xA10 노드에서 Qwen2.5의 LoRA 미세 조정 하이퍼파라미터를 검색합니다. 부트스트랩 명령이 노드를 아우르는 Ray 클러스터를 시작하고, 드라이버는 Ray Tune에 각 trial마다 GPU 1개를 요청합니다. 클러스터는 한 번에 4개의 트라이얼을 실행하고, 나머지는 GPU가 무료가 되면 시작됩니다.

검색은 ASHA 스케줄러(Asynchronous Successive Halving)를 사용합니다. 각 시험은 eval_loss 일정한 단계 간격으로 진행되며, ASHA는 모든 후보자를 완수할 때까지 훈련시키는 대신 뒤처진 시험을 중단합니다.

이 예시는 공개 모델(Qwen2.5-0.5B)을 사용하므로 Hugging Face 토큰 없이 as-is 실행됩니다.

워크로드는 다음을 수행합니다.

  • code_source: snapshot를 사용하여 로컬 프로젝트를 업로드합니다.
  • 데이터셋을 드라이버에서 한 번만 토큰화한 후 텐서 형태로 각 실행에 전달합니다.
  • 8개의 LoRA 구성을 샘플링하고 한 번에 4개씩 돌립니다.
  • 스윕 설정, 최적의 구성, 그리고 시행당 손실을 MLflow에 기록합니다.

사전 요구 사항

프로젝트 배치

다음 파일을 사용하여 디렉터리를 만듭니다.

ray_tune_lora/
├── tune.yaml           # air workload config (inline dependencies + Ray bootstrap)
└── tune_lora.py        # Ray Tune driver + per-trial LoRA fine-tuning

1단계: 워크로드 YAML 작성

tune.yaml는 4개의 GPU_1xA10 노드를 요청하고 version 아래에 종속성을 인라인으로 선언합니다(런타임 environment 포함). 워크로드는 command 노드 간 Ray 클러스터를 시작한 후 드라이버를 실행하므로 예시에는 별도의 의존성 파일이나 런처 스크립트가 필요하지 않습니다:

experiment_name: air-ray-tune-lora

environment:
  version: 'databricks_ai_v5'
  dependencies:
    # databricks_ai_v5 ships ray, transformers, and datasets. It does not ship peft
    # and needs a newer fsspec for huggingface_hub.
    - peft>=0.13
    - fsspec>=2024.6.1

# 4 1xA10 nodes. Ray Tune runs one trial per GPU.
compute:
  num_accelerators: 4
  accelerator_type: GPU_1xA10

code_source:
  type: snapshot
  snapshot:
    root_path: .

command: |
  cd $CODE_SOURCE_PATH
  set -e
  RAY_HEAD_PORT=6379
  GPUS_PER_NODE=${LOCAL_WORLD_SIZE:-1}

  if [ "${NODE_RANK:-0}" = "0" ]; then
    echo "NODE_RANK=0: starting Ray head with $GPUS_PER_NODE GPU(s)..."
    ray start --head --port=$RAY_HEAD_PORT --num-gpus="$GPUS_PER_NODE" --dashboard-host=0.0.0.0
    # Stop the cluster on exit, even if the driver fails, so workers don't wait out the timeout.
    trap 'ray stop --grace-period 5' EXIT
    python tune_lora.py
  else
    echo "NODE_RANK=$NODE_RANK: connecting to Ray head at $MASTER_ADDR:$RAY_HEAD_PORT..."
    # `ray start` returns as soon as this node joins, so the worker controls its own exit.
    joined=""
    for i in $(seq 1 12); do
      if ray start --address="$MASTER_ADDR:$RAY_HEAD_PORT" --num-gpus="$GPUS_PER_NODE" 2>/dev/null; then
        joined=1
        break
      fi
      echo "Attempt $i failed, retrying in 5s..."
      sleep 5
    done
    if [ -z "$joined" ]; then
      echo "Worker failed to join the Ray head after all retries; aborting." >&2
      exit 1
    fi

    # health-check exits non-zero once the head runs `ray stop`, which is this worker's cue
    # to exit. The timeout keeps each probe short so the job finishes promptly; the counter
    # caps the total wait.
    echo "Worker joined; waiting for the head to finish the sweep..."
    for _ in $(seq 1 360); do
      if ! timeout 5 ray health-check --address "$MASTER_ADDR:$RAY_HEAD_PORT" 2>/dev/null; then
        break
      fi
      sleep 5
    done
    echo "Head is no longer healthy; stopping local Ray and exiting."
    ray stop --grace-period 5
  fi

max_retries: 0
timeout_minutes: 45

env_variables:
  NCCL_SOCKET_IFNAME: eth0
  HF_HOME: /tmp/hf

2단계: 탐색 공간과 스케줄러 정의

드라이버 함수는 main 데이터를 한 번 토큰화하고, 검색 공간을 정의한 후 ASHA를 구성합니다:

tuner = tune.Tuner(
    # with_resources gives each trial a whole GPU so trials never share a device.
    tune.with_resources(
        tune.with_parameters(train_fn, train_data=train_data, eval_data=eval_data),
        resources={"gpu": 1},
    ),
    param_space={
        "lr": tune.loguniform(1e-5, 1e-3),
        "lora_r": tune.choice([8, 16, 32]),
        "lora_alpha_ratio": tune.choice([1, 2]),
        "lora_dropout": tune.uniform(0.0, 0.1),
        "weight_decay": tune.choice([0.0, 0.01]),
        "batch_size": tune.choice([4, 8]),
    },
    tune_config=tune.TuneConfig(
        metric="eval_loss",
        mode="min",
        scheduler=ASHAScheduler(
            max_t=MAX_ITERATIONS, grace_period=GRACE_PERIOD, reduction_factor=2
        ),
        num_samples=NUM_SAMPLES,
    ),
)
results = tuner.fit()

tune.with_resources(..., resources={"gpu": 1}) 수색 결과를 클러스터에 매핑합니다. Ray Tune은 클러스터에 GPU가 4개 있으므로 trial 4개를 동시에 실행합니다. 따라서 스윕 범위를 넓히려면 코드를 변경하는 대신 YAML에서 num_accelerators 값을 늘리세요.

각 시험은 모든 EVAL_STEPS 최적화 단계를 보고합니다. grace_period는 시도가 중단되기 전에 받는 보고 횟수를 정하고, max_t는 살아남은 시도가 받을 수 있는 보고 횟수의 상한을 두며, reduction_factor=2는 각 단계에서 대략 하위 절반을 중단시킵니다.

3단계: 각 시험에서 가지치기 지표를 보고하세요

train_fn는 하나의 시도입니다. tune.report 이 통화에서 ASHA는 시험을 중단할지 계속할지 결정합니다:

def train_fn(config, train_data=None, eval_data=None):
    # Ray Tune pins one GPU per trial via CUDA_VISIBLE_DEVICES, so cuda:0 is this trial's.
    device = torch.device("cuda")

    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, dtype=torch.bfloat16)
    model.config.use_cache = False
    lora = LoraConfig(
        r=config["lora_r"],
        lora_alpha=config["lora_r"] * config["lora_alpha_ratio"],
        lora_dropout=config["lora_dropout"],
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
        task_type="CAUSAL_LM",
    )
    model = get_peft_model(model, lora).to(device)
    ...
    if step % EVAL_STEPS == 0:
        tune.report({
            "eval_loss": evaluate(model, eval_loader, device),
            "train_loss": out.loss.item(),
            "step": step,
        })

ASHA는 훈련 손실이 아니라 홀드아웃 분할의 eval_loss 에서 시도들을 비교하는데, 훈련 손실을 기준으로 하면 가장 빨리 과적합되는 구성에 유리해지기 때문이다. build_datasets 데이터를 드라이버에서 한 번 토큰화하고 객체를 반환 TensorDataset 합니다. tune.with_parameters 이를 다른 노드에서의 시도로 전송합니다. 텐서는 값 자체로 직렬화되는 반면, Hugging Face 데이터셋은 다른 노드에서 열 수 없는 메모리 매핑된 파일의 경로로 전달됩니다.

전체 스크립트는 이 페이지 끝의 풀 튜닝 스크립트 로 나와 있습니다.

4단계: 런 제출하기

air run -f tune.yaml --dry-run
air run -f tune.yaml --watch

5단계: 실행 확인

air get run <run-id>
air logs <run-id>

드라이버는 노드 0에서 실행되므로 Ray Tune 상태 테이블은 해당 노드의 로그에서 스트리밍되며, 각 시도마다 샘플링된 구성, 반복 횟수, 그리고 최신 eval_loss.이 표시됩니다. ASHA가 중단한 시도는 max_t보다 반복 횟수가 적은 것으로 표시됩니다TERMINATED.

결과가 표시되는 위치

실행이 끝나면 드라이버는 최적의 구성과 그 eval_loss를 출력하고, 이 둘과 함께 스윕 설정 및 각 시도의 최종 experiment_name를 eval_loss에 지정된 MLflow 실험에 기록합니다.

운전자는 어떤 시도든 실패하면 오류를 제기합니다.

예시는 어댑터 가중치를 유지하지 않습니다. 최상의 어댑터를 유지하려면 모든 노드가 접근할 수 있는 Unity 카탈로그 볼륨에 어댑터 RunConfig(storage_path=...) 를 설정 tune.Tuner 하세요.

스윕 크기를 조정하세요

상단 tune_lora.py 의 상수는 스윕의 크기를 제어합니다. 몇 분 안에 변경 사항을 빠르게 검증하려면 그것들을 더 작게 설정하되, 그러면 eval_loss 수치의 변동성이 너무 커서 구성의 순위를 매기기 어렵습니다.

실제 경과 시간은 NUM_SAMPLES / num_accelerators에 따라 달라지므로, 스윕이 너무 오래 걸리면 탐색 범위를 줄이기보다는 num_accelerators를 높이세요. 더 큰 모델이라면 더 큰 GPU로 업그레이드 accelerator_type 하세요. 무작위로 샘플링하는 대신 구성을 선택하려면 Optuna와 같은 것을 search_alg 전달 TuneConfig 합니다.

풀 튜닝 스크립트

복사-붙여넣기를 위한 전체 tune_lora.py:

#!/usr/bin/env python3
"""LoRA hyperparameter search for Qwen2.5-0.5B with Ray Tune + ASHA on 4 1xA10 nodes.

The workload's `command` starts a Ray head on node 0 and joins the other nodes as workers,
then runs this script on the head. Ray Tune requests one GPU per trial, so every node runs
one trial at a time. ASHA concentrates GPU time on the promising configurations by stopping
trials that fall behind at each rung.

Uses a public model (no Hugging Face token required) so the example runs as-is.
"""

import os

import mlflow
import ray
import torch
from datasets import load_dataset
from peft import LoraConfig, get_peft_model
from ray import tune
from ray.tune.schedulers import ASHAScheduler
from torch.utils.data import DataLoader, TensorDataset
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_NAME = "Qwen/Qwen2.5-0.5B"
DATASET_NAME = "tatsu-lab/alpaca"
MAX_SEQ_LEN = 512

# Trials report every EVAL_STEPS optimizer steps, so ASHA sees at most MAX_ITERATIONS
# reports per trial and can start pruning once a trial has sent GRACE_PERIOD of them.
EVAL_STEPS = 25
MAX_ITERATIONS = 12
GRACE_PERIOD = 3

NUM_SAMPLES = 8
TRAIN_EXAMPLES = 2000
EVAL_EXAMPLES = 200


def build_datasets():
    """Tokenizes the SFT data once on the driver.

    Returns TensorDatasets so the tokenized splits serialize by value, which is what lets
    tune.with_parameters hand them to trials on any node in the cluster.
    """
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    raw = load_dataset(DATASET_NAME, split=f"train[:{TRAIN_EXAMPLES + EVAL_EXAMPLES}]")

    def format_example(row):
        prompt = f"### Instruction:\n{row['instruction']}\n\n"
        if row.get("input"):
            prompt += f"### Input:\n{row['input']}\n\n"
        text = f"{prompt}### Response:\n{row['output']}{tokenizer.eos_token}"
        out = tokenizer(text, truncation=True, max_length=MAX_SEQ_LEN, padding="max_length")
        # -100 is cross-entropy's ignore_index, so the loss covers only real tokens and
        # eval_loss stays a meaningful signal for ASHA to rank trials by.
        out["labels"] = [token if mask == 1 else -100 for token, mask in zip(out["input_ids"], out["attention_mask"])]
        return out

    tokenized = raw.map(format_example, remove_columns=raw.column_names)
    split = tokenized.train_test_split(test_size=EVAL_EXAMPLES, shuffle=True, seed=0)

    def to_tensors(ds):
        return TensorDataset(
            torch.tensor(ds["input_ids"], dtype=torch.long),
            torch.tensor(ds["attention_mask"], dtype=torch.long),
            torch.tensor(ds["labels"], dtype=torch.long),
        )

    return to_tensors(split["train"]), to_tensors(split["test"])


def evaluate(model, loader, device):
    """Mean cross-entropy over the held-out split. This is the metric ASHA prunes on."""
    model.eval()
    total, batches = 0.0, 0
    with torch.no_grad():
        for input_ids, attention_mask, labels in loader:
            out = model(
                input_ids=input_ids.to(device),
                attention_mask=attention_mask.to(device),
                labels=labels.to(device),
            )
            total += out.loss.item()
            batches += 1
    model.train()
    return total / max(batches, 1)


def train_fn(config, train_data=None, eval_data=None):
    """One trial: LoRA fine-tunes Qwen on a single GPU and reports eval_loss to ASHA."""
    # Ray Tune pins one GPU per trial via CUDA_VISIBLE_DEVICES, so cuda:0 is this trial's.
    device = torch.device("cuda")

    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, dtype=torch.bfloat16)
    model.config.use_cache = False
    lora = LoraConfig(
        r=config["lora_r"],
        lora_alpha=config["lora_r"] * config["lora_alpha_ratio"],
        lora_dropout=config["lora_dropout"],
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
        task_type="CAUSAL_LM",
    )
    model = get_peft_model(model, lora).to(device)

    train_loader = DataLoader(train_data, batch_size=config["batch_size"], shuffle=True, drop_last=True)
    eval_loader = DataLoader(eval_data, batch_size=config["batch_size"])

    optimizer = torch.optim.AdamW(
        (p for p in model.parameters() if p.requires_grad),
        lr=config["lr"],
        weight_decay=config["weight_decay"],
    )

    model.train()
    step = 0
    max_steps = EVAL_STEPS * MAX_ITERATIONS
    # Cycle the loader over multiple epochs until the step budget is spent.
    while step < max_steps:
        for input_ids, attention_mask, labels in train_loader:
            out = model(
                input_ids=input_ids.to(device),
                attention_mask=attention_mask.to(device),
                labels=labels.to(device),
            )
            out.loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            optimizer.zero_grad()
            step += 1

            if step % EVAL_STEPS == 0:
                # ASHA stops or continues the trial based on this report.
                tune.report(
                    {
                        "eval_loss": evaluate(model, eval_loader, device),
                        "train_loss": out.loss.item(),
                        "step": step,
                    }
                )
            if step >= max_steps:
                break


def main():
    ray.init(address="auto")

    num_nodes = int(os.environ.get("NUM_NODES", 1))
    total_gpus = int(ray.cluster_resources().get("GPU", 0))
    if total_gpus < 1:
        raise SystemExit("No GPUs registered with Ray; check GPU discovery on the cluster.")
    print(f"Cluster ready: {num_nodes} node(s), {total_gpus} GPU(s)", flush=True)
    print(f"Running {NUM_SAMPLES} trials, up to {total_gpus} concurrently\n", flush=True)

    train_data, eval_data = build_datasets()

    param_space = {
        "lr": tune.loguniform(1e-5, 1e-3),
        "lora_r": tune.choice([8, 16, 32]),
        "lora_alpha_ratio": tune.choice([1, 2]),
        "lora_dropout": tune.uniform(0.0, 0.1),
        "weight_decay": tune.choice([0.0, 0.01]),
        "batch_size": tune.choice([4, 8]),
    }

    tuner = tune.Tuner(
        # with_resources gives each trial a whole GPU so trials never share a device.
        tune.with_resources(
            tune.with_parameters(train_fn, train_data=train_data, eval_data=eval_data),
            resources={"gpu": 1},
        ),
        param_space=param_space,
        tune_config=tune.TuneConfig(
            metric="eval_loss",
            mode="min",
            scheduler=ASHAScheduler(
                max_t=MAX_ITERATIONS,
                grace_period=GRACE_PERIOD,
                reduction_factor=2,
            ),
            num_samples=NUM_SAMPLES,
        ),
    )

    results = tuner.fit()

    # Surface trial failures: a best result is only meaningful when the whole sweep ran.
    if results.num_errors:
        raise RuntimeError(
            f"{results.num_errors} of {len(results)} trials errored; see the per-trial error files above."
        )

    best = results.get_best_result("eval_loss", "min")
    print(f"\nBest config:    {best.config}", flush=True)
    print(f"Best eval_loss: {best.metrics['eval_loss']:.4f}", flush=True)

    # AI Runtime injects MLFLOW_RUN_ID and configures the databricks tracking URI on the
    # node, so logging needs no credentials here. Gating on the variable keeps the script
    # runnable off-platform, where it is unset.
    if os.environ.get("MLFLOW_RUN_ID"):
        with mlflow.start_run(run_id=os.environ["MLFLOW_RUN_ID"]):
            mlflow.log_params(
                {
                    "model": MODEL_NAME,
                    "dataset": DATASET_NAME,
                    "num_samples": NUM_SAMPLES,
                    "scheduler": "ASHA",
                    "asha_max_t": MAX_ITERATIONS,
                    "asha_grace_period": GRACE_PERIOD,
                    **{f"best_{k}": v for k, v in best.config.items()},
                }
            )
            mlflow.log_metric("best_eval_loss", best.metrics["eval_loss"])
            for i, result in enumerate(results):
                if result.metrics and "eval_loss" in result.metrics:
                    mlflow.log_metric("trial_eval_loss", result.metrics["eval_loss"], step=i)

    ray.shutdown()


if __name__ == "__main__":
    main()

추가 리소스

  • Last updated on