Supervised fine-tuning (Full) and serving of Qwen3.5-0.8B

Open in Databricks

Fine-tune the compact open-weight Qwen3.5-0.8B-Base large language model on AI Runtime (serverless GPU), then deploy it behind a Model Serving endpoint. This example runs end-to-end on a single H100 GPU and shows you how to:

Run supervised fine-tuning (SFT) with TRL's SFTTrainer on an instruction-following dataset
Compare model responses before and after fine-tuning to see the effect of SFT
Register the fine-tuned model in Unity Catalog for governance and deployment
Serve the model behind a Custom Foundation Model endpoint running a vLLM OpenAI-compatible server

Key concepts:

Supervised fine-tuning (SFT): Continues training a base model on curated instruction/response pairs so it follows instructions in the target style
TRL: A library for supervised fine-tuning and reinforcement learning of language models
Custom Foundation Model serving: Serves your own fine-tuned LLM weights on GPU-backed Model Serving with an OpenAI-compatible API

note

This example requires the AI Runtime environment version 6 or above (the serving step uses vLLM and flashinfer, which ship in v6).

Connect to serverless GPU compute

This notebook requires serverless GPU compute. To connect:

Click the notebook's compute selector in the top right and select Serverless GPU.
On the right side, click the environment button.
Select H100 as the Accelerator.
Choose AI v6 from the base environment.
Click Apply.

Configuration

The next cell defines widgets for the Unity Catalog location where the fine-tuned model is registered. The model is registered as {uc_catalog}.{uc_schema}.{uc_model_name}, and the serving endpoint is named {uc_model_name}-endpoint.

Set deploy_endpoint to false to stop after model registration and the local vLLM test, without deploying a managed serving endpoint.

Python
dbutils.widgets.text("uc_catalog", "main")
dbutils.widgets.text("uc_schema", "default")
dbutils.widgets.text("uc_model_name", "qwen3_5_0_8b_sft")
dbutils.widgets.dropdown("deploy_endpoint", "true", ["true", "false"])

UC_CATALOG = dbutils.widgets.get("uc_catalog")
UC_SCHEMA = dbutils.widgets.get("uc_schema")
UC_MODEL_NAME_BASE = dbutils.widgets.get("uc_model_name")
# Whether to deploy the managed serving endpoint (Steps 8-9). Set to "false" to stop after
# registration and the local vLLM test (for example, on workspaces where entrypoint-based
# Custom Foundation Model serving is not enabled).
DEPLOY_ENDPOINT = dbutils.widgets.get("deploy_endpoint").lower() == "true"

print(f"UC_CATALOG: {UC_CATALOG}")
print(f"UC_SCHEMA: {UC_SCHEMA}")
print(f"UC_MODEL_NAME: {UC_MODEL_NAME_BASE}")
print(f"DEPLOY_ENDPOINT: {DEPLOY_ENDPOINT}")

Import libraries

Load the libraries used throughout the notebook. The AI Runtime v6 environment already includes torch, transformers, trl, and datasets, so no installation is required.

Python
import torch
import pandas as pd
from datasets import load_dataset, Dataset
import transformers
from transformers import TrainingArguments, AutoTokenizer, AutoModelForCausalLM
from trl import SFTTrainer, SFTConfig

Step 1: Load base model and tokenizer

Load the pre-trained Qwen3.5-0.8B-Base checkpoint from Hugging Face Hub.

Architecture: Qwen3 (decoder-only transformer, ~0.8B parameters)
"Base" checkpoint: no instruction-tuning applied yet. This is the model to fine-tune below.
The model is moved to GPU immediately after loading for faster inference and training

Python
# Qwen3.5-0.8B-Base: ~0.8B parameter decoder-only model.
# "Base" = no instruction-tuning yet; this is the checkpoint fine-tuned below.
model_name = "Qwen/Qwen3.5-0.8B-Base"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

Python
# Move all model weights to the GPU for faster inference and training.
model.to("cuda")

Configure tokenizer

Base checkpoints ship without a chat template. Define a minimal System / User / Assistant Jinja template so the tokenizer can correctly format prompts for single-turn and multi-turn conversations. Set pad_token = eos_token because the base vocabulary has no dedicated padding token.

