Swiftorial Logo
Home
Swift Lessons
AI Tools
Learn More
Career
Resources
Deploying Fine-Tuned LLMs with FastAPI and Docker

Deploying Fine-Tuned LLMs with FastAPI and Docker

A practical guide for developers on taking fine-tuned Large Language Models from development to production, leveraging FastAPI for high-performance APIs and Docker for consistent, scalable deployments.

1. Introduction: Bridging the Gap to Production

Fine-tuning Large Language Models (LLMs) allows for powerful specialization, but the journey from a successful training run in a notebook to a robust, scalable, and reliable production service is a significant leap. To truly unlock the value of your fine-tuned LLM, you need to expose it as an API that can handle real-time inference requests efficiently. This guide focuses on a popular and effective combination for this task: **FastAPI** for building the API endpoint and **Docker** for packaging and deploying your application consistently. Together, they provide a solid foundation for operationalizing your specialized LLMs.

2. Why FastAPI and Docker for LLM Deployment?

a. FastAPI: High Performance and Developer Experience

  • **Speed:** FastAPI is built on Starlette (for web parts) and Pydantic (for data validation), making it incredibly fast and efficient, crucial for low-latency LLM inference.
  • **Asynchronous Support:** Natively supports `async`/`await`, allowing it to handle many concurrent requests efficiently, which is vital for high-throughput AI services.
  • **Automatic Docs:** Generates interactive API documentation (Swagger UI / OpenAPI) automatically, simplifying API testing and integration for other developers.
  • **Type Hinting:** Leverages Python type hints for data validation, auto-completion, and clear code.

b. Docker: Consistency and Portability

  • **Environment Isolation:** Docker containers package your application and all its dependencies (Python versions, libraries, CUDA drivers) into a single, isolated unit. This eliminates "it works on my machine" problems.
  • **Consistency:** Ensures that your application runs the same way, regardless of the underlying infrastructure (development, staging, production).
  • **Portability:** A Docker image can be easily moved and run on any system that supports Docker (local machine, cloud VMs, Kubernetes).
  • **Scalability:** Containers are the building blocks for modern orchestration systems like Kubernetes, enabling easy horizontal scaling.
# FastAPI + Docker = Fast, Reliable, Scalable LLM APIs
# FastAPI: The "engine" for your API logic.
# Docker: The "shipping container" for your entire application.

3. The Deployment Workflow: Step-by-Step

Let's outline the practical steps to deploy your fine-tuned LLM.

Step 1: Save Your Fine-Tuned Model and Tokenizer

After fine-tuning (e.g., using Hugging Face Transformers and PEFT/LoRA), save your model weights and tokenizer. For LoRA, you'll typically save only the small adapter weights, which are then loaded on top of the base model during inference.

# Conceptual: Saving your fine-tuned LoRA adapter
# lora_model.save_pretrained("./my_fine_tuned_lora_adapter")
# tokenizer.save_pretrained("./my_tokenizer")

Step 2: Create a FastAPI Application (`app.py`)

This Python script will define your API endpoint(s) and handle model loading and inference. It's crucial to load the model once when the application starts, not per request, for efficiency.

# app.py: Your FastAPI application
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel # Required if you fine-tuned with LoRA

app = FastAPI(
    title="Fine-Tuned LLM Inference API",
    description="API for a specialized Large Language Model.",
    version="1.0.0"
)

# Global variables to store model and tokenizer
# These will be loaded once when the app starts
tokenizer = None
model = None
device = "cuda" if torch.cuda.is_available() else "cpu"

# Define your base model ID and LoRA adapter path
BASE_MODEL_ID = "mistralai/Mistral-7B-v0.1" # Example: Replace with your base model
LORA_ADAPTER_PATH = "./my_fine_tuned_lora_adapter" # Path where your LoRA adapter is saved

# Pydantic model for request body validation
class InferenceRequest(BaseModel):
    text: str
    max_new_tokens: int = 100
    temperature: float = 0.7

# Pydantic model for response body
class InferenceResponse(BaseModel):
    generated_text: str

