Edge Device Deployment: Running Tiny LLMs on Raspberry Pi, Mobile, and IoT

# Base dependencies (Ubuntu/Debian)
sudo apt update && sudo apt upgrade -y
sudo apt install -y \
    python3-pip \
    python3-dev \
    cmake \
    build-essential \
    git
 
# ARM-specific (Raspberry Pi, Jetson)
sudo apt install -y libopenblas-dev
 
# NVIDIA Jetson-specific
sudo apt install -y nvidia-jetpack tensorrt
 
# Google Coral-specific
curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -
echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list
sudo apt update
sudo apt install -y libedgetpu1-std python3-pycoral
 
# Python packages
pip install --upgrade pip
pip install \
    torch==2.1.0 \
    transformers==4.36.0 \
    onnxruntime==1.16.3 \
    onnx==1.15.0 \
    llama-cpp-python==0.2.20
 
# For NVIDIA devices with GPU support
pip install onnxruntime-gpu==1.16.3
pip install tensorrt  # Jetson only
 
# For quantization tools
pip install bitsandbytes optimum

# test_edge_setup.py
import sys
import torch
import onnxruntime as ort
from transformers import AutoTokenizer
 
print("=== Edge Deployment Environment Check ===\n")
 
# Python version
print(f"Python: {sys.version}")
 
