cake

加入我们的项目社区服务器！ <a href="https://discord.gg/https://discord.gg/btZpkp45gQ" target="_blank" title="加入我们的社区！"> <img src="https://dcbadge.limes.pink/api/server/https://discord.gg/btZpkp45gQ"/> </a> <hr/>

Cake是一个基于Candle的Rust框架，用于分布式推理大型模型，如LLama3和Stable Diffusion。该项目的目标是通过将消费级硬件重新用作iOS、Android、macOS、Linux和Windows设备的异构集群，来运行大型（70B+）模型，有效地利用计划报废作为工具，使人工智能更加accessible和民主化。

⚠ 这是正在积极开发和快速变化的实验性代码，请预期会有bug ⚠

这个想法是将transformer块分片到多个设备上，以便能够对通常无法装入单个设备GPU内存的模型进行推理。同一工作节点上连续transformer块的推理会被批处理，以最小化数据传输造成的延迟。

支持

操作系统	架构	加速	状态
GNU/Linux	arm, arm64, x86_64	-	✅
GNU/Linux	arm, arm64, x86_64	CUDA	✅
GNU/Linux	arm, arm64, x86_64	BLAS	✅
Windows	x86_64	BLAS	未测试
Windows	x86_64	CUDA	✅
macOS	x86_64	-	✅
macOS	aarch64	-	✅
macOS	aarch64	Metal	✅
Android	arm, arm64, x86_64	-	✅
Android	arm, arm64, x86_64	CUDA	未测试
iOS / iPadOS	aarch64	-	✅
iOS / iPadOS	aarch64	Metal	🛠️ 90% 完成，进行中
Web	-	WebGPU	理论上可行，尚未完成

CUDA加速系统需要CUDA >= 12.2。

编译

安装Rust后，您可以使用不同的加速方式构建核心库和CLI工具。

无加速（将使用CPU）：

cargo build --release

Apple Silicon的Metal加速：

cargo build --release --features metal

CUDA加速：

cargo build --release --features cuda

生成可通过XCode编译和部署的iOS绑定：

make ios

使用方法

运行工作节点：

cake-cli --model /path/to/Meta-Llama-3-8B \ # 模型路径，阅读下文了解如何优化工作节点的模型大小
         --mode worker \                    # 以工作节点模式运行
         --name worker0 \                   # 拓扑文件中的工作节点名称
         --topology topology.yml \          # 拓扑
         --address 0.0.0.0:10128            # 绑定地址

运行带有OpenAI兼容REST API的主节点：

cake-cli --model /path/to/Meta-Llama-3-8B \ # 模型路径
         --api 0.0.0.0:8080               \ # API绑定地址
         --topology topology.yml            # 拓扑文件

topology.yml 决定了哪些层由哪个工作节点提供服务（你可以在模型的张量索引文件中找到所有层的列表）：

linux_server_1:
  host: 'linux_server.host:10128'
  description: 'NVIDIA Titan X Pascal (12GB)'
  layers:
    - 'model.layers.0-5'

linux_server_2:
  host: 'linux_server2.host:10128'
  description: 'NVIDIA GeForce 3080 (10GB)'
  layers:
    - 'model.layers.6-16'

iphone:
  host: 'iphone.host:10128'
  description: 'iPhone 15 Pro Max'
  layers:
    - 'model.layers.17'

ipad:
  host: 'ipad.host:10128'
  description: 'iPad'
  layers:
    - 'model.layers.18-19'

macbook:
  host: 'macbook.host:10128'
  description: 'M1 Max'
  layers:
    - 'model.layers.20-31'

现在你可以通过以下方式与集群交互：

curl http://master-ip:8080/api/v1/chat/completions \                                                                                                                           ~
  -H "Content-Type: application/json" \
  -d '{
    "messages": [
        {
            "role": "system",
            "content": "你是一个乐于助人的AI助手。"
        },
        {
            "role": "user",
            "content": "为什么天空是蓝色的？"
        }
    ]
}'

拆分模型

作为内存和磁盘空间优化，你可能希望只给工作节点提供它实际需要的模型数据，而不是整个文件夹。在这种情况下，你可以使用 cake-split-model 工具。例如，要生成 llama3 safetensors 的较小版本，你可以：

cake-split-model --model-path path/to/Meta-Llama-3-8B \ # 要拆分的源模型
                 --topology path/to/topology.yml \      # 拓扑文件
                 --output output-folder-name            # 保存所有工作节点数据包的输出文件夹

这将创建一个较小的文件夹，其中只包含所需层的张量和特定工作节点的拓扑文件。在部署工作节点之前，记得也要将其他模型内容（config.json、tokenizer.json 等）复制到工作节点包中。

Stable Diffusion 图像生成

在 topology.yml 中定义模型部分：

wsl2_on_windows:
  host: 192.168.1.2:10128
  description: NVIDIA RTX 4090 24GB
  layers:
  - unet

macbook:
  host: 192.168.1.3:10128
  description: Macbook M2
  layers:
  - clip
  - vae

运行工作节点：

cake-cli --model /path/to/hf/cache \        # Huggingface 模型的缓存目录
         --mode worker \                    # 以工作节点模式运行
         --name wsl2_on_windows \           # 拓扑文件中的工作节点名称
         --model-type image-model \         # 对 SD 使用 image-model，对 LLM 使用 text-model 或跳过
         --topology topology.yml \          # 拓扑文件
         --address 0.0.0.0:10128            # 绑定地址

可以通过指定更多命令行参数在 SD1.5、SD2.1、SDXL 和 SDXL Turbo 之间切换模型。

如果在本地缓存目录中找不到模型文件，将自动从 Huggingface 下载。

运行带有 REST API 的主节点：

cake-cli --model /path/to/hf/cache \        # Huggingface 模型的缓存目录
         --api 0.0.0.0:8080 \               # API 绑定地址
         --model-type image-model \         # 对 SD 使用 image-model，对 LLM 使用 text-model 或跳过
         --topology topology.yml            # 拓扑文件

使用集群生成图像：

curl http://master-ip:8080/api/v1/image \                                                                                                                           ~
  -H "Content-Type: application/json" \
  -d '{
    "image_args": {
      "sd-image-prompt": "一个老人坐在海边的椅子上",
      "sd-num-samples": 1,
      "sd-image-seed": 2439383
    }
}'

更多控制参数可以在代码中找到。