Deploying DeepSeek R1 on Databricks: A Step-by-Step Guide

Registering DeepSeek R1 Model on Databricks

The next step is to register the model. 

• Once the workspace is created, click the “+ New” button and select the notebook to launch it.

• Create a Jupyter cell and run the following code:

1. Install the required Python packages:

%%capture
!pip install torch transformers mlflow accelerate torchvision
%restart_python

2. Load the tokens, configurations, and model from the Hugging Face repository:

import pandas as pd
import mlflow
import mlflow.transformers
import torch
from mlflow.models.signature import infer_signature
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig, pipeline


# Specify the model from HuggingFace transformers
model_name = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"


# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(model_name)
config = AutoConfig.from_pretrained(model_name)


model = AutoModelForCausalLM.from_pretrained(
    model_name,
    config=config,
    torch_dtype=torch.float16  
)

3. Test the loaded model by providing it with a sample prompt. 

We will also create a signature that we will add during the model registration:

text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
example_prompt = "How does a computer work?"
example_inputs = pd.DataFrame({"inputs": [example_prompt]})
example_outputs = text_generator(example_prompt, max_length=200)
signature = infer_signature(example_inputs, example_outputs)


print(example_outputs)

The model should be working perfectly!

The output it showed for me is:

[{'generated_text': "How does a computer work? What is the computer? What is the computer used for? What is the computer used for in real life?\n\nI need to answer this question, but I need to do it step by step. I need to start with the very basic level and build up from there. I need to make sure I understand each concept before moving on. I need to use a lot of examples to explain each idea. I need to write my thoughts as if I'm explaining them to someone else, but I need to make sure I understand how to structure the answer properly.\n\nOkay, let's start with the basic level. What is a computer? It's an electronic device, right? And it has a central processing unit (CPU) that does the processing. But I think the central processing unit is more efficient, so maybe it's the CPU. Then, it has memory and storage. I remember that memory is like RAM and storage is like ROM. But wait, I think"}]

4. Set up the conda environment with the correct Python version and the required Python packages. 

This will help us create the container that has all the necessary tools to run the model:

conda_env = {
    "name": "mlflow-env",
    "channels": ["defaults", "conda-forge"],
    "dependencies": [
        "python=3.11",
        "pip",
        {
            "pip": [
                "mlflow",
                "transformers",
                "accelerate",
                "torch",
                "torchvision"
            ]
        }
    ]
}

5. Register the model 

Provide the mlflow.transformers.log_model function with the text generation pipeline, artifact path, signature, example input, model name, and conda environment:

with mlflow.start_run() as run:
    mlflow.transformers.log_model(
        transformers_model=text_generator,
        artifact_path="deepseek_model",
        signature=signature,
        input_example=example_inputs,
        registered_model_name="deepseek_r1_llama_8b",
        conda_env=conda_env
    )

It will take a few minutes to register the model.

Learn about the power of Databricks Lakehouse and enhance your data engineering and machine learning skills by taking the Data Management in Databricks course.

Deploying Deepseek on Databricks

It’s time to deploy the model!

• Once the model is registered, click the “Models” tab at the bottom left of the Databricks dashboard. 
• Select your model and click on the “Serve this model” button.

• Enter the endpoint serving name, select the compute options, and set the concurrency before clicking the “Create” button.

It will take a few minutes to set up the serving endpoint. The status will turn green once the setup is complete and display “Ready.”

Additionally, you can fine-tune the DeepSeek R1 on a custom dataset before registering the model. Follow the Fine-Tuning DeepSeek R1 tutorial to learn everything you need to know about this process.

Accessing the Deployed Model

There are many ways to access or use this model. 

1. You can write your question in the Databricks playground, which will generate a response in the browser. 
2. Additionally, you can access the model using CURL, the Python requests package, or SQL commands.

To get started, we will write the question using the browser and generate the response.

To access the model locally or integrate it into your application, you first need to generate a Databricks API key. 

• Go to the settings.
• Click on "Developer" and then generate a new token. 
• Copy this token and create a local environment variable.

• Use the CURL command and provide it with the API token, ensuring that the prompt and server endpoint URL are in the proper JSON format. Here’s an example:

curl \
  -u token:$DATABRICKS_TOKEN \
  -X POST \
  -H "Content-Type: application/json" \
  -d '{"dataframe_split": {"columns": ["inputs"], "data": [["How does a computer work?"]]}}' \
  https://dbc-85e34089.cloud.databricks.com/serving-endpoints/deepseek_r1/invocations

When you run the command, generating the response will take a few seconds. It's that simple!

Check out the DeepSeek R1 vs V3 blog to learn about the best large language models available from DeepSeek.

Conclusion

Registering and deploying the DeepSeek R1 model on Databricks is straightforward. You can even register and deploy the large model using a CPU cluster or a local CPU machine, all without incurring any costs. However, running the model on a CPU can be slow and requires patience, especially when building the docker image.

In this tutorial, we covered the entire model deployment process step by step. We started by setting up Databricks and registering the pre-trained DeepSeek Distilled R1 model in the Databricks Model Registry. Then, we used the Databricks dashboard to deploy the model. Finally, we tested the deployed model and demonstrated how to use it locally with a simple CURL command.

If you are new to AI and large language models, I recommend taking the Introduction to LLMs in Python course. This will help you build a solid foundation, understand key terminologies, and get started with working on advanced models like DeepSeek R1!