Gemma 3 - Zerfoo Implementation

This repository contains a Gemma 3 language model implementation using the Zerfoo ML framework. The implementation has been updated to use the new Zerfoo architecture with ZMF (Zerfoo Model Format) support.

Features

• Modern Zerfoo Architecture: Uses the latest Zerfoo framework with generic tensor types and graph-based computation
• ZMF Model Loading: Supports loading models from ZMF format instead of legacy ONNX
• Tokenizer Integration: Includes SentencePiece tokenizer support via github.com/sugarme/tokenizer
• Comprehensive Testing: Unit tests for all components plus integration tests
• CPU Engine Support: Optimized for CPU inference using Zerfoo's CPU engine

Architecture

Core Components

1. Model (gemma/gemma.go): Main model combining embedding, transformer stack, and LM head
2. GemmaStack (gemma/gemma_stack.go): Multi-layer transformer implementation with local attention
3. Tokenizer (tokenizer/tokenizer.go): SentencePiece tokenizer wrapper for text processing

Key Features

• Generic Tensor Types: Uses TensorNumeric[T] for type-safe numeric operations
• Local Attention: Implements sliding window attention with configurable window size
• Weight Sharing: Shares weights between embedding and LM head layers
• Flexible Architecture: Configurable number of layers, heads, dimensions, etc.

Usage

Basic Inference

package main
import (
"context""github.com/zerfoo/gemma/tokenizer""github.com/zerfoo/zerfoo/compute""github.com/zerfoo/zerfoo/layers/registry""github.com/zerfoo/zerfoo/model""github.com/zerfoo/zerfoo/numeric"
)
funcmain() {
// Initialize layer registryregistry.RegisterAll()
// Load ZMF modelzmfModel, err:=model.LoadZMF("data/model_with_weights.zmf")
iferr!=nil {
panic(err)
}
// Build Zerfoo graphops:= numeric.Float32Ops{}
engine:= compute.NewCPUEngine[float32](ops)
graph, err:=model.BuildFromZMF[float32](engine, ops, zmfModel)
iferr!=nil {
panic(err)
}
// Initialize tokenizertokenizer, err:=tokenizer.NewGemmaTokenizer("data/tokenizer.json")
iferr!=nil {
panic(err)
}
// Tokenize inputtokens, err:=tokenizer.Encode("What is the meaning of life?")
iferr!=nil {
panic(err)
}
// Run inferenceoutput, err:=graph.Forward(context.Background(), /* inputs */)
iferr!=nil {
panic(err)
}
}

Model Creation (Programmatic)

import"github.com/zerfoo/gemma/gemma"// Create a Gemma model programmaticallymodel, err:=gemma.New[float32](
engine, // Compute engineops, // Numeric operationsvocabSize, // Vocabulary sizehiddenSize, // Hidden dimensionnumHeads, // Number of attention headsnumKeyValueHeads, // Number of key-value headsffnDim, // Feed-forward dimensionepsilon, // RMS norm epsilonbase, // Rotary embedding basemaxSeqLen, // Maximum sequence lengthnumLayers, // Number of layerslocalWindowSize, // Local attention window sizeglobalInterval, // Global attention interval
)

Testing

The project includes comprehensive tests:

Unit Tests

# Test model components
go test ./gemma/...
# Test tokenizer
go test ./tokenizer/...

Integration Tests

# API integration (fast)
go test -run TestAPIIntegration
# Model architecture compatibility
go test -run TestModelArchitectureCompatibility
# End-to-end test (requires model files)
go test -run TestGemma3EndToEnd

Test Coverage

• Gemma Stack: Tests forward pass with mock data
• Gemma Model: Tests full model forward pass
• Tokenizer API: Tests special token handling and encoding
• API Integration: Tests component compatibility and tensor operations
• Architecture Tests: Validates model creation with various parameters

Model Files

The implementation expects these files in the data/ directory:

• model_with_weights.zmf: ZMF model file with full weights
• tokenizer.json: SentencePiece tokenizer configuration
• Additional config files (optional)

Dependencies

require (
github.com/sugarme/tokenizerv0.2.2github.com/zerfoo/zerfoov0.2.0
)

Performance

• CPU Optimized: Uses Zerfoo's optimized CPU engine
• Memory Efficient: Supports various quantization formats
• Scalable: Configurable model size and attention patterns

Known Issues

1. Tokenizer Compatibility: Some Gemma tokenizer.json files may not be fully compatible with the sugarme/tokenizer library
2. Large Model Loading: Full-size models may require significant memory and loading time
3. Test Timeouts: Integration tests with large models may timeout

Contributing

1. Follow the existing code style
2. Add tests for new features
3. Update documentation
4. Ensure all tests pass before submitting PRs

License

See LICENSE file for details.