Skip to content

Training your own models on Databricks AI Runtime

Preetham Reddy8 min read

Too often, teams that want a model trained on their own data spend more time securing GPU capacity and waiting for experiments to launch than improving the model. Running GPUs separately from the lakehouse also adds governance gaps and extra copies of data that someone then has to account for.

At TechFabric we're a Databricks partner. Databricks announced AI Runtime as generally available on 21 October 2026, and it's the piece we'd use to help clients keep that training inside the workspace where their data and permissions already live.

With it, training runs on serverless GPUs inside the Unity Catalog governance perimeter, and usage is tracked per workload in billing system tables. Whether training your own model is worth it still depends on the use case, so start with what Databricks shipped and what it changes for each person who signs off on the project.

What Databricks shipped

AI Runtime (AIR) is serverless GPU compute for training and fine-tuning on Databricks, with GPU procurement, cluster configuration and networking handled by Databricks. The documentation lists four accelerator shapes, and choosing between them is mostly a cost decision:

Accelerator GPU memory Intended use
1xA10 24 GB Classic ML and smaller LLM fine-tuning
1xH100 80 GB Medium-size LLM fine-tuning
8xH100 80 GB per GPU Large-scale training and fine-tuning
8xB300 288 GB per GPU The largest models and longest contexts (Public Preview on AWS, by waitlist)

For a first fine-tuning experiment we'd pick a single-GPU shape, 1xA10 or 1xH100, and move to 8xH100 only when the evaluation shows a larger model earns its cost. Only the 8-GPU shapes support multi-GPU distributed training, so that move is also the point where the training code needs to be written for it.

What it changes for each stakeholder

AI Runtime is an infrastructure release, and its value shows up differently for each person who signs off on a training project.

  • CTO or chief data officer. A first experiment no longer waits on a GPU contract or a Kubernetes cluster, and a training job becomes a scheduled Lakeflow job with a few lines of YAML. A use case can be tested on real data before anyone funds an infrastructure project.
  • Head of ML or data science. PyTorch DDP, FSDP, DeepSpeed and Ray run on it, as do custom containers, so engineers keep their training code and stop maintaining GPU clusters. Teams already on Kubernetes or SLURM can use Databricks' AIR Skills to translate existing jobs.
  • CISO and risk. The Unity Catalog grants your security team already reviewed for the source tables also govern the training data, checkpoints and logs, so training doesn't need a separate access model. Ask your account team about residency, because under high demand AIR may use GPUs in another region, which the documentation says can incur egress costs.
  • CFO and FinOps. Usage lands in billing system tables per workload, which you can tie to a project and a budget, instead of a reserved cluster that bills whether anyone uses it or not. Lakeflow's priority and pre-emption reduce idle GPU time on scheduled jobs.

When to train your own model, and when not to

Our position is to fine-tune only after prompting and retrieval have missed the quality bar on a measured test. The flow below is the one we'd follow with any client.

Yes

No

No

Yes

Yes

No

Define the task and an evaluation set

Prompting and retrieval meet the bar?

Ship with a hosted model

Enough labelled examples?

Collect and label data first

Fine-tune an open model on AI Runtime

Beats production on the evaluation set?

Register and serve it

We only fine-tune or train when prompting and retrieval miss the quality bar on a fixed evaluation set and there is labelled data to learn from.

In our view, fine-tuning a smaller open model is worth considering when the task is narrow and repeated at high volume, so inference cost per call matters. We'd also consider it when the vocabulary or document formats are your own, when latency needs a smaller model, or when the data can't be sent to an outside model provider. Training from scratch is rarer, and we'd only recommend it when you have the data volume and a reason no existing base model fits.

Two constraints shape how we'd plan the work. GPU capacity can be constrained in some regions, so we wouldn't put a training run on a release's critical path without a fallback date. Running a scheduled job that uses a workspace base environment, including the Databricks AI environment, is in Beta, so for production retraining we'd pin dependencies in a custom container instead.

AI Runtime is the wrong tool when the honest answer is better retrieval, a clearer prompt or cleaner data. It's also not a fit yet for workspaces that need the compliance security profile under FedRAMP High or DoD IL5, which the documentation lists as unsupported. Long runs need checkpointing too: CLI jobs can run up to 14 days, and notebook sessions up to two.

A lending document model, end to end

Financial services shows the pattern well. A lender receives thousands of loan files a day, among them pay stubs, bank statements, titles and payoff letters. A general-purpose model reads them adequately, but a fine-tuned model that knows the lender's own document types and exception codes could classify them at lower cost per page, and the files can't leave the lender's governed environment.

Better

Worse

Loan files in a Unity Catalog volume