Python
# Base checkpoints ship without a chat template.
# Define a minimal System / User / Assistant Jinja template so the tokenizer
# can format both single-turn and multi-turn conversations correctly.
if not tokenizer.chat_template:
    print("No chat template — applying default template")
    tokenizer.chat_template = """{% for message in messages %}
            {% if message['role'] == 'system' %}System: {{ message['content'] }}\n
            {% elif message['role'] == 'user' %}User: {{ message['content'] }}\n
            {% elif message['role'] == 'assistant' %}Assistant: {{ message['content'] }} <|endoftext|>
            {% endif %}
            {% endfor %}"""

# Set pad_token = eos_token because the base vocabulary has no dedicated pad token.
if not tokenizer.pad_token:
    print("No pad token — using eos_token as pad_token")
    tokenizer.pad_token = tokenizer.eos_token

Python
print(tokenizer.chat_template[0:100])
print(tokenizer.pad_token)

Step 2: Baseline inference (pre-SFT)

Run a quick sanity-check inference with the base (un-finetuned) model. The response here serves as a reference point. Compare it against the SFT model output in Step 5.

Python
# Build a single-turn chat in OpenAI-style message format.
# The tokenizer's chat template will convert this list into a formatted prompt string.
messages = []
user_message = "Give me a one-sentence introduction to LLMs."
messages.append({"role": "user", "content": user_message})
messages

Python
# Render the message list into a raw text string using the chat template.
# add_generation_prompt=True appends the "Assistant:" prefix to trigger generation.
# enable_thinking=False disables Qwen3's chain-of-thought reasoning mode.
prompt = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
    enable_thinking=False,
)
prompt

Python
# Tokenize the prompt string and move tensors to the same device as the model.
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
inputs

Python
max_new_tokens = 100
with torch.no_grad():          # no gradient tracking needed during inference
    outputs = model.generate(
        **inputs,              # pass tokenized prompt (input_ids + attention_mask)
        max_new_tokens=max_new_tokens,
        do_sample=False,       # greedy decoding — deterministic output
        pad_token_id=tokenizer.eos_token_id,
        eos_token_id=tokenizer.eos_token_id,
    )

outputs

Python
# Slice off the prompt tokens; decode only the newly generated portion.
input_len = inputs["input_ids"].shape[1]
generated_ids = outputs[0][input_len:]
response = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
response

Helper functions

Reusable utilities defined for this notebook:

generate_responses: formats a prompt with the chat template and runs greedy decoding
test_model_with_questions: benchmarks a list of questions and prints model outputs side-by-side
load_model_and_tokenizer: loads model + tokenizer with optional GPU placement and template patching
display_dataset: renders the first 3 rows of a chat-format dataset as a readable table

Python
def generate_responses(model, tokenizer, user_message, system_message=None,
                       max_new_tokens=100):
    # Format chat using tokenizer's chat template
    messages = []
    if system_message:
        messages.append({"role": "system", "content": system_message})

    # Assume the data are all single-turn conversations
    messages.append({"role": "user", "content": user_message})

    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,
    )

    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    # Recommended to use vllm, sglang or TensorRT
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    input_len = inputs["input_ids"].shape[1]
    generated_ids = outputs[0][input_len:]
    response = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()

    return response

Python
def test_model_with_questions(model, tokenizer, questions,
                              system_message=None, title="Model Output"):
    print(f"\n=== {title} ===")
    for i, question in enumerate(questions, 1):
        response = generate_responses(model, tokenizer, question,
                                      system_message)
        print(f"\nModel Input {i}:\n{question}\nModel Output {i}:\n{response}\n")

Python
def load_model_and_tokenizer(model_name, use_gpu = False):

    # Load base model and tokenizer
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)

    if use_gpu:
        model.to("cuda")

    if not tokenizer.chat_template:
        tokenizer.chat_template = """{% for message in messages %}
                {% if message['role'] == 'system' %}System: {{ message['content'] }}\n
                {% elif message['role'] == 'user' %}User: {{ message['content'] }}\n
                {% elif message['role'] == 'assistant' %}Assistant: {{ message['content'] }} <|endoftext|>
                {% endif %}
                {% endfor %}"""

    # Tokenizer config
    if not tokenizer.pad_token:
        tokenizer.pad_token = tokenizer.eos_token

    return model, tokenizer

