Skip to main content
Last updated on

Hyperparameter search with Ray Tune

Preview

This feature is in Public Preview.

This example uses Ray Tune to search LoRA fine-tuning hyperparameters for Qwen2.5 across 4 1xA10 nodes. A bootstrap command starts a Ray cluster that spans the nodes, and the driver asks Ray Tune for one GPU per trial. The cluster runs 4 trials at a time, and the rest start as GPUs become free.

The search uses the ASHA scheduler (Asynchronous Successive Halving). Each trial reports held-out eval_loss at a fixed step interval, and ASHA stops the trials that fall behind instead of training every candidate to completion.

The example uses a public model (Qwen2.5-0.5B), so it runs as-is without a Hugging Face token.

The workload does the following:

  • Uploads the local project with code_source: snapshot.
  • Tokenizes the dataset once on the driver and passes it to the trials as tensors.
  • Samples 8 LoRA configurations and runs 4 at a time.
  • Logs the sweep settings, the best configuration, and the per-trial losses to MLflow.

Prerequisites​

Project layout​

Create a directory with the following files.

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

Step 1: Write the workload YAML​

tune.yaml requests 4 GPU_1xA10 nodes and declares its dependencies inline under environment (with the runtime version). The workload's command starts a Ray cluster across the nodes, then runs the driver, so the example needs no separate dependency file or launcher script:

YAML
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

Step 2: Define the search space and the scheduler​

The driver's main function tokenizes the data once, defines the search space, then configures ASHA:

Python
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}) maps the search onto the cluster. Ray Tune keeps 4 trials in flight because the cluster has 4 GPUs, so to widen the sweep, raise num_accelerators in the YAML rather than changing the code.

Each trial reports every EVAL_STEPS optimizer steps. grace_period sets how many reports a trial gets before it can be stopped, max_t caps how many a surviving trial gets, and reduction_factor=2 stops roughly the bottom half at every rung.

Step 3: Report the pruning metric from each trial​

train_fn is one trial. The tune.report call is where ASHA stops or continues the trial:

Python
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 compares trials on eval_loss from a held-out split rather than on training loss, which would favor the configurations that overfit fastest. build_datasets tokenizes the data once on the driver and returns TensorDataset objects. tune.with_parameters ships them to trials on other nodes. Tensors serialize by value, whereas a Hugging Face dataset would arrive as a path to a memory-mapped file that the other nodes cannot open.

The complete script is listed in Full tuning script at the end of this page.

Step 4: Submit the run​

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

Step 5: Inspect the run​

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

The driver runs on node 0, so the Ray Tune status table streams from that node's logs, with one row per trial showing its sampled configuration, iteration count, and latest eval_loss. Trials that ASHA stopped appear as TERMINATED with fewer iterations than max_t.

Where results land​

At the end of the run, the driver prints the best configuration and its eval_loss, and logs both to the MLflow experiment named in experiment_name, along with the sweep settings and each trial's final eval_loss.

The driver raises an error if any trial failed.

The example does not persist adapter weights. To keep the best adapter, give tune.Tuner a RunConfig(storage_path=...) on a Unity Catalog volume that every node can reach.

Adjust the size of the sweep​

The constants at the top of tune_lora.py control the size of the sweep. Set them smaller to smoke-test a change in a couple of minutes, though the eval_loss figures are then too noisy to rank configurations.

Wall-clock time tracks NUM_SAMPLES / num_accelerators, so raise num_accelerators rather than shrinking the search when a sweep takes too long. For a larger model, raise accelerator_type to a bigger GPU. To choose configurations instead of sampling them randomly, pass TuneConfig a search_alg such as Optuna.

Full tuning script​

The complete tune_lora.py for copy-paste:

Python
#!/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()

Additional resources​

On this page