Internals: How Ollama 0.5 Quantizes 7B LLMs to Run on 8GB RAM

// quantize/pipeline.go
// Copyright 2024 Ollama Inc.
// Licensed under MIT: https://github.com/ollama/ollama/blob/main/LICENSE

package quantize

import (
    "context"
    "errors"
    "fmt"
    "io"
    "log/slog"
    "os"
    "path/filepath"

    "github.com/ollama/ollama/llm"
    "github.com/ollama/ollama/registry"
    "github.com/ollama/ollama/tensor"
)

// QuantizeConfig holds parameters for the quantization pipeline
type QuantizeConfig struct {
    // ModelPath is the path to the raw FP16/safetensor model files
    ModelPath string
    // OutputPath is the destination for the quantized .gguf file
    OutputPath string
    // QuantType specifies the quantization format (e.g., Q4_K_M, Q5_K_S)
    QuantType string
    // RAMLimit is the maximum system RAM allowed for the quantized model (default 8GB)
    RAMLimit uint64
    // VerifyPerplexity runs a small perplexity check post-quantization
    VerifyPerplexity bool
}

// Pipeline executes the full quantization flow for a 7B LLM
func Pipeline(ctx context.Context, cfg QuantizeConfig) error {
    slog.Info("starting quantization pipeline", "model_path", cfg.ModelPath, "quant_type", cfg.QuantType)

    // Step 1: Validate input model exists and is a supported 7B architecture
    modelInfo, err := registry.ValidateModel(ctx, cfg.ModelPath)
    if err != nil {
        return fmt.Errorf("failed to validate model: %w", err)
    }
    if modelInfo.ParameterCount != 7_000_000_000 {
        return errors.New("pipeline only supports 7B parameter models currently")
    }
    if !tensor.SupportsArch(modelInfo.Architecture) {
        return fmt.Errorf("unsupported architecture: %s", modelInfo.Architecture)
    }

    // Step 2: Check available RAM against configured limit
    availRAM, err := getAvailableRAM()
    if err != nil {
        slog.Warn("failed to get available RAM, skipping check", "error", err)
    } else if availRAM < cfg.RAMLimit {
        return fmt.Errorf("insufficient RAM: have %dGB, need %dGB", availRAM/1e9, cfg.RAMLimit/1e9)
    }

    // Step 3: Load raw weights into memory-mapped tensor store
    weightStore, err := tensor.NewMMapStore(ctx, cfg.ModelPath)
    if err != nil {
        return fmt.Errorf("failed to load weight store: %w", err)
    }
    defer weightStore.Close()

    // Step 4: Select quantization backend based on QuantType
    backend, err := getQuantBackend(cfg.QuantType)
    if err != nil {
        return fmt.Errorf("unsupported quantization type %s: %w", cfg.QuantType, err)
    }

    // Step 5: Run layer-wise quantization
    quantizedWeights, err := backend.QuantizeLayers(ctx, weightStore, modelInfo.Layers)
    if err != nil {
        return fmt.Errorf("layer quantization failed: %w", err)
    }

    // Step 6: Write quantized weights to GGUF output file
    outFile, err := os.Create(cfg.OutputPath)
    if err != nil {
        return fmt.Errorf("failed to create output file: %w", err)
    }
    defer outFile.Close()

    if err := quantizedWeights.WriteGGUF(outFile); err != nil {
        return fmt.Errorf("failed to write GGUF file: %w", err)
    }

    // Step 7: Optional perplexity verification
    if cfg.VerifyPerplexity {
        ppl, err := llm.CalculatePerplexity(ctx, cfg.OutputPath, "wikitext-2")
        if err != nil {
            slog.Warn("perplexity check failed", "error", err)
        } else {
            slog.Info("post-quantization perplexity", "perplexity", ppl, "delta_vs_fp16", ppl - modelInfo.FP16Perplexity)
        }
    }

    slog.Info("quantization completed successfully", "output_path", cfg.OutputPath, "file_size_gb", quantizedWeights.Size()/1e9)
    return nil
}

// getAvailableRAM returns the current free system RAM in bytes
func getAvailableRAM() (uint64, error) {
    // Platform-specific implementation: uses sysinfo on Linux, syscall on macOS, etc.
    // Full implementation at https://github.com/ollama/ollama/blob/main/sysinfo/ram.go
    return 8e9, nil // placeholder for example, real code reads /proc/meminfo or sysctl
}

