Zeta 是一个模块化的 PyTorch 框架，旨在通过提供可重用、高性能的构建模块来简化 AI 模型的开发。你可以把它想象成一套用于 AI 的乐高积木——每个组件都经过精心设计、测试和优化，使你能够快速搭建最先进的模型，而无需从头开始重新发明轮子。

概述

Zeta 提供了一个全面的模块化组件库，这些组件广泛应用于现代 AI 架构中，包括：

• 注意力机制：多查询注意力、sigmoid 注意力、FlashAttention 等
• 专家混合（MoE）：高效的专家路由和门控机制
• 神经网络模块：前馈网络、激活函数、归一化层
• 量化：BitLinear、动态量化等优化技术
• 架构：Transformer、编码器、解码器、视觉 Transformer，以及完整的模型实现
• 训练工具：优化算法、日志记录和性能监控

每个组件的设计特点如下：

• 模块化：即插即用，与 PyTorch 无缝兼容
• 高性能：在适用情况下采用融合内核进行优化实现
• 经过充分测试：拥有全面的测试覆盖率，确保可靠性
• 生产就绪：已在多个领域的数百个模型中得到应用

安装

pip3 install -U zetascale

快速入门

多查询注意力

多查询注意力通过在不同注意力头之间共享键和值投影，从而在保持模型质量的同时降低内存占用。

import torch
from zeta import MultiQueryAttention

# 初始化模型
model = MultiQueryAttention(
    dim=512,
    heads=8,
)

# 前向传播
text = torch.randn(2, 4, 512)
output, _, _ = model(text)
print(output.shape)  # torch.Size([2, 4, 512])

SwiGLU 激活函数

SwiGLU 激活函数通过门控机制有选择地在网络中传递信息。

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])

相对位置偏置

相对位置偏置将位置之间的距离量化为若干桶，并利用嵌入向量提供位置感知的注意力偏置。

import torch
from torch import nn
from zeta.nn import RelativePositionBias

# 初始化模块
rel_pos_bias = RelativePositionBias()

# 计算注意力机制的偏置矩阵
bias_matrix = rel_pos_bias(1, 10, 10)

# 在自定义注意力中使用
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))
        # 在注意力计算中使用偏置
        return None

前馈网络

一个灵活的前馈模块，带有可选的 GLU 激活函数和 LayerNorm，常用于 Transformer 架构中。

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 量化

BitLinear 通过量化和反量化执行线性变换，能够在保持性能的同时减少内存占用。其灵感来源于 BitNet: 面向大型语言模型的 1 位 Transformer 扩展。

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：多模态架构

PalmE 是一种完整的多模态模型架构实现，结合了 ViT 图像编码器和 Transformer 解码器，用于视觉-语言任务。

import torch
from zeta.structs import (
    AutoRegressiveWrapper,
    Decoder,
    Encoder,
    Transformer,
    ViTransformerWrapper,
)

class PalmE(torch.nn.Module):
    """
    PalmE 是一种使用 ViT 编码器和 Transformer 解码器的 Transformer 架构。
    
    该实现展示了如何结合 Zeta 的模块化组件来构建一个完整的多模态模型架构。
    """
    
    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__()
        
        # 视觉编码器
        self.encoder = ViTransformerWrapper(
            image_size=image_size,
            patch_size=patch_size,
            attn_layers=Encoder(
                dim=encoder_dim, 
                depth=encoder_depth, 
                heads=encoder_heads
            ),
        )
        
        # 语言解码器
        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,
            ),
        )
        
        # 启用自回归生成
        self.decoder = AutoRegressiveWrapper(self.decoder)
    
    def forward(self, img: torch.Tensor, text: torch.Tensor):
        """模型前向传播。"""
        encoded = self.encoder(img, return_embeddings=True)
        return self.decoder(text, context=encoded)

# 使用示例
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 架构

用于图像分割和生成任务的完整 U-Net 实现。

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"输入形状: {x.shape}")
print(f"输出形状: {y.shape}")

视觉嵌入

将图像转换为适合基于 Transformer 的视觉模型的补丁嵌入。

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)

使用 Niva 进行动态量化

Niva 提供针对特定层类型的动态量化功能，非常适合具有可变运行时激活的模型。

import torch
from torch import nn
from zeta import niva

# 加载预训练模型
model = YourModelClass()

# 对模型进行动态量化
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,
)

融合操作

Zeta 包含多种融合操作，可以将多个操作合并为单个内核，从而提升性能。

FusedDenseGELUDense

