Fine-tune Mistral 7B with LoRA on Kubernetes

Fine-tuning Mistral 7B with LoRA adapts selected parts of a model for a task while retaining its base weights. On Kubernetes, the surrounding workflow also needs to identify the training data, schedule suitable GPU resources, retain checkpoints, and evaluate the exported adapter. Polyaxon can organize those executions and their evidence.

This guide uses Mistral-7B-v0.3 as a concrete base-model choice. It is not a claim that this model is the best option for every workload. Establish a baseline for your task before spending time on fine-tuning.

Choose a behavior you can evaluate

Start with a narrow output requirement. For example, a support application may need to convert a case summary into a structured routing decision. Define the allowed fields and labels, the information the model may use, and the expected handling of incomplete cases.

Measure the base model on a held-out set using the intended inference prompt. This distinguishes a training problem from a prompt or data-quality problem. If the application mainly lacks changing factual information, consider whether retrieval addresses that requirement more directly.

Choose acceptance criteria before training: valid output structure, task correctness, relevant error slices, latency, and any security checks required by the application. A lower training loss alone does not establish an improvement in those outcomes.

Freeze the data and serialization

Create separate training, validation, and final test splits. Split related records together where necessary: two summaries from the same support case should not appear on opposite sides of the evaluation boundary. Retain the split manifest and preprocessing revision.

For a plain-text prompt-completion task, one synthetic training record could look like this:

{
  "prompt": "Route this support case. Return JSON with a team field.\nCase: The customer was charged twice for the same renewal.\nDecision:\n",
  "completion": "{\"team\":\"billing\"}"
}

This is an illustration of the format, not a sufficient training dataset. Real examples need representative language, difficult cases, and reviewed targets. Keep the separators, end-of-sequence handling, and inference formatting consistent with training.

TRL's SFT trainer documentation distinguishes language-modeling, prompt-completion, and conversational datasets. For this prompt-completion format, completion-only loss trains on the target tokens. For a conversational task, explicitly select and test a compatible chat template rather than assuming that a base-model tokenizer has the format your application expects.

Inspect tokenized examples before training. Check truncation, target-token counts, padding, and whether the desired response remains inside the sequence limit. A pipeline can execute successfully while training on the wrong text.

Choose LoRA or QLoRA deliberately

LoRA adds a relatively small set of trainable adapter parameters. QLoRA combines adapters with a quantized base model to reduce the memory required to hold its weights. Hugging Face's PEFT quantization guide documents the model-loading and preparation steps for that approach.

Quantized weights do not eliminate activation memory, optimizer state, temporary buffers, or checkpoint storage. Profile the actual sequence length and batch configuration on the intended GPU before increasing the training budget. Support for the chosen compute dtype and quantization backend must match the hardware and library build.

The following configuration fragment illustrates a small LoRA trial using PEFT and TRL. It defines settings only; your trainer must load the pinned model and dataset and pass these objects to the training API. The values are starting points to validate, not tested quality or memory guarantees.

from peft import LoraConfig
from trl import SFTConfig

adapter_config = LoraConfig(
    task_type="CAUSAL_LM",
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
)

training_config = SFTConfig(
    output_dir="outputs/mistral-lora",
    max_length=1024,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=8,
    max_steps=100,
    learning_rate=1e-4,
    completion_only_loss=True,
    logging_steps=10,
    save_steps=50,
    report_to="none",
)

Pin a compatible set of PyTorch, Transformers, PEFT, TRL, datasets, and accelerator dependencies in the training image. For QLoRA, also configure quantized loading and prepare the model for low-bit training; the adapter configuration above does not enable quantization by itself.

Package training as a Polyaxon job

Keep the trainer, dependency lockfiles, and container build in your code repository. Express the execution through a Polyaxon component, with parameters for the base-model revision, data revision, training configuration, and output location.

For a cluster exposing NVIDIA GPUs through a device plugin, this container resource fragment requests one GPU:

resources:
  limits:
    nvidia.com/gpu: 1

Place it under run.container in the component and add CPU, memory, storage, and node-placement settings appropriate to your measured workload. The fragment is not a complete job definition. It also does not select a GPU with a particular amount of memory. Kubernetes documents the device-plugin resource rules; shared resources and MIG profiles need the corresponding cluster configuration.

Connect model downloads and dataset access through configured storage and credentials. Keep a cache of reusable model files where appropriate, while recording the immutable model revision each run loads. Verify the cache and checkpoint locations have enough capacity for the whole execution.

Use Polyaxon's Hugging Face tracking integration to capture trainer metrics, and log the relevant configuration and exported artifacts. Useful records include sequence length, effective batch configuration, adapter settings, token counts, loss, elapsed time, and peak memory. Preserve the exact evaluation setup alongside these training records.

Test recovery before a long run

Run a short training budget first, save a checkpoint, stop the process, and resume it in a fresh execution. Verify that it restores the intended training step and the required optimizer, scheduler, and random-number state.

An exported adapter is not necessarily a complete resumable training checkpoint. Keep those two output types distinct. Put checkpoints in storage that survives the expected failure, and make the resume location explicit. A Kubernetes restart or Polyaxon retry does not automatically teach the training program which checkpoint to load.

Evaluate the package that will actually be served

Load the exported adapter in a fresh inference process using the recorded base-model revision. Compare it with the baseline on the frozen evaluation suite, using equivalent generation settings. Preserve per-case outputs so aggregate improvements do not hide regressions on an important category.

Evaluate task correctness separately from formatting. A valid JSON object can still route a case to the wrong team. For applications exposed to untrusted text, add the relevant prompt injection tests and check the effect of fine-tuning on existing controls.

If deployment requires merged weights, evaluate that exported form too. A successful training run should not bypass verification of the package and inference stack that will be used downstream.

Register the adapter with its dependencies

The PEFT checkpoint format separates adapter weights and configuration from the base model. Preserve the base-model identifier and revision, tokenizer files, adapter configuration, training-data manifest, environment, and evaluation report with the handoff.

Use Polyaxon model versions to connect the reviewed package to its originating run. Record whether the artifact is an adapter or a merged model and provide the corresponding loading instructions. Promote it only after the task-specific comparison supports that decision; keep the previous version available for a controlled rollback.