Fine-Tuning Mistral 7B using QLoRA with PyTorch pt. 2: K8s & GKE | ML Engineering & MLOps

 

Hi All

Continuing from Part 1, this post details the Kubernetes and Observability configs of the project. The full source is available here. Let's break the code down shall we. 

1. K3's Server Config (infra/server-config.yaml) 


write-kubeconfig-mode: "0644" 

*    Sets file permissions for kubeconfig file (readable by all users in the group)

*    0644 means owner can read/write, group and others can only read


disable:
  - traefik
  - servicelb
  - local-storage
  - metrics-server

*    Disables default k3s components that we'll replace with better alternatives.

Components Disabled:

*    _traefik: Replaced with ingress-nginx for better control

*    _servicelb: Replaced with MetalLB or cloud load balancer

*    _local-storage: Replaced with Longhon for dynamic provisioning

*    _metrics-server: Replaced with Prometheus for better monitoring.


kube-apiserver-arg:
  - "service-node-port-range=8000-65535"
  - "enable-admission-plugins=NodeRestriction,LimitRanger,ResourceQuota"

* Configures the k8s API server

Key Parameters:

* _service-node-port-range: Expands port range for GPU workloads.

* _enable-admission-plugins: Enables security policies (NodeResitriction prevents nodes from modifying themsleves. LimitRanger enforces resource limits, ResourceQuota manages resource quotas)


kube-controller-manager-arg:
  - "terminated-pod-gc-threshold=10"

*    Configures the controller manager

*    Details: Sets how many terminated pods to keep before garbage collection (default is 5000, wereduce to 10 for better cleanups). 


kubelet-arg:
  - "eviction-hard=memory.available=500Mi"
  - "image-gc-high-threshold=85"
  - "image-gc-low-threshold=75"
  - "kube-reserved=cpu=500m,memory=1Gi,ephemeral-storage=1Gi"
  - "system-reserved=cpu=500m,memory=2Gi,ephemeral-storage=1Gi"

*    Configures the kubelet (node agent)

Key Parameters:

*    eviction-hard: More aggressive pod eviction when memory is low

*    image-gc-*: More frequent image garbage collection (keeps disk usage in check)

*    kube-reserved: Resources reserved for k8s system components

*    system-reserved: Resources reserved for non-Kubernetes system processes.


node-ip: "auto"
flannel-backend: "wireguard-native"
disable-network-policy: false

 

*    Network configuration

Details:

*    node-ip: "auto": Automatically detects and uses the nodes primary IP

*    flannel-backend: "wireguard-native": Uses WireGuard for better performance than VXLAN.

*    disable-network-policy: false: Enables K8s network policies for security.

 

2. k3s Agent Config for GPU Nodes


server: "https://${k3s-server-ip}:6443"
token: "${your-k3s-token}"

* Connects the agent to the server.

Details:

* Replace <k3s-server-ip> with your server's IP address

* Replace <your-k3s-token> with the token from /var/lib/rancher/k3s/server/node-token


node-label:
  - "node-role.kubernetes.io/gpu=true"
  - "nvidia.com/gpu=true"

*    Labels the node for GPU workloads

Details:

*    node-role.kubernetes.io/gpu=true: Marks the node as a GPU node

*    nvidia.com/gpu=true: indicate NVIDIA GPU support


node-taint:
  - "nvidia.com/gpu=true:NoSchedule"

*    Prevents non-GPU workloads from running on GPU nodes

*    The taint prevents scheduling unless a pod has matching tolerations.


kubelet-arg:
  - "eviction-hard=memory.available=1Gi"
  - "image-gc-high-threshold=90"
  - "image-gc-low-threshold=80"
  - "kube-reserved=cpu=1,memory=4Gi,ephemeral-storage=2Gi"
  - "system-reserved=cpu=1,memory=4Gi,ephemeral-storage=2Gi"

*    More generous resource reservations for GPU nodes

*    Details:

*    Higher thresholds for image garbage collection

*    More resources reserved for system components (GPU nodes typically have more resources)


container-runtime-endpoint: "unix:///run/containerd/containerd.sock"

*    Explicitly sets the container runtime socket

*    Details: Ensures k3s uses containerd (the default runtime) with the correct socket path.

 

3. NVIDIA GPU Operator Helm Values


operator:
  defaultRuntime: containerd
  deploy:
    daemonset: true
    driver: true
    toolkit: true
    devicePlugin: true
    dcgmExporter: true
    gfd: true
    migManager: false

*    Configures which components of the GPU operator to deploy

Components:

*    daemonset: Deploys GPU operator components as daemonsets

*    driver: Installs NVIDIA drivers

*    toolset: Installs CUDA toolkit

*    devicePlugin: Manages GPU resources in K8s

*    dcgmExporter: Exports GPU metrics

*    gfd: GPU Feature Discovery (labels nodes with GPU capabilities)