Python
def display_dataset(dataset):
    # Visualize the dataset
    rows = []
    for i in range(3):
        example = dataset[i]
        user_msg = next(m['content'] for m in example['messages']
                        if m['role'] == 'user')
        assistant_msg = next(m['content'] for m in example['messages']
                             if m['role'] == 'assistant')
        rows.append({
            'User Prompt': user_msg,
            'Assistant Response': assistant_msg
        })

    # Display as table
    df = pd.DataFrame(rows)
    pd.set_option('display.max_colwidth', None)  # Avoid truncating long strings
    display(df)

Step 3: Load the training dataset

Load banghua/DL-SFT-Dataset from the Hugging Face Hub, an instruction-following dataset from the DeepLearning.AI SFT course. Each example is a messages list with user and assistant turns.

This example uses a 100-example subset to keep training time short. Increase the subset size for real fine-tuning.

Python
train_dataset = load_dataset("banghua/DL-SFT-Dataset")['train']

train_dataset=train_dataset.select(range(100))

display_dataset(train_dataset)

Step 4: SFT fine-tuning

Use TRL's SFTTrainer to run Supervised Fine-Tuning. Key hyperparameter choices for this demo:

Parameter	Value	Notes
`learning_rate`	`8e-5`	Standard starting point for SFT on small models
`num_train_epochs`	`1`	Single pass for demo; increase for real training
`per_device_train_batch_size`	`1`	Tune with `gradient_accumulation_steps`
`gradient_accumulation_steps`	`8`	Effective batch size = 1 × 8 = 8
`gradient_checkpointing`	`False`	Disabled for speed; enable to reduce VRAM on larger models

Parameter	Value	Notes
`learning_rate`	`8e-5`	Standard starting point for SFT on small models
`num_train_epochs`	`1`	Single pass for demo; increase for real training
`per_device_train_batch_size`	`1`	Tune with `gradient_accumulation_steps`
`gradient_accumulation_steps`	`8`	Effective batch size = 1 × 8 = 8
`gradient_checkpointing`	`False`	Disabled for speed; enable to reduce VRAM on larger models

Python
# SFTConfig is a superset of HuggingFace TrainingArguments with SFT-specific defaults.
sft_config = SFTConfig(
    # --- Training hyperparameters ---
    learning_rate=8e-5,               # standard starting point for SFT on small models
    num_train_epochs=1,               # single pass for demo; increase for real training
    per_device_train_batch_size=1,    # fits H100 VRAM; tune together with gradient_accumulation_steps
    gradient_accumulation_steps=8,    # effective batch size = 1 × 8 = 8
    gradient_checkpointing=False,     # disable for speed; enable to reduce VRAM on larger models
    logging_steps=2,

    # --- Logging ---
    report_to=[],                     # disable W&B / MLflow / etc.
    logging_strategy="steps",
    logging_first_step=True,

    # --- Checkpointing ---
    output_dir="./checkpoints/tiny-finetune-exp1",
    run_name="tiny-finetune-exp1-run",   # must differ from output_dir to avoid W&B conflicts
    save_strategy="no",                  # set to "epoch" to persist a final checkpoint
    save_total_limit=1
)

Python
# SFTTrainer handles dataset formatting, tokenization, and response-label masking
# automatically based on the tokenizer's chat template.
sft_trainer = SFTTrainer(
    model=model,
    args=sft_config,
    train_dataset=train_dataset,
    processing_class=tokenizer    # replaces the deprecated tokenizer= arg in TRL 0.12+
)

# Run training — loss should decrease over the single epoch on this 100-example subset.
sft_trainer.train()

Step 5: Post-SFT evaluation

Compare the fine-tuned model's responses on the same questions used in Step 2. Look for improved formatting and instruction-following style that reflects the training data distribution.

Python
questions = [
    "Calculate 1+1-1",
    "What's the difference between thread and process?"
]
test_model_with_questions(sft_trainer.model, tokenizer, questions,
                          title="Fine-tuned model output")

Step 6: Serve the fine-tuned model with vLLM

