Skip to main content

Can I fine-tune and serve open source models on Databricks?

Summary

  • Yes. Databricks supports the whole path in one governed platform: take an open-weight model, fine-tune it on your data, register the result in Unity Catalog through MLflow, and deploy it on Model Serving for production inference.
  • You can fine-tune with a managed service (Model Training) that provisions the GPUs and orchestrates the run, or use Databricks AI Runtime with open-source frameworks like LLM Foundry, Unsloth, TRL, and Axolotl for full or parameter-efficient (LoRA/QLoRA) fine-tuning.
  • Fine-tuned models are registered in Unity Catalog through MLflow, so you keep ownership of the weights and get permissions, versioning, and lineage from the training data through to the serving endpoint.
  • Model Serving deploys open-weight and fine-tuned models on GPU infrastructure, with pay-per-token endpoints for experimentation and provisioned throughput for production.
  • Governance runs through Unity Catalog and Unity Gateway: access control, rate limits, budgets, guardrails, and request logging apply to the models you serve.

Can I fine-tune and serve open source models on Databricks?

Yes. Databricks supports the whole path in one governed platform: take an open-weight model, fine-tune it on your data, register the result in Unity Catalog through MLflow, and deploy it on Model Serving for production inference. Because training and serving sit next to your data, you keep full ownership of the fine-tuned weights and get one set of permissions, versioning, and lineage from the training dataset through to the live endpoint. That matters when a model encodes proprietary data and has to meet enterprise access and audit requirements.

What fine-tuning an open source model involves

Open-weight models, such as the Llama and Mistral families, ship with downloadable weights you can adapt to your own data. Two broad approaches exist. Full fine-tuning updates all of a model's parameters and gives the most control, at higher GPU cost. Parameter-efficient methods like LoRA and QLoRA train a small set of added weights, which cuts memory use and lets you fine-tune larger models on smaller GPU footprints. The right choice depends on your quality targets, dataset size, and available compute.

Capabilities to evaluate

Capability Why it matters
Fine-tuning options A managed service is fastest to start; open-source frameworks give full control over the training loop.
GPU training compute Fine-tuning open-weight models needs GPUs, ideally without you managing clusters by hand.
Model registry and versioning Each fine-tuned version should be tracked, comparable, and reproducible.
Weight ownership You should own and be able to move the fine-tuned weights, not rent access to a black box.
Serving modes Support for both low-commitment endpoints and production-grade throughput covers experimentation and scale.
Governance Access control, rate limits, budgets, guardrails, and audit should apply to every served model.
Data-to-endpoint lineage Tracing which data and code produced a deployed model is essential for compliance and debugging.

How Databricks approaches this

For fine-tuning, you have two routes. Model Training is a managed service: you pick a base model and a training dataset, and Databricks handles the GPU provisioning and training orchestration, returning weights you own. For more control, Databricks AI Runtime lets you fine-tune with open-source frameworks including LLM Foundry, Unsloth, TRL, and Axolotl, using full fine-tuning or LoRA and QLoRA. The fine-tuning guide walks through the tradeoffs, and the AI Runtime tutorials show end-to-end runs on models like Llama 3.
Once training finishes, you register the model and its tokenizer in Unity Catalog with MLflow:

mlflow.set_registry_uri("databricks-uc")
with mlflow.start_run(run_id=run_id):
    mlflow.transformers.log_model(
        transformers_model={"model": model, "tokenizer": tokenizer},
        name="model",
        task="llm/v1/chat",
        registered_model_name=REGISTERED_MODEL_NAME,
    )

Registration gives the model a governed home with permissions and lineage that tracks the training data and downstream consumers. From there, Model Serving deploys it on GPU infrastructure. You can serve a registered fine-tuned model as a custom model, or serve open-weight foundation models directly. Two serving styles cover different needs: pay-per-token endpoints for experimentation without infrastructure commitments, and provisioned throughput for production workloads that need steady performance. Governance runs through Unity Catalog and Unity Gateway, which add access control, rate limits, budgets, guardrails, and request logging across the models you serve. The foundation model serving overview covers the deployment options in detail.

Use cases

Common patterns include adapting an open-weight model to a company's domain language or support history, then serving it behind an internal app or agent; standing up a private endpoint for a fine-tuned model so proprietary data stays inside the governed platform; and running batch or real-time inference on a fine-tuned model as part of a larger LLMOps workflow with monitoring and retraining.

FAQs

Which open source models can I fine-tune?

Open-weight model families such as Llama and Mistral are supported, along with Databricks' own DBRX. The managed Model Training service and the AI Runtime frameworks both work with popular open-weight models.

Do I own the fine-tuned model?

Yes. Fine-tuning returns weights you own, registered in your Unity Catalog. You control access and can serve or move the model as your governance allows.

Can I serve a model without fine-tuning it first?

Yes. Model Serving can deploy open-weight foundation models directly, so you can start serving and add fine-tuning later.

How is a served model governed?

Serving endpoints are governed through Unity Catalog and Unity Gateway, which provide access control, rate limits, budgets, guardrails, and request logging.

What is the difference between pay-per-token and provisioned throughput?

Pay-per-token endpoints suit experimentation and variable traffic without infrastructure commitments. Provisioned throughput reserves capacity for production workloads that need consistent performance.

The information provided herein is for general informational purposes only and may not reflect the most current product capabilities or configurations.