ML and GPU Compute
This page covers model development and serving — not agent integration. For Genie, MCP, and Unity Gateway, start at AI Platform.
Choose your track
| Track | When to use | Compute |
|---|---|---|
| GPU compute for LLMs | Fine-tune, continue pretraining, or serve a custom or open-weight LLM | Serverless GPU / AI Runtime |
| Classic ML | sklearn, XGBoost, forecasting, ranking, and similar models | CPU or GPU Model Serving |
Hosted foundation models through Unity Gateway are pay-per-token and do not require partner GPU sizing.
Train and fine-tune
GPU / AI Runtime
LoRA, QLoRA, and full fine-tuning run on AI Runtime serverless GPU compute. Use this for custom and open-weight LLMs.
Documentation: AI Runtime
Classic ML training and experiment tracking
Log classic ML experiments (sklearn, XGBoost, forecasting) with MLflow.
Documentation: MLflow
Register and govern
Log runs and checkpoints with MLflow; register models in Unity Catalog for access control, lineage, and serving. This applies to both GPU and classic ML models.
Documentation: Manage model lifecycle in Unity Catalog | Experiment tracking and observability | MLflow
Serve and infer
Custom Model Serving
Deploy a Python or custom model as a production API on CPU or GPU. Match serving GPU or CPU to how the model was trained. For hosted foundation models, use Foundation Model APIs through Unity Gateway instead.
Documentation: Custom models
Batch inference and AI Functions
Apply models to data in SQL or Python with AI Functions (ai_query and task-specific functions), including batch inference.
Documentation: AI Functions
Classic batch and online serving
Serve registered classic ML models through Model Serving or batch jobs.
Documentation: Model Serving
Features and training data
Use Unity Catalog–governed features for training and inference. Partners can read feature tables for real-time or batch inference and write engineered features back into Unity Catalog for shared use across models and teams.
Documentation: Feature Store
Partner patterns
Data annotation and labeling
Annotation partners read data from Databricks, apply labels, and write results back. Labeled sets can also feed LLM evaluation and GPU fine-tuning.
Integration principles:
- Read input data from Unity Catalog: Tables for structured data, Volumes for unstructured files (PDFs, images, audio)
- Access via SQL, Databricks SQL REST APIs, SDKs, JDBC/ODBC, or connectors
- Write labeled output back into Unity Catalog Delta tables using SQL (
PUT,MERGE,COPY INTO) or APIs/SDKs - Store raw files into Volumes, extracted text/metadata/labels into Tables
Extraction and classification
Document AI partners that extract metadata and classify documents (no agent loop) store raw files in Volumes and extracted text, metadata, and classifications in Unity Catalog tables.
Documentation: Model Serving
Forecasting, ranking, and predictive models
Train, register, and serve predictive models through the lifecycle above.
Documentation: MLflow | Model Serving
What's next
- Return to AI Platform
- Unity Gateway for hosted foundation models (no GPU to manage)
- Review the integration requirements for foundational guidance
- Learn about telemetry and attribution for usage tracking
- Explore other Partner product categories for additional integration patterns