With training done, deploy the fine-tuned model behind a Databricks Custom Foundation Model serving endpoint that runs a vLLM OpenAI-compatible server.

A wrinkle worth knowing: Qwen/Qwen3.5-0.8B-Base is really a multimodal model, and AutoModelForCausalLM above loaded only its text backbone (Qwen3_5ForCausalLM). vLLM (AI env v6) can serve that text backbone natively, but a few accommodations are needed because it is the text half of a vision-language model:

Save the fine-tuned weights with the composite Qwen3_5Config (the one that carries vision_config) plus the model's original processor files, otherwise vLLM's processor rejects the checkpoint.
At launch, tell vLLM there are zero images/videos (--limit-mm-per-prompt) so it never touches the (absent) vision tower, run the linear-attention (GDN) kernels through Triton (--gdn-prefill-backend triton) so nothing has to JIT-compile with ninja/nvcc, and route sampling through native torch (VLLM_USE_FLASHINFER_SAMPLER=0).

Everything below runs on the same Serverless GPU (H100) + AI environment v6 session.

Save the fine-tuned model for serving

The trainer keeps the model in memory only (save_strategy="no"), so persist it before restarting Python:

save_pretrained the fine-tuned text weights (named model.*).
Save the model's original processor (preprocessor_config.json, etc.). The composite config declares a vision component, so vLLM insists on it. Then re-save the new tokenizer on top so the custom chat template wins.
Overwrite the flat text config with the composite Qwen3_5Config, with architectures pinned to Qwen3_5ForCausalLM so vLLM loads the text model (not the full VL model).

Run this while sft_trainer, tokenizer, and model_name are still in scope.

Python
import os, tempfile
from transformers import AutoConfig, AutoProcessor

# Local-disk working dir. Use a fixed path (not a random tmpdir) so it survives %restart_python below;
# ARTIFACTS_PATH is a relative basename because the vLLM entrypoint's --model must match it both here
# and inside the packaged model's artifacts/ dir at serving time.
WORKDIR = os.path.join(tempfile.gettempdir(), "sft_serve")
ARTIFACTS_PATH = "qwen3_sft"
SAVE_DIR = os.path.join(WORKDIR, ARTIFACTS_PATH)
os.makedirs(SAVE_DIR, exist_ok=True)

# 1. Fine-tuned text backbone.
sft_trainer.model.save_pretrained(SAVE_DIR)

# 2. Original processor files, then the new tokenizer (with the custom chat template) on top.
try:
    AutoProcessor.from_pretrained(model_name).save_pretrained(SAVE_DIR)
except Exception as e:
    print("processor save skipped:", e)
tokenizer.save_pretrained(SAVE_DIR)

# 3. Composite Qwen3_5Config (text_config + vision_config), arch pinned to the text model.
cfg = AutoConfig.from_pretrained(model_name)
cfg.architectures = ["Qwen3_5ForCausalLM"]
cfg.save_pretrained(SAVE_DIR)

print("saved ->", SAVE_DIR)
print("config architectures:", cfg.architectures, "| model_type:", cfg.model_type,
      "| has vision_config:", hasattr(cfg, "vision_config"))

Restart Python to free the H100

Nothing needs installing: AI env v6 already ships flashinfer and its prebuilt flashinfer-cubin kernels, the GDN (linear-attention) prefill runs on Triton, and sampling runs on native torch (both set in the launch flags below), so no kernel JIT-compiles at startup.

The one thing needed is to release the GPU memory the trainer is holding so vLLM can claim it. %restart_python does that; the saved checkpoint on local disk survives the restart (same driver node).

Python
%restart_python

Serving configuration

The top of the next cell holds the values you set for your workspace: the model/endpoint names and endpoint sizing. The rest is internal wiring you can leave as-is. At minimum, set UC_MODEL_NAME to a Unity Catalog path you can write to before running the register/deploy cells.

Python
from databricks.sdk.service.serving import ServingModelWorkloadType

# Re-read the widgets: %restart_python reset the Python process, but the widget
# values set at the top of the notebook persist and can be read again here.
UC_CATALOG = dbutils.widgets.get("uc_catalog")
UC_SCHEMA = dbutils.widgets.get("uc_schema")
UC_MODEL_NAME_BASE = dbutils.widgets.get("uc_model_name")
DEPLOY_ENDPOINT = dbutils.widgets.get("deploy_endpoint").lower() == "true"

