使用模块化、高性能和可扩展的构建块,构建最先进的AI模型速度提高80%!
![在Hacker News上分享](https://img.shields.io/badge/-在Hacker News上分享-orange)
在构建了数千个神经网络并面临相同的令人烦恼的瓶颈(混乱的代码库、缺乏模块化和低性能模块)之后,Zeta需要诞生,以使我和其他人能够快速原型化、训练和优化最新的SOTA神经网络并将其部署到生产环境中。
Zeta极度强调可用性、模块化和性能。Zeta现在已经在我的GitHub和其他人的GitHub上的数百个模型中得到应用。 从下面开始,如果你需要我的帮助来构建任何模型,我随时为你效劳😊💜
安装
$ pip3 install -U zetascale
使用
开始你的旅程
创建一个具有上述突破性研究特性的模型非常简单。以下是如何快速实现著名的Flash Attention
import torch
from zeta.nn import FlashAttention
q = torch.randn(2, 4, 6, 8)
k = torch.randn(2, 4, 10, 8)
v = torch.randn(2, 4, 10, 8)
attention = FlashAttention(causal=False, dropout=0.1, flash=True)
output = attention(q, k, v)
print(output.shape)
SwiGLU
SwiGLU激活函数接受一个输入张量并应用门控机制来选择性地传递信息。它由两部分组成: "switch"门和"glu"门。switch门控制信息流,而glu门对输入执行非线性变换。
import torch
from zeta.nn import SwiGLUStacked
x = torch.randn(5, 10)
swiglu = SwiGLUStacked(10, 20)
swiglu(x).shape
在这个例子中,我们首先导入必要的模块,包括torch用于张量操作,以及zeta.nn中的SwiGLUStacked用于SwiGLU激活函数。
然后我们创建一个形状为(5, 10)的随机输入张量x。接下来,我们实例化一个SwiGLUStacked,输入大小为10,输出大小为20。
最后,我们将输入张量x传递给swiglu模块,它对其应用SwiGLU激活函数。结果输出张量存储在output变量中。我们打印输出张量的形状以查看结果。
RelativePositionBias
RelativePositionBias
将两个位置之间的距离量化为一定数量的桶,然后使用嵌入来获得相对位置偏差。这种机制通过提供基于查询和键之间相对位置的偏差来辅助注意力机制,而不是仅仅依赖它们的绝对位置。
import torch
from torch import nn
from zeta.nn import RelativePositionBias
# 初始化RelativePositionBias模块
rel_pos_bias = RelativePositionBias()
# 示例1:计算单个批次的偏差
bias_matrix = rel_pos_bias(1, 10, 10)
# 示例2:与注意力机制结合使用
# 注意:这是一个模拟示例,可能不代表实际注意力机制的完整实现。
class MockAttention(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 # 占位符
# 示例3:修改默认配置
custom_rel_pos_bias = RelativePositionBias(
bidirectional=False, num_buckets=64, max_distance=256, num_heads=8
)
FeedForward
FeedForward模块对输入张量x执行前馈操作。它由一个多层感知器(MLP)组成,可选择性地包含激活函数和LayerNorm。 在大多数语言、多模态和现代神经网络中使用。
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)
BitLinear
- BitLinear模块对输入数据执行线性变换,然后进行量化和反量化。量化过程使用absmax_quantize函数,该函数根据绝对最大值对输入张量进行量化,来自论文
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()
# 创建一个大小为(128, 10)的随机张量
input = torch.randn(128, 10)
# 执行前向传播
output = model(input)
# 打印输出的大小
print(output.size()) # torch.Size([128, 20])
PalmE
- 这是使用解码器LLM作为主干和VIT图像编码器处理视觉的多模态Palm-E模型的实现,它与GPT4、Kosmos、RTX2和许多其他多模态模型架构非常相似
import torch
from zeta.structs import (
AutoRegressiveWrapper,
Decoder,
Encoder,
Transformer,
ViTransformerWrapper,
)
class PalmE(torch.nn.Module):
"""
PalmE是一个使用ViT编码器和transformer解码器的transformer架构。
参数:
image_size (int):图像大小。
patch_size (int):补丁大小。
encoder_dim (int):编码器维度。
encoder_depth (int):编码器深度。
encoder_heads (int):编码器头数。
num_tokens (int):令牌数量。
max_seq_len (int):最大序列长度。
decoder_dim (int):解码器维度。
decoder_depth (int):解码器深度。
decoder_heads (int):解码器头数。
alibi_num_heads (int):alibi注意力中的头数。
attn_kv_heads (int):注意力键值投影中的头数。
use_abs_pos_emb (bool):是否使用绝对位置嵌入。
cross_attend (bool):解码器中是否进行交叉注意力。
alibi_pos_bias (bool):是否在alibi注意力中使用位置偏差。
rotary_xpos (bool):是否使用旋转位置嵌入。
attn_flash (bool):是否使用注意力闪存。
qk_norm (bool):是否在注意力层中对查询和键进行归一化。
返回:
torch.Tensor:模型的输出。
用法:
img = torch.randn(1, 3, 256, 256)
text = torch.randint(0, 20000, (1, 1024))
model = PalmE()
output = model(img, text)
print(output)
"""
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__()
# vit架构
self.encoder = ViTransformerWrapper(
image_size=image_size,
patch_size=patch_size,
attn_layers=Encoder(
dim=
# 动态量化模型,指定要量化的层
niva(
model=model,
model_path="pretrained_model_weights.pt的路径",
output_path="量化后的模型.pt",
quant_type="dynamic",
quantize_layers=[nn.Linear, nn.Conv2d],
dtype=torch.qint8,
)
FusedDenseGELUDense
- 使用此模块将模型速度提高2倍,它将来自bits and bytes的2个超优化密集操作和gelu融合在一起!
import torch
from zeta.nn import FusedDenseGELUDense
x = torch.randn(1, 512)
model = FusedDenseGELUDense(512, 1024)
out = model(x)
out.shape
FusedDropoutLayerNorm
- FusedDropoutLayerNorm是dropout和layernorm的融合内核,可以将FFN或MLP的速度提高2倍
import torch
from torch import nn
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
- 新的SSM模型架构Mamba的Pytorch实现
import torch
from zeta.nn import MambaBlock
# 初始化Mamba
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层
film_layer = Film(dim=128, hidden_dim=64, expanse_ratio=4)
# 为条件和隐藏层创建一些虚拟数据
conditions = torch.randn(10, 128) # 批次大小为10,特征大小为128
hiddens = torch.randn(
10, 1, 128
) # 批次大小为10,序列长度为1,特征大小为128
# 将数据通过Film层
modulated_features = film_layer(conditions, hiddens)
# 打印输出的形状
print(modulated_features.shape) # 应该是[10, 1, 128]
hyper_optimize
- 一个包含torch.fx、torch.script、torch.compile、动态量化、通过torch.amp进行混合精度以及执行时间指标的单一包装器!
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 - 直接策略优化
直接策略优化用于许多LLM的RLHF应用。
import torch
from torch import nn
from zeta.rl import DPO
# 定义一个简单的策略模型
class PolicyModel(nn.Module):
def __init__(self, input_dim, output_dim):
super().__init__()
self.fc = nn.Linear(input_dim, output_dim)
def forward(self, x):
return self.fc(x)
input_dim = 10
output_dim = 5
policy_model = PolicyModel(input_dim, output_dim)
# 用策略模型初始化DPO
dpo_model = DPO(model=policy_model, beta=0.1)
# 采样首选和非首选序列
preferred_seq = torch.randint(0, output_dim, (3, input_dim))
unpreferred_seq = torch.randint(0, output_dim, (3, input_dim))
# 计算损失
loss = dpo_model(preferred_seq, unpreferred_seq)
print(loss)
文档
所有类都必须有文档,如果你看到没有文档的类或函数,请向我报告,邮箱是kye@apac.ai
文档位于zeta.apac.ai
运行测试
你应该用pre-commit install安装预提交钩子。这将在每次提交时运行linter、mypy和部分测试。
关于如何运行完整测试套件的更多例子,请参考CI工作流。
在本地运行测试的一些例子:
python3 -m pip install -e '.[testing]' # 安装测试所需的额外依赖
python3 -m pytest tests/ # 整个测试套件
社区
加入我们遍布全球的不断增长的社区,获得实时支持、想法,并讨论如何构建更好的模型😊
🤝 安排一对一会话
想要为现实世界的任务训练自定义AI模型,比如通用多模态模型、人脸识别、药物发现、人形机器人吗?我会帮你创建模型架构,然后训练模型并优化它以满足你的质量保证标准。
在这里预约与创始人Kye的一对一会话,讨论任何问题,提供反馈,或探讨我们如何改进Zeta以满足你的需求或帮助你构建自己的定制模型!
🫶 贡献:
最简单的贡献方式是选择任何带有good first issue
标签的问题💪。在这里阅读贡献指南。发现Bug?在这里提交 | 功能请求?在这里提交
Zeta是一个开源项目,非常欢迎贡献。如果你想贡献,你可以创建新功能、修复bug或改进基础设施。请参考CONTRIBUTING.md和我们的贡献板来参与路线图讨论!
加速积压工作
通过财务支持帮助我们加速积压工作!注意,我们是一个开源公司,所以目前我们产生的所有收入都是通过捐赠;)
许可证
- Apache
引用
@misc{zetascale,
title = {Zetascale Framework},
author = {Kye Gomez},
year = {2024},
howpublished = {\url{https://github.com/kyegomez/zeta}},
}