**Zeta** is a modular PyTorch framework designed to simplify the development of AI models by providing reusable, high-performance building blocks. Think of it as a collection of LEGO blocks for AI each component is carefully crafted, tested, and optimized, allowing you to quickly assemble state-of-the-art models without reinventing the wheel.
## Overview
Zeta provides a comprehensive library of modular components commonly used in modern AI architectures, including:
- **Attention Mechanisms**: Multi-query attention, sigmoid attention, flash attention, and more
- **Mixture of Experts (MoE)**: Efficient expert routing and gating mechanisms
- **Neural Network Modules**: Feedforward networks, activation functions, normalization layers
- **Quantization**: BitLinear, dynamic quantization, and other optimization techniques
- **Architectures**: Transformers, encoders, decoders, vision transformers, and complete model implementations
- **Training Utilities**: Optimization algorithms, logging, and performance monitoring
Each component is designed to be:
- **Modular**: Drop-in replacements that work seamlessly with PyTorch
- **High-Performance**: Optimized implementations with fused kernels where applicable
- **Well-Tested**: Comprehensive test coverage ensuring reliability
- **Production-Ready**: Used in hundreds of models across various domains
## Installation
```bash
pip3 install -U zetascale
```
## Quick Start
### Multi-Query Attention
Multi-query attention reduces memory usage while maintaining model quality by sharing key and value projections across attention heads.
```python
import torch
from zeta import MultiQueryAttention
# Initialize the model
model = MultiQueryAttention(
dim=512,
heads=8,
)
# Forward pass
text = torch.randn(2, 4, 512)
output, _, _ = model(text)
print(output.shape) # torch.Size([2, 4, 512])
```
### SwiGLU Activation
The SwiGLU activation function applies a gating mechanism to selectively pass information through the network.
```python
import torch
from zeta.nn import SwiGLUStacked
x = torch.randn(5, 10)
swiglu = SwiGLUStacked(10, 20)
output = swiglu(x)
print(output.shape) # torch.Size([5, 20])
```
### Relative Position Bias
Relative position bias quantizes the distance between positions into buckets and uses embeddings to provide position-aware attention biases.
```python
import torch
from torch import nn
from zeta.nn import RelativePositionBias
# Initialize the module
rel_pos_bias = RelativePositionBias()
# Compute bias for attention mechanism
bias_matrix = rel_pos_bias(1, 10, 10)
# Use in custom attention
class CustomAttention(nn.Module):
def __init__(self):
super().__init__()
self.rel_pos_bias = RelativePositionBias()
def forward(self, queries, keys):
bias = self.rel_pos_bias(queries.size(0), queries.size(1), keys.size(1))
# Use bias in attention computation
return None
```
### FeedForward Network
A flexible feedforward module with optional GLU activation and LayerNorm, commonly used in transformer architectures.
```python
import torch
from zeta.nn import FeedForward
model = FeedForward(256, 512, glu=True, post_act_ln=True, dropout=0.2)
x = torch.randn(1, 256)
output = model(x)
print(output.shape) # torch.Size([1, 512])
```
### BitLinear Quantization
BitLinear performs linear transformation with quantization and dequantization, reducing memory usage while maintaining performance. Based on [BitNet: Scaling 1-bit Transformers for Large Language Models](https://arxiv.org/abs/2310.11453).
```python
import torch
from torch import nn
import zeta.quant as qt
class MyModel(nn.Module):
def __init__(self):
super().__init__()
self.linear = qt.BitLinear(10, 20)
def forward(self, x):
return self.linear(x)
model = MyModel()
input = torch.randn(128, 10)
output = model(input)
print(output.size()) # torch.Size([128, 20])
```
### PalmE: Multi-Modal Architecture
A complete implementation of the PalmE multi-modal model architecture, combining a ViT image encoder with a transformer decoder for vision-language tasks.
```python
import torch
from zeta.structs import (
AutoRegressiveWrapper,
Decoder,
Encoder,
Transformer,
ViTransformerWrapper,
)
class PalmE(torch.nn.Module):
"""
PalmE is a transformer architecture that uses a ViT encoder and a transformer decoder.
This implementation demonstrates how to combine Zeta's modular components to build
a complete multi-modal model architecture.
"""
def __init__(
self,
image_size=256,
patch_size=32,
encoder_dim=512,
encoder_depth=6,
encoder_heads=8,
num_tokens=20000,
max_seq_len=1024,
decoder_dim=512,
decoder_depth=6,
decoder_heads=8,
alibi_num_heads=4,
attn_kv_heads=2,
use_abs_pos_emb=False,
cross_attend=True,
alibi_pos_bias=True,
rotary_xpos=True,
attn_flash=True,
qk_norm=True,
):
super().__init__()
# Vision encoder
self.encoder = ViTransformerWrapper(
image_size=image_size,
patch_size=patch_size,
attn_layers=Encoder(
dim=encoder_dim,
depth=encoder_depth,
heads=encoder_heads
),
)
# Language decoder
self.decoder = Transformer(
num_tokens=num_tokens,
max_seq_len=max_seq_len,
use_abs_pos_emb=use_abs_pos_emb,
attn_layers=Decoder(
dim=decoder_dim,
depth=decoder_depth,
heads=decoder_heads,
cross_attend=cross_attend,
alibi_pos_bias=alibi_pos_bias,
alibi_num_heads=alibi_num_heads,
rotary_xpos=rotary_xpos,
attn_kv_heads=attn_kv_heads,
attn_flash=attn_flash,
qk_norm=qk_norm,
),
)
# Enable autoregressive generation
self.decoder = AutoRegressiveWrapper(self.decoder)
def forward(self, img: torch.Tensor, text: torch.Tensor):
"""Forward pass of the model."""
encoded = self.encoder(img, return_embeddings=True)
return self.decoder(text, context=encoded)
# Usage
img = torch.randn(1, 3, 256, 256)
text = torch.randint(0, 20000, (1, 1024))
model = PalmE()
output = model(img, text)
print(output.shape)
```
### U-Net Architecture
A complete U-Net implementation for image segmentation and generative tasks.
```python
import torch
from zeta.nn import Unet
model = Unet(n_channels=1, n_classes=2)
x = torch.randn(1, 1, 572, 572)
y = model(x)
print(f"Input shape: {x.shape}")
print(f"Output shape: {y.shape}")
```
### Vision Embeddings
Convert images into patch embeddings suitable for transformer-based vision models.
```python
import torch
from zeta.nn import VisionEmbedding
vision_embedding = VisionEmbedding(
img_size=224,
patch_size=16,
in_chans=3,
embed_dim=768,
contain_mask_token=True,
prepend_cls_token=True,
)
input_image = torch.rand(1, 3, 224, 224)
output = vision_embedding(input_image)
print(output.shape)
```
### Dynamic Quantization with Niva
Niva provides dynamic quantization for specific layer types, ideal for models with variable runtime activations.
```python
import torch
from torch import nn
from zeta import niva
# Load a pre-trained model
model = YourModelClass()
# Quantize the model dynamically
niva(
model=model,
model_path="path_to_pretrained_weights.pt",
output_path="quantized_model.pt",
quant_type="dynamic",
quantize_layers=[nn.Linear, nn.Conv2d],
dtype=torch.qint8,
)
```
### Fused Operations
Zeta includes several fused operations that combine multiple operations into single kernels for improved performance.
#### FusedDenseGELUDense
Fuses two dense operations with GELU activation for up to 2x speedup.
```python
import torch
from zeta.nn import FusedDenseGELUDense
x = torch.randn(1, 512)
model = FusedDenseGELUDense(512, 1024)
out = model(x)
print(out.shape) # torch.Size([1, 1024])
```
#### FusedDropoutLayerNorm
Fuses dropout and layer normalization for faster feedforward networks.
```python
import torch
from zeta.nn import FusedDropoutLayerNorm
model = FusedDropoutLayerNorm(dim=512)
x = torch.randn(1, 512)
output = model(x)
print(output.shape) # torch.Size([1, 512])
```
### Mamba: State Space Model
PyTorch implementation of the Mamba state space model architecture.
```python
import torch
from zeta.nn import MambaBlock
block = MambaBlock(dim=64, depth=1)
x = torch.randn(1, 10, 64)
y = block(x)
print(y.shape) # torch.Size([1, 10, 64])
```
### FiLM: Feature-wise Linear Modulation
Feature-wise Linear Modulation for conditional feature transformation.
```python
import torch
from zeta.nn import Film
film_layer = Film(dim=128, hidden_dim=64, expanse_ratio=4)
conditions = torch.randn(10, 128)
hiddens = torch.randn(10, 1, 128)
modulated_features = film_layer(conditions, hiddens)
print(modulated_features.shape) # torch.Size([10, 1, 128])
```
### Model Optimization
The `hyper_optimize` decorator` provides a unified interface for multiple optimization techniques.
```python
import torch
from zeta.nn import hyper_optimize
@hyper_optimize(
torch_fx=False,
torch_script=False,
torch_compile=True,
quantize=True,
mixed_precision=True,
enable_metrics=True,
)
def model(x):
return x @ x
out = model(torch.randn(1, 3, 32, 32))
print(out)
```
### Direct Policy Optimization (DPO)
DPO implementation for reinforcement learning from human feedback (RLHF) applications.
```python
import torch
from torch import nn
from zeta.rl import DPO
class PolicyModel(nn.Module):
def __init__(self, dim, output_dim):
super().__init__()
self.fc = nn.Linear(dim, output_dim)
def forward(self, x):
return self.fc(x)
dim = 10
output_dim = 5
policy_model = PolicyModel(dim, output_dim)
dpo_model = DPO(model=policy_model, beta=0.1)
preferred_seq = torch.randint(0, output_dim, (3, dim))
unpreferred_seq = torch.randint(0, output_dim, (3, dim))
loss = dpo_model(preferred_seq, unpreferred_seq)
print(loss)
```
### PyTorch Model Logging
A decorator for comprehensive model execution logging, including parameters, gradients, and memory usage.
```python
import torch
from torch import nn
from zeta.utils.verbose_execution import verbose_execution
@verbose_execution(log_params=True, log_gradients=True, log_memory=True)
class YourPyTorchModel(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 64, 3)
self.relu = nn.ReLU()
self.flatten = nn.Flatten()
self.fc = nn.Linear(64 * 222 * 222, 10)
def forward(self, x):
x = self.conv1(x)
x = self.relu(x)
x = self.flatten(x)
x = self.fc(x)
return x
model = YourPyTorchModel()
input_tensor = torch.randn(1, 3, 224, 224)
output = model(input_tensor)
# Gradient information requires backward pass
loss = output.sum()
loss.backward()
```
### Sigmoid Attention
An attention mechanism that replaces softmax with sigmoid, providing up to 18% speedup while maintaining performance.
```python
import torch
from zeta import SigmoidAttention
batch_size = 32
seq_len = 128
dim = 512
heads = 8
x = torch.rand(batch_size, seq_len, dim)
mask = torch.ones(batch_size, seq_len, seq_len)
sigmoid_attn = SigmoidAttention(dim, heads, seq_len)
output = sigmoid_attn(x, mask)
print(output.shape) # torch.Size([32, 128, 512])
```
## Documentation
Comprehensive documentation is available at [zeta.apac.ai](https://zeta.apac.ai/).
## Quick Examples
There are various examples that you can try out in the [examples folder](examples/README.md)
## Running Tests
To run the full test suite:
```bash
python3 -m pip install -e '.[testing]' # Install extra dependencies for testing
python3 -m pytest tests/ # Run the entire test suite
```
For more details, refer to the CI workflow configuration.
## Community
Join our growing community for real-time support, ideas, and discussions on building better AI models.
| Platform | Link | Description |
|-------------|------------------------------------------------------------------------------|-----------------------------|
| Docs | [zeta.apac.ai](https://zeta.apac.ai) | Official documentation |
| Discord | [Join our Discord](https://discord.gg/EamjgSaEQf) | Live chat & community |
| Twitter | [@kyegomez](https://twitter.com/kyegomez) | Follow for updates |
| LinkedIn | [The Swarm Corporation](https://www.linkedin.com/company/the-swarm-corporation) | Connect professionally |
| YouTube | [YouTube Channel](https://www.youtube.com/channel/UC9yXyitkbU_WSy7bd_41SqQ) | Watch our videos |
## Contributing
Zeta is an open-source project, and contributions are welcome! If you want to create new features, fix bugs, or improve the infrastructure, we'd love to have you contribute.
**Getting Started:**
- Pick any issue with the `good first issue` tag to get started
- Read our [Contributing Guidelines](CONTRIBUTING.md)
- Check out our [contributing board](https://github.com/users/kyegomez/projects/1) for roadmap discussions
**Report Issues:**
- [Bug Report](https://github.com/kyegomez/zeta/issues/new/choose)
- [Feature Request](https://github.com/kyegomez/zeta/issues/new/choose)
## Our Contributors
Thank you to all of our contributors who have built this great framework 🙌
---
## Citation
If you use Zeta in your research or projects, please cite it:
```bibtex
@misc{zetascale,
title = {Zetascale Framework},
author = {Kye Gomez},
year = {2024},
howpublished = {\url{https://github.com/kyegomez/zeta}},
}
```
## License
Apache 2.0 License