@app.on_event("startup")
async def load_model():
    """
    Load the fine-tuned model and tokenizer when the FastAPI application starts.
    This ensures the model is loaded only once.
    """
    global tokenizer, model

    print(f"Loading tokenizer from {BASE_MODEL_ID}...")
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)
    if tokenizer.pad_token is None:
        tokenizer.add_special_tokens({'pad_token': tokenizer.eos_token})
    print("Tokenizer loaded.")

    print(f"Loading base model {BASE_MODEL_ID} on {device}...")
    try:
        # Load base model, potentially with quantization (e.g., load_in_4bit for QLoRA)
        base_model = AutoModelForCausalLM.from_pretrained(
            BASE_MODEL_ID,
            torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32, # Use bfloat16 for GPU, float32 for CPU
            device_map="auto" if device == "cuda" else None, # Auto-distribute on GPU
            # For Flash Attention 2, if supported by model and installed:
            # attention_implementation="flash_attention_2"
        )

        # Load LoRA adapter weights on top of the base model
        print(f"Loading LoRA adapter from {LORA_ADAPTER_PATH}...")
        model = PeftModel.from_pretrained(base_model, LORA_ADAPTER_PATH)
        model.eval() # Set model to evaluation mode
        print("Model and LoRA adapter loaded successfully.")

    except Exception as e:
        print(f"Error loading model: {e}")
        # In a real app, you might want to raise an exception or set a flag
        # to indicate the service is not ready.
        raise HTTPException(status_code=500, detail=f"Failed to load model: {e}")


@app.post("/generate", response_model=InferenceResponse)
async def generate_text(request: InferenceRequest):
    """
    API endpoint to generate text using the fine-tuned LLM.
    """
    if model is None or tokenizer is None:
        raise HTTPException(status_code=503, detail="Model not loaded yet. Please try again.")

    try:
        # Prepare input for the model
        inputs = tokenizer(request.text, return_tensors="pt").to(device)

        # Generate output
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=request.max_new_tokens,
                temperature=request.temperature,
                pad_token_id=tokenizer.eos_token_id # Important for consistent generation
            )

        # Decode the generated text
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

        # You might want to post-process the generated text here
        # E.g., remove the input prompt from the output if it's a completion task
        # For instruction tuning, the model often generates the full response including the instruction
        # A common pattern is to find the instruction in the output and return only the response part.
        # Example for instruction tuning:
        # if "### Response:" in generated_text:
        #     generated_text = generated_text.split("### Response:", 1)[1].strip()

        return InferenceResponse(generated_text=generated_text)

    except Exception as e:
        print(f"Error during inference: {e}")
        raise HTTPException(status_code=500, detail=f"Inference failed: {e}")

# To run this locally (for testing):
# 1. Save the above code as `app.py`
# 2. Ensure you have your LoRA adapter saved at `./my_fine_tuned_lora_adapter`
# 3. Create a `requirements.txt` with:
#    fastapi
#    uvicorn[standard]
#    torch
#    transformers
#    peft
#    bitsandbytes # If using QLoRA
#    accelerate # If using accelerate for multi-GPU
#    xformers # If using Flash Attention 2
# 4. pip install -r requirements.txt
# 5. Run: `uvicorn app:app --host 0.0.0.0 --port 8000`
# 6. Access docs at http://localhost:8000/docs

Step 3: Dockerize Your Application (Dockerfile)

Create a `Dockerfile` in the same directory as `app.py` and `requirements.txt`. This file instructs Docker on how to build your container image.

# Dockerfile
# Use a Python base image with CUDA support for GPU inference
FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04

# Set environment variables
ENV PYTHONUNBUFFERED 1
ENV LC_ALL C.UTF-8
ENV LANG C.UTF-8

