Tune LLM hyperparameters within a compute budget

Six fine-tuning runs can produce six lower training losses without answering which model should handle your application. The useful comparison connects each configuration to a task score, the resources it consumed, and the conditions under which that score was measured.

Start with a bounded study: three learning rates, two LoRA ranks, one training recipe, and one validation procedure. Polyaxon can schedule that grid and retain its experiment records. Your training library implements the adapter, and your evaluator decides whether its outputs meet the task requirements.

Separate training, decoding, and serving settings

The phrase “LLM hyperparameters” covers several different experiments:

Changing decoding parameters does not train model weights. Hugging Face's generation reference documents how these settings affect token generation. Keep them fixed while comparing fine-tuning candidates; otherwise, a different generation policy could explain the apparent improvement.

This article focuses on fine-tuning. Use controlled LLM performance and cost experiments for the subsequent serving study.

Define what earns another training run

Suppose a model routes support cases to a fixed set of teams. Evaluate the base model with the intended prompt before tuning. Record macro F1 across routing labels, invalid responses, and errors on rare but important categories. Count invalid outputs in the evaluation rather than dropping difficult cases from the denominator.

Create separate training, validation, and final test splits. Group related messages from one case together to prevent leakage. Use validation results for hyperparameter and checkpoint selection; keep the final test set untouched until those choices are frozen.

Decide the minimum acceptable task quality and resource allowance before launching the sweep. Reserve budget for the baseline, evaluation, failed attempts, and repeated finalists. Spending the entire allowance on the initial candidates leaves no room to check whether the apparent winner is repeatable.

If you still need a training setup, the Mistral LoRA guide covers data preparation, GPU scheduling, checkpoints, and adapter packaging. Here, we build on that workflow to compare configurations.

Choose a small, meaningful search space

LoRA keeps the base weights frozen and trains low-rank updates. The rank controls the size of those updates; it is a candidate to evaluate, not a quality score. QLoRA combines adapter training with a quantized base model. Choose the loading and quantization recipe before this study, and keep it consistent across candidates.

For an initial pilot, consider learning rates of 0.00005, 0.0001, and 0.0002, crossed with ranks 8 and 16. These are illustrative values to review against your model's recipe, not established optima.

Keep the target modules, dropout, optimizer, scheduler, warmup, and loss masking fixed. Make the adapter scaling rule explicit: for standard LoRA, this study uses lora_alpha = 2 × r, keeping alpha/r constant as rank changes. Rank-stabilized LoRA uses a different scaling rule; see the PEFT implementation.

Freeze the base-model and tokenizer revisions, split manifest, prompt format, dependency versions, GPU type, and evaluator. TRL's SFT documentation explains how dataset format and completion-only loss determine which tokens contribute to training. Changing that policy between trials changes the experiment.

In ordinary data-parallel training, effective batch size is the per-device batch multiplied by the number of data-parallel devices and gradient accumulation steps. Transformers documents these controls. With one GPU, a microbatch of 1, and accumulation of 8, each complete optimizer update covers eight sequences. Record actual non-padding and target-token counts too: sequence length, packing, and masking affect how much training signal those sequences contain.

Express the six candidates as a Polyaxon matrix

Package your trainer and validation evaluator in a component called finetune-component.yaml. It must declare the following inputs, pass them to the training program, and use the frozen recipe above. Declare learning_rate as a float and the remaining values as integers; the trainer derives lora_alpha from lora_rank.

This operation describes the search. It assumes that component already provides its container image, training code, one-GPU resource allocation, credentials, and dataset access; it is not a standalone model-training implementation.

Use Polyaxon CLI 2.13 or newer for the short-form parameter values below. Save this operation as lora-study.yaml beside finetune-component.yaml; the CLI resolves pathRef when submitting the study.

version: 1.1
kind: operation
name: support-routing-lora-study
tags: [llm-tuning, support-routing]
pathRef: ./finetune-component.yaml
params:
  max_steps: 300
  max_length: 1024
  per_device_train_batch_size: 1
  gradient_accumulation_steps: 8
  seed: 23
matrix:
  kind: grid
  concurrency: 2
  params:
    learning_rate:
      kind: choice
      value: [0.00005, 0.0001, 0.0002]
    lora_rank:
      kind: choice
      value: [8, 16]

Submit the study from that directory:

polyaxon run -f lora-study.yaml

The Cartesian product contains six configurations. concurrency: 2 permits at most two child operations at once; it does not allocate GPUs or impose a financial budget. Each candidate starts from the same base revision with a fresh adapter and optimizer state. Do not resume one candidate from another candidate's checkpoint.

The 300 steps are an illustrative pilot budget. Keep the training schedule and data order comparable, and record the actual completed steps. If later exploring a broader space, random search uses numRuns to bound the number of sampled configurations. Its matrix seed controls search sampling; the training program still needs its own seed.

Bound resources without hiding failed candidates

Profile one short run before setting the full study budget. Include validation and checkpoint writing in the measurement. For a fixed allocation, allocated GPU-hours equal GPU count multiplied by elapsed allocation hours. Keep that quantity separate from GPU utilization and monetary cost: prices, idle nodes, storage, and other overhead need their own accounting.

Set a termination timeout and deliberate retry policy on the training component. Record every attempt; retries consume resources even when the final run succeeds. A timeout is a runtime boundary, not evidence that the model is poor.

Use resource monitoring to inspect memory and utilization where enabled. An out-of-memory candidate is useful evidence about feasibility. Preserve its configuration and failure reason instead of quietly rerunning it with a smaller batch.

Add early stopping only after deciding what comparable progress means. A slow warmup can make a useful configuration look weak early. Polyaxon's early-stopping reference distinguishes stopping individual runs through a policy from a metric threshold without a policy, which can stop the whole pipeline. The initial six-run matrix intentionally leaves metric-based stopping unset.

Keep the evidence beside each configuration

Use the Hugging Face callback for trainer-reported metrics. Add application scores and resource measurements through the tracking API. A useful run record contains:

• The resolved training configuration, data and model revisions, seed, and exported adapter identity.
• Training and validation loss at recorded steps, task macro F1, invalid-output count, and slice-level results.
• Completed training tokens, measured elapsed time, GPU allocation, peak memory, and any interrupted attempts.
• A case-level prediction report, scorer version, and checkpoint-selection rule.

Save reports and adapters under the run's outputs location using artifact tracking. Reference-logging methods such as log_file_ref record lineage; the application must actually save the referenced file. A metric callback also does not automatically calculate task F1, allocated GPU-hours, or a cost estimate.

Repeat finalists, then evaluate the frozen choice

In the comparison dashboard, filter to the same dataset, evaluator, hardware, and training budget. Compare quality against measured resource use, then inspect the error reports. A configuration that improves common labels while degrading an essential category may fail the selection rule.

Repeat shortlisted configurations with additional training seeds on the same validation procedure. Report variation and evaluation sample size; do not select a model on a tiny score difference without checking the cases behind it. If increasing the training budget, evaluate finalists at that budget again rather than assuming their short-run ranking will hold.

Freeze the chosen recipe and checkpoint rule before running the final test. Evaluate the actual exported adapter with its pinned base model and generation settings. Retain that report with a model registry version so the result can be traced back to its run.

Start with the six-candidate study and a written selection rule. Expand the search when its evidence identifies a worthwhile next question, whether that is another learning-rate range, a different adapter target, or a data issue that tuning cannot resolve.