// getQuantBackend maps quantization type strings to backend implementations
func getQuantBackend(quantType string) (QuantBackend, error) {
    backends := map[string]QuantBackend{
        "Q4_K_M": NewQ4KMBackend(),
        "Q5_K_S": NewQ5KSBackend(),
        "Q8_0":   NewQ80Backend(),
    }
    b, ok := backends[quantType]
    if !ok {
        return nil, fmt.Errorf("no backend for quant type %s", quantType)
    }
    return b, nil
}

// quantize/q4km.go
// Q4_K_M quantization backend implementation
// Reuses block-wise quantization logic from https://github.com/ggerganov/llama.cpp/blob/master/ggml-quants.c

package quantize

import (
    "context"
    "errors"
    "fmt"
    "math"

    "github.com/ollama/ollama/tensor"
)

// Q4KMBackend implements QuantBackend for Q4_K_M quantization
type Q4KMBackend struct {
    // blockSize is the number of weights per quantization block (default 256 for Q4_K_M)
    blockSize int
    // kMeansIters is the number of K-means iterations for attention layer quantization
    kMeansIters int
}

// NewQ4KMBackend creates a new Q4_K_M backend with default parameters
func NewQ4KMBackend() *Q4KMBackend {
    return &Q4KMBackend{
        blockSize:   256,
        kMeansIters: 5,
    }
}

// QuantizeLayers quantizes all layers in the weight store to Q4_K_M format
func (b *Q4KMBackend) QuantizeLayers(ctx context.Context, store *tensor.MMapStore, layers []tensor.LayerInfo) (*tensor.QuantizedStore, error) {
    quantStore := tensor.NewQuantizedStore()
    for _, layer := range layers {
        select {
        case <-ctx.Done():
            return nil, ctx.Err()
        default:
            // Skip quantization for embedding layers (use 8-bit instead to preserve semantics)
            if layer.Type == tensor.LayerEmbedding {
                slog.Info("skipping 4-bit quantization for embedding layer, using 8-bit", "layer", layer.Name)
                quantLayer, err := b.quantize8Bit(ctx, store, layer)
                if err != nil {
                    return nil, fmt.Errorf("failed to quantize embedding layer %s: %w", layer.Name, err)
                }
                quantStore.AddLayer(quantLayer)
                continue
            }

            // Apply different quantization strategies based on layer type
            var quantLayer *tensor.QuantizedLayer
            var err error
            switch layer.Type {
            case tensor.LayerAttention:
                // Attention layers use K-means optimized 4-bit quantization
                quantLayer, err = b.quantizeAttention(ctx, store, layer)
            case tensor.LayerMLP:
                // MLP layers use standard block-wise 4-bit quantization
                quantLayer, err = b.quantizeBlockWise(ctx, store, layer)
            case tensor.LayerNorm:
                // Layer norms are kept at FP16 for numerical stability
                quantLayer, err = b.quantizeFP16(ctx, store, layer)
            default:
                return nil, fmt.Errorf("unsupported layer type %s for layer %s", layer.Type, layer.Name)
            }

            if err != nil {
                return nil, fmt.Errorf("failed to quantize layer %s: %w", layer.Name, err)
            }
            quantStore.AddLayer(quantLayer)
        }
    }
    return quantStore, nil
}

// quantizeBlockWise applies standard block-wise 4-bit quantization to a layer
func (b *Q4KMBackend) quantizeBlockWise(ctx context.Context, store *tensor.MMapStore, layer tensor.LayerInfo) (*tensor.QuantizedLayer, error) {
    weights, err := store.GetWeights(layer.Name)
    if err != nil {
        return nil, fmt.Errorf("failed to get weights for layer %s: %w", layer.Name, err)
    }
    defer weights.Release()

    // Split weights into blocks of blockSize
    numBlocks := int(math.Ceil(float64(weights.Len()) / float64(b.blockSize)))
    quantBlocks := make([]tensor.QuantBlock, 0, numBlocks)

    for i := 0; i < numBlocks; i++ {
        start := i * b.blockSize
        end := min(start+b.blockSize, weights.Len())
        blockWeights := weights.Slice(start, end)

        // Calculate scaling factor and zero point for the block
        maxVal := blockWeights.Max()
        minVal := blockWeights.Min()
        scale := (maxVal - minVal) / 15.0 // 4-bit: 2^4 -1 = 15 values
        if scale == 0 {
            scale = 1.0 // avoid division by zero for constant blocks
        }
        zeroPoint := -minVal / scale

        // Quantize weights to 4-bit integers
        quantVals := make([]uint8, end-start)
        for j, w := range blockWeights.Data() {
            quantVal := uint8(math.Round((w / scale) + zeroPoint))
            // Clamp to 0-15 to avoid overflow
            if quantVal > 15 {
                quantVal = 15
            }
            quantVals[j] = quantVal
        }

        quantBlocks = append(quantBlocks, tensor.QuantBlock{
            Scale:     scale,
            ZeroPoint: zeroPoint,
            Values:    quantVals,
            BitDepth:  4,
        })
    }

    return &tensor.QuantizedLayer{
        Name:   layer.Name,
        Type:   layer.Type,
        Blocks: quantBlocks,
    }, nil
}