# Install Python and pip
RUN apt-get update && apt-get install -y --no-install-recommends \
    python3.10 \
    python3-pip \
    && rm -rf /var/lib/apt/lists/*

# Set working directory in the container
WORKDIR /app

# Copy requirements.txt and install Python dependencies
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy your application code and model weights (LoRA adapter)
COPY app.py .
# Assuming your LoRA adapter is in a folder named 'my_fine_tuned_lora_adapter'
COPY my_fine_tuned_lora_adapter ./my_fine_tuned_lora_adapter
# If you need the full base model weights in the container (less common for LoRA)
# COPY path/to/base_model ./base_model

# Expose the port FastAPI will run on
EXPOSE 8000

# Command to run the FastAPI application using Uvicorn
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]

Step 4: Build Your Docker Image

Navigate to your project directory in the terminal and build the Docker image:

# Build the Docker image
docker build -t my-llm-api .

Step 5: Run Your Docker Container

Once the image is built, you can run your container. For GPU inference, ensure Docker is configured with NVIDIA Container Toolkit.

# Run the Docker container (with GPU support)
docker run --gpus all -p 8000:8000 my-llm-api

Your API should now be accessible at `http://localhost:8000` (or your server's IP) and the interactive docs at `http://localhost:8000/docs`.

4. Key Considerations for Production Deployment

a. General Performance Optimization

  • **Flash Attention:** Ensure your model and environment support Flash Attention for faster attention computation, especially with long context windows.
  • **Batching:** Implement dynamic batching in your FastAPI application (e.g., using a queue) to group multiple incoming requests and process them in a single inference pass on the GPU, maximizing throughput.
  • **Model Compilation:** Explore tools like ONNX Runtime, TensorRT, or PyTorch Compile to optimize the model graph for faster inference on specific hardware.

b. Quantization for Edge Deployment

**Quantization** is a technique that reduces the precision of the numbers (weights and activations) used in a neural network, typically from 32-bit floating point (FP32) to lower precision formats like 16-bit floating point (FP16), 8-bit integer (INT8), or even lower (e.g., 4-bit integer). This process is crucial for deploying LLMs to **edge devices** (e.g., mobile phones, IoT devices, embedded systems) or environments with very limited computational resources.

Why Quantize for Edge?

  • **Reduced Memory Footprint:** Lower precision numbers require less memory to store the model weights, allowing larger models to fit on devices with limited RAM.
  • **Faster Inference:** Operations on lower precision data are often significantly faster on specialized hardware (e.g., mobile NPUs, embedded GPUs) that are optimized for integer arithmetic.
  • **Lower Power Consumption:** Less memory access and simpler computations lead to reduced power consumption, extending battery life for mobile or IoT devices.

How to Apply Quantization with Fine-Tuning:

The most common approach for LLMs is **Quantization-Aware Training (QAT)** or **Post-Training Quantization (PTQ)**. For fine-tuning, **QLoRA** (Quantized LoRA) has become a popular method, where the base LLM is loaded in a quantized format (e.g., 4-bit), and only the small LoRA adapters are trained in a higher precision. This allows fine-tuning huge models on consumer GPUs.

  • **During Fine-Tuning (QLoRA):**
    • **Action:** When loading your base model for fine-tuning, specify `load_in_4bit=True` (or `load_in_8bit=True`) using the `transformers` library. The `bitsandbytes` library handles the actual quantization. This means your fine-tuning process itself is memory-efficient.
    • **Example (from `app.py` startup):** 
      # Load base model with 4-bit quantization (QLoRA)
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL_ID,
    load_in_4bit=True, # This is the key for 4-bit quantization
    torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
    device_map="auto" if device == "cuda" else None,
)
# Ensure `prepare_model_for_kbit_training` is used if fine-tuning with QLoRA
# model = prepare_model_for_kbit_training(model) # This would be done during training
  • **For Inference on Edge (Post-Training Quantization - PTQ):**
    • **Action:** After fine-tuning, you might export the entire fine-tuned model (base model + LoRA adapter merged) to a format like ONNX or OpenVINO, and then apply further quantization (e.g., INT8) using specialized tools like ONNX Runtime, OpenVINO Toolkit, or TensorFlow Lite. This step typically requires a representative calibration dataset to minimize accuracy loss.
    • **Considerations:**
      • **Accuracy-Performance Trade-off:** Lower precision often comes with a slight reduction in model accuracy. Thoroughly evaluate your quantized model.
      • **Hardware Support:** Ensure your target edge device's hardware (e.g., NPU, DSP) supports the chosen quantization format.
      • **Tooling:** Different frameworks (PyTorch, TensorFlow) and deployment targets have specific quantization tools and workflows.

c. Resource Management

Carefully select the appropriate GPU instance type (e.g., NVIDIA A100, V100, T4) based on your model size, expected traffic, and budget. Monitor GPU utilization, memory, and CPU usage.

d. Scalability

For high traffic, deploy multiple instances of your Docker container behind a load balancer. Use container orchestration platforms like Kubernetes to manage auto-scaling based on demand.

e. Monitoring and Logging

Integrate logging (e.g., using Python's `logging` module) and push logs to a centralized logging system. Monitor API request rates, latency, error rates, and model-specific metrics (e.g., inference time per request, token generation speed).

f. Security

Secure your API endpoints with authentication (e.g., API keys, OAuth2). Use HTTPS for all traffic. Regularly update base images and dependencies to patch vulnerabilities.

g. CI/CD Pipelines

Automate the entire process: building Docker images, running tests, and deploying new versions of your API. This ensures consistent and reliable updates.

5. Scaling Fine-Tuning Workflows Across Multiple GPUs

While PEFT methods like LoRA significantly reduce the memory footprint for fine-tuning, training very large LLMs or processing massive datasets still often requires leveraging multiple GPUs. Scaling fine-tuning workflows across multiple GPUs can dramatically reduce training time and enable you to tackle more ambitious projects. There are two primary strategies for multi-GPU training:

a. Data Parallelism

This is the most common and straightforward approach. In data parallelism, each GPU gets a copy of the entire model, but the training data is split across the GPUs. Each GPU processes a different mini-batch of data, computes its local gradients, and then these gradients are averaged across all GPUs before a single weight update is performed. This effectively increases the overall batch size, leading to faster convergence and more stable training.

  • **How it works:**
    1. The dataset is divided into smaller chunks.
    2. Each GPU loads a copy of the model and processes a different chunk of data.
    3. Gradients are computed independently on each GPU.
    4. Gradients are then synchronized (averaged) across all GPUs.
    5. Model weights are updated based on the averaged gradients.
  • **Benefits:** Relatively easy to implement, good scaling for smaller to medium-sized models.
  • **Limitations:** Each GPU still needs enough memory to hold a full copy of the model.
# Conceptual Data Parallelism with Hugging Face Accelerate
# from accelerate import Accelerator
# from transformers import Trainer, TrainingArguments
#
# accelerator = Accelerator()
# model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
#     model, optimizer, train_dataloader, eval_dataloader
# )
#
# # Then use the accelerator's `backward()` method instead of `loss.backward()`
# # and `accelerator.wait_for_everyone()` for synchronization.
# # The Hugging Face Trainer handles this automatically when `accelerate` is used.

b. Model Parallelism (or Tensor Parallelism / Pipeline Parallelism)

When a single GPU doesn't have enough memory to even fit the entire model, **model parallelism** becomes necessary. In this approach, the model itself is split across multiple GPUs. Each GPU holds only a portion of the model's layers or parameters.

  • **Tensor Parallelism:** Splits individual model layers (e.g., attention matrices) across GPUs. This is very fine-grained.
  • **Pipeline Parallelism:** Splits the model vertically, with different GPUs responsible for different layers or stages of the model's computation pipeline. Data flows sequentially through the GPUs.
  • **Benefits:** Enables training of models that are too large for a single GPU.
  • **Limitations:** More complex to implement, can suffer from "bubble" (idle time) in pipeline parallelism if stages aren't perfectly balanced.
# Conceptual Model Parallelism (often handled by libraries like DeepSpeed)
# from transformers import AutoModelForCausalLM
#
# # For very large models, device_map="auto" can try to do some basic model parallelism
# model = AutoModelForCausalLM.from_pretrained(
#     "meta-llama/Llama-2-70b-hf",
#     device_map="auto" # Hugging Face will try to distribute layers
# )
#
# # For more advanced model parallelism, DeepSpeed is commonly used:
# # deepspeed --num_gpus=8 your_training_script.py --deepspeed_config ds_config.json

c. Orchestration Libraries: `accelerate` and `DeepSpeed`

Implementing multi-GPU training from scratch is complex. Libraries like Hugging Face's `accelerate` and Microsoft's `DeepSpeed` abstract away much of this complexity:

  • **Hugging Face `accelerate`:** Provides a simple API to run PyTorch training scripts across any distributed setup (multiple GPUs, TPUs, CPUs) with minimal code changes. It primarily focuses on data parallelism but can integrate with DeepSpeed for more advanced techniques.
  • **Microsoft `DeepSpeed`:** A powerful optimization library that supports various forms of parallelism (data, model, pipeline), mixed precision training, and ZeRO (Zero Redundancy Optimizer) for extreme memory efficiency. It's essential for training truly massive models.
# Installing DeepSpeed
# pip install deepspeed
#
# Example of a DeepSpeed configuration file (ds_config.json)
# {
#   "train_batch_size": 16,
#   "gradient_accumulation_steps": 2,
#   "optimizer": {
#     "type": "AdamW",
#     "params": {
#       "lr": 2e-5
#     }
#   },
#   "fp16": {
#     "enabled": true
#   },
#   "zero_optimization": {
#     "stage": 2 # ZeRO-2 for memory optimization
#   }
# }

6. Conclusion: Robust AI Services in Production

Deploying fine-tuned LLMs into production is a critical step in realizing their value. By combining the high performance and developer-friendly features of **FastAPI** with the consistency and portability of **Docker**, you can build robust, scalable, and efficient API services for your specialized models. Remember to optimize for performance, manage resources effectively, and implement continuous monitoring to ensure your AI applications deliver reliable and impactful results in the real world.

← Back to Articles