Cactus-Compute/needle · Hugging Face 来源：https://huggingface.co/Cactus-Compute/needle 我们将 Gemini 3.1 蒸馏成一个 2600 万参数的“简单注意力网络（Simple Attention Network，https://huggingface.co/Cactus-Compute/needle/blob/main/docs/simple_attention_networks.md）”，你甚至可以在本地的 Mac/PC 上对其进行微调。在生产环境中，Needle 运行于 Cactus（https://github.com/cactus-compute/cactus）之上，预填充速度达 6000 tok/s，解码速度达 1200 tok/s。权重完全开放于 Cactus-Compute/needle（https://huggingface.co/Cactus-Compute/needle），数据集生成也同样开放。 参数26M架构Encoder-decoder，纯注意力（无 FFN）Encoder12 层，GQA（8H/4KV），RoPE，门控残差Decoder8 层，自注意力 + 交叉注意力，门控残差d_model512词表8192（SentencePiece BPE）归一化ZCRMSNorm（零中心化，init=0）精度bfloat16（训练期间使用 INT4 QAT）预训练200B tokens，在 16×TPU v6e 上训练（27 小时）后训练2B tokens 的函数调用数据（45 分钟） d=512, 8H/4KV, BPE=8192 ┌──────────────┐ │ Tool Call │ └──────┬───────┘ ┌┴──────────┐ │ Softmax │ └─────┬─────┘ ┌─────┴─────┐ │ Linear (T)│ <- tied └─────┬─────┘ ┌─────┴─────┐ │ ZCRMSNorm │ └─────┬─────┘ ┌────────┴────────┐ │ Decoder x 8 │ │┌───────────────┐│ ││ ZCRMSNorm ││ ││ Masked Self ││ ││ Attn + RoPE ││ ││ Gated Residual││ │├───────────────┤│ ┌──────────────┐ ││ ZCRMSNorm ││ │ Encoder x 12 │─────────────────────>Cross Attn ││ │ │ ││ Gated Residual││ │ ┌──────────┐ │ │└───────────────┘│ │ │ZCRMSNorm │ │ └────────┬────────┘ │ │Self Attn │ │ ┌─────┴─────┐ │ │ GQA+RoPE │ │ │ Embedding │ <- shared │ │Gated Res │ │ └─────┬─────┘ │ │ │ │ ┌───────┴────────┐ │ │ (no FFN) │ │ │[EOS]│ │ └──────────┘ │ │ + answer │ │ │ └────────────────┘ └──────┬───────┘ │ ┌────┴──────┐ │ Embedding │ └────┬──────┘ │ ┌────┴──────┐ │ Text │ │ query │ └───────────┘ ## 快速开始 git clone https://github.com/cactus-compute/needle.git cd needle && source ./setup needle playground 打开一个 Web 界面，地址为 http://127.0.0.1:7860/，你可以在其中测试和在自己的工具上微调模型。权重会自动下载。 ## Python 用法 from needle import load_checkpoint, generate, SimpleAttentionNetwork, get_tokenizer params, config = load_checkpoint("checkpoints/needle.pkl") model = SimpleAttentionNetwork(config) tokenizer = get_tokenizer() result = generate( model, params, tokenizer, query="What's the weather in San Francisco?", tools='[{"name":"get_weather","parameters":{"location":"string"}}]', stream=False, ) print(result) # [{"name":"get_weather","arguments":{"location":"San Francisco"}}] ## 微调 通过 Web 界面或 CLI 在自己的工具上进行微调： # Web 界面（通过 Gemini 生成数据、训练、评估、打包结果） needle playground # CLI（若本地没有权重则自动下载） needle finetune data.jsonl ## 链接 - Needle（https://github.com/cactus-compute/needle）——训练、微调和推理代码 - Cactus（https://github.com/cactus-compute/cactus）——设备端运行时（6000 tok/s 预填充，1200 tok/s 解码） - Simple Attention Networks（https://github.com/cactus-compute/needle/blob/main/docs/simple_attention_networks.md）——架构详情 ## 许可证 MIT ## 引用 @misc{ndubuaku2026needle, title={Needle}, author={Henry Ndubuaku and Jakub Mroz and Karen Mosoyan and Roman Shemet and Parkirat Sandhu and Satyajit Kumar and Noah Cylich and Justin H. Lee}, year={2026}, url={https://github.com/cactus-compute/needle} }