将两个全连接层与 GELU 激活函数融合，速度最高可提升 2 倍。

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

将 Dropout 和层归一化融合，以加快前馈网络的速度。

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：状态空间模型

Mamba 状态空间模型架构的 PyTorch 实现。

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：特征线性调制

用于条件特征变换的特征线性调制。

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])

模型优化

hyper_optimize 装饰器提供了一个统一的接口，用于多种优化技术。

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)

直接策略优化（DPO）

用于从人类反馈中学习（RLHF）应用的 DPO 实现。

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 模型日志记录

一个用于全面模型执行日志记录的装饰器，包括参数、梯度和内存使用情况。

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)

# 梯度信息需要反向传播
loss = output.sum()
loss.backward()

Sigmoid 注意力机制

一种用 sigmoid 替代 softmax 的注意力机制，可在保持性能的同时提供高达 18% 的速度提升。

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])

文档

完整的文档可在 zeta.apac.ai 上找到。

快速示例

您可以在 examples 文件夹 中尝试各种示例。

运行测试

要运行完整的测试套件：

python3 -m pip install -e '.[testing]'  # 安装测试所需的额外依赖
python3 -m pytest tests/                # 运行整个测试套件

更多详细信息，请参阅 CI 工作流配置。

社区

加入我们不断壮大的社区，获取实时支持、创意以及关于构建更优 AI 模型的讨论。

平台	链接	描述
文档	zeta.apac.ai	官方文档
Discord	加入我们的 Discord	实时聊天与社区
Twitter	@kyegomez	关注以获取最新动态
LinkedIn	The Swarm Corporation	专业交流
YouTube	YouTube 频道	观看我们的视频

贡献

Zeta 是一个开源项目，欢迎各位贡献！如果您想添加新功能、修复 bug 或改进基础设施，我们都十分期待您的参与。

开始贡献：

• 选择带有 good first issue 标签的问题作为起点
• 阅读我们的 贡献指南
• 查看我们的 贡献看板 以了解路线图讨论

报告问题：

• Bug 报告
• 功能请求

我们的贡献者

感谢所有为我们构建这一优秀框架的贡献者 🙌

---

引用

如果您在研究或项目中使用 Zeta，请引用如下：

@misc{zetascale,
    title = {Zetascale 框架},
    author = {Kye Gomez},
    year = {2024},
    howpublished = {\url{https://github.com/kyegomez/zeta}},
}

许可证

Apache 2.0 许可证

Zeta 快速上手指南

Zeta 是一个模块化的 PyTorch 框架，旨在通过提供可复用、高性能的构建块来简化 AI 模型的开发。你可以将其视为 AI 领域的"LEGO 积木”，快速组装出最先进的模型架构。

环境准备

在开始之前，请确保你的开发环境满足以下要求：

• 操作系统: Linux, macOS, 或 Windows
• Python 版本: Python 3.8 及以上
• 核心依赖: PyTorch (建议安装与你的 CUDA 版本匹配的最新版)
• 包管理工具: pip

安装步骤

使用 pip 直接安装最新版本的 zetascale 包：

pip3 install -U zetascale

提示：国内开发者若遇到下载速度慢的问题，可使用清华源或阿里源加速安装：

pip3 install -U zetascale -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

Zeta 的核心优势在于其模块化组件。以下是一个最基础的示例，展示如何使用 Multi-Query Attention（多查询注意力机制）模块。该机制能在保持模型质量的同时显著降低显存占用。

示例：使用 Multi-Query Attention

import torch
from zeta import MultiQueryAttention

# 1. 初始化模型组件
# dim: 嵌入维度，heads: 注意力头数
model = MultiQueryAttention(
    dim=512,
    heads=8,
)

# 2. 准备输入数据 (Batch_Size, Sequence_Length, Dimension)
text = torch.randn(2, 4, 512)

# 3. 执行前向传播
output, _, _ = model(text)

# 4. 查看输出形状
print(output.shape)  # 输出: torch.Size([2, 4, 512])

更多常用组件

Zeta 提供了丰富的即插即用模块，以下是几个高频使用的场景：

1. SwiGLU 激活函数

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])

2. 量化线性层 (BitLinear) 基于 BitNet 技术，减少内存占用并保持性能。

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])

3. 完整架构示例：U-Net 直接调用完整的 U-Net 架构用于图像分割任务。

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"Output shape: {y.shape}")

通过以上示例，你可以快速将 Zeta 的高性能模块集成到自己的 PyTorch 项目中，无需重复造轮子。