Lakeflow builds a labelled training table

AI Runtime fine-tuning job

MLflow records run and checkpoints

Evaluation gate vs production model

Register model in Unity Catalog

Model Serving endpoint

Drift monitoring triggers retraining

Loan files, the training run, the evaluation and the served model all stay inside Unity Catalog, so the permissions on the source data carry through to the model.

Here is how we'd build it:

  1. Data. Loan documents land in a Unity Catalog volume. A Lakeflow pipeline joins them to the outcomes underwriters already recorded, which become labels, and writes a versioned training table with lineage back to the source files.
  2. Training. An AI Runtime job fine-tunes an open model on 1xH100 for the first experiments, and moves to 8xH100 only if the evaluation says a larger model is worth it. The job streams the volume through Databricks' built-in data loaders instead of exporting it, and MLflow records each run's parameters, metrics and checkpoints.
  3. Evaluation. Every run is compared with the production model on a fixed, held-out set of labelled files, including awkward ones such as scanned documents, mismatched names and expired payoffs. A model that doesn't beat production doesn't ship.
  4. Serving and retraining. The winning model is registered in Unity Catalog and served from Model Serving. When drift monitoring shows accuracy falling on new document types, the same Lakeflow job retrains it on a schedule, with a rollback to the previous version.

Because each step runs under Unity Catalog grants, an underwriter who can't see a borrower's file today can't read it through the training table, the checkpoints or the model's logs either.

How TechFabric helps you train your own models

Choosing the use case, labelling the data, agreeing what "better" means and keeping the model accurate after launch are still your team's work, and AI Runtime doesn't do any of it. This is how we'd work alongside your engineers on those parts:

  • Pick the use case worth training. We'd start with an AI Readiness Assessment of your data, permissions and platform, and use it to decide which use case is worth training and which ones prompting and retrieval already solve.
  • Get the data ready in Unity Catalog. We'd build the Lakeflow pipelines that turn operational data into labelled, versioned training tables, with the Unity Catalog grants your security team will sign off.
  • Build the training and evaluation harness. We'd write the AI Runtime jobs, MLflow tracking and an evaluation gate with TechFabric Experiments, so each candidate model has to beat the one in production before it ships.
  • Keep GPU spend predictable. We'd tie each workload's usage in system tables to a project and a budget, and size runs before they scale up, as part of our Databricks cost optimization work.
  • Take it to production and hand it over. Our Databricks MLOps practice covers MLflow, Model Serving, drift monitoring and retraining, and the goal is that your engineers can run and extend it without us.

Where to start

Pick one model. Choose a task with a clear owner, a labelled dataset you already have, and a held-out evaluation set the business agrees on. Measure what prompting and retrieval achieve on that set first, then run a single fine-tuning experiment on 1xH100 with a fixed GPU budget, and compare.

Bring the person who owns the use case and the person who approves access to its data into the first conversation, because together they decide whether the model can ship. If you'd like help scoping that first model, our machine learning team can run it with you.

Frequently asked questions

What is Databricks AI Runtime?

AI Runtime is Databricks' serverless GPU compute for training and fine-tuning models, generally available since 21 October 2026. Compared with running your own GPU cluster, the capacity, cluster setup and governance come from Databricks and Unity Catalog, and your team keeps the training code, the data and the evaluation.

Does our data leave Unity Catalog when we train on AI Runtime?

Training data, checkpoints and logs stay under the Unity Catalog grants you already use, with PrivateLink and serverless egress controls available. The exception to plan for is capacity. During high demand, AI Runtime may use GPUs in another region, which the documentation says can incur egress costs, so we'd check the availability-by-geography rules before training on data with residency requirements.

How do we control what GPU training costs on Databricks?

Usage and cost are recorded in Databricks billing system tables at the workload level, so each training workload can be tied to a project and a budget. In practice we'd start on the smallest accelerator that fits, scale up only when the evaluation justifies it, checkpoint long runs, and use Lakeflow's priority and pre-emption for scheduled jobs.

How long can a training run on AI Runtime last?

According to the documentation, CLI jobs can run up to 14 days, and connected notebook sessions and notebook jobs up to two days. For longer training, Databricks recommends checkpointing and restarting the run when it reaches the limit, and we'd checkpoint any run long enough that losing it would cost a day of GPU time.

Do we need to fine-tune a model, or is retrieval enough?

Often retrieval is enough. We'd measure prompting and retrieval against a fixed evaluation set first, and fine-tune only when they miss the quality bar and you have labelled examples to learn from. In our view, fine-tuning is worth considering for narrow, high-volume tasks, your own document formats, tight latency targets, or data that can't go to an outside model provider.