Ray Tune ile hiperparametre arama

Important

Bu özellik Genel Önizleme aşamasındadır.

Bu örnek, 4 adet 1xA10 düğümünde Qwen2.5 için LoRA ince ayar hiperparametrelerini aramak için Ray Tune kullanmaktadır. Bir önyükleme komutu, düğümleri kapsayan bir Ray kümesi başlatır ve sürücü, her deneme için Ray Tune'dan bir GPU ister. Küme aynı anda 4 deneme çalıştırıyor ve geri kalanlar GPU'ların serbest hale gelmesiyle başlıyor.

Arama ASHA zamanlayıcısını (Asenkron Ardışık Yarılama) kullanır. Her deneme, sabit bir adım aralıkla gerçekleştirildiğini eval_loss bildirir ve ASHA, her adayı tamamlamak için eğitim vermek yerine geride kalan denemeleri durdurur.

Örnek, kamuya açık bir model (Qwen2.5-0.5B) kullanıyor, yani Hugging Face belirtekosu olmadan as-is çalışıyor.

İş yükü aşağıdakileri yapar:

  • Yerel projeyi code_source: snapshot ile karşıya yükler.
  • Veri kümesini driver üzerinde bir kez tokenize eder ve denemelere tensörler olarak iletir.
  • 8 LoRA konfigürasyonunu örnekler ve aynı anda 4'ünü çalıştırır.
  • Sweep ayarlarını, en iyi yapılandırmayı ve deneme başına MLflow kayıplarını kaydediyor.

Prerequisites

Proje düzeni

Aşağıdaki dosyaları içeren bir dizin oluşturun.

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

1. Adım: İş yükü YAML'sini yazma

tune.yaml 4 GPU_1xA10 düğüm ister ve bağımlılıklarını environment altında satır içinde bildirir (version çalışma zamanı ile birlikte). İş yükü command düğümler arasında bir Ray kümesi başlatır, ardından sürücü çalıştırılır, bu yüzden örnek ayrı bir bağımlılık dosyası veya başlatıcı betiği gerektirmez:

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

Adım 2: Arama alanını ve zamanlayıcıyı tanımlayın

Sürücünün main fonksiyonu veriyi bir kez tokenize eder, arama alanını tanımlar, ardından ASHA'yı yapılandırır:

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}) Aramayı küme ile haritalandırır. Ray Tune, kümede 4 GPU olduğu için aynı anda 4 denemeyi yürütür; bu nedenle arama alanını genişletmek için kodu değiştirmek yerine YAML'deki num_accelerators değerini artırın.

Her deneme, her EVAL_STEPS optimizasyon adımında rapor verir. grace_period bir denemenin durdurulmadan önce kaç rapor alacağını ayarlar, max_t elenmeyen bir denemenin alacağı rapor sayısını sınırlar ve reduction_factor=2 her basamakta yaklaşık olarak en alttaki yarıyı durdurur.

Adım 3: Her denemedeki budama metriğini raporlayın

train_fn bir denemedir. tune.report çağrısı, ASHA'nın denemeyi durdurduğu veya sürdürdüğü yerdir:

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, en hızlı şekilde aşırı uyum gösteren konfigürasyonları avantajlı hâle getirecek olan eğitim kaybına göre değil, ayrılmış bir veri bölümündeki eval_loss denemelerini karşılaştırır. build_datasets verileri sürücüde bir kez tokenize eder ve nesneleri döndürür TensorDataset . tune.with_parameters onları diğer düğümlerdeki denemelere gönderir. Tensorlar değere göre serileştirirken, Hugging Face veri seti diğer düğümlerin açamadığı bellek haritalı bir dosyaya giden yol olarak gelir.

Betiğin tamamı, bu sayfanın sonunda yer alan Tam ince ayar betiği başlığında listelenmiştir.

Adım 4: Koşuyu gönderin

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

Adım 5: Uçuşu kontrol et

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

Sürücü, 0 numaralı düğümde çalışır; bu nedenle Ray Tune durum tablosu o düğümün günlüklerinden aktarılır ve deneme başına bir satırda örneklenen yapılandırması, yineleme sayısı ve en son eval_loss değerini gösterir. ASHA'nın durdurduğu denemeler, TERMINATED'den daha az yinelemeye sahip max_t olarak görüntülenir.

Sonuçların geldiği yer

Çalıştırma sonunda sürücü, en iyi yapılandırmayı ve onun eval_loss değerini yazdırır ve her ikisini de, tarama ayarları ve her denemenin son eval_loss değeriyle birlikte, adı experiment_name içinde belirtilen MLflow deneyine kaydeder.

Sürücü, herhangi bir deneme başarısız olursa hata yapar.

Örnek, adaptör ağırlıklarını kalıcı olarak saklamaz. En iyi adaptörü saklamak için, her düğümün erişebileceği bir Unity Catalog biriminde tune.Tuner için bir RunConfig(storage_path=...) sağlayın.

Süpürme boyutunu ayarlayın

En üstteki tune_lora.py sabitler süpürmenin boyutunu kontrol eder. Bunları daha küçük ayarlarsanız, bir değişikliği birkaç dakika içinde hızlıca sınayabilirsiniz; ancak bu durumda eval_loss değerleri, yapılandırmaları sıralamak için fazla gürültülü kalır.

Duvar saati zaman izlerini NUM_SAMPLES / num_acceleratorstakip eder, bu yüzden tarama çok uzun sürdüğünde aramayı küçültmek yerine artırın num_accelerators . Daha büyük bir model için daha büyük bir GPU'ya yükseltin accelerator_type . Konfigürasyonları rastgele örneklemek yerine seçmek için, TuneConfig gibi bir search_alg geçin.

Tam ayar komut dosyası

Kopyala-yapıştır için eksiksiz 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()

Ek kaynaklar

  • Last updated on