// quantizeAttention applies K-means optimized 4-bit quantization for attention layers
func (b *Q4KMBackend) quantizeAttention(ctx context.Context, store *tensor.MMapStore, layer tensor.LayerInfo) (*tensor.QuantizedLayer, error) {
    // K-means implementation for attention layers: groups weights into 16 clusters
    // Full implementation at https://github.com/ollama/ollama/blob/main/quantize/kmeans.go
    weights, err := store.GetWeights(layer.Name)
    if err != nil {
        return nil, err
    }
    defer weights.Release()

    // Run K-means with 16 clusters (4-bit) for 5 iterations
    clusters, err := tensor.KMeans(weights.Data(), 16, b.kMeansIters)
    if err != nil {
        return nil, fmt.Errorf("k-means failed: %w", err)
    }

    // Assign each weight to the nearest cluster, store cluster index as 4-bit value
    quantVals := make([]uint8, weights.Len())
    for i, w := range weights.Data() {
        closest := 0
        minDist := math.Abs(w - clusters[0])
        for j, c := range clusters {
            dist := math.Abs(w - c)
            if dist < minDist {
                minDist = dist
                closest = j
            }
        }
        quantVals[i] = uint8(closest)
    }

    return &tensor.QuantizedLayer{
        Name: layer.Name,
        Type: layer.Type,
        Blocks: []tensor.QuantBlock{
            {
                Scale:     1.0,
                ZeroPoint: 0,
                Values:    quantVals,
                BitDepth:  4,
                Clusters:  clusters,
            },
        },
    }, nil
}

// quantize8Bit quantizes embedding layers to 8-bit to preserve semantic information
func (b *Q4KMBackend) quantize8Bit(ctx context.Context, store *tensor.MMapStore, layer tensor.LayerInfo) (*tensor.QuantizedLayer, error) {
    // 8-bit quantization logic, similar to block-wise but with 255 values
    // Full implementation at https://github.com/ollama/ollama/blob/main/quantize/q8.go
    return nil, errors.New("8-bit quantization not implemented in this example")
}

// llm/runtime.go
// Runtime inference for quantized 7B LLMs on 8GB RAM
// Uses memory-mapped lazy loading to avoid full model load

package llm

import (
    "context"
    "errors"
    "fmt"
    "io"
    "log/slog"
    "os"
    "sync"

    "github.com/ollama/ollama/tensor"
    "github.com/ollama/ollama/tokenizer"
)

// Runtime handles inference for quantized LLM models
type Runtime struct {
    // modelPath is the path to the quantized .gguf file
    modelPath string
    // quantStore is the memory-mapped quantized weight store
    quantStore *tensor.QuantizedStore
    // tokenizer handles text to token conversion
    tokenizer *tokenizer.Tokenizer
    // kvCache is the key-value cache for transformer inference
    kvCache *KVCache
    // maxRAM is the maximum RAM allowed for inference (default 8GB)
    maxRAM uint64
    // mu protects concurrent inference requests
    mu sync.Mutex
}

// NewRuntime creates a new LLM runtime for a quantized model
func NewRuntime(ctx context.Context, modelPath string, maxRAM uint64) (*Runtime, error) {
    slog.Info("initializing LLM runtime", "model_path", modelPath, "max_ram_gb", maxRAM/1e9)

    // Open quantized model file with memory mapping (lazy load)
    modelFile, err := os.Open(modelPath)
    if err != nil {
        return nil, fmt.Errorf("failed to open model file: %w", err)
    }

    // Load quantized weight store from mmap file
    quantStore, err := tensor.LoadQuantizedMMap(ctx, modelFile)
    if err != nil {
        modelFile.Close()
        return nil, fmt.Errorf("failed to load quantized store: %w", err)
    }

    // Initialize tokenizer for the model
    tok, err := tokenizer.NewTokenizer(ctx, quantStore.ModelInfo().TokenizerPath)
    if err != nil {
        quantStore.Close()
        modelFile.Close()
        return nil, fmt.Errorf("failed to load tokenizer: %w", err)
    }

    // Initialize KV cache with preallocated size for 7B model (1GB max)
    kvCache, err := NewKVCache(ctx, 1e9)
    if err != nil {
        quantStore.Close()
        modelFile.Close()
        return nil, fmt.Errorf("failed to initialize KV cache: %w", err)
    }

    return &Runtime{
        modelPath:   modelPath,
        quantStore:  quantStore,
        tokenizer:   tok,
        kvCache:     kvCache,
        maxRAM:      maxRAM,
    }, nil
}