# --- Model / endpoint names ---
UC_MODEL_NAME = f"{UC_CATALOG}.{UC_SCHEMA}.{UC_MODEL_NAME_BASE}"   # Unity Catalog catalog.schema.model
ENDPOINT_NAME = f"{UC_MODEL_NAME_BASE}-endpoint"                   # serving endpoint name; unique per workspace
SERVED_MODEL_NAME = UC_MODEL_NAME_BASE                            # name vLLM exposes the model under

# --- Endpoint sizing (adjust if needed) ---
# --gdn-prefill-backend triton (see the entrypoint) JITs the GDN kernels for whatever GPU the pod
# lands on, so this is not pinned to Hopper. GPU_MEDIUM fits the 0.8B model; verify at deploy.
WORKLOAD_TYPE = ServingModelWorkloadType.GPU_MEDIUM
WORKLOAD_SIZE = "Small"
SCALE_TO_ZERO_ENABLED = True

# --- Internal wiring: leave as-is ---
import os, tempfile

# Local-disk working dir; must match the save cell (survives %restart_python; same driver node).
WORKDIR = os.path.join(tempfile.gettempdir(), "sft_serve")
ARTIFACTS_PATH = "qwen3_sft"          # relative basename; entrypoint --model resolves to ./qwen3_sft
os.chdir(WORKDIR)                     # so `--model qwen3_sft` works locally and matches artifacts/ at serving

# Allowlisted ports for serverless GPU notebooks are 3000-3999. Model Serving requires 8080.
LOCAL_PORT = 3080
SERVING_PORT = 8080

# vLLM tuning.
DTYPE = "float16"                     # model is bf16-native; this matches what shipped
MAX_MODEL_LEN = 8192                  # keep <= the model's max_position_embeddings (config.json)
GPU_MEMORY_UTILIZATION = 0.85

Test the fine-tuned model locally with vLLM

Boot a vLLM OpenAI server against the saved checkpoint and smoke-test it before spending ~40 minutes on an endpoint deploy. A bad checkpoint or flag fails here in seconds instead.

The entrypoint() command is defined once and reused for both the local test (LOCAL_PORT) and the serving endpoint (SERVING_PORT); the exact same string is stored in the model's metadata and re-run inside the serving container.

Python
def entrypoint(port: int) -> str:
    args = [
        "VLLM_USE_FLASHINFER_SAMPLER=0",   # native torch sampling; avoids the flashinfer sampler JIT
        "python", "-u", "-m", "vllm.entrypoints.openai.api_server",
        "--model", ARTIFACTS_PATH,
        "--served-model-name", SERVED_MODEL_NAME,
        "--host", "0.0.0.0",
        "--port", str(port),
        "--dtype", DTYPE,
        "--max-model-len", str(MAX_MODEL_LEN),
        "--gpu-memory-utilization", str(GPU_MEMORY_UTILIZATION),
        # Run GDN (linear-attention) prefill through Triton (its own bundled compiler) instead of
        # flashinfer's JIT kernel, so no ninja/nvcc is needed here or in the serving pod.
        "--gdn-prefill-backend", "triton",
        # Text-only serving of a multimodal-config model: allow zero images/videos so vLLM never
        # profiles or exercises the vision tower (absent from the text weights).
        "--limit-mm-per-prompt", "'{\"image\": 0, \"video\": 0}'",
    ]
    return " ".join(args)

Python
# Start the vLLM server in the background; logs stream to process.log.
import subprocess

log = open("process.log", "w")
subprocess.Popen(
    ["bash", "-lc", entrypoint(LOCAL_PORT)],
    stdout=log,
    stderr=subprocess.STDOUT,
    text=True,
    start_new_session=True,
)

Python
%sh
# Tail logs until vLLM is ready. If this hangs, vLLM startup probably hit an error (read process.log).
tail -f process.log | sed -u '/Application startup complete/q'

Python
# Smoke test (sync): the endpoint speaks the OpenAI chat schema at /invocations.
import requests

