Skip to content

Fine-tuning with Kubeflow

Fine-tune a language model on your research data using a Kubeflow Pipeline. This example uses LoRA (Low-Rank Adaptation) to efficiently fine-tune Llama 3.1 8B for domain-specific tasks.

Platform access required

This example requires Kubeflow access — you need a research group namespace with GPU quota. See Requesting Access.

What you'll build

A Kubeflow Pipeline that: 1. Loads and preprocesses your training data from VAST storage 2. Runs LoRA fine-tuning with configurable hyperparameters 3. Evaluates the fine-tuned model 4. Logs all metrics and the adapter weights to MLflow

Prerequisites

pip install kfp kfp-kubernetes transformers peft accelerate mlflow datasets

Pipeline definition

# finetune_pipeline.py
import kfp
from kfp import dsl
from kfp.dsl import Dataset, Model, Input, Output, Artifact
import os


@dsl.component(
    base_image="python:3.11",
    packages_to_install=["datasets", "transformers", "boto3"],
)
def prepare_dataset(
    data_path: str,
    output_dataset: Output[Dataset],
    val_split: float = 0.1,
):
    """Load JSONL data from VAST storage and split into train/val."""
    import json
    from pathlib import Path

    records = []
    for line in Path(data_path).read_text().splitlines():
        records.append(json.loads(line))

    split_idx = int(len(records) * (1 - val_split))
    train_records = records[:split_idx]
    val_records = records[split_idx:]

    import json
    with open(output_dataset.path, "w") as f:
        json.dump({"train": train_records, "val": val_records}, f)

    print(f"Train: {len(train_records)}, Val: {len(val_records)}")


@dsl.component(
    base_image="nvcr.io/nvidia/pytorch:24.08-py3",
    packages_to_install=["transformers", "peft", "accelerate", "mlflow", "trl"],
)
def finetune_lora(
    dataset: Input[Dataset],
    output_model: Output[Model],
    base_model: str = "meta-llama/Llama-3.1-8B-Instruct",
    lora_r: int = 16,
    lora_alpha: int = 32,
    learning_rate: float = 2e-4,
    num_epochs: int = 3,
    batch_size: int = 4,
    mlflow_tracking_uri: str = "",
):
    """LoRA fine-tuning using PEFT + TRL SFTTrainer."""
    import json
    import mlflow
    from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments
    from peft import LoraConfig, get_peft_model
    from trl import SFTTrainer

    data = json.load(open(dataset.path))

    if mlflow_tracking_uri:
        mlflow.set_tracking_uri(mlflow_tracking_uri)
        mlflow.set_experiment("lora_finetune")

    tokenizer = AutoTokenizer.from_pretrained(base_model)
    model = AutoModelForCausalLM.from_pretrained(base_model, load_in_4bit=True)

    lora_config = LoraConfig(
        r=lora_r,
        lora_alpha=lora_alpha,
        target_modules=["q_proj", "v_proj"],
        lora_dropout=0.05,
        task_type="CAUSAL_LM",
    )
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()

    training_args = TrainingArguments(
        output_dir=output_model.path,
        num_train_epochs=num_epochs,
        per_device_train_batch_size=batch_size,
        learning_rate=learning_rate,
        logging_steps=10,
        save_strategy="epoch",
        report_to="mlflow" if mlflow_tracking_uri else "none",
    )

    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=data["train"],
        eval_dataset=data["val"],
    )
    trainer.train()
    trainer.save_model(output_model.path)
    print(f"Model saved to {output_model.path}")


@dsl.pipeline(
    name="lora-finetune-pipeline",
    description="LoRA fine-tuning pipeline for domain adaptation",
)
def finetune_pipeline(
    data_path: str = "/vast/rg-compsci/datasets/finetune_data.jsonl",
    lora_r: int = 16,
    learning_rate: float = 2e-4,
    num_epochs: int = 3,
    mlflow_uri: str = "https://mlflow.pais.auckland.ac.nz",
):
    prep = prepare_dataset(data_path=data_path)
    prep.set_caching_options(enable_caching=True)

    train = finetune_lora(
        dataset=prep.outputs["output_dataset"],
        base_model="meta-llama/Llama-3.1-8B-Instruct",
        lora_r=lora_r,
        learning_rate=learning_rate,
        num_epochs=num_epochs,
        mlflow_tracking_uri=mlflow_uri,
    )
    train.set_accelerator_type("GPU")
    train.set_gpu_limit(1)


if __name__ == "__main__":
    kfp.compiler.Compiler().compile(
        finetune_pipeline,
        package_path="finetune_pipeline.yaml",
    )
    print("Pipeline compiled to finetune_pipeline.yaml")

Compiling and submitting

# Compile
python finetune_pipeline.py

# Submit via Kubeflow Pipelines UI
# 1. Open https://kubeflow.pais.auckland.ac.nz
# 2. Pipelines → Upload Pipeline → select finetune_pipeline.yaml
# 3. Create Run with your parameters

Or submit programmatically:

import kfp

client = kfp.Client(host="https://kubeflow.pais.auckland.ac.nz")

run = client.create_run_from_pipeline_func(
    finetune_pipeline,
    arguments={
        "data_path": "/vast/rg-compsci/datasets/my_training_data.jsonl",
        "lora_r": 32,
        "learning_rate": 1e-4,
        "num_epochs": 5,
    },
    run_name="lora-r32-lr1e-4",
    experiment_name="domain-adaptation-v2",
)
print(f"Run URL: https://kubeflow.pais.auckland.ac.nz/pipeline/#/runs/details/{run.run_id}")

See also