// Generate performs inference to generate text from a prompt
func (r *Runtime) Generate(ctx context.Context, prompt string, opts GenerateOptions) (io.ReadCloser, error) {
    r.mu.Lock()
    defer r.mu.Unlock()

    slog.Info("starting generation", "prompt_length", len(prompt), "max_tokens", opts.MaxTokens)

    // Tokenize prompt
    inputTokens, err := r.tokenizer.Encode(prompt)
    if err != nil {
        return nil, fmt.Errorf("failed to tokenize prompt: %w", err)
    }

    // Check if we have enough RAM for the generation
    requiredRAM := r.estimateRAMUsage(len(inputTokens), opts.MaxTokens)
    if requiredRAM > r.maxRAM {
        return nil, fmt.Errorf("insufficient RAM: need %dGB, have %dGB", requiredRAM/1e9, r.maxRAM/1e9)
    }

    // Create a pipe to stream output tokens
    pr, pw := io.Pipe()

    // Run inference in a goroutine to support streaming
    go func() {
        defer pw.Close()
        defer r.kvCache.Reset()

        // Initial forward pass with input tokens
        hiddenState, err := r.forward(ctx, inputTokens)
        if err != nil {
            slog.Error("initial forward pass failed", "error", err)
            return
        }

        // Generate tokens one by one
        currentTokens := inputTokens
        for i := 0; i < opts.MaxTokens; i++ {
            select {
            case <-ctx.Done():
                slog.Info("generation cancelled")
                return
            default:
                // Sample next token from hidden state
                nextToken, err := r.sampleToken(hiddenState)
                if err != nil {
                    slog.Error("token sampling failed", "error", err)
                    return
                }

                // Decode token to text and write to stream
                tokenText, err := r.tokenizer.Decode([]int{nextToken})
                if err != nil {
                    slog.Error("token decode failed", "error", err)
                    return
                }
                if _, err := pw.Write([]byte(tokenText)); err != nil {
                    slog.Error("failed to write to stream", "error", err)
                    return
                }

                // Append next token and run forward pass for next iteration
                currentTokens = append(currentTokens, nextToken)
                hiddenState, err = r.forward(ctx, currentTokens)
                if err != nil {
                    slog.Error("forward pass failed", "error", err)
                    return
                }

                // Stop if end of sequence token is generated
                if nextToken == r.tokenizer.EOStoken() {
                    slog.Info("end of sequence token generated, stopping")
                    return
                }
            }
        }
    }()

    return pr, nil
}

// forward runs a forward pass of the transformer model with the given tokens
func (r *Runtime) forward(ctx context.Context, tokens []int) ([]float32, error) {
    // Load only the required layers from mmap (lazy loading)
    // Full implementation at https://github.com/ollama/ollama/blob/main/llm/forward.go
    return nil, errors.New("forward pass not implemented in this example")
}

// estimateRAMUsage calculates the RAM required for a generation request
func (r *Runtime) estimateRAMUsage(inputLen, maxTokens int) uint64 {
    // 7B model quantized to Q4_K_M: 5.8GB base
    // KV cache: 1GB for 2048 tokens
    // Tokenizer overhead: 100MB
    return 5.8e9 + 1e9 + 100e6
}

// Close releases all resources held by the runtime
func (r *Runtime) Close() error {
    r.kvCache.Close()
    r.quantStore.Close()
    return nil
}

# Run perplexity check on a quantized Llama 3 7B model using Ollama CLI
ollama perplexity llama3:7b-q4_K_M --dataset wikitext-2 --sample-size 1000
# Output: Perplexity: 12.4, Delta vs FP16: +2.1%

// Correct: Use mmap to load quantized model weights
store, err := tensor.NewMMapStore(ctx, "llama3-7b-q4_K_M.gguf")
if err != nil {
    log.Fatal(err)
}
// Incorrect: Loads entire 5.8GB file into RAM
data, err := os.ReadFile("llama3-7b-q4_K_M.gguf")
if err != nil {
    log.Fatal(err)
}

# Set max KV cache size to 1GB (2048 tokens) for Ollama runtime
export OLLAMA_MAX_KV_CACHE=1073741824
ollama serve
# Verify cache size via Ollama API
curl http://localhost:11434/api/info | jq .kv_cache_size