resp = requests.post(f"http://localhost:{LOCAL_PORT}/invocations", json={"messages": [{"role": "user", "content": "Hello"}]})
resp.json()["choices"][0]["message"]["content"]

Python
# Smoke test (streaming): vLLM streams completions as Server-Sent Events — one JSON chunk per
# `data: ` line, terminated by `data: [DONE]`.
import requests
import json

resp = requests.post(
    f"http://localhost:{LOCAL_PORT}/invocations",
    json={"messages": [{"role": "user", "content": "Tell me a story that is about 300 words!"}], "stream": True},
    stream=True,
)

for line in resp.iter_lines():
    if not line:
        continue
    if line == b"data: [DONE]":
        break
    if line.startswith(b"data: "):
        data = json.loads(line[6:])
        delta = data["choices"][0].get("delta", {})
        if "content" in delta:
            print(delta["content"], end="", flush=True)

Python
%sh
# Stop the local server before logging/registering (the endpoint runs its own copy).
pkill -f vllm.entrypoints.openai.api_server

Step 7: Log and register the model

Log an MLflow ChatModel whose metadata carries task = llm/v1/chat and the vLLM entrypoint command. Serving runs that entrypoint, not python_model.predict, so the class body is just a required placeholder.

Registering with env_pack="databricks_model_serving" produces the Serverless Optimized Deployment (SOD) artifacts (weights + packaged environment) that Custom LLM Serving requires. Logging must happen from this Serverless GPU runtime so the correct GPU dependencies are packaged.

Python
import os
import mlflow
from mlflow.pyfunc.model import ChatModel, ChatCompletionResponse

# Required placeholder. Serving runs the entrypoint, not python_model.predict.
class LLMModel(ChatModel):
    def predict(self, context, messages, params):
        return ChatCompletionResponse.from_dict({"choices": []})

# You must log and register from a Serverless GPU runtime, otherwise the model is packaged
# with CPU deps and the GPU serving endpoint fails to start.
if not os.environ.get("DATABRICKS_ACCELERATOR"):
    raise RuntimeError(
        "This model MUST be logged+registered from a serverless GPU runtime, otherwise the correct dependencies will not be packaged for serving."
    )

model_info = mlflow.pyfunc.log_model(
    name=SERVED_MODEL_NAME,
    python_model=LLMModel(),
    artifacts={
        "model_dir": ARTIFACTS_PATH,
    },
    metadata={
        "task": "llm/v1/chat",
        "entrypoint": entrypoint(SERVING_PORT),
    },
    # Pin whatever mlflow ships in AI env v6 (rather than a hardcoded version).
    extra_pip_requirements=[f"mlflow=={mlflow.__version__}"],
)
model_info.model_uri

Python
import mlflow

# env_pack is required. Custom LLM Serving depends on Serverless Optimized Deployments (SOD).
# The endpoint will not work without it.
# https://docs.databricks.com/aws/en/machine-learning/model-serving/serverless-optimized-deployments
model_version = mlflow.register_model(model_info.model_uri, UC_MODEL_NAME, env_pack="databricks_model_serving")

Step 8: Create the serving endpoint

Create a Custom Foundation Model endpoint that serves the registered version. create_and_wait blocks until the endpoint is ready (up to 40 minutes for the first deploy; the container downloads the SOD artifacts and boots vLLM).

Deploy the endpoint (optional)

The remaining steps deploy a managed serving endpoint. They require entrypoint-based Custom Foundation Model serving (Serverless Optimized Deployment). If deploy_endpoint is false, the notebook stops here.

Python
if not DEPLOY_ENDPOINT:
    dbutils.notebook.exit("deploy_endpoint=false: skipping managed serving endpoint deployment")

Python
from databricks.sdk import WorkspaceClient
from datetime import timedelta
from databricks.sdk.service.serving import EndpointCoreConfigInput, ServedEntityInput

served_entities = [
    ServedEntityInput(
        entity_name=UC_MODEL_NAME,
        entity_version=str(model_version.version),
        workload_type=WORKLOAD_TYPE,
        workload_size=WORKLOAD_SIZE,
        scale_to_zero_enabled=SCALE_TO_ZERO_ENABLED,
    )
]

w = WorkspaceClient()