*    migManager: Disabled for simplicity (Multi-Instance GPU management)


driver:
  enabled: true
  version: "535.86.10"

*    Configures NVIDIA driver installation

*    Details: Specifies the exact driver version to install (should match your GPU's supported version)


toolkit:
  enabled: true

Installs CUDA toolkit components needed for GPU workloads 


devicePlugin:
  enabled: true

 

*    Manages CPU resources in K8s (exposes GPU counts to the scheduler)


dcgmExporter:
  enabled: true
  serviceMonitor:
    enabled: true

*    Exports GPU metrics for monitoring

*    Details: serviceMonitor enables Prometheus to scrape these metrics.


gfd:
  enabled: true

*    GPU Feature Discovery

*    Details: Automatically labels nodes with GPU capabilities (eg. nvidia.com/gpu: "true")

 

4. Persistent Storage Configuration (Longhorn)


# Install Longhorn
kubectl apply -f https://raw.githubusercontent.com/longhorn/longhorn/v1.5.1/deploy/longhorn.yaml

*    Installs Longhorn storage system

*    Details: Longhorn is distributed block storage system for K8s 


apiVersion: storage.k8s.io/v1
kind: StorageClass
metadata:
  name: longhorn-gpu
provisioner: driver.longhorn.io
allowVolumeExpansion: true
parameters:
  numberOfReplicas: "2"
  staleReplicaTimeout: "30"
  fsType: "ext4"
volumeBindingMode: WaitForFirstConsumer

*    Creates a custom Storage for GPU workloads

 

Key Parameters:

*    NumberOfReplicas: "2": Creates 2 copies of each volume for redundancy

*    staleReplicaTimeout: "30": Time (in minutes) before considering a replica stale

*    fsType: "ext4": Filesystem type for volumes

*    volumeBindingMode: WaitForFirstConsumer: Delays volume binding until a pod is scheduled (ensures volume is created on the same node as the pod)

 

5. GPU Workload Configuration Example


apiVersion: v1
kind: PersistentVolumeClaim
metadata:
  name: mistral-data
spec:
  storageClassName: longhorn-gpu
  accessModes:
    - ReadWriteOnce
  resources:
    requests:
      storage: 100Gi


* Creates a PersistentVolumeClaim for the Mistral fine-tuning job

Details:

* Uses the longhorn-gpu StorageClass

* ReadWriteOnce: Only one pod can mount the volume in read-write mode

* Requests 100GB of storage


apiVersion: batch/v1
kind: Job
metadata:
  name: mistral-qlora-finetune
spec:
  backoffLimit: 0
  template:

* Defines a K8s job for the fine-tuning workload

Details:

* backoffLimit: Disables retries on failure (job will fail immediately if it fails)

* templates: Defines the pod template


spec:
  nodeSelector:
    node-role.kubernetes.io/gpu: "true"
  tolerations:
    - key: "nvidia.com/gpu"
      operator: "Exists"
      effect: "NoSchedule"

* Ensures the pod only runs on GPU nodes

Details:

* nodeSelector: Selects nodes with the GPU label

* tolerations: Allow the pod to tolerate the GPU taint


containers:
- name: finetune
  image: mistral-qlora-finetune:latest
  command: ["python", "finetune_mistral_qlora.py"]

* Defines the container specification

Details:

* Uses the custom Docker image with the fine-tuning code

* Runs the Python script when the container starts


env:
- name: NCCL_DEBUG
  value: "INFO"
- name: NCCL_SOCKET_IFNAME
  value: "eth0"

* Sets environment variables for NCCL

Details:

* NCCL_DEBUG: "INFO": Enables NCCL debugging output

* NCCL_SOCKET_IFNAME: "eth0": Specifies the network interface to use for NCCL communication.


resources:
  limits:
    nvidia.com/gpu: 1
    memory: "40Gi"
    cpu: "8"
  requests:
    nvidia.com/gpu: 1
    memory: "32Gi"
    cpu: "6"

 

* Defines resources the pod can use (prevents OOM kills)

* Details:

* limits: Maximum resources the pod can use (prevents OOM kills)

* requests: Resources guaranteed to the pod (used for scheduling)

* Requests 1 GPU, 32GB memory, and 6 CPUs (with limits of 40GB memory and 8 CPUs)


volumeMounts:
- name: data
  mountPath: /workspace
volumes:
- name: data
  persistentVolumeClaim:
    claimName: mistral-data

 

* Mounts the persistent volume

Details:

* Mounts the PVC at /workspace in the container

* The PVC is named mistral-data 


restartPolicy: Never

 

* Configures pod restart behavior

* The pod will not restart if it fails (appropriate for a batch job)

 

6. Ingress Configuration


# Install ingress-nginx
kubectl apply -f https://raw.githubusercontent.com/kubernetes/ingress-nginx/controller-v1.8.2/deploy/static/provider/cloud/deploy.yaml

* Installs the ingress-nginx controller

* Ingress-nginx is a popular ingress controller for K8s


apiVersion: networking.k8s.io/v1
kind: Ingress
metadata:
  name: mistral-ingress
  annotations:
    nginx.ingress.kubernetes.io/proxy-body-size: "0"
    nginx.ingress.kubernetes.io/proxy-read-timeout: "3600"
    nginx.ingress.kubernetes.io/proxy-send-timeout: "3600"
spec:
  ingressClassName: nginx
  rules:
  - host: mistral.example.com
    http:
      paths:
      - path: /
        pathType: Prefix
        backend:
          service:
            name: mistral-service
            port:
              number: 80


* Defines an Ingress resource for the Mistral service

Key Annotations:

* proxy-body-size: "9": Disables body size limits (useful for large model files)

* proxy-read-timeout: "3600": Sets read timeout to 1 hour

* proxy-send-timeout: "3600": Sets send timeout to 1 hour

Details:

* Routes traffic from mistral.example.com to the mistral-service on port 80

* Uses the nginx ingress class 

 

7. Monitoring Stack (Prometheus & Grafana) 


# Install kube-prometheus-stack
helm repo add prometheus-community https://prometheus-community.github.io/helm-charts
helm repo update
helm install prometheus prometheus-community/kube-prometheus-stack \
  --namespace monitoring --create-namespace \
  --set grafana.ingress.enabled=true \
  --set grafana.ingress.hosts[0]=grafana.example.com

 

* Installs Prometheus, Grafana, and related monitoring components

Details:

* Deploys Prometheus for metrics collection

* Deploys Grafana for visualization

* Enables Grafana ingress with hostname grafana.example.com

* Creates a dedicated monitoring namespace

The above is runnable production grade inference example using an edge capable and quantized popular small model Mistral 7B. The infra code is available here, and the full source here. That's all for now folks, much more ML Backend Engineering on the way though.


Later


- Ed (@lightspeed001)

 

 

 


 

 

 


 


 

 

 

Comments

Post a Comment

Popular posts from this blog

Fine-Tuning Mistral 7B using QLoRA with PyTorch pt. 1: The Model | ML Engineering

Image
     Hi All Today we're working with a popular and slightly bigger model than our previous example. Mistral 7B is capable of chat and light coding tasks, for older hardware it's a winner for sure.  Here's a complete, runnable example of fine-tuning Mistral 7B using QLoRA with the peft , transformers , and bitsandbytes libraries. This example assumes you're working with a single GPU (eg. an A100 or similar). First install the required packages: pip install -q bitsandbytes datasets accelerate peft transformers trl View full script below, also available here :   Full breakdown of the script above, block-by-block. 1.      Dataset Loading dataset = load_dataset("timdettmers/openassistant-guanaco", split="train") *      Loads a preprocessed instruction-following dataset (Guanco, derived from OpenAssistant). *      split="train" selects the training portion *      The dataset is in a conversational ...
Read more

Deploying LoRA Optimised BERT as a FastApi service on GKE | ML Engineering & MLOps

Image
  Hi All  Alright today on "Bored MLE" we're doing some inference, using a LoRA optimized BERT model on GKE . Now there's a lot more to MLOps than just this (like cluster level observability for example), but for this example I'm keeping it brief.  The full code is available on GitHub with added minikube deployment instructions (I only cover the Cloud deployment here).    I assume familiarity with Python, PyTorch, K8s, FastApi, Minikube and GKE. Let's get on with it.  View the full FastApi service source below, also available here : Let's breakdown the above code, block by block. 0. Install Dependencies :  pip install fastapi uvicorn torch transformers peft accelerate  1. Imports from fastapi import FastAPI, HTTPException from pydantic import BaseModel from typing import List import torch from transformers import AutoModelForSequenceClassification, AutoTokenizer from peft import PeftModel, PeftConfig  *      fastapi : F...
Read more

Advanced Rust ML: Custom Modules with Tch-rs | ML Engineering

Image
  Hi All Another day, another Rust ML backend application. Today we're looking at defining custom modules in Rust in the ML backend. We're going to be creating a custom linear layer in Rust using Tch-rs ( Rust PyTorch bindings ).   View full source below, also available here :   Let's break the above code down. Block by block. 1. Imports use tch::{nn, nn::Module, Tensor}; What it does:  *        tch::nn - Neural network module containing layer definitions. *      nn:Module - The trait that all neural network modules must implement. *      Tensor - The Tensor type used throughout tch-rs   2. Struct Definition   struct CustomLayer { weight: Tensor, bias: Tensor, } What it does: *      weight - A tensor holding a layer's weights (matrix) *      bias -  A tensor holding the layer's bias (vector) This is essentially a linear layer (fully connected laye...
Read more