Fine-Tuning LLMs at Scale With Databricks MLflow and Spark

# Databricks Runtime ML 14.x+ recommended (ships CUDA, PyTorch, Transformers)
# Install additional packages in your cluster init script or notebook

%pip install \
    transformers==4.40.0 \
    peft==0.10.0 \
    trl==0.8.6 \
    accelerate==0.29.3 \
    horovod[spark]==0.28.1 \
    datasets==2.19.0 \
    evaluate==0.4.1 \
    --quiet

dbutils.library.restartPython()

import os
import mlflow
import mlflow.transformers
import torch
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling,
)
from peft import LoraConfig, get_peft_model, TaskType
from pyspark.sql import functions as F
from datasets import Dataset

# ── MLflow setup ──────────────────────────────────────────────────────────────
# On Databricks, MLflow tracking URI is pre-configured to the workspace
# mlflow.set_tracking_uri("databricks")   # uncomment for external clusters

EXPERIMENT_NAME = "/Users/[email protected]/llm-finetuning/support-classifier"
mlflow.set_experiment(EXPERIMENT_NAME)

BASE_MODEL   = "mistralai/Mistral-7B-Instruct-v0.2"
CATALOG      = "prod"
GOLD_DB      = f"{CATALOG}.gold"
MODEL_NAME   = f"{CATALOG}.ml.support_intent_classifier"   # Unity Catalog model path

print(f"GPU available: {torch.cuda.is_available()}")
print(f"Device count:  {torch.cuda.device_count()}")

# ── Spark Data Preparation ────────────────────────────────────────────────────

def build_prompt(row):
    """
    Format a support conversation into an instruction-following prompt.
    Uses the Mistral instruct template: [INST] ... [/INST]
    """
    return f"[INST] Classify the intent of this support message:\n\n{row['message']} [/INST] {row['intent_label']}"


# Load from Delta Gold table
raw_df = (
    spark.table(f"{GOLD_DB}.support_conversations")
        .filter(F.col("quality_score") >= 0.85)           # keep high-quality labels only
        .filter(F.col("intent_label").isNotNull())
        .filter(F.length("message") > 20)                 # filter empty/stub messages
        .filter(F.length("message") < 2048)               # filter messages too long to tokenize
        .dropDuplicates(["message_hash"])                  # remove exact duplicates
        .select("message", "intent_label", "created_date")
        .limit(500_000)                                    # cap for this training run
)

print(f"Training candidates: {raw_df.count():,}")

# Build prompt strings using Spark — parallelized across all workers
prompt_udf = F.udf(
    lambda msg, label: f"[INST] Classify the intent of this support message:\n\n{msg} [/INST] {label}",
    returnType="string"
)

prepared_df = (
    raw_df
        .withColumn("prompt", prompt_udf(F.col("message"), F.col("intent_label")))
        .withColumn("token_count",
            F.size(F.split(F.col("prompt"), r"\s+")))         # rough word count proxy
        .filter(F.col("token_count") < 512)                   # stay within model context
        .select("prompt", "token_count", "created_date")
)

# Stratified split using Spark (reproducible with seed)
train_df, val_df, test_df = prepared_df.randomSplit([0.80, 0.10, 0.10], seed=42)

# Persist splits to Delta for lineage + reproducibility
train_df.write.format("delta").mode("overwrite").saveAsTable(f"{GOLD_DB}.llm_train_split")
val_df.write.format("delta").mode("overwrite").saveAsTable(f"{GOLD_DB}.llm_val_split")
test_df.write.format("delta").mode("overwrite").saveAsTable(f"{GOLD_DB}.llm_test_split")

print(f"Train: {train_df.count():,} | Val: {val_df.count():,} | Test: {test_df.count():,}")

# ── Distributed Fine-Tuning with Horovod on Spark ────────────────────────────
# Best for: datasets > 5M tokens, models > 13B params, or when you need
# to reduce wall-clock training time below a business SLA.

import horovod.torch as hvd
from sparkdl import HorovodRunner