# Create the endpoint, or update it in place if an endpoint of this name already exists,
# so the notebook is safe to re-run. The first deploy can take up to ~40 minutes while the
# container downloads the Serverless Optimized Deployment artifacts and boots vLLM.
existing = {e.name for e in w.serving_endpoints.list()}
if ENDPOINT_NAME in existing:
    print(f"Updating existing endpoint: {ENDPOINT_NAME}")
    w.serving_endpoints.update_config_and_wait(
        name=ENDPOINT_NAME, served_entities=served_entities, timeout=timedelta(minutes=40)
    )
else:
    print(f"Creating endpoint: {ENDPOINT_NAME}")
    config = EndpointCoreConfigInput(name=ENDPOINT_NAME, served_entities=served_entities)
    w.serving_endpoints.create_and_wait(
        name=ENDPOINT_NAME, config=config, timeout=timedelta(minutes=40)
    )

Step 9: Query the endpoint

Once the endpoint is ready, query it three ways: the Databricks SDK, the OpenAI client, and the OpenAI client with streaming.

Python
# Query using the Databricks SDK.
from databricks.sdk import WorkspaceClient
from databricks.sdk.service.serving import ChatMessage, ChatMessageRole

w = WorkspaceClient()

resp = w.serving_endpoints.query(
    name=ENDPOINT_NAME,
    messages=[ChatMessage(role=ChatMessageRole.USER, content="Hi, what model are you?")],
)

print(resp.choices[0].message.content)

Python
# Query using the OpenAI client.
from openai import OpenAI

DATABRICKS_HOST = dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiUrl().get()
DATABRICKS_TOKEN = dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiToken().get()

client = OpenAI(
    api_key=DATABRICKS_TOKEN,
    base_url=f"{DATABRICKS_HOST}/serving-endpoints",
)

response = client.chat.completions.create(
    model=ENDPOINT_NAME,
    messages=[
        {"role": "user", "content": "Hello"},
    ],
)
print(response.choices[0].message.content)

Python
# Query using the OpenAI client (streaming).
stream = client.chat.completions.create(
    model=ENDPOINT_NAME,
    messages=[
        {"role": "user", "content": "Hello, tell me a 200 word story"},
    ],
    stream=True,
)

for event in stream:
    delta = event.choices[0].delta
    print(delta.content, end="|")

Next steps

Now that you've fine-tuned, registered, and served your model, you can:

Query the endpoint from your applications: Query foundation models and external models
Learn more about custom LLM serving: Serve custom LLMs
Optimize serverless GPU usage: Best practices for AI Runtime
Troubleshoot issues: Troubleshoot issues on AI Runtime

Example notebook

Supervised fine-tuning (Full) and serving of Qwen3.5-0.8B

Open notebook in new tab Open in Databricks

Connect to serverless GPU compute​

Configuration​

Import libraries​

Step 1: Load base model and tokenizer​

Configure tokenizer​

Step 2: Baseline inference (pre-SFT)​

Helper functions​

Step 3: Load the training dataset​

Step 4: SFT fine-tuning​

Step 5: Post-SFT evaluation​

Step 6: Serve the fine-tuned model with vLLM​

Save the fine-tuned model for serving​

Restart Python to free the H100​

Serving configuration​

Test the fine-tuned model locally with vLLM​

Step 7: Log and register the model​

Step 8: Create the serving endpoint​

Deploy the endpoint (optional)​

Step 9: Query the endpoint​

Next steps​

Example notebook​

Supervised fine-tuning (Full) and serving of Qwen3.5-0.8B

Connect to serverless GPU compute

Configuration

Import libraries

Step 1: Load base model and tokenizer

Configure tokenizer

Step 2: Baseline inference (pre-SFT)

Helper functions

Step 3: Load the training dataset

Step 4: SFT fine-tuning

Step 5: Post-SFT evaluation

Step 6: Serve the fine-tuned model with vLLM

Save the fine-tuned model for serving

Restart Python to free the H100

Serving configuration

Test the fine-tuned model locally with vLLM

Step 7: Log and register the model

Step 8: Create the serving endpoint

Deploy the endpoint (optional)

Step 9: Query the endpoint

Next steps

Example notebook