# PyTorch
print(f"PyTorch: {torch.__version__}")
print(f"CUDA Available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"CUDA Version: {torch.version.cuda}")
 
# ONNX Runtime
print(f"\nONNX Runtime: {ort.__version__}")
print(f"Available Providers: {ort.get_available_providers()}")
 
# Check for hardware-specific acceleration
if "CUDAExecutionProvider" in ort.get_available_providers():
    print("✅ CUDA acceleration available")
elif "TensorrtExecutionProvider" in ort.get_available_providers():
    print("✅ TensorRT acceleration available")
elif "CoreMLExecutionProvider" in ort.get_available_providers():
    print("✅ CoreML acceleration available (Apple)")
else:
    print("ℹ️  CPU-only mode (expected for Raspberry Pi)")
 
print("\n✅ Environment ready for edge deployment!")

# Quantize with ONNX Runtime
import onnxruntime as ort
from onnxruntime.quantization import quantize_dynamic, QuantType
 
def quantize_model_for_edge(
    model_path="model.onnx",
    output_path="model_int8.onnx",
    weight_type=QuantType.QUInt8
):
    """
    Quantize ONNX model for edge deployment.
    
    Args:
        model_path: Path to FP32/FP16 ONNX model
        output_path: Output path for quantized model
        weight_type: QUInt8 (unsigned) or QInt8 (signed)
    """
    quantize_dynamic(
        model_path,
        output_path,
        weight_type=weight_type,
        optimize_model=True,  # Apply graph optimizations
    )
    
    print(f"Quantized model saved to {output_path}")
    
    # Compare sizes
    import os
    original_size = os.path.getsize(model_path) / 1e9
    quantized_size = os.path.getsize(output_path) / 1e9
    
    print(f"Original: {original_size:.2f} GB")
    print(f"Quantized: {quantized_size:.2f} GB")
    print(f"Reduction: {original_size / quantized_size:.2f}×")
 
# Usage
quantize_model_for_edge("tinyllama.onnx", "tinyllama_int8.onnx")
# Original: 2.20 GB
# Quantized: 0.55 GB
# Reduction: 4.00×

import onnx
from onnxruntime.transformers import optimizer
 
def optimize_onnx_graph(input_path, output_path, model_type="gpt2"):
    """
    Optimize ONNX graph for inference.
    
    Optimizations:
    - Fuse multi-head attention
    - Fuse layer normalization
    - Constant folding
    - Remove identity operations
    """
    from onnxruntime.transformers.optimizer import optimize_model
    
    optimized_model = optimize_model(
        input_path,
        model_type=model_type,
        num_heads=32,  # TinyLlama config
        hidden_size=2048,
        optimization_options=None  # Use default optimizations
    )
    
    optimized_model.save_model_to_file(output_path)
    print(f"Optimized model saved to {output_path}")
 
optimize_onnx_graph("tinyllama_int8.onnx", "tinyllama_optimized.onnx")

# Install llama.cpp (optimized for ARM)
git clone https://github.com/ggerganov/llama.cpp
cd llama.cpp
 
# Compile with ARM NEON
make LLAMA_NO_ACCELERATE=1
 
# Convert model to GGUF format (optimized binary format)
python convert.py ../tinyllama.pth --outfile tinyllama.gguf --outtype q8_0
 
# Result: 550 MB INT8 model with NEON-optimized kernels

# Update system
sudo apt update && sudo apt upgrade -y
 
# Install dependencies
sudo apt install -y python3-pip cmake build-essential
 
# Install optimized BLAS for ARM
sudo apt install -y libopenblas-dev

# Clone and compile
git clone https://github.com/ggerganov/llama.cpp
cd llama.cpp
make -j4  # Use all 4 cores
 
# Download or convert model
./convert.py /path/to/tinyllama --outfile models/tinyllama-q8.gguf --outtype q8_0

# Run inference
./main -m models/tinyllama-q8.gguf \
       -p "The capital of France is" \
       -n 50 \
       -t 4  # Use 4 threads
 
# Monitor performance
./main -m models/tinyllama-q8.gguf -p "Hello" -n 128 --timing
# Output:
# llama_print_timings:        load time =  1245.32 ms
# llama_print_timings:      sample time =    12.45 ms
# llama_print_timings: prompt eval time =   156.78 ms
# llama_print_timings:        eval time =  4567.89 ms /   128 tokens ( 35.69 ms per token)
# llama_print_timings:       total time =  6023.45 ms
 
# ~28 tokens/second

from llama_cpp import Llama
 
# Load model
llm = Llama(
    model_path="models/tinyllama-q8.gguf",
    n_ctx=2048,  # Context window
    n_threads=4,  # Use all cores
    n_batch=512,  # Batch size
    use_mlock=True,  # Lock model in RAM (faster, no swapping)
)
 
# Generate
output = llm(
    "The capital of France is",
    max_tokens=50,
    temperature=0.7,
    stop=["\n"]
)
 
print(output['choices'][0]['text'])

# Disable swap (prevent slowdowns)
sudo dphys-swapfile swapoff
sudo systemctl disable dphys-swapfile
 
# CPU governor to performance mode
echo performance | sudo tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor
 
# Increase GPU memory allocation (if using for display)
# In /boot/config.txt: gpu_mem=16 (minimum)

# Install JetPack SDK (includes CUDA, cuDNN)
sudo apt install nvidia-jetpack
 
# Install TensorRT for optimized inference
sudo apt install tensorrt

import tensorrt as trt
import onnx
 
def convert_to_tensorrt(onnx_path, engine_path, precision="fp16"):
    """
    Convert ONNX model to TensorRT engine.
    
    Args:
        onnx_path: Path to ONNX model
        engine_path: Output path for TensorRT engine
        precision: "fp32", "fp16", or "int8"
    """
    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    parser = trt.OnnxParser(network, logger)
    
    # Parse ONNX
    with open(onnx_path, 'rb') as f:
        if not parser.parse(f.read()):
            print("ERROR: Failed to parse ONNX model")
            return
    
    # Build config
    config = builder.create_builder_config()
    config.max_workspace_size = 1 << 30  # 1GB
    
    if precision == "fp16":
        config.set_flag(trt.BuilderFlag.FP16)
    elif precision == "int8":
        config.set_flag(trt.BuilderFlag.INT8)
    
    # Build engine
    engine = builder.build_engine(network, config)
    
    # Save
    with open(engine_path, 'wb') as f:
        f.write(engine.serialize())
    
    print(f"TensorRT engine saved to {engine_path}")
 
convert_to_tensorrt("tinyllama.onnx", "tinyllama.trt", precision="fp16")

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
 
class TensorRTInference:
    def __init__(self, engine_path):
        self.logger = trt.Logger(trt.Logger.WARNING)
        with open(engine_path, 'rb') as f:
            self.engine = trt.Runtime(self.logger).deserialize_cuda_engine(f.read())
        self.context = self.engine.create_execution_context()
    
    def infer(self, input_ids):
        # Allocate device memory
        d_input = cuda.mem_alloc(input_ids.nbytes)
        d_output = cuda.mem_alloc(output_shape.nbytes)
        
        # Copy input to device
        cuda.memcpy_htod(d_input, input_ids)
        
        # Run inference
        self.context.execute_v2([int(d_input), int(d_output)])
        
        # Copy output to host
        output = np.empty(output_shape, dtype=np.float32)
        cuda.memcpy_dtoh(output, d_output)
        
        return output
 
# Usage
trt_model = TensorRTInference("tinyllama.trt")
output = trt_model.infer(input_ids)

# Install Edge TPU compiler
curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -
echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list
sudo apt update
sudo apt install edgetpu-compiler
 
# Compile TFLite model for TPU
edgetpu_compiler tinyllama_int8.tflite
 
# Output: tinyllama_int8_edgetpu.tflite

from pycoral.utils import edgetpu
from pycoral.adapters import common
import numpy as np
 
class CoralInference:
    def __init__(self, model_path):
        self.interpreter = edgetpu.make_interpreter(model_path)
        self.interpreter.allocate_tensors()
    
    def infer(self, input_data):
        # Set input
        common.set_input(self.interpreter, input_data)
        
        # Run inference
        self.interpreter.invoke()
        
        # Get output
        return common.output_tensor(self.interpreter, 0)
 
# Usage
coral = CoralInference("tinyllama_int8_edgetpu.tflite")
output = coral.infer(input_ids)

FROM balenalib/raspberry-pi-debian:latest
 
# Install dependencies
RUN apt-get update && apt-get install -y \
    python3-pip \
    cmake \
    build-essential \
    libopenblas-dev \
    && rm -rf /var/lib/apt/lists/*
 
# Copy llama.cpp
WORKDIR /app
COPY llama.cpp /app/llama.cpp
WORKDIR /app/llama.cpp
RUN make -j4
 
# Copy model
COPY models/tinyllama-q8.gguf /app/models/
 
# Python API server
COPY server.py /app/
RUN pip3 install flask llama-cpp-python
 
# Expose port
EXPOSE 8000
 
# Run server
CMD ["python3", "/app/server.py"]

# server.py
from flask import Flask, request, jsonify
from llama_cpp import Llama
 
app = Flask(__name__)
 
# Load model once at startup
llm = Llama(
    model_path="/app/models/tinyllama-q8.gguf",
    n_ctx=2048,
    n_threads=4,
    use_mlock=True
)
 
@app.route('/generate', methods=['POST'])
def generate():
    data = request.json
    prompt = data.get('prompt', '')
    max_tokens = data.get('max_tokens', 100)
    
    try:
        output = llm(prompt, max_tokens=max_tokens)
        return jsonify({
            'success': True,
            'text': output['choices'][0]['text']
        })
    except Exception as e:
        return jsonify({'success': False, 'error': str(e)}), 500
 
@app.route('/health', methods=['GET'])
def health():
    return jsonify({'status': 'healthy'})
 
if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8000)

# Build for ARM
docker buildx build --platform linux/arm64 -t tinyllm-pi:latest .
 
# Run
docker run -d -p 8000:8000 --name tinyllm tinyllm-pi:latest
 
# Test
curl -X POST http://localhost:8000/generate \
  -H "Content-Type: application/json" \
  -d '{"prompt": "Hello, world!", "max_tokens": 50}'

import psutil
import logging
from prometheus_client import Counter, Histogram, Gauge, start_http_server
 
# Metrics
REQUEST_COUNT = Counter('tinyllm_requests_total', 'Total requests')
REQUEST_LATENCY = Histogram('tinyllm_request_latency_seconds', 'Request latency')
CPU_USAGE = Gauge('tinyllm_cpu_percent', 'CPU usage')
MEMORY_USAGE = Gauge('tinyllm_memory_mb', 'Memory usage MB')
 
# Start metrics server
start_http_server(9090)
 
# Monitor system
def monitor_system():
    while True:
        CPU_USAGE.set(psutil.cpu_percent())
        MEMORY_USAGE.set(psutil.virtual_memory().used / 1e6)
        time.sleep(5)
 
# In Flask app
@app.route('/generate', methods=['POST'])
def generate():
    REQUEST_COUNT.inc()
    
    start_time = time.time()
    # ... inference ...
    REQUEST_LATENCY.observe(time.time() - start_time)