def train_fn(hparams):
    """
    Training function executed on each Horovod worker.
    Each worker trains on a data shard; gradients are averaged across workers.
    """
    import horovod.torch as hvd
    from transformers import AutoModelForCausalLM, Trainer, TrainingArguments
    from datasets import load_from_disk

    hvd.init()

    # Each worker loads only its shard
    local_rank = hvd.local_rank()
    world_size = hvd.size()

    torch.cuda.set_device(local_rank)

    # Load dataset shard for this worker
    dataset = load_from_disk(f"/dbfs/tmp/llm-finetune/train_shards/shard_{local_rank}")

    model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        torch_dtype=torch.bfloat16,
    ).to(f"cuda:{local_rank}")

    # Wrap optimizer with Horovod DistributedOptimizer
    optimizer = torch.optim.AdamW(model.parameters(), lr=hparams["lr"])
    optimizer = hvd.DistributedOptimizer(
        optimizer,
        named_parameters=model.named_parameters(),
        compression=hvd.Compression.fp16,    # compress gradient communication
    )

    # Broadcast initial model weights from rank 0 to all workers
    hvd.broadcast_parameters(model.state_dict(), root_rank=0)
    hvd.broadcast_optimizer_state(optimizer, root_rank=0)

    training_args = TrainingArguments(
        output_dir=f"/dbfs/tmp/llm-finetune/hvd_output",
        num_train_epochs=hparams["epochs"],
        per_device_train_batch_size=hparams["batch_size"],
        bf16=True,
        no_cuda=False,
        dataloader_num_workers=2,
        # Only rank 0 logs and saves — avoids duplicated artifacts
        report_to="mlflow" if hvd.rank() == 0 else "none",
        save_strategy="epoch" if hvd.rank() == 0 else "no",
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
        optimizers=(optimizer, None),
    )
    trainer.train()

    # Only rank 0 registers the model
    if hvd.rank() == 0:
        mlflow.transformers.log_model(
            transformers_model={"model": model, "tokenizer": tokenizer},
            artifact_path="model",
            registered_model_name=MODEL_NAME,
        )


# Launch distributed training across N GPU workers
# np = number of processes = number of GPUs across all nodes
hr = HorovodRunner(np=8, driver_log_verbosity="all")   # 8 GPUs (e.g., 2 × 4-GPU nodes)
hr.run(train_fn, hparams={
    "lr":         2e-5,
    "epochs":     3,
    "batch_size": 2,       # per GPU; effective = 2 × 8 = 16
})

# ── Model Registry Promotion Workflow ─────────────────────────────────────────

from mlflow.tracking import MlflowClient

client = MlflowClient()

# Get the latest registered version from the training run
latest_version = client.get_registered_model(MODEL_NAME).latest_versions[0].version

# Tag the new version as a candidate for review
client.set_registered_model_alias(
    name=MODEL_NAME,
    alias="candidate",
    version=latest_version,
)

client.set_model_version_tag(
    name=MODEL_NAME,
    version=latest_version,
    key="fine_tuned_on",
    value="gold.support_conversations",
)

client.set_model_version_tag(
    name=MODEL_NAME,
    version=latest_version,
    key="eval_loss",
    value=str(round(eval_results["eval_loss"], 4)),
)

# After human review / automated eval gates pass → promote to staging
client.set_registered_model_alias(
    name=MODEL_NAME,
    alias="staging",
    version=latest_version,
)

# After integration tests pass → promote to champion (production)
client.set_registered_model_alias(
    name=MODEL_NAME,
    alias="champion",
    version=latest_version,
)

# Load model by alias — decouples code from version numbers
champion_model = mlflow.transformers.load_model(f"models:/{MODEL_NAME}@champion")

# ── Deploy to Databricks Model Serving ────────────────────────────────────────
# Can also be done via the UI: Models > Serving > Create Endpoint

import requests, json

WORKSPACE_URL = "https://<your-workspace>.azuredatabricks.net"
TOKEN         = dbutils.secrets.get("prod-scope", "databricks-token")

endpoint_config = {
    "name": "support-intent-classifier",
    "config": {
        "served_models": [
            {
                "name":                    "mistral-7b-lora-champion",
                "model_name":              MODEL_NAME,
                "model_version":           latest_version,
                "workload_size":           "Small",      # 1 GPU
                "scale_to_zero_enabled":   True,
                "workload_type":           "GPU_LARGE",  # A10G
            }
        ],
        "traffic_config": {
            "routes": [
                {"served_model_name": "mistral-7b-lora-champion", "traffic_percentage": 100}
            ]
        },
        "auto_capture_config": {
            "catalog_name":  CATALOG,
            "schema_name":   "ml",
            "table_name":    "support_classifier_inference_log",
            "enabled":       True,                        # log all requests/responses to Delta
        }
    }
}

response = requests.post(
    f"{WORKSPACE_URL}/api/2.0/serving-endpoints",
    headers={"Authorization": f"Bearer {TOKEN}", "Content-Type": "application/json"},
    data=json.dumps(endpoint_config),
)
print(response.json())

# ── Query the endpoint ────────────────────────────────────────────────────────
def classify_intent(message: str) -> str:
    payload = {
        "inputs": {"prompt": f"[INST] Classify the intent of this support message:\n\n{message} [/INST]"},
        "params": {"max_new_tokens": 50, "temperature": 0.1},
    }
    resp = requests.post(
        f"{WORKSPACE_URL}/serving-endpoints/support-intent-classifier/invocations",
        headers={"Authorization": f"Bearer {TOKEN}", "Content-Type": "application/json"},
        data=json.dumps(payload),
    )
    return resp.json()["predictions"][0]

print(classify_intent("My order hasn't arrived and it's been 10 days"))
# → "shipping_delay"