diff --git a/Cargo.toml b/Cargo.toml index 4705281b..1fd51bb7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "candle-vllm" -version = "0.2.1" +version = "0.3.0" edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html diff --git a/README-CN.md b/README-CN.md index 599d1081..49f3c1a1 100644 --- a/README-CN.md +++ b/README-CN.md @@ -88,29 +88,25 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #同时包含flash att ### 构建/运行参数 -- [`ENV_PARAM`] cargo run [`BUILD_PARAM`] -- [`PROGRAM_PARAM`] [`MODEL_ID/MODEL_WEIGHT_PATH`] [`MODEL_TYPE`] [`MODEL_PARAM`] +- [`ENV_PARAM`] cargo run [`BUILD_PARAM`] -- [`PROGRAM_PARAM`] [`MODEL_ID/MODEL_WEIGHT_PATH`]
显示详情 **示例:** ```shell - [RUST_LOG=warn] cargo run [--release --features cuda,nccl] -- [--multi-process --log --dtype bf16 --p 2000 --d "0,1" --mem 8192] [--w /home/weights/Qwen3-27B-GPTQ-4Bit] [qwen3] [--quant gptq --temperature 0.7 --penalty 1.0 --top-k 32 --top-p 0.95 --thinking] + [RUST_LOG=warn] cargo run [--release --features cuda,nccl] -- [--log --dtype bf16 --p 2000 --d 0,1 --mem 8192] [--w /home/weights/Qwen3-27B-GPTQ-4Bit] ``` `ENV_PARAM`: RUST_LOG=warn `BUILD_PARAM`: --release --features cuda,nccl - `PROGRAM_PARAM`：--multi-process --log --dtype bf16 --p 2000 --d "0,1" --mem 8192 + `PROGRAM_PARAM`：--log --dtype bf16 --p 2000 --d 0,1 --mem 8192 `MODEL_WEIGHT_PATH`: --w /home/weights/Qwen3-27B-GPTQ-4Bit - `MODEL_TYPE`: qwen3 - - `MODEL_PARAM`: --quant gptq --temperature 0.7 --penalty 1.0 --top-k 32 --top-p 0.95 --thinking - - 其中，`mem` (`kvcache-mem-gpu`) 参数控制KV Cache缓存，长文本或批量推理量请增大缓存；`MODEL_TYPE`可选值为：["llama", "llama3", "mistral", "phi2", "phi3", "qwen2", "qwen3", "glm4", "gemma", "gemma3", "yi", "stable-lm", "deep-seek"] + 其中，`mem` (`kvcache-mem-gpu`) 参数控制KV Cache缓存，长文本或批量推理量请增大缓存；支持的模型架构有：["llama", "llama3", "mistral", "phi2", "phi3", "qwen2", "qwen3", "glm4", "gemma", "gemma3", "yi", "stable-lm", "deep-seek"]
## 如何运行？ @@ -122,13 +118,13 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #同时包含flash att **本地路径** ```shell - target/release/candle-vllm --p 2000 --w /home/DeepSeek-R1-Distill-Llama-8B/ llama3 --temperature 0. --penalty 1.0 + target/release/candle-vllm --w /home/DeepSeek-R1-Distill-Llama-8B/ ``` **模型ID（从Huggingface下载）** ```shell - target/release/candle-vllm --model-id deepseek-ai/DeepSeek-R1-0528-Qwen3-8B qwen3 + target/release/candle-vllm --m deepseek-ai/DeepSeek-R1-0528-Qwen3-8B ``` @@ -140,13 +136,13 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #同时包含flash att **本地路径（指定端口、数据类型、采样参数）** ```shell - target/release/candle-vllm --p 2000 --dtype bf16 --f /home/data/DeepSeek-R1-0528-Qwen3-8B-Q2_K.gguf qwen3 --quant gguf --temperature 0.7 --penalty 1.1 + target/release/candle-vllm --f /home/data/DeepSeek-R1-0528-Qwen3-8B-Q2_K.gguf ``` **模型ID（从Huggingface下载）** ```shell - target/release/candle-vllm --model-id unsloth/DeepSeek-R1-0528-Qwen3-8B-GGUF --f DeepSeek-R1-0528-Qwen3-8B-Q2_K.gguf qwen3 --quant gguf + target/release/candle-vllm --m unsloth/DeepSeek-R1-0528-Qwen3-8B-GGUF --f DeepSeek-R1-0528-Qwen3-8B-Q2_K.gguf ``` @@ -158,13 +154,13 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #同时包含flash att **本地路径（假设模型已下载到/home）** ```shell - cargo run --release --features metal -- --p 2000 --dtype bf16 --f /home/qwq-32b-q4_k_m.gguf qwen2 --quant gguf --temperature 0. --penalty 1.0 + cargo run --release --features metal -- --f /home/qwq-32b-q4_k_m.gguf qwen2 ``` **模型ID（从Huggingface下载）** ```shell - cargo run --release --features metal -- --p 2000 --dtype bf16 --model-id Qwen/QwQ-32B-GGUF --f qwq-32b-q4_k_m.gguf qwen2 --quant gguf --temperature 0. --penalty 1.0 + cargo run --release --features metal -- --m Qwen/QwQ-32B-GGUF --f qwq-32b-q4_k_m.gguf ``` @@ -173,13 +169,13 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #同时包含flash att
显示命令 - **只需在运行未量化模型时添加`quant`参数** + **只需在运行未量化模型时添加`isq`参数** ```shell - target/release/candle-vllm --p 2000 --w /home/DeepSeek-R1-Distill-Llama-8B/ llama3 --quant q4k --temperature 0. --penalty 1.0 + target/release/candle-vllm --w /home/DeepSeek-R1-Distill-Llama-8B/ llama3 --isq q4k ``` - 注：原位量化加载可能需要更长的加载时间，原位`quant`参数选项：["q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "q2k", "q3k","q4k","q5k","q6k"] + 注：原位量化加载可能需要更长的加载时间，原位`isq`参数选项：["q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "q2k", "q3k","q4k","q5k","q6k"]
- 运行**Marlin兼容的GPTQ模型**（4位GPTQ，128分组，desc_act=False） @@ -189,19 +185,19 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #同时包含flash att **本地路径** ```shell - target/release/candle-vllm --dtype bf16 --p 2000 --w /home/DeepSeek-R1-Distill-Qwen-14B-GPTQ_4bit-128g qwen2 --quant gptq --temperature 0. --penalty 1.0 + target/release/candle-vllm --w /home/DeepSeek-R1-Distill-Qwen-14B-GPTQ_4bit-128g ``` **模型ID（从Huggingface下载）** ```shell - target/release/candle-vllm --model-id thesven/Llama-3-8B-GPTQ-4bit llama3 --quant gptq + target/release/candle-vllm --m thesven/Llama-3-8B-GPTQ-4bit llama3 ``` **将未压缩模型转换为Marlin兼容格式** ```shell python3 examples/convert_marlin.py --src /home/DeepSeek-R1-Distill-Qwen-14B/ --dst /home/DeepSeek-R1-Distill-Qwen-14B-GPTQ_4bit-128g - target/release/candle-vllm --dtype bf16 --p 2000 --w /home/DeepSeek-R1-Distill-Qwen-14B-GPTQ_4bit-128g qwen2 --quant gptq --temperature 0. --penalty 1.0 + target/release/candle-vllm --w /home/DeepSeek-R1-Distill-Qwen-14B-GPTQ_4bit-128g ``` @@ -217,7 +213,7 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #同时包含flash att **运行转换后的AWQ模型** ```shell - target/release/candle-vllm --multi-process --dtype f16 --p 2000 --d "0" --w /home/Meta-Llama-3.1-8B-Instruct-AWQ-INT4-Marlin/ llama3 --quant awq --temperature 0. --penalty 1.0 + target/release/candle-vllm --d 0 --w /home/Meta-Llama-3.1-8B-Instruct-AWQ-INT4-Marlin/ ``` @@ -227,7 +223,7 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #同时包含flash att 显示命令 ```shell - target/release/candle-vllm --dtype bf16 --p 2000 --w /home/DeepSeek-R1-Distill-Qwen-14B-GPTQ-Marlin/ qwen2 --quant marlin --penalty 1.0 --temperature 0. + target/release/candle-vllm --w /home/DeepSeek-R1-Distill-Qwen-14B-GPTQ-Marlin/ ``` @@ -238,7 +234,7 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #同时包含flash att **在两块GPU上运行QwQ-32B BF16模型** ```shell - cargo run --release --features cuda,nccl -- --multi-process --dtype bf16 --p 2000 --d "0,1" --w /home/QwQ-32B/ qwen2 --penalty 1.0 --temperature 0. + cargo run --release --features cuda,nccl -- --d 0,1 --w /home/QwQ-32B/ ``` **在两块GPU上运行QwQ-32B 4位AWQ模型** @@ -250,7 +246,7 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #同时包含flash att 2) 运行转换后的AWQ模型 ```shell - cargo run --release --features cuda,nccl -- --multi-process --dtype bf16 --p 2000 --d "0,1" --w /home/QwQ-32B-AWQ-Marlin/ qwen2 --quant awq --penalty 1.0 --temperature 0. + cargo run --release --features cuda,nccl -- --d 0,1 --w /home/QwQ-32B-AWQ-Marlin/ ``` **注意**：使用的GPU数量（`--d`）必须为2的幂次方（例如2、4或8）。 @@ -260,11 +256,11 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #同时包含flash att
显示命令 - 只需移除`--multi-process`参数。 + 只需添加`--multithread`参数。 **在两块GPU上运行QwQ-32B BF16模型** ```shell - cargo run --release --features cuda,nccl -- --dtype bf16 --p 2000 --d "0,1" --w /home/QwQ-32B/ qwen2 --penalty 1.0 --temperature 0. + cargo run --release --features cuda,nccl -- --multithread --d 0,1 --w /home/QwQ-32B/ ``` 如果在多线程多GPU模式下遇到问题，可以尝试： @@ -285,7 +281,7 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #同时包含flash att **2. 在8块A100（40GB）上运行DeepSeek-R1模型** ```shell - cargo run --release --features cuda,nccl -- --log --multi-process --dtype bf16 --p 2000 --d "0,1,2,3,4,5,6,7" --w /data/DeepSeek-R1-AWQ-Marlin/ deep-seek --quant awq --temperature 0. --penalty 1.0 --num-experts-offload-per-rank 15 + cargo run --release --features cuda,nccl -- --log --d 0,1,2,3,4,5,6,7 --w /data/DeepSeek-R1-AWQ-Marlin/ --num-experts-offload-per-rank 15 ``` **注意**：此设置将每个rank的15个专家（总共256个专家中的120个）卸载到CPU（需要约150GB的额外主机内存）。在推理过程中，这些卸载的专家会根据需要交换回GPU内存。如果GPU内存更少，可以增加`--num-experts-offload-per-rank`参数（最大支持每个rank卸载32个专家）。 @@ -322,7 +318,7 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #同时包含flash att **4. 使用MPI运行器在两个节点上运行模型** ```shell - sudo mpirun -np 16 -x RUST_LOG=info -hostfile ./hostfile --allow-run-as-root -bind-to none -map-by slot --mca plm_rsh_args "-p 22" --mca btl_tcp_if_include %NET_INTERFACE% target/release/candle-vllm --log --multi-process --dtype bf16 --p 2000 --d "0,1,2,3,4,5,6,7" --w /data/DeepSeek-R1-AWQ-Marlin/ deep-seek --quant awq --temperature 0. --penalty 1.0 + sudo mpirun -np 16 -x RUST_LOG=info -hostfile ./hostfile --allow-run-as-root -bind-to none -map-by slot --mca plm_rsh_args "-p 22" --mca btl_tcp_if_include %NET_INTERFACE% target/release/candle-vllm --log --d 0,1,2,3,4,5,6,7 --w /data/DeepSeek-R1-AWQ-Marlin/ ```
@@ -341,13 +337,13 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #同时包含flash att 如果你想使用全部 GPU 进行推理，下面的 NUMA 绑定配置可以获得最佳性能： ```shell - MAP_NUMA_NODE=0,0,0,0,1,1,1,1 numactl --cpunodebind=0 --membind=0 cargo run --release --features cuda,nccl -- --multi-process --dtype bf16 --p 2000 --d "0,1,2,3,4,5,6,7" --w /home/data/DeepSeek-V2-Chat-AWQ-Marlin deep-seek --quant awq --temperature 0. --penalty 1.0 + MAP_NUMA_NODE=0,0,0,0,1,1,1,1 numactl --cpunodebind=0 --membind=0 cargo run --release --features cuda,nccl -- --d 0,1,2,3,4,5,6,7 --w /home/data/DeepSeek-V2-Chat-AWQ-Marlin ``` 如果你只使用 4 张 GPU，可以使用如下的 NUMA 绑定方式： ```shell - MAP_NUMA_NODE=0,0,0,0 numactl --cpunodebind=0 --membind=0 cargo run --release --features cuda,nccl -- --multi-process --dtype bf16 --p 2000 --d "0,1,2,3" --w /home/data/DeepSeek-V2-Chat-AWQ-Marlin deep-seek --quant awq --temperature 0. --penalty 1.0 + MAP_NUMA_NODE=0,0,0,0 numactl --cpunodebind=0 --membind=0 cargo run --release --features cuda,nccl -- --d 0,1,2,3 --w /home/data/DeepSeek-V2-Chat-AWQ-Marlin ``` 以上命令中 `numactl --cpunodebind=0 --membind=0`指定了master进程（master rank）绑定的NUMA node，其必须与 `MAP_NUMA_NODE`相匹配。 @@ -361,7 +357,7 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #同时包含flash att 1) 启动`Qwen3-Reranker`模型服务 ```shell - target/release/candle-vllm --p 2000 --multi-process --f /home/data/Qwen3-Reranker-4B-q4_k_m.gguf qwen3 --quant gguf + target/release/candle-vllm --f /home/data/Qwen3-Reranker-4B-q4_k_m.gguf ``` 2) 启动迷你聊天机器人并传入`system prompt` @@ -574,12 +570,12 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #同时包含flash att
显示量化配置 - Candle-vllm支持在模型加载时将默认权重（F32/F16/BF16）转换为任何GGML/GGUF格式，或将`4位GPTQ/AWQ`权重转换为`Marlin`格式进行加速。此功能有助于节省GPU内存（或通过Marlin内核加速推理性能），使其更适合消费级GPU（例如RTX 4090）。要使用此功能，只需在运行candle-vllm时传递相应`quant`参数。 + Candle-vllm支持在模型加载时将默认权重（F32/F16/BF16）转换为任何GGML/GGUF格式，或将`4位GPTQ/AWQ`权重转换为`Marlin`格式进行加速。此功能有助于节省GPU内存（或通过Marlin内核加速推理性能），使其更适合消费级GPU（例如RTX 4090）。要使用此功能，只需在运行candle-vllm时传递相应`isq`参数。 **对于未量化模型：** ``` - cargo run --release --features cuda -- --p 2000 --w /home/Meta-Llama-3.1-8B-Instruct/ llama3 --quant q4k + cargo run --release --features cuda -- --w /home/Meta-Llama-3.1-8B-Instruct/ llama3 --isq q4k ``` `quant`参数选项：["q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "q2k", "q3k","q4k","q5k","q6k"] @@ -587,7 +583,7 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #同时包含flash att **对于4位GPTQ量化模型：** ``` - cargo run --release --features cuda -- --p 2000 --w /home/mistral_7b-int4/ mistral --quant marlin + cargo run --release --features cuda -- --w /home/mistral_7b-int4/ mistral --isq marlin ``` **关于Marlin的注意事项**： @@ -613,7 +609,7 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #同时包含flash att 你可以传递`penalty`和`temperature`参数给模型以**防止潜在的重复**，例如： ``` - cargo run --release --features cuda -- --p 2000 --w /home/mistral_7b/ mistral --repeat-last-n 64 --penalty 1.1 --temperature 0.7 + cargo run --release --features cuda -- --w /home/mistral_7b/ ``` `--max-gen-tokens`参数用于控制每次聊天响应的最大输出令牌数。默认值将设置为`max_sequence_len`的1/5。 @@ -621,10 +617,10 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #同时包含flash att 对于`消费级GPU`，建议以GGML格式（或Marlin格式）运行模型，例如： ``` - cargo run --release --features cuda -- --p 2000 --w /home/Meta-Llama-3.1-8B-Instruct/ llama3 --quant q4k + cargo run --release --features cuda -- --w /home/Meta-Llama-3.1-8B-Instruct/ llama3 --isq q4k ``` - 其中`quant`可选值为：["q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "q2k", "q3k","q4k","q5k","q6k", "awq", "gptq", "marlin", "gguf", "ggml"]。 + 其中`isq`可选值为：["q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "q2k", "q3k","q4k","q5k","q6k", "awq", "gptq", "marlin", "gguf", "ggml"]。
- **GPTQ/AWQ模型通过Marlin Kernel加速** @@ -634,13 +630,13 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #同时包含flash att Candle-vllm现在支持GPTQ/AWQ（Marlin内核），如果你有`Marlin`格式的量化权重，可以传递`quant`（marlin）参数，例如： ```shell - cargo run --release --features cuda -- --p 2000 --dtype f16 --w /home/Meta-Llama-3.1-8B-Instruct-GPTQ-INT4-Marlin/ llama3 --quant marlin --temperature 0. --penalty 1. + cargo run --release --features cuda -- --w /home/Meta-Llama-3.1-8B-Instruct-GPTQ-INT4-Marlin/ ``` 或者，将现有的AWQ 4位模型转换为Marlin兼容格式： ```shell python3 examples/convert_awq_marlin.py --src /home/Meta-Llama-3.1-8B-Instruct-AWQ-INT4/ --dst /home/Meta-Llama-3.1-8B-Instruct-AWQ-INT4-Marlin/ --bits 4 --method awq --group 128 --nk False - cargo run --release --features cuda,nccl -- --multi-process --dtype f16 --p 2000 --d "0" --w /home/Meta-Llama-3.1-8B-Instruct-AWQ-INT4-Marlin/ llama3 --quant awq --temperature 0. --penalty 1.0 + cargo run --release --features cuda,nccl -- --d 0 --w /home/Meta-Llama-3.1-8B-Instruct-AWQ-INT4-Marlin/ ``` 你也可以使用`GPTQModel`通过脚本`examples/convert_marlin.py`将模型转换为Marlin兼容格式。 diff --git a/README.md b/README.md index a1d1d4d8..1042513b 100644 --- a/README.md +++ b/README.md @@ -95,22 +95,18 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #build with flash-attn **Example:** ```shell - [RUST_LOG=warn] cargo run [--release --features cuda,nccl] -- [--multi-process --log --dtype bf16 --p 2000 --d "0,1" --mem 8192] [--w /home/weights/Qwen3-27B-GPTQ-4Bit] [qwen3] [--quant gptq --temperature 0.7 --penalty 1.0 --top-k 32 --top-p 0.95 --thinking] + [RUST_LOG=warn] cargo run [--release --features cuda,nccl] -- [--log --dtype bf16 --p 2000 --d 0,1 --mem 8192] [--w /home/weights/Qwen3-27B-GPTQ-4Bit] ``` `ENV_PARAM`: RUST_LOG=warn `BUILD_PARAM`: --release --features cuda,nccl - `PROGRAM_PARAM`：--multi-process --log --dtype bf16 --p 2000 --d "0,1" --mem 8192 + `PROGRAM_PARAM`：--log --dtype bf16 --p 2000 --d 0,1 --mem 8192 - `MODEL_WEIGHT_PATH`: --w /home/weights/Qwen3-27B-GPTQ-4Bit + `MODEL_WEIGHT_PATH`: --w /home/weights/Qwen3-27B-GPTQ-4Bit (or `--m` specify model-id) - `MODEL_TYPE`: qwen3 - - `MODEL_PARAM`: --quant gptq --temperature 0.7 --penalty 1.0 --top-k 32 --top-p 0.95 --thinking - - where, `--mem` (`kvcache-mem-gpu`) is the key parameter to control KV cache usage (increase this for large batch); `MODEL_TYPE` in ["llama", "llama3", "mistral", "phi2", "phi3", "qwen2", "qwen3", "glm4", "gemma", "gemma3", "yi", "stable-lm", "deep-seek"] + where, `--mem` (`kvcache-mem-gpu`) is the key parameter to control KV cache usage (increase this for large batch); supported model archs include ["llama", "llama3", "mistral", "phi2", "phi3", "qwen2", "qwen3", "glm4", "gemma", "gemma3", "yi", "stable-lm", "deep-seek"] @@ -124,13 +120,13 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #build with flash-attn **Local Path** ```shell - target/release/candle-vllm --p 2000 --w /home/DeepSeek-R1-Distill-Llama-8B/ llama3 --temperature 0. --penalty 1.0 + target/release/candle-vllm --p 2000 --w /home/DeepSeek-R1-Distill-Llama-8B/ ``` **Model-ID (download from Huggingface)** ```shell - target/release/candle-vllm --m deepseek-ai/DeepSeek-R1-0528-Qwen3-8B qwen3 + target/release/candle-vllm --m deepseek-ai/DeepSeek-R1-0528-Qwen3-8B ``` @@ -142,13 +138,13 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #build with flash-attn **Local Path (with port, dtype, sampling parameter specified)** ```shell - target/release/candle-vllm --p 2000 --dtype bf16 --f /home/data/DeepSeek-R1-0528-Qwen3-8B-Q2_K.gguf qwen3 --quant gguf --temperature 0.7 --penalty 1.1 + target/release/candle-vllm --f /home/data/DeepSeek-R1-0528-Qwen3-8B-Q2_K.gguf ``` **Model-ID (download from Huggingface)** ```shell - target/release/candle-vllm --m unsloth/DeepSeek-R1-0528-Qwen3-8B-GGUF --f DeepSeek-R1-0528-Qwen3-8B-Q2_K.gguf qwen3 --quant gguf + target/release/candle-vllm --m unsloth/DeepSeek-R1-0528-Qwen3-8B-GGUF --f DeepSeek-R1-0528-Qwen3-8B-Q2_K.gguf ``` @@ -160,13 +156,13 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #build with flash-attn **Local Path (assume model downloaded in /home)** ```shell - cargo run --release --features metal -- --p 2000 --dtype bf16 --f /home/qwq-32b-q4_k_m.gguf qwen2 --quant gguf --temperature 0. --penalty 1.0 + cargo run --release --features metal -- --f /home/qwq-32b-q4_k_m.gguf ``` **Model-ID (download from Huggingface)** ```shell - cargo run --release --features metal -- --p 2000 --dtype bf16 --m Qwen/QwQ-32B-GGUF --f qwq-32b-q4_k_m.gguf qwen2 --quant gguf --temperature 0. --penalty 1.0 + cargo run --release --features metal -- --m Qwen/QwQ-32B-GGUF --f qwq-32b-q4_k_m.gguf ``` @@ -175,13 +171,13 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #build with flash-attn
Show command - **Simply add `quant` parameter when running unquantized models** + **Simply add `isq` parameter when running unquantized models** ```shell - target/release/candle-vllm --p 2000 --w /home/DeepSeek-R1-Distill-Llama-8B/ llama3 --quant q4k --temperature 0. --penalty 1.0 + target/release/candle-vllm --p 2000 --w /home/DeepSeek-R1-Distill-Llama-8B/ llama3 --isq q4k ``` - Options for in-site `quant` parameters: ["q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "q2k", "q3k","q4k","q5k","q6k"] + Options for in-site `isq` parameters: ["q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "q2k", "q3k","q4k","q5k","q6k"]
@@ -192,19 +188,19 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #build with flash-attn **Local Path** ```shell - target/release/candle-vllm --dtype bf16 --p 2000 --w /home/DeepSeek-R1-Distill-Qwen-14B-GPTQ_4bit-128g qwen2 --quant gptq --temperature 0. --penalty 1.0 + target/release/candle-vllm --w /home/DeepSeek-R1-Distill-Qwen-14B-GPTQ_4bit-128g ``` **Model-ID (download from Huggingface)** ```shell - target/release/candle-vllm --m thesven/Llama-3-8B-GPTQ-4bit llama3 --quant gptq + target/release/candle-vllm --m thesven/Llama-3-8B-GPTQ-4bit ``` **Convert Any uncompressed model to marlin-compatible format** ```shell python3 examples/convert_marlin.py --src /home/DeepSeek-R1-Distill-Qwen-14B/ --dst /home/DeepSeek-R1-Distill-Qwen-14B-GPTQ_4bit-128g - target/release/candle-vllm --dtype bf16 --p 2000 --w /home/DeepSeek-R1-Distill-Qwen-14B-GPTQ_4bit-128g qwen2 --quant gptq --temperature 0. --penalty 1.0 + target/release/candle-vllm --w /home/DeepSeek-R1-Distill-Qwen-14B-GPTQ_4bit-128g ``` @@ -220,7 +216,7 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #build with flash-attn **Run the converted AWQ model** ```shell - target/release/candle-vllm --multi-process --dtype f16 --p 2000 --d "0" --w /home/Meta-Llama-3.1-8B-Instruct-AWQ-INT4-Marlin/ llama3 --quant awq --temperature 0. --penalty 1.0 + target/release/candle-vllm --d 0 --w /home/Meta-Llama-3.1-8B-Instruct-AWQ-INT4-Marlin/ ``` @@ -230,7 +226,7 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #build with flash-attn Show command ```shell - target/release/candle-vllm --dtype bf16 --p 2000 --w /home/DeepSeek-R1-Distill-Qwen-14B-GPTQ-Marlin/ qwen2 --quant marlin --penalty 1.0 --temperature 0. + target/release/candle-vllm --w /home/DeepSeek-R1-Distill-Qwen-14B-GPTQ-Marlin/ ``` @@ -242,7 +238,7 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #build with flash-attn **QwQ-32B BF16 model on two GPUs** ```shell - cargo run --release --features cuda,nccl -- --multi-process --dtype bf16 --p 2000 --d "0,1" --w /home/QwQ-32B/ qwen2 --penalty 1.0 --temperature 0. + cargo run --release --features cuda,nccl -- --d 0,1 --w /home/QwQ-32B/ ``` **QwQ-32B 4-bit AWQ model on two GPUs** @@ -254,7 +250,7 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #build with flash-attn 2) Run the converted AWQ model ```shell - cargo run --release --features cuda,nccl -- --multi-process --dtype bf16 --p 2000 --d "0,1" --w /home/QwQ-32B-AWQ-Marlin/ qwen2 --quant awq --penalty 1.0 --temperature 0. + cargo run --release --features cuda,nccl -- --d 0,1 --w /home/QwQ-32B-AWQ-Marlin/ ``` **Note:** number of GPUs (`--d`) used must be aligned to 2^n (e.g., 2, 4, or 8). @@ -264,11 +260,11 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #build with flash-attn
Show command - Simply remove the `--multi-process` parameter + Simply add the `--multithread` parameter **QwQ-32B BF16 model on two GPUs** ```shell - cargo run --release --features cuda,nccl -- --dtype bf16 --p 2000 --d "0,1" --w /home/QwQ-32B/ qwen2 --penalty 1.0 --temperature 0. + cargo run --release --features cuda,nccl -- --multithread --d 0,1 --w /home/QwQ-32B/ ``` If you encountered problems under Multi-threaded Multi-GPU mode, you may: @@ -289,7 +285,7 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #build with flash-attn **2. Run DeepSeek-R1 model on 8 x A100(40GB)** ```shell - cargo run --release --features cuda,nccl -- --log --multi-process --dtype bf16 --p 2000 --d "0,1,2,3,4,5,6,7" --w /data/DeepSeek-R1-AWQ-Marlin/ deep-seek --quant awq --temperature 0. --penalty 1.0 --num-experts-offload-per-rank 15 + cargo run --release --features cuda,nccl -- --log --d 0,1,2,3,4,5,6,7 --w /data/DeepSeek-R1-AWQ-Marlin/--num-experts-offload-per-rank 15 ``` **Note:** This setup offloads 15 experts per rank (a total of 120 out of 256 experts) to the CPU (around 150GB additional host memory required). During inference, these offloaded experts are swapped back into GPU memory as needed. If you have even less GPU memory, consider increasing the `--num-experts-offload-per-rank` parameter (up to a maximum of 32 experts per rank in this case). @@ -326,7 +322,7 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #build with flash-attn **4. Run the model on two nodes with MPI runner** ```shell - sudo mpirun -np 16 -x RUST_LOG=info -hostfile ./hostfile --allow-run-as-root -bind-to none -map-by slot --mca plm_rsh_args "-p 22" --mca btl_tcp_if_include %NET_INTERFACE% target/release/candle-vllm --log --multi-process --dtype bf16 --p 2000 --d "0,1,2,3,4,5,6,7" --w /data/DeepSeek-R1-AWQ-Marlin/ deep-seek --quant awq --temperature 0. --penalty 1.0 + sudo mpirun -np 16 -x RUST_LOG=info -hostfile ./hostfile --allow-run-as-root -bind-to none -map-by slot --mca plm_rsh_args "-p 22" --mca btl_tcp_if_include %NET_INTERFACE% target/release/candle-vllm --log --d 0,1,2,3,4,5,6,7 --w /data/DeepSeek-R1-AWQ-Marlin/ deep-seek ```
@@ -344,13 +340,13 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #build with flash-attn To achieve optimal performance during inference using all GPUs, use the following NUMA binding: ```shell - MAP_NUMA_NODE=0,0,0,0,1,1,1,1 numactl --cpunodebind=0 --membind=0 target/release/candle-vllm --multi-process --dtype bf16 --p 2000 --d "0,1,2,3,4,5,6,7" --w /home/data/DeepSeek-V2-Chat-AWQ-Marlin deep-seek --quant awq --temperature 0. --penalty 1.0 + MAP_NUMA_NODE=0,0,0,0,1,1,1,1 numactl --cpunodebind=0 --membind=0 target/release/candle-vllm --d 0,1,2,3,4,5,6,7 --w /home/data/DeepSeek-V2-Chat-AWQ-Marlin ``` To use only 4 GPUs, you can apply this NUMA binding: ```shell - MAP_NUMA_NODE=0,0,0,0 numactl --cpunodebind=0 --membind=0 target/release/candle-vllm --multi-process --dtype bf16 --p 2000 --d "0,1,2,3" --w /home/data/DeepSeek-V2-Chat-AWQ-Marlin deep-seek --quant awq --temperature 0. --penalty 1.0 + MAP_NUMA_NODE=0,0,0,0 numactl --cpunodebind=0 --membind=0 target/release/candle-vllm --d 0,1,2,3 --w /home/data/DeepSeek-V2-Chat-AWQ-Marlin ``` *where* `numactl --cpunodebind=0 --membind=0` above indicates NUMA binding for the master rank (master process) which should be matched to `MAP_NUMA_NODE`. @@ -363,7 +359,7 @@ cargo build --release --features cuda,nccl,flash-attn,mpi #build with flash-attn 1) Start the backend service for `Qwen3-Reranker` model ```shell - target/release/candle-vllm --p 2000 --multi-process --f /home/data/Qwen3-Reranker-4B-q4_k_m.gguf qwen3 --quant gguf + target/release/candle-vllm --p 2000 --f /home/data/Qwen3-Reranker-4B-q4_k_m.gguf ``` 2) Start the chatbot with `system prompt` for qwen3-reranker @@ -573,20 +569,20 @@ Chat frontend (any frontend compatible with openai API, simple options available
Show quantization config - Candle-vllm supports in-situ quantization, allowing the transformation of default weights (F32/F16/BF16) into any GGML/GGUF format, or `4-bit GPTQ/AWQ` weights into `marlin format` during model loading. This feature helps conserve GPU memory and speedup inference performance, making it more efficient for consumer-grade GPUs (e.g., RTX 4090). To use this feature, simply supply the `quant` parameter when running candle-vllm. + Candle-vllm supports in-situ quantization, allowing the transformation of default weights (F32/F16/BF16) into any GGML/GGUF format, or `4-bit GPTQ/AWQ` weights into `marlin format` during model loading. This feature helps conserve GPU memory and speedup inference performance, making it more efficient for consumer-grade GPUs (e.g., RTX 4090). To use this feature, simply supply the `isq` parameter when running candle-vllm. **For unquantized models:** ``` - cargo run --release --features cuda -- --p 2000 --w /home/Meta-Llama-3.1-8B-Instruct/ llama3 --quant q4k + cargo run --release --features cuda -- --p 2000 --w /home/Meta-Llama-3.1-8B-Instruct/ llama3 --isq q4k ``` - Options for `quant` parameters: ["q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "q2k", "q3k","q4k","q5k","q6k"] + Options for `isq` parameters: ["q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "q2k", "q3k","q4k","q5k","q6k"] **For quantized 4-bit GPTQ model:** ``` - cargo run --release --features cuda -- --p 2000 --w /home/mistral_7b-int4/ mistral --quant marlin + cargo run --release --features cuda -- --p 2000 --w /home/mistral_7b-int4/ ``` **Please note for marlin**: @@ -608,10 +604,8 @@ Chat frontend (any frontend compatible with openai API, simple options available For chat streaming, the `stream` flag in chat request need to be set to `True`. - You may supply `penalty` and `temperature` to the model to **prevent potential repetitions**, for example: - ``` - cargo run --release --features cuda -- --p 2000 --w /home/mistral_7b/ mistral --repeat-last-n 64 --penalty 1.1 --temperature 0.7 + cargo run --release --features cuda -- --p 2000 --w /home/mistral_7b/ ``` `--max-gen-tokens` parameter is used to control the maximum output tokens per chat response. The value will be set to 1/5 of max_sequence_len by default. @@ -619,27 +613,27 @@ Chat frontend (any frontend compatible with openai API, simple options available For `consumer GPUs`, it is suggested to run the models under GGML formats (or Marlin format), e.g., ``` - cargo run --release --features cuda -- --p 2000 --w /home/Meta-Llama-3.1-8B-Instruct/ llama3 --quant q4k + cargo run --release --features cuda -- --p 2000 --w /home/Meta-Llama-3.1-8B-Instruct/ llama3 --isq q4k ``` - where `quant` is one of ["q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "q2k", "q3k","q4k","q5k","q6k", "awq", "gptq", "marlin", "gguf", "ggml"]. + where `isq` is one of ["q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "q2k", "q3k","q4k","q5k","q6k", "awq", "gptq", "marlin", "gguf", "ggml"].
- **Use Marlin kernel to speedup GPTQ/AWQ models**
Show details - Candle-vllm now supports GPTQ/AWQ Marlin kernel, you may supply the `quant` (marlin) parameter if you have `Marlin` format quantized weights, such as: + Candle-vllm now supports GPTQ/AWQ Marlin kernel, you can run these models directly, such as: ```shell - cargo run --release --features cuda -- --p 2000 --dtype f16 --w /home/Meta-Llama-3.1-8B-Instruct-GPTQ-INT4-Marlin/ llama3 --quant marlin --temperature 0. --penalty 1. + cargo run --release --features cuda -- --dtype f16 --w /home/Meta-Llama-3.1-8B-Instruct-GPTQ-INT4-Marlin/ ``` or, convert existing AWQ 4bit model to marlin compatible format ```shell python3 examples/convert_awq_marlin.py --src /home/Meta-Llama-3.1-8B-Instruct-AWQ-INT4/ --dst /home/Meta-Llama-3.1-8B-Instruct-AWQ-INT4-Marlin/ --bits 4 --method awq --group 128 --nk False - cargo run --release --features cuda,nccl -- --multi-process --dtype f16 --p 2000 --d "0" --w /home/Meta-Llama-3.1-8B-Instruct-AWQ-INT4-Marlin/ llama3 --quant awq --temperature 0. --penalty 1.0 + cargo run --release --features cuda,nccl -- --dtype f16 --d 0 --w /home/Meta-Llama-3.1-8B-Instruct-AWQ-INT4-Marlin/ ``` You may also use `GPTQModel` to transform a model to marlin-compatible format using the given script `examples/convert_marlin.py`. diff --git a/src/backend/cache.rs b/src/backend/cache.rs index d6dde5e2..4fdee730 100644 --- a/src/backend/cache.rs +++ b/src/backend/cache.rs @@ -1,5 +1,3 @@ -#[cfg(feature = "cuda")] -use crate::{openai::responses::APIError, try_api}; #[cfg(feature = "metal")] use candle_core::{ backend::BackendStorage, CpuStorage, Device, IndexOp, Layout, MetalDevice, MetalStorage, @@ -9,7 +7,7 @@ use candle_core::{ use candle_core::{ cuda_backend::cudarc::driver::{CudaSlice, DevicePtr}, cuda_backend::CudaStorageSlice, - Device, IndexOp, Storage, Tensor, + Device, IndexOp, Result, Storage, Tensor, }; #[cfg(feature = "cuda")] use kernels::ffi::{copy_blocks_bf16, copy_blocks_f16, copy_blocks_f32}; @@ -22,7 +20,7 @@ pub unsafe fn copy_blocks( key_caches: Vec<&mut Tensor>, value_caches: Vec<&mut Tensor>, block_mapping: HashMap>, -) -> Result<(), APIError> { +) -> Result<()> { use candle_core::DType; let cache_dev = key_caches.first().unwrap().device(); @@ -30,18 +28,18 @@ pub unsafe fn copy_blocks( panic!("Expected the key caches to be on a CUDA device.") }; if !cache_dev.same_device(value_caches.first().unwrap().device()) { - return Err(APIError::new(format!( + candle_core::bail!( "`key` and `value` caches have different devices, got {:?} and {:?} respectively.", cache_dev, value_caches.first().unwrap().device() - ))); + ) } if key_caches.first().unwrap().dtype() != value_caches.first().unwrap().dtype() { - return Err(APIError::new(format!( + candle_core::bail!( "Key and value caches have different types, got {:?} and {:?}.", key_caches.first().unwrap().dtype(), value_caches.first().unwrap().dtype() - ))); + ) } let num_layers: u32 = key_caches.len().try_into().unwrap(); if num_layers == 0 { @@ -55,8 +53,8 @@ pub unsafe fn copy_blocks( let mut dtype = DType::F32; for (key_cache, value_cache) in zip(&key_caches, &value_caches) { - try_api!(key_cache.to_device(cache_dev)); - try_api!(value_cache.to_device(cache_dev)); + key_cache.to_device(cache_dev)?; + value_cache.to_device(cache_dev)?; let key_offset: u64 = key_cache .storage_and_layout() @@ -99,9 +97,7 @@ pub unsafe fn copy_blocks( (ptr_key, ptr_value) } _ => { - return Err(APIError::from( - "only f32, f16 and bf16 input data type supported!", - )); + candle_core::bail!("only f32, f16 and bf16 input data type supported!") } }; key_cache_ptrs.push(key_ptr + key_offset); @@ -121,7 +117,10 @@ pub unsafe fn copy_blocks( let value_cache_ptr = value_cache_ptrs.as_mut_ptr() as *mut core::ffi::c_void; let block_mapping_ptr = block_mapping_vec.as_mut_ptr() as *const core::ffi::c_void; - let numel_per_block: u32 = try_api!(key_caches.first().unwrap().i(0)) + let numel_per_block: u32 = key_caches + .first() + .unwrap() + .i(0)? .shape() .dims() .iter() @@ -174,19 +173,23 @@ pub fn swap_blocks( src: Tensor, dst: &mut Tensor, block_mapping: HashMap, -) -> Result<(), APIError> { +) -> Result<()> { let block_size_in_bytes = src.dtype().size_in_bytes() * src.dims()[0]; match (src.device(), dst.device()) { (Device::Cuda(src_dev), Device::Cuda(dst_dev)) => { if src_dev.ordinal() != dst_dev.ordinal() { - return Err(APIError::new(format!("Tensors must be on the same device to copy, got ordinals {} (src) and {} (dst).", src_dev.ordinal(), dst_dev.ordinal()))) + candle_core::bail!("Tensors must be on the same device to copy, got ordinals {} (src) and {} (dst).", src_dev.ordinal(), dst_dev.ordinal()) } let (src_storage, src_layout) = src.storage_and_layout(); let (dst_storage, dst_layout) = dst.storage_and_layout(); assert!(matches!(&*src_storage, Storage::Cuda(_))); assert!(matches!(&*dst_storage, Storage::Cuda(_))); - let Storage::Cuda(src_storage) = &*src_storage else { unreachable!() }; - let Storage::Cuda(dst_storage) = &*dst_storage else { unreachable!() }; + let Storage::Cuda(src_storage) = &*src_storage else { + unreachable!() + }; + let Storage::Cuda(dst_storage) = &*dst_storage else { + unreachable!() + }; let (src_ptr, dst_ptr) = match (&src_storage.slice, &dst_storage.slice) { (CudaStorageSlice::BF16(slice_src), CudaStorageSlice::BF16(slice_dst)) => { let ptr_src = *slice_src.slice(src_layout.start_offset()..).device_ptr(); @@ -204,7 +207,7 @@ pub fn swap_blocks( (ptr_src, ptr_dst) } _ => { - return Err(APIError::from("only f32, f16 and bf16 input data type supported!")); + candle_core::bail!("only f32, f16 and bf16 input data type supported!"); } }; // let src_ptr = src_storage.as_cuda_slice::().map_err(APIError::from)?.device_ptr() + TryInto::::try_into(src_layout.start_offset()).unwrap(); @@ -214,10 +217,16 @@ pub fn swap_blocks( let src_offset: u64 = (src_block_number * block_size_in_bytes).try_into().unwrap(); let dst_offset: u64 = (dst_block_number * block_size_in_bytes).try_into().unwrap(); // u8s because we copy by bytes - let src_slice: CudaSlice = unsafe { src_dev.upgrade_device_ptr(src_ptr+src_offset, block_size_in_bytes) }; - let mut dst_slice = unsafe { dst_dev.upgrade_device_ptr(dst_ptr+dst_offset, block_size_in_bytes) }; - - try_api!(src_dev.dtod_copy(&src_slice, &mut dst_slice)); + let src_slice: CudaSlice = unsafe { + src_dev.upgrade_device_ptr(src_ptr + src_offset, block_size_in_bytes) + }; + let mut dst_slice = unsafe { + dst_dev.upgrade_device_ptr(dst_ptr + dst_offset, block_size_in_bytes) + }; + + src_dev + .dtod_copy(&src_slice, &mut dst_slice) + .map_err(candle_core::Error::wrap)?; } } (Device::Cpu, Device::Cuda(dst_dev)) => { @@ -225,22 +234,34 @@ pub fn swap_blocks( let (dst_storage, dst_layout) = dst.storage_and_layout(); assert!(matches!(&*src_storage, Storage::Cpu(_))); assert!(matches!(&*dst_storage, Storage::Cuda(_))); - let Storage::Cpu(src_storage) = &*src_storage else { unreachable!() }; - let Storage::Cuda(dst_storage) = &*dst_storage else { unreachable!() }; - let dst_ptr = dst_storage.as_cuda_slice::().map_err(APIError::from)?.device_ptr() + TryInto::::try_into(dst_layout.start_offset()).unwrap(); - let src_slice = try_api!(src_storage.as_slice()); + let Storage::Cpu(src_storage) = &*src_storage else { + unreachable!() + }; + let Storage::Cuda(dst_storage) = &*dst_storage else { + unreachable!() + }; + let dst_ptr = dst_storage.as_cuda_slice::()?.device_ptr() + + TryInto::::try_into(dst_layout.start_offset()).unwrap(); + let src_slice = src_storage.as_slice()?; for (src_block_number, dst_block_number) in block_mapping { let src_offset = src_block_number * block_size_in_bytes; let dst_offset: u64 = (dst_block_number * block_size_in_bytes).try_into().unwrap(); // u8s because we copy by bytes - let mut dst_slice: CudaSlice = unsafe { dst_dev.upgrade_device_ptr(dst_ptr+dst_offset, block_size_in_bytes) }; + let mut dst_slice: CudaSlice = unsafe { + dst_dev.upgrade_device_ptr(dst_ptr + dst_offset, block_size_in_bytes) + }; - try_api!(dst_dev.htod_sync_copy_into(&src_slice[src_offset..src_offset+block_size_in_bytes], &mut dst_slice)); + dst_dev + .htod_sync_copy_into( + &src_slice[src_offset..src_offset + block_size_in_bytes], + &mut dst_slice, + ) + .map_err(candle_core::Error::wrap)?; } } (src, dst) => { - return Err(APIError::new(format!("Tensors must be on either the GPU or CPU to swap,, got {src:?} (src) and {dst:?} (dst)."))) + candle_core::bail!("Tensors must be on either the GPU or CPU to swap,, got {src:?} (src) and {dst:?} (dst).") } } diff --git a/src/backend/gguf.rs b/src/backend/gguf.rs index 5a3e4e02..c197ce05 100644 --- a/src/backend/gguf.rs +++ b/src/backend/gguf.rs @@ -618,3 +618,15 @@ impl TryFrom> for NormalizerWrapper { Ok(value) } } + +pub fn get_arch_and_num_of_layers(ct: gguf_file::Content) -> Result<(String, usize)> { + let md_get = |s: &str| match ct.metadata.get(s) { + None => candle_core::bail!("cannot find {s} in metadata"), + Some(v) => Ok(v), + }; + let architecture = md_get("general.architecture")?.to_string()?; + + let nlayers = + md_get(format!("{}.block_count", architecture.as_str()).as_str())?.to_u32()? as usize; + Ok((architecture.clone(), nlayers)) +} diff --git a/src/backend/gptq.rs b/src/backend/gptq.rs index c940ef1f..2c530eb8 100644 --- a/src/backend/gptq.rs +++ b/src/backend/gptq.rs @@ -1,15 +1,17 @@ +#[allow(unused_imports)] use candle::backend::BackendStorage; #[cfg(feature = "cuda")] use candle::CudaStorage; +#[allow(unused_imports)] use candle::{CpuStorage, DType, Layout, Result, Shape, Storage, Tensor}; use candle_core as candle; -use half::{bf16, f16}; #[cfg(feature = "cuda")] use kernels::ffi::{ awq_repack, gemm_half_q_half_alt, gptq_repack, marlin_4bit_bf16, marlin_4bit_f16, marlin_awq_4bit_bf16, marlin_awq_4bit_f16, }; +#[allow(unused)] struct GPTQMatMul { qzeros: Option, g_idx: Option, @@ -228,8 +230,10 @@ impl candle::CustomOp3 for GPTQMatMul { scale_l: &Layout, ) -> Result<(CudaStorage, Shape)> { match x.dtype() { - DType::F16 => self.cuda_fwd_t::(x, x_l, qweight, qweight_l, scale, scale_l), - DType::BF16 => self.cuda_fwd_t::(x, x_l, qweight, qweight_l, scale, scale_l), + DType::F16 => self.cuda_fwd_t::(x, x_l, qweight, qweight_l, scale, scale_l), + DType::BF16 => { + self.cuda_fwd_t::(x, x_l, qweight, qweight_l, scale, scale_l) + } dt => candle::bail!("GPTQMatMul is only supported for f16 and bf16 ({dt:?})"), } } diff --git a/src/backend/mod.rs b/src/backend/mod.rs index aa2b33f1..0a14563e 100644 --- a/src/backend/mod.rs +++ b/src/backend/mod.rs @@ -6,10 +6,11 @@ mod paged_attention; pub fn get_or_load_func( ptx_file: &'static str, kernel_base: &str, - dtype: DType, + dtype: candle_core::DType, suffix: Option<&str>, device: &CudaDevice, ) -> Result { + use candle_core::DType; let spec = match dtype { DType::U8 => "_u8", DType::U32 => "_u32", @@ -30,9 +31,9 @@ pub fn get_or_load_func( .map_err(APIError::from) } +#[allow(unused_imports)] use crate::openai::responses::APIError; pub use cache::*; -use candle_core::DType; #[cfg(feature = "cuda")] use candle_core::{cuda_backend::cudarc::driver::CudaFunction, CudaDevice}; pub use gptq::*; diff --git a/src/backend/paged_attention.rs b/src/backend/paged_attention.rs index a004b050..2be3682e 100644 --- a/src/backend/paged_attention.rs +++ b/src/backend/paged_attention.rs @@ -5,8 +5,6 @@ use candle::CudaStorage; use candle::MetalStorage; use candle::{CpuStorage, DType, Layout, Result, Shape, Storage, Tensor}; use candle_core as candle; -use half::{bf16, f16}; -use std::ffi::c_int; #[allow(dead_code)] struct PagedAttention { softmax_scale: f32, @@ -30,6 +28,7 @@ impl PagedAttention { ) -> Result<(CudaStorage, Shape)> { use candle::cuda_backend::cudarc::driver::DevicePtr; use candle::cuda_backend::WrapErr; + use core::ffi::c_int; let dtype = q.dtype(); let internal_type = match dtype { DType::F16 => 0, @@ -480,8 +479,8 @@ impl candle::CustomOp1 for PagedAttention { fn cuda_fwd(&self, q: &CudaStorage, q_l: &Layout) -> Result<(CudaStorage, Shape)> { match q.dtype() { DType::F32 => self.cuda_fwd_t::(q, q_l), - DType::F16 => self.cuda_fwd_t::(q, q_l), - DType::BF16 => self.cuda_fwd_t::(q, q_l), + DType::F16 => self.cuda_fwd_t::(q, q_l), + DType::BF16 => self.cuda_fwd_t::(q, q_l), dt => candle::bail!("paged-attention is only supported for f32/f16/bf16 ({dt:?})"), } } @@ -557,6 +556,7 @@ impl ReshapeCache { slot_mapping: &Tensor, ) -> Result<()> { use candle::cuda_backend::cudarc::driver::DevicePtr; + use core::ffi::c_int; let dtype = k.dtype(); let dev = k.device(); let internal_type = match dtype { @@ -842,7 +842,7 @@ impl candle::InplaceOp1 for ReshapeCache { &self.value_cache, &self.slot_mapping, ), - DType::F16 => self.cuda_fwd_t::( + DType::F16 => self.cuda_fwd_t::( k, k_l, &self.value, @@ -850,7 +850,7 @@ impl candle::InplaceOp1 for ReshapeCache { &self.value_cache, &self.slot_mapping, ), - DType::BF16 => self.cuda_fwd_t::( + DType::BF16 => self.cuda_fwd_t::( k, k_l, &self.value, diff --git a/src/lib.rs b/src/lib.rs index c1ecd4ec..312f69f8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,839 +2,12 @@ use candle::utils::{cuda_is_available, metal_is_available}; use candle::{Device, Result}; use candle_core as candle; -use clap::Subcommand; -use openai::pipelines::pipeline::DefaultLoader; -use std::fmt::Display; use std::path::Path; use tracing::warn; pub mod backend; pub mod openai; pub mod paged_attention; pub mod scheduler; -#[derive(Debug, Subcommand)] -pub enum ModelSelected { - /// Select the llama model (default llama2-7b). - Llama { - /// Control the application of repeat penalty for the last n tokens - #[arg(long)] - repeat_last_n: Option, - - #[arg(long)] - temperature: Option, - - #[arg(long)] - top_p: Option, - - #[arg(long)] - top_k: Option, - - #[arg(long)] - penalty: Option, - - #[arg(long)] - max_gen_tokens: Option, - - #[arg(long)] - quant: Option, - }, - - /// Select the llama3 model (default llama3.1-8b). - Llama3 { - /// Control the application of repeat penalty for the last n tokens - #[arg(long)] - repeat_last_n: Option, - - #[arg(long)] - temperature: Option, - - #[arg(long)] - top_p: Option, - - #[arg(long)] - top_k: Option, - - #[arg(long)] - penalty: Option, - - #[arg(long)] - max_gen_tokens: Option, - - #[arg(long)] - quant: Option, - - #[arg(long, default_value_t = false)] - thinking: bool, - }, - - /// Select the phi2 model (default 2.7b). - Phi2 { - /// Control the application of repeat penalty for the last n tokens - #[arg(long)] - repeat_last_n: Option, - - #[arg(long)] - temperature: Option, - - #[arg(long)] - top_p: Option, - - #[arg(long)] - top_k: Option, - - #[arg(long)] - penalty: Option, - - #[arg(long)] - max_gen_tokens: Option, - - #[arg(long)] - quant: Option, - }, - - /// Select the phi3 model (default 3.8b). - Phi3 { - /// Control the application of repeat penalty for the last n tokens - #[arg(long)] - repeat_last_n: Option, - - #[arg(long)] - temperature: Option, - - #[arg(long)] - top_p: Option, - - #[arg(long)] - top_k: Option, - - #[arg(long)] - penalty: Option, - - #[arg(long)] - max_gen_tokens: Option, - - #[arg(long)] - quant: Option, - }, - - /// Select the qwen model (default 1.8b). - Qwen2 { - /// Control the application of repeat penalty for the last n tokens - #[arg(long)] - repeat_last_n: Option, - - #[arg(long)] - temperature: Option, - - #[arg(long)] - top_p: Option, - - #[arg(long)] - top_k: Option, - - #[arg(long)] - penalty: Option, - - #[arg(long)] - max_gen_tokens: Option, - - #[arg(long)] - quant: Option, - }, - - Qwen3 { - /// Control the application of repeat penalty for the last n tokens - #[arg(long)] - repeat_last_n: Option, - - #[arg(long)] - temperature: Option, - - #[arg(long)] - top_p: Option, - - #[arg(long)] - top_k: Option, - - #[arg(long)] - penalty: Option, - - #[arg(long)] - max_gen_tokens: Option, - - #[arg(long)] - quant: Option, - - #[arg(long, default_value_t = false)] - thinking: bool, - }, - - /// Select the gemma model (default 2b). - Gemma { - /// Control the application of repeat penalty for the last n tokens - #[arg(long)] - repeat_last_n: Option, - - #[arg(long)] - temperature: Option, - - #[arg(long)] - top_p: Option, - - #[arg(long)] - top_k: Option, - - #[arg(long)] - penalty: Option, - - #[arg(long)] - max_gen_tokens: Option, - - #[arg(long)] - quant: Option, - }, - - Gemma3 { - /// Control the application of repeat penalty for the last n tokens - #[arg(long)] - repeat_last_n: Option, - - #[arg(long)] - temperature: Option, - - #[arg(long)] - top_p: Option, - - #[arg(long)] - top_k: Option, - - #[arg(long)] - penalty: Option, - - #[arg(long)] - max_gen_tokens: Option, - - #[arg(long)] - quant: Option, - - #[arg(long, default_value_t = false)] - thinking: bool, - }, - - /// Select the mistral model (default 7b). - Mistral { - /// Control the application of repeat penalty for the last n tokens - #[arg(long)] - repeat_last_n: Option, - - #[arg(long)] - temperature: Option, - - #[arg(long)] - top_p: Option, - - #[arg(long)] - top_k: Option, - - #[arg(long)] - penalty: Option, - - #[arg(long)] - max_gen_tokens: Option, - - #[arg(long)] - quant: Option, - }, - - /// Select the Yi model (default 6b). - Yi { - /// Control the application of repeat penalty for the last n tokens - #[arg(long)] - repeat_last_n: Option, - - #[arg(long)] - temperature: Option, - - #[arg(long)] - top_p: Option, - - #[arg(long)] - top_k: Option, - - #[arg(long)] - penalty: Option, - - #[arg(long)] - max_gen_tokens: Option, - - #[arg(long)] - quant: Option, - }, - - /// Select the stable-lm model (default zephyr-3b). - StableLM { - /// Control the application of repeat penalty for the last n tokens - #[arg(long)] - repeat_last_n: Option, - - #[arg(long)] - temperature: Option, - - #[arg(long)] - penalty: Option, - - #[arg(long)] - max_gen_tokens: Option, - - #[arg(long)] - quant: Option, - }, - - GLM4 { - /// Control the application of repeat penalty for the last n tokens - #[arg(long)] - repeat_last_n: Option, - - #[arg(long)] - temperature: Option, - - #[arg(long)] - top_p: Option, - - #[arg(long)] - top_k: Option, - - #[arg(long)] - penalty: Option, - - #[arg(long)] - max_gen_tokens: Option, - - #[arg(long)] - quant: Option, - }, - - /// Select the deepseek model (default deepseek-v2-lite-chat). - DeepSeek { - /// Control the application of repeat penalty for the last n tokens - #[arg(long)] - repeat_last_n: Option, - - #[arg(long)] - temperature: Option, - - #[arg(long)] - top_p: Option, - - #[arg(long)] - top_k: Option, - - #[arg(long)] - penalty: Option, - - #[arg(long)] - max_gen_tokens: Option, - - #[arg(long)] - quant: Option, - - //the number of experts offloaded per rank, - //suppose there are 256 experts in total which split into 8 devices (rank 8), - //each rank has 32 experts, num-experts-offload-per-rank=16 means - //half of the experts can be offloaded to cpu memory - #[arg(long)] - num_experts_offload_per_rank: Option, - - #[arg(long, default_value_t = false)] - thinking: bool, - }, -} - -impl Display for ModelSelected { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - ModelSelected::Llama { .. } => write!(f, "llama"), - ModelSelected::Llama3 { .. } => write!(f, "llama3"), - ModelSelected::Phi2 { .. } => write!(f, "phi2"), - ModelSelected::Phi3 { - repeat_last_n: _, - temperature: _, - top_k: _, - top_p: _, - penalty: _, - max_gen_tokens: _, - quant: _, - } => write!(f, "phi3"), - ModelSelected::Qwen2 { - repeat_last_n: _, - temperature: _, - top_k: _, - top_p: _, - penalty: _, - max_gen_tokens: _, - quant: _, - } => write!(f, "qwen2"), - ModelSelected::Qwen3 { - repeat_last_n: _, - temperature: _, - top_k: _, - top_p: _, - penalty: _, - max_gen_tokens: _, - quant: _, - thinking: _, - } => write!(f, "qwen3"), - ModelSelected::Gemma { .. } => write!(f, "gemma"), - ModelSelected::Gemma3 { .. } => write!(f, "gemma3"), - ModelSelected::Mistral { .. } => write!(f, "mistral"), - ModelSelected::Yi { .. } => write!(f, "yi"), - ModelSelected::StableLM { .. } => write!(f, "stablelm"), - ModelSelected::GLM4 { .. } => write!(f, "glm4"), - ModelSelected::DeepSeek { .. } => write!(f, "deepseek"), - } - } -} - -#[derive(Debug, Clone)] -pub struct SpecificConfig { - repeat_last_n: Option, - temperature: Option, - top_k: Option, - top_p: Option, - penalty: Option, - max_gen_tokens: Option, - quant: Option, - num_experts_offload_per_rank: Option, - thinking: bool, -} - -impl SpecificConfig { - pub fn new( - repeat_last_n: Option, - temperature: Option, - top_k: Option, - top_p: Option, - penalty: Option, - max_gen_tokens: Option, - quant: Option, - num_experts_offload_per_rank: Option, - thinking: bool, - ) -> Self { - Self { - repeat_last_n, - temperature, - top_k, - top_p, - penalty, - max_gen_tokens, - quant, - num_experts_offload_per_rank, - thinking, - } - } -} - -pub fn get_model_loader( - selected_model: ModelSelected, - model_id: Option, -) -> (Box, String, Option) { - match selected_model { - ModelSelected::Llama { - repeat_last_n, - temperature, - top_k, - top_p, - penalty, - max_gen_tokens, - quant, - } => ( - Box::new(DefaultLoader::new( - SpecificConfig::new( - repeat_last_n, - temperature, - top_k, - top_p, - penalty, - max_gen_tokens, - quant.clone(), - None, - false, - ), - "llama".to_string(), - )), - if let Some(model_id) = model_id { - model_id - } else { - "meta-llama/Llama-2-7b-chat-hf".to_string() - }, - quant, - ), - ModelSelected::Llama3 { - repeat_last_n, - temperature, - top_k, - top_p, - penalty, - max_gen_tokens, - quant, - thinking, - } => ( - Box::new(DefaultLoader::new( - SpecificConfig::new( - repeat_last_n, - temperature, - top_k, - top_p, - penalty, - max_gen_tokens, - quant.clone(), - None, - thinking, - ), - "llama3".to_string(), - )), - if let Some(model_id) = model_id { - model_id - } else { - "meta-llama/Meta-Llama-3.1-8B-Instruct".to_string() - }, - quant, - ), - ModelSelected::Phi2 { - repeat_last_n, - temperature, - top_k, - top_p, - penalty, - max_gen_tokens, - quant, - } => ( - Box::new(DefaultLoader::new( - SpecificConfig::new( - repeat_last_n, - temperature, - top_k, - top_p, - penalty, - max_gen_tokens, - quant.clone(), - None, - false, - ), - "phi2".to_string(), - )), - if let Some(model_id) = model_id { - model_id - } else { - "microsoft/phi-2".to_string() - }, - quant, - ), - ModelSelected::Phi3 { - repeat_last_n, - temperature, - top_k, - top_p, - penalty, - max_gen_tokens, - quant, - } => ( - Box::new(DefaultLoader::new( - SpecificConfig::new( - repeat_last_n, - temperature, - top_k, - top_p, - penalty, - max_gen_tokens, - quant.clone(), - None, - false, - ), - "phi3".to_string(), - )), - if let Some(model_id) = model_id { - model_id - } else { - "microsoft/Phi-3-mini-4k-instruct".to_string() - }, - quant, - ), - ModelSelected::Qwen2 { - repeat_last_n, - temperature, - top_k, - top_p, - penalty, - max_gen_tokens, - quant, - } => ( - Box::new(DefaultLoader::new( - SpecificConfig::new( - repeat_last_n, - temperature, - top_k, - top_p, - penalty, - max_gen_tokens, - quant.clone(), - None, - false, - ), - "qwen2".to_string(), - )), - if let Some(model_id) = model_id { - model_id - } else { - "Qwen/Qwen1.5-1.8B-Chat".to_string() - }, - quant, - ), - ModelSelected::Qwen3 { - repeat_last_n, - temperature, - top_k, - top_p, - penalty, - max_gen_tokens, - quant, - thinking, - } => ( - Box::new(DefaultLoader::new( - SpecificConfig::new( - repeat_last_n, - temperature, - top_k, - top_p, - penalty, - max_gen_tokens, - quant.clone(), - None, - thinking, - ), - "qwen3".to_string(), - )), - if let Some(model_id) = model_id { - model_id - } else { - "Qwen/Qwen3-8B".to_string() - }, - quant, - ), - ModelSelected::Gemma { - repeat_last_n, - temperature, - top_k, - top_p, - penalty, - max_gen_tokens, - quant, - } => ( - Box::new(DefaultLoader::new( - SpecificConfig::new( - repeat_last_n, - temperature, - top_k, - top_p, - penalty, - max_gen_tokens, - quant.clone(), - None, - false, - ), - "gemma".to_string(), - )), - if let Some(model_id) = model_id { - model_id - } else { - "google/gemma-2b-it".to_string() - }, - quant, - ), - ModelSelected::Gemma3 { - repeat_last_n, - temperature, - top_k, - top_p, - penalty, - max_gen_tokens, - quant, - thinking, - } => ( - Box::new(DefaultLoader::new( - SpecificConfig::new( - repeat_last_n, - temperature, - top_k, - top_p, - penalty, - max_gen_tokens, - quant.clone(), - None, - thinking, - ), - "gemma3".to_string(), - )), - if let Some(model_id) = model_id { - model_id - } else { - "google/gemma-3-4b-it".to_string() - }, - quant, - ), - ModelSelected::Mistral { - repeat_last_n, - temperature, - top_k, - top_p, - penalty, - max_gen_tokens, - quant, - } => ( - Box::new(DefaultLoader::new( - SpecificConfig::new( - repeat_last_n, - temperature, - top_k, - top_p, - penalty, - max_gen_tokens, - quant.clone(), - None, - false, - ), - "mistral".to_string(), - )), - if let Some(model_id) = model_id { - model_id - } else { - "mistralai/Mistral-7B-Instruct-v0.3".to_string() - }, - quant, - ), - - ModelSelected::Yi { - repeat_last_n, - temperature, - top_k, - top_p, - penalty, - max_gen_tokens, - quant, - } => ( - Box::new(DefaultLoader::new( - SpecificConfig::new( - repeat_last_n, - temperature, - top_k, - top_p, - penalty, - max_gen_tokens, - quant.clone(), - None, - false, - ), - "yi".to_string(), - )), - if let Some(model_id) = model_id { - model_id - } else { - "01-ai/Yi-6B-Chat".to_string() - }, - quant, - ), - - ModelSelected::StableLM { - repeat_last_n, - temperature, - penalty, - max_gen_tokens, - quant, - } => ( - Box::new(DefaultLoader::new( - SpecificConfig::new( - repeat_last_n, - temperature, - None, - None, - penalty, - max_gen_tokens, - quant.clone(), - None, - false, - ), - "stablelm".to_string(), - )), - if let Some(model_id) = model_id { - model_id - } else { - "stabilityai/stablelm-zephyr-3b".to_string() - }, - quant, - ), - ModelSelected::GLM4 { - repeat_last_n, - temperature, - top_k, - top_p, - penalty, - max_gen_tokens, - quant, - } => ( - Box::new(DefaultLoader::new( - SpecificConfig::new( - repeat_last_n, - temperature, - top_k, - top_p, - penalty, - max_gen_tokens, - quant.clone(), - None, - false, - ), - "glm4".to_string(), - )), - if let Some(model_id) = model_id { - model_id - } else { - "ZhipuAI/GLM-4-9B-0414".to_string() - }, - quant, - ), - ModelSelected::DeepSeek { - repeat_last_n, - temperature, - top_k, - top_p, - penalty, - max_gen_tokens, - quant, - num_experts_offload_per_rank, - thinking, - } => ( - Box::new(DefaultLoader::new( - SpecificConfig::new( - repeat_last_n, - temperature, - top_k, - top_p, - penalty, - max_gen_tokens, - quant.clone(), - num_experts_offload_per_rank, - thinking, - ), - "deepseek".to_string(), - )), - if let Some(model_id) = model_id { - model_id - } else { - "deepseek-ai/DeepSeek-V2-Lite-Chat".to_string() - }, - quant, - ), - } -} pub fn hub_load_local_safetensors( path: &String, diff --git a/src/main.rs b/src/main.rs index aead678c..08307f1a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,23 +3,20 @@ use axum::{ routing::post, Router, }; -use candle_core::{DType, Device}; +use candle_core::{DType, Device, Result}; #[cfg(feature = "nccl")] use candle_vllm::backend::heartbeat; use candle_vllm::openai::openai_server::chat_completions; use candle_vllm::openai::pipelines::llm_engine::LLMEngine; -use candle_vllm::openai::pipelines::pipeline::DefaultModelPaths; -use candle_vllm::openai::responses::APIError; +use candle_vllm::openai::pipelines::pipeline::DefaultLoader; use candle_vllm::openai::OpenAIServerData; use candle_vllm::scheduler::cache_engine::{CacheConfig, CacheEngine}; use candle_vllm::scheduler::SchedulerConfig; -use candle_vllm::{get_model_loader, hub_load_local_safetensors, ModelSelected}; use clap::Parser; -use std::{path::PathBuf, sync::Arc}; +use std::sync::Arc; use tracing::{info, warn}; const SIZE_IN_MB: usize = 1024 * 1024; use candle_vllm::openai::models::Config; -use std::path::Path; use tokio::sync::Notify; use tower_http::cors::{AllowOrigin, CorsLayer}; #[derive(Parser, Debug)] @@ -35,16 +32,13 @@ struct Args { hf_token_path: Option, /// Port to serve on (localhost:port) - #[arg(long = "p")] + #[arg(long = "p", default_value_t = 2000)] port: u16, /// Set verbose mode (print all requests) #[arg(long)] verbose: bool, - #[clap(subcommand)] - command: ModelSelected, - /// Maximum number of sequences to allow #[arg(long, default_value_t = 256)] max_num_seqs: usize, @@ -69,6 +63,9 @@ struct Args { #[arg(long)] dtype: Option, + #[arg(long)] + isq: Option, + #[arg(long, default_value_t = false)] cpu: bool, @@ -92,9 +89,9 @@ struct Args { #[arg(long, default_value_t = 500)] holding_time: usize, - //Whether the program running in multiprocess or multithread model for parallel inference + //Whether the program is forced running in multithread model for parallel inference (for debug) #[arg(long, default_value_t = false)] - multi_process: bool, + multithread: bool, #[arg(long, default_value_t = false)] log: bool, @@ -105,20 +102,21 @@ fn get_cache_config( kvcache_mem_cpu: usize, block_size: usize, config: &Config, + kv_dtype: DType, num_shards: usize, ) -> CacheConfig { - let dsize = config.kv_cache_dtype.size_in_bytes(); + let dsize = kv_dtype.size_in_bytes(); let num_gpu_blocks = kvcache_mem_gpu * SIZE_IN_MB / dsize / block_size - / (config.num_key_value_heads / num_shards) + / (config.num_key_value_heads.unwrap() / num_shards) / config.k_head_dim() / config.num_hidden_layers / 2; let num_cpu_blocks = kvcache_mem_cpu * SIZE_IN_MB / dsize / block_size - / (config.num_key_value_heads / num_shards) + / (config.num_key_value_heads.unwrap() / num_shards) / config.k_head_dim() / config.num_hidden_layers / 2; @@ -127,15 +125,11 @@ fn get_cache_config( num_gpu_blocks: Some(num_gpu_blocks), num_cpu_blocks: Some(num_cpu_blocks), fully_init: true, - dtype: config.kv_cache_dtype, + dtype: kv_dtype, } } -fn config_log( - logger: ftail::Ftail, - log_enable: bool, - log_file: String, -) -> Result<(), ftail::error::FtailError> { +fn config_log(logger: ftail::Ftail, log_enable: bool, log_file: String) -> Result<()> { if !log_enable { return Ok(()); } @@ -165,10 +159,12 @@ fn config_log( .console(cfg_filter) .single_file(log_file.as_str(), true, cfg_filter) .init() + .map_err(candle_core::Error::wrap) } #[tokio::main] -async fn main() -> Result<(), APIError> { +#[allow(unused_mut)] +async fn main() -> Result<()> { let args = Args::parse(); if !args.log { tracing_subscriber::fmt() @@ -176,90 +172,13 @@ async fn main() -> Result<(), APIError> { .init(); } - let (loader, model_id, quant) = get_model_loader(args.command, args.model_id.clone()); - if args.model_id.is_none() && args.weight_path.is_none() && args.weight_file.is_none() { - info!("No model id specified, using the default model_id or specified in the weight_path to retrieve config files!"); - } + let loader = Box::new(DefaultLoader::new( + args.model_id, + args.weight_path, + args.weight_file, + )); - let paths = match (&args.weight_path, &args.weight_file) { - //model in a folder (safetensor format, huggingface folder structure) - (Some(path), None) => DefaultModelPaths { - tokenizer_filename: Path::new(path).join("tokenizer.json"), - tokenizer_config_filename: Path::new(path).join("tokenizer_config.json"), - config_filename: Path::new(path).join("config.json"), - filenames: if Path::new(path) - .join("model.safetensors.index.json") - .exists() - { - hub_load_local_safetensors(path, "model.safetensors.index.json").unwrap() - } else { - //a single weight file case - let mut safetensors_files = Vec::::new(); - safetensors_files.insert(0, Path::new(path).join("model.safetensors")); - safetensors_files - }, - }, - //model in a quantized file (gguf/ggml format) - (path, Some(file)) => DefaultModelPaths { - tokenizer_filename: PathBuf::new(), - tokenizer_config_filename: PathBuf::new(), - config_filename: PathBuf::new(), - filenames: { - let path = path.clone().unwrap_or("".to_string()); - if Path::new(&path).join(file).exists() { - vec![Path::new(&path).join(file)] - } else { - panic!("Model file not found {file}"); - } - }, - }, - _ => { - //try download model anonymously - let loaded = loader.download_model( - model_id.clone(), - args.weight_file.clone(), - quant.clone(), - None, - args.hf_token.clone(), - args.hf_token_path.clone(), - ); - if loaded.is_ok() { - loaded.unwrap() - } else { - //if it's failed, try using huggingface token - info!("Try request model using cached huggingface token..."); - if args.hf_token.is_none() && args.hf_token_path.is_none() { - //no token provided - let token_path = format!( - "{}/.cache/huggingface/token", - dirs::home_dir() - .ok_or(APIError::new_str("No home directory"))? - .display() - ); - if !Path::new(&token_path).exists() { - //also no token cache - use std::io::Write; - let mut input_token = String::new(); - warn!("Unable to request model, please provide your huggingface token to download model:\n"); - std::io::stdin() - .read_line(&mut input_token) - .expect("Failed to read token!"); - std::fs::create_dir_all(Path::new(&token_path).parent().unwrap()).unwrap(); - let mut output = std::fs::File::create(token_path).unwrap(); - write!(output, "{}", input_token.trim()).expect("Failed to save token!"); - } - } - loader.download_model( - model_id, - args.weight_file, - quant.clone(), - None, - args.hf_token, - args.hf_token_path, - )? - } - } - }; + let (paths, gguf) = loader.prepare_model_weights(args.hf_token, args.hf_token_path)?; let dtype = match args.dtype.as_deref() { Some("f16") => DType::F16, @@ -281,17 +200,29 @@ async fn main() -> Result<(), APIError> { "More than one shard was given, but NCCL is not enabled for parallel inference!" ); - if num_shards > 1 - && quant.is_some() - && matches!(quant.as_ref().unwrap().as_str(), "ggml" | "gguf") - { + if gguf && num_shards > 1 { panic!("Multiple device-ids detected: ggml/gguf model is not supported for multi-rank inference!"); } - let logger = ftail::Ftail::new(); + if gguf && args.isq.is_some() { + panic!("Quantized gguf/ggml model does not support isq option!"); + } + + let multi_process = if num_shards > 1 { + if args.multithread { + tracing::warn!("The program is forced running under multithread mode (for debug purpose), which may not stable!"); + false + } else { + tracing::warn!("Multi-process mode is automatically enabled for multi-rank inference!"); + true + } + } else { + !args.multithread + }; + let logger: ftail::Ftail = ftail::Ftail::new(); let mut port = args.port; #[cfg(feature = "nccl")] - let (pipelines, global_rank, daemon_manager) = if args.multi_process { + let (pipelines, global_rank, daemon_manager) = if multi_process { use candle_vllm::openai::communicator::init_subprocess; let (id, local_rank, global_rank, global_world_size, daemon_manager) = init_subprocess(device_ids.clone()).unwrap(); @@ -310,7 +241,8 @@ async fn main() -> Result<(), APIError> { .load_model( paths, dtype, - &quant, + gguf, + args.isq.clone(), vec![device_ids[local_rank]], Some(id), Some(local_rank), @@ -330,7 +262,16 @@ async fn main() -> Result<(), APIError> { ( loader .load_model( - paths, dtype, &quant, device_ids, None, None, None, None, None, + paths, + dtype, + gguf, + args.isq.clone(), + device_ids, + None, + None, + None, + None, + None, ) .await, 0, @@ -341,7 +282,7 @@ async fn main() -> Result<(), APIError> { #[cfg(feature = "nccl")] info!( "parallel model: {}!", - if args.multi_process { + if multi_process { "multiprocess" } else { "multithread" @@ -354,7 +295,7 @@ async fn main() -> Result<(), APIError> { let _ = config_log(logger, args.log, log_file); ( loader - .load_model(paths, dtype, &quant, device_ids, None, None) + .load_model(paths, dtype, gguf, args.isq.clone(), device_ids, None, None) .await, 0, ) @@ -371,11 +312,17 @@ async fn main() -> Result<(), APIError> { .into_iter() .map(|pipeline| { let cfg = pipeline.get_model_config(); + let kv_dtype = if matches!(pipeline.name(), "phi2" | "PhiForCausalLM") { + DType::F32 + } else { + dtype + }; let cache_cfg = get_cache_config( args.kvcache_mem_gpu, args.kvcache_mem_cpu, //dummy 512MB for cpu args.block_size, &cfg, + kv_dtype, num_shards, ); let cache_engine = CacheEngine::new( @@ -410,7 +357,7 @@ async fn main() -> Result<(), APIError> { Arc::new(Notify::new()), args.holding_time, num_shards, - args.multi_process, + multi_process, #[cfg(feature = "nccl")] daemon_manager, )?; @@ -429,7 +376,7 @@ async fn main() -> Result<(), APIError> { } #[cfg(feature = "nccl")] - if args.multi_process { + if multi_process { let e = server_data.model.read(); let mut daemon_manager = e.daemon_manager.write(); daemon_manager.as_mut().unwrap().mpi_sync(); @@ -462,10 +409,10 @@ async fn main() -> Result<(), APIError> { let listener = tokio::net::TcpListener::bind(format!("0.0.0.0:{port}")) .await - .map_err(|e| APIError::new(e.to_string()))?; + .map_err(candle_core::Error::wrap)?; axum::serve(listener, app) .await - .map_err(|e| APIError::new(e.to_string()))?; + .map_err(candle_core::Error::wrap)?; Ok(()) } diff --git a/src/openai/mod.rs b/src/openai/mod.rs index e2f57475..812ffb05 100644 --- a/src/openai/mod.rs +++ b/src/openai/mod.rs @@ -39,12 +39,6 @@ where pub struct PipelineConfig { pub max_model_len: usize, pub default_max_tokens: usize, - pub penalty: f32, - pub repeat_last_n: usize, - pub temperature: Option, - pub top_k: Option, - pub top_p: Option, - pub thinking: Option, } #[derive(Deserialize, Debug, Clone)] diff --git a/src/openai/models/deepseek.rs b/src/openai/models/deepseek.rs index 99c56e5f..f6d9afb9 100644 --- a/src/openai/models/deepseek.rs +++ b/src/openai/models/deepseek.rs @@ -1,7 +1,6 @@ #![allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)] use super::{ - Config, DeepSeekRopeScaling, MoEConfig, QuantConfig, ScoringFunc, SpecificConfig, TokenID, - TopkMethod, + Config, DeepSeekRopeScaling, MoEConfig, QuantConfig, ScoringFunc, TokenID, TopkMethod, }; use crate::backend::custom_ops::moe::{masked_fill, NonZeroOp, SplitOp, TopKLastDimOp, TopKOutput}; use crate::backend::progress::{ProgressLike, ProgressReporter}; @@ -19,6 +18,7 @@ use std::cell::{Cell, RefCell}; use std::collections::HashSet; use std::f32::consts::PI; use std::iter::{zip, FromIterator}; +use std::path::PathBuf; use std::rc::Rc; use std::sync::{Arc, RwLock}; @@ -39,7 +39,7 @@ serde_default_fn!(usize, first_k_dense_replace, 0); serde_default_fn!(bool, norm_topk_prob, false); serde_default_fn!(ScoringFunc, scoring_func, ScoringFunc::Softmax); serde_default_fn!(Activation, hidden_act, Activation::Silu); -// serde_default_fn!(bool, tie_word_embeddings, false); +serde_default_fn!(bool, tie_word_embeddings, false); #[derive(Deserialize, Clone, Debug)] pub struct DeepSeekConfig { @@ -70,8 +70,8 @@ pub struct DeepSeekConfig { pub(crate) hidden_act: Activation, pub(crate) max_position_embeddings: usize, pub(crate) rms_norm_eps: f64, - // #[serde(default = "tie_word_embeddings")] - // pub(crate) tie_word_embeddings: bool, + #[serde(default = "tie_word_embeddings")] + pub(crate) tie_word_embeddings: bool, pub(crate) rope_theta: f32, pub(crate) rope_scaling: Option, // pub(crate) attention_bias: bool, @@ -94,68 +94,105 @@ pub struct DeepSeekV2RotaryEmbedding { cos: Tensor, } -impl DeepSeekConfig { - pub fn into_config( - self, - use_flash_attn: bool, - kv_cache_dtype: DType, - scfg: &SpecificConfig, - ) -> Config { +impl DeepSeek { + pub fn load_config(filename: &PathBuf, isq: Option) -> Result { + let config = match std::fs::read(filename.clone()) { + Ok(f) => { + let config: DeepSeekConfig = + serde_json::from_slice(&f).map_err(candle_core::Error::wrap)?; + config + } + Err(e) => panic!("Unable to load config file {:?}", e), + }; + + let num_experts_offload_per_rank = if let Ok(num_experts_offload_per_rank) = + std::env::var("NUM_EXPERTS_OFFLOAD_PER_RANK") + { + let number: usize = num_experts_offload_per_rank + .trim() + .parse::() + .expect("Failed to parse num_experts_offload_per_rank to number"); + Some(number) + } else { + None + }; + let moe_config = MoEConfig { - num_experts_per_tok: self.num_experts_per_tok, - n_routed_experts: self.n_routed_experts.unwrap_or(0), - moe_intermediate_size: self.moe_intermediate_size, - scoring_func: self.scoring_func, - topk_method: self.topk_method, - norm_topk_prob: self.norm_topk_prob, - routed_scaling_factor: self.routed_scaling_factor, - n_shared_experts: self.n_shared_experts, - qk_nope_head_dim: self.qk_nope_head_dim, - qk_rope_head_dim: self.qk_rope_head_dim, - v_head_dim: self.v_head_dim, - kv_lora_rank: self.kv_lora_rank, - first_k_dense_replace: self.first_k_dense_replace, - moe_layer_freq: self.moe_layer_freq, - rope_scaling: self.rope_scaling, - q_lora_rank: self.q_lora_rank, - n_group: self.n_group, - topk_group: self.topk_group, - num_experts_offload_per_rank: scfg.num_experts_offload_per_rank, + num_experts_per_tok: config.num_experts_per_tok, + n_routed_experts: config.n_routed_experts.unwrap_or(0), + moe_intermediate_size: config.moe_intermediate_size, + scoring_func: config.scoring_func, + topk_method: config.topk_method, + norm_topk_prob: config.norm_topk_prob, + routed_scaling_factor: config.routed_scaling_factor, + n_shared_experts: config.n_shared_experts, + qk_nope_head_dim: config.qk_nope_head_dim, + qk_rope_head_dim: config.qk_rope_head_dim, + v_head_dim: config.v_head_dim, + kv_lora_rank: config.kv_lora_rank, + first_k_dense_replace: config.first_k_dense_replace, + moe_layer_freq: config.moe_layer_freq, + rope_scaling: config.rope_scaling, + q_lora_rank: config.q_lora_rank, + n_group: config.n_group, + topk_group: config.topk_group, + num_experts_offload_per_rank, }; - Config { - hidden_size: self.hidden_size, - head_dim: Some(self.hidden_size / self.num_attention_heads), - intermediate_size: self.intermediate_size, - vocab_size: self.vocab_size, - num_hidden_layers: self.num_hidden_layers, - num_attention_heads: self.num_attention_heads, - num_key_value_heads: self.num_key_value_heads.unwrap_or(self.num_attention_heads), - rms_norm_eps: self.rms_norm_eps, - rope_theta: f64::from(self.rope_theta), + let quant = if config.quantization_config.is_some() { + Some( + config + .quantization_config + .as_ref() + .unwrap() + .quant_method + .clone(), + ) + } else if isq.is_some() { + panic!("DeepSeek does not support isq quantization yet!"); + } else { + None + }; + + let config = Config { + architectures: Some(vec!["DeepseekV3ForCausalLM".to_string()]), + hidden_size: config.hidden_size, + head_dim: Some(config.hidden_size / config.num_attention_heads), + intermediate_size: config.intermediate_size, + vocab_size: config.vocab_size, + num_hidden_layers: config.num_hidden_layers, + num_attention_heads: config.num_attention_heads, + num_key_value_heads: Some( + config + .num_key_value_heads + .unwrap_or(config.num_attention_heads), + ), + rms_norm_eps: config.rms_norm_eps, + rope_theta: f64::from(config.rope_theta), rope_local_base_freq: None, - use_flash_attn, - bos_token_id: self.bos_token_id, - eos_token_id: self.eos_token_id, - max_seq_len: self.max_position_embeddings, - sliding_window: self.sliding_window, + bos_token_id: Some(config.bos_token_id), + eos_token_id: config.eos_token_id, + max_seq_len: config.max_position_embeddings, + sliding_window: config.sliding_window, sliding_window_pattern: None, - hidden_act: Some(self.hidden_act), - tie_word_embeddings: false, + hidden_act: Some(config.hidden_act), + hidden_activation: None, + tie_word_embeddings: config.tie_word_embeddings, rope_scaling: None, - original_max_position_embeddings: None, - attention_bias: false, + max_position_embeddings: Some(config.max_position_embeddings), + original_max_position_embeddings: config.max_position_embeddings, + attention_bias: Some(false), partial_rotary_factor: None, - qk_layer_rms_norm: None, - kv_cache_dtype, + qk_layernorm: false, use_qkv_bias: None, custom_stop_tokens: None, - specific_config: scfg.clone(), attn_logit_softcapping: None, final_logit_softcapping: None, - quantization_config: self.quantization_config, + quantization_config: config.quantization_config.clone(), moe_config: Some(moe_config), - } + quant, + }; + Ok(config) } } @@ -391,14 +428,15 @@ impl Attention { ) -> Result { let q_head_dim = cfg.q_head_dim(); let moe_cfg = cfg.moe_config.as_ref().unwrap(); + let attention_bias = cfg.attention_bias.unwrap(); let q = match moe_cfg.q_lora_rank { Some(lora_rank) => { let a = ReplicatedLinear::load_b( cfg.hidden_size, lora_rank, - cfg.attention_bias, + attention_bias, vb.pp("q_a_proj"), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let norm = rms_norm(lora_rank, cfg.rms_norm_eps, vb.pp("q_a_layernorm"))?; @@ -408,7 +446,7 @@ impl Attention { false, vb.pp("q_b_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; QProj::Lora { a, norm, b } @@ -419,7 +457,7 @@ impl Attention { false, vb.pp("q_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?), }; @@ -427,9 +465,9 @@ impl Attention { let kv_a_proj_with_mqa = ReplicatedLinear::load_b( cfg.hidden_size, moe_cfg.kv_lora_rank + moe_cfg.qk_rope_head_dim, - cfg.attention_bias, + attention_bias, vb.pp("kv_a_proj_with_mqa"), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let kv_a_layernorm = rms_norm( @@ -443,22 +481,22 @@ impl Attention { false, vb.pp("kv_b_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let o_proj = TensorParallelRowLinear::load_with_hints( cfg.num_attention_heads * moe_cfg.v_head_dim, cfg.hidden_size, - cfg.attention_bias, + attention_bias, vb.pp("o_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let num_attention_heads = cfg.num_attention_heads / comm.world_size(); - let num_kv_heads = cfg.num_key_value_heads / comm.world_size(); + let num_kv_heads = cfg.num_key_value_heads.unwrap() / comm.world_size(); Ok(Self { q, @@ -588,7 +626,7 @@ impl Mlp { hidden_size, intermediate_size, vb.pp("gate_proj"), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; @@ -596,7 +634,7 @@ impl Mlp { hidden_size, intermediate_size, vb.pp("up_proj"), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; @@ -604,7 +642,7 @@ impl Mlp { intermediate_size, hidden_size, vb.pp("down_proj"), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; diff --git a/src/openai/models/gemma.rs b/src/openai/models/gemma.rs index 00924075..3fe9c856 100644 --- a/src/openai/models/gemma.rs +++ b/src/openai/models/gemma.rs @@ -1,95 +1,46 @@ -use super::{Config, QuantConfig}; +use super::Config; use crate::backend::progress::{ProgressLike, ProgressReporter}; use crate::openai::distributed::{ embedding, Comm, ReplicatedLinear, TensorParallelColumnLinear, TensorParallelRowLinear, VarBuilder, }; -use crate::openai::models::TokenID; use crate::paged_attention::input_metadata::InputMetadata; use crate::paged_attention::PagedAttention; -use crate::SpecificConfig; use candle::{DType, Device, IndexOp, Module, Result, Tensor}; use candle_core as candle; -use candle_nn::{Activation, RmsNorm}; +use candle_nn::RmsNorm; use std::iter::zip; +use std::path::PathBuf; use std::rc::Rc; use std::sync::{Arc, RwLock}; -use tracing::warn; -#[derive(serde::Deserialize, Debug, Clone)] -pub struct GemmaConfig { - pub attention_bias: bool, - pub head_dim: Option, - // The code gemma configs include both hidden_act and hidden_activation. - pub hidden_act: Option, - pub hidden_activation: Option, - pub hidden_size: usize, - pub intermediate_size: usize, - pub num_attention_heads: usize, - pub num_hidden_layers: usize, - pub num_key_value_heads: usize, - pub rms_norm_eps: f64, - pub rope_theta: f64, - pub vocab_size: usize, - pub bos_token_id: TokenID, - pub eos_token_id: TokenID, - pub sliding_window: Option, - pub max_position_embeddings: Option, - pub attn_logit_softcapping: Option, - pub final_logit_softcapping: Option, - pub quantization_config: Option, -} -impl GemmaConfig { - pub fn into_config( - self, - use_flash_attn: bool, - kv_cache_dtype: DType, - scfg: &SpecificConfig, - ) -> Config { - let hidden_act = match (self.hidden_act, self.hidden_activation) { - (None, Some(act)) | (Some(act), None) => Some(act), - (Some(act), Some(_)) => { - warn!("both hidden_act and hidden_activation are set"); - Some(act) - } - (None, None) => panic!("none of hidden_act and hidden_activation are set"), - }; - Config { - hidden_size: self.hidden_size, - head_dim: Some( - self.head_dim - .unwrap_or(self.hidden_size / self.num_attention_heads), - ), - intermediate_size: self.intermediate_size, - vocab_size: self.vocab_size, - num_hidden_layers: self.num_hidden_layers, - num_attention_heads: self.num_attention_heads, - num_key_value_heads: self.num_key_value_heads, - rms_norm_eps: self.rms_norm_eps, - rope_theta: self.rope_theta, - rope_local_base_freq: None, - use_flash_attn, - bos_token_id: self.bos_token_id, - eos_token_id: self.eos_token_id, - max_seq_len: self.max_position_embeddings.unwrap_or(4096), - sliding_window: self.sliding_window, - sliding_window_pattern: None, - hidden_act, - tie_word_embeddings: false, - rope_scaling: None, - original_max_position_embeddings: None, - attention_bias: self.attention_bias, - partial_rotary_factor: None, - qk_layer_rms_norm: None, - kv_cache_dtype, - use_qkv_bias: None, - custom_stop_tokens: None, - specific_config: scfg.clone(), - attn_logit_softcapping: self.attn_logit_softcapping, - final_logit_softcapping: self.final_logit_softcapping, - quantization_config: self.quantization_config, - moe_config: None, +impl Gemma { + pub fn load_config(filename: &PathBuf, isq: Option) -> Result { + let mut config = Config::load_config(filename.clone())?; + config.head_dim = Some( + config + .head_dim + .unwrap_or(config.hidden_size / config.num_attention_heads), + ); + config.num_key_value_heads = Some( + config + .num_key_value_heads + .unwrap_or(config.num_attention_heads), + ); + config.max_seq_len = config.max_position_embeddings.unwrap_or(config.max_seq_len); + if config.quantization_config.is_some() { + config.quant = Some( + config + .quantization_config + .as_ref() + .unwrap() + .quant_method + .clone(), + ); + } else if isq.is_some() { + config.quant = Some(isq.unwrap().to_string()); } + Ok(config) } } @@ -164,7 +115,7 @@ impl Mlp { false, vb.pp("gate_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let up_proj = TensorParallelColumnLinear::load_with_hints( @@ -173,7 +124,7 @@ impl Mlp { false, vb.pp("up_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let down_proj = TensorParallelRowLinear::load_with_hints( @@ -182,7 +133,7 @@ impl Mlp { false, vb.pp("down_proj"), comm, - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; Ok(Self { @@ -223,16 +174,16 @@ impl Attention { ) -> Result { let hidden_sz = cfg.hidden_size; let num_heads = cfg.num_attention_heads; - let num_kv_heads = cfg.num_key_value_heads; + let num_kv_heads = cfg.num_key_value_heads.unwrap(); let head_dim = cfg.head_dim.unwrap(); - let bias = cfg.attention_bias; + let bias = cfg.attention_bias.unwrap(); let q_proj = TensorParallelColumnLinear::load_with_hints( hidden_sz, num_heads * head_dim, bias, vb.pp("q_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let k_proj = TensorParallelColumnLinear::load_with_hints( @@ -241,7 +192,7 @@ impl Attention { bias, vb.pp("k_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let v_proj = TensorParallelColumnLinear::load_with_hints( @@ -250,7 +201,7 @@ impl Attention { bias, vb.pp("v_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; @@ -260,11 +211,11 @@ impl Attention { bias, vb.pp("o_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let attention_heads = cfg.num_attention_heads / comm.world_size(); - let kv_heads = cfg.num_key_value_heads / comm.world_size(); + let kv_heads = cfg.num_key_value_heads.unwrap() / comm.world_size(); Ok(Self { q_proj, k_proj, diff --git a/src/openai/models/gemma3.rs b/src/openai/models/gemma3.rs index d0e3b6b6..77be8aa1 100644 --- a/src/openai/models/gemma3.rs +++ b/src/openai/models/gemma3.rs @@ -4,16 +4,16 @@ use crate::openai::distributed::{ embedding, Comm, ReplicatedLinear, TensorParallelColumnLinear, TensorParallelRowLinear, VarBuilder, }; -use crate::openai::models::RopeScaling; use crate::openai::models::TokenID; +use crate::openai::models::{RopeScaling, ScalingValue}; use crate::paged_attention::input_metadata::InputMetadata; -use crate::SpecificConfig; use candle::{DType, Device, IndexOp, Module, Result, Tensor}; use candle_core as candle; use candle_nn::{Activation, RmsNorm}; use either::Either; use std::collections::HashMap; use std::iter::zip; +use std::path::PathBuf; use std::rc::Rc; use std::sync::{Arc, RwLock}; @@ -98,31 +98,38 @@ pub struct Gemma3Config { pub text_config: GemmaTextConfig, } -impl Gemma3Config { - pub fn into_config( - self, - use_flash_attn: bool, - kv_cache_dtype: DType, - scfg: &SpecificConfig, - ) -> Config { - let bos_token_id = self +impl Gemma3 { + pub fn load_config(filename: &PathBuf, isq: Option) -> Result { + let config = match std::fs::read(filename.clone()) { + Ok(f) => { + let config: Gemma3Config = + serde_json::from_slice(&f).map_err(candle_core::Error::wrap)?; + config + } + Err(e) => panic!("Unable to load config file {:?}", e), + }; + + let bos_token_id = config .text_config .bos_token_id - .or(self.bos_token_id) + .or(config.bos_token_id) .unwrap_or(super::TokenID(Either::Left(Some(2)))); - let eos_token_id = self + let eos_token_id = config .text_config .eos_token_id - .or(self.eos_token_id) + .or(config.eos_token_id) .unwrap_or(super::TokenID(Either::Left(Some(1)))); - let ropescaling = if self.text_config.rope_scaling.is_some() { + let ropescaling = if config.text_config.rope_scaling.is_some() { let mut ropescaling = HashMap::::new(); - for (key, value) in self.text_config.rope_scaling.as_ref().unwrap() { + for (key, value) in config.text_config.rope_scaling.as_ref().unwrap() { match value { Gemma3RopeScaling(Either::Left(l)) => { - ropescaling.insert(key.to_string(), RopeScaling(Either::Left(vec![*l]))); + ropescaling.insert( + key.to_string(), + RopeScaling(Either::Left(ScalingValue(Either::Left(*l)))), + ); } Gemma3RopeScaling(Either::Right(r)) => { ropescaling @@ -135,39 +142,57 @@ impl Gemma3Config { None }; - Config { - hidden_size: self.text_config.hidden_size, - head_dim: Some(self.text_config.head_dim), - intermediate_size: self.text_config.intermediate_size, - vocab_size: self.text_config.vocab_size, - num_hidden_layers: self.text_config.num_hidden_layers, - num_attention_heads: self.text_config.num_attention_heads, - num_key_value_heads: self.text_config.num_key_value_heads, - rms_norm_eps: self.text_config.rms_norm_eps, - rope_theta: self.text_config.rope_theta, - rope_local_base_freq: Some(self.text_config.rope_local_base_freq), - use_flash_attn, - bos_token_id, + let quant = if config.text_config.quantization_config.is_some() { + Some( + config + .text_config + .quantization_config + .as_ref() + .unwrap() + .quant_method + .clone(), + ) + } else if isq.is_some() { + Some(isq.unwrap().to_string()) + } else { + None + }; + + let config = Config { + architectures: Some(vec!["Gemma3ForConditionalGeneration".to_string()]), + hidden_size: config.text_config.hidden_size, + head_dim: Some(config.text_config.head_dim), + intermediate_size: config.text_config.intermediate_size, + vocab_size: config.text_config.vocab_size, + num_hidden_layers: config.text_config.num_hidden_layers, + num_attention_heads: config.text_config.num_attention_heads, + num_key_value_heads: Some(config.text_config.num_key_value_heads), + rms_norm_eps: config.text_config.rms_norm_eps, + rope_theta: config.text_config.rope_theta, + rope_local_base_freq: Some(config.text_config.rope_local_base_freq), + bos_token_id: Some(bos_token_id), eos_token_id, - max_seq_len: self.text_config.max_position_embeddings, - sliding_window: self.text_config.sliding_window, - sliding_window_pattern: Some(self.text_config.sliding_window_pattern), - hidden_act: Some(self.text_config.hidden_activation), - tie_word_embeddings: self.text_config.tie_word_embeddings, + max_seq_len: config.text_config.max_position_embeddings, + sliding_window: config.text_config.sliding_window, + sliding_window_pattern: Some(config.text_config.sliding_window_pattern), + hidden_act: Some(config.text_config.hidden_activation), + hidden_activation: None, + tie_word_embeddings: config.text_config.tie_word_embeddings, rope_scaling: ropescaling, - original_max_position_embeddings: None, - attention_bias: self.text_config.attention_bias, + max_position_embeddings: Some(config.text_config.max_position_embeddings), + original_max_position_embeddings: config.text_config.max_position_embeddings, + attention_bias: Some(config.text_config.attention_bias), partial_rotary_factor: None, - qk_layer_rms_norm: None, - kv_cache_dtype, + qk_layernorm: false, use_qkv_bias: None, custom_stop_tokens: None, - specific_config: scfg.clone(), - attn_logit_softcapping: self.text_config.attn_logit_softcapping, - final_logit_softcapping: self.text_config.final_logit_softcapping, - quantization_config: self.text_config.quantization_config, + attn_logit_softcapping: config.text_config.attn_logit_softcapping, + final_logit_softcapping: config.text_config.final_logit_softcapping, + quantization_config: config.text_config.quantization_config.clone(), moe_config: None, - } + quant, + }; + Ok(config) } } @@ -305,7 +330,7 @@ impl Mlp { false, vb.pp("gate_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let up_proj = TensorParallelColumnLinear::load_with_hints( @@ -314,7 +339,7 @@ impl Mlp { false, vb.pp("up_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let down_proj = TensorParallelRowLinear::load_with_hints( @@ -323,7 +348,7 @@ impl Mlp { false, vb.pp("down_proj"), comm, - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; Ok(Self { @@ -368,9 +393,9 @@ impl Attention { ) -> Result { let hidden_sz = cfg.hidden_size; let num_heads = cfg.num_attention_heads; - let num_kv_heads = cfg.num_key_value_heads; + let num_kv_heads = cfg.num_key_value_heads.unwrap(); let head_dim = cfg.head_dim.unwrap(); - let bias = cfg.attention_bias; + let bias = cfg.attention_bias.unwrap(); let q_proj = TensorParallelColumnLinear::load_with_hints( hidden_sz, @@ -378,7 +403,7 @@ impl Attention { bias, vb.pp("q_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let k_proj = TensorParallelColumnLinear::load_with_hints( @@ -387,7 +412,7 @@ impl Attention { bias, vb.pp("k_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let v_proj = TensorParallelColumnLinear::load_with_hints( @@ -396,7 +421,7 @@ impl Attention { bias, vb.pp("v_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; @@ -406,7 +431,7 @@ impl Attention { bias, vb.pp("o_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; @@ -414,7 +439,7 @@ impl Attention { let k_norm = rms_norm(head_dim, cfg.rms_norm_eps, vb.pp("k_norm"))?; let attention_heads = cfg.num_attention_heads / comm.world_size(); - let kv_heads = cfg.num_key_value_heads / comm.world_size(); + let kv_heads = cfg.num_key_value_heads.unwrap() / comm.world_size(); Ok(Self { q_proj, k_proj, diff --git a/src/openai/models/glm4.rs b/src/openai/models/glm4.rs index 87213fec..1975c6e9 100644 --- a/src/openai/models/glm4.rs +++ b/src/openai/models/glm4.rs @@ -1,83 +1,52 @@ -use super::{Config, QuantConfig}; +use super::Config; +use crate::backend::progress::{ProgressLike, ProgressReporter}; use crate::openai::distributed::{ embedding, rms_norm, Comm, MergedParallelColumnLinear, ReplicatedLinear, TensorParallelColumnLinear, TensorParallelRowLinear, VarBuilder, }; use crate::paged_attention::input_metadata::InputMetadata; -use crate::SpecificConfig; use candle::{DType, Device, IndexOp, Result, Tensor, D}; use candle_core as candle; use candle_nn::{Embedding, Module, RmsNorm}; -pub const MAX_SEQ_LEN: usize = 4096; -use crate::backend::progress::{ProgressLike, ProgressReporter}; -use crate::openai::models::TokenID; -use either::Either; use std::iter::zip; +use std::path::PathBuf; pub use std::rc::Rc; use std::sync::{Arc, RwLock}; -#[derive(Debug, Clone, serde::Deserialize)] -pub struct GLMConfig { - pub num_hidden_layers: Option, - pub vocab_size: usize, - pub hidden_size: usize, - pub head_dim: Option, - pub num_attention_heads: usize, - pub num_key_value_heads: usize, - pub intermediate_size: usize, - pub rms_norm_eps: f64, - pub rope_theta: Option, - pub partial_rotary_factor: Option, - pub hidden_act: candle_nn::Activation, - pub attention_bias: Option, - pub sliding_window: Option, - pub eos_token_id: TokenID, - pub max_position_embeddings: Option, - pub quantization_config: Option, -} -impl GLMConfig { - pub fn into_config( - self, - use_flash_attn: bool, - kv_cache_dtype: DType, - scfg: &SpecificConfig, - ) -> Config { - Config { - hidden_size: self.hidden_size, - head_dim: Some( - self.head_dim - .unwrap_or(self.hidden_size / self.num_attention_heads), - ), - intermediate_size: self.intermediate_size, - vocab_size: self.vocab_size, - num_hidden_layers: self.num_hidden_layers.unwrap_or(40), - num_attention_heads: self.num_attention_heads, - num_key_value_heads: self.num_key_value_heads, - rms_norm_eps: self.rms_norm_eps, - rope_theta: self.rope_theta.unwrap_or(10_000f64), - use_flash_attn, - bos_token_id: super::TokenID(Either::Left(Some(128256))), - eos_token_id: self.eos_token_id, - max_seq_len: self.max_position_embeddings.unwrap_or(32768), - sliding_window: self.sliding_window, - hidden_act: Some(self.hidden_act), - tie_word_embeddings: false, - rope_local_base_freq: None, - sliding_window_pattern: None, - rope_scaling: None, - original_max_position_embeddings: None, - attention_bias: self.attention_bias.unwrap_or(false), - partial_rotary_factor: self.partial_rotary_factor, - qk_layer_rms_norm: None, - use_qkv_bias: None, - kv_cache_dtype, - custom_stop_tokens: None, - specific_config: scfg.clone(), - attn_logit_softcapping: None, - final_logit_softcapping: None, - quantization_config: self.quantization_config, - moe_config: None, +impl GLM4 { + pub fn load_config(filename: &PathBuf, isq: Option) -> Result { + let mut config = Config::load_config(filename.clone())?; + config.head_dim = Some( + config + .head_dim + .unwrap_or(config.hidden_size / config.num_attention_heads), + ); + config.num_key_value_heads = Some( + config + .num_key_value_heads + .unwrap_or(config.num_attention_heads), + ); + config.num_hidden_layers = config.num_hidden_layers; + config.max_seq_len = config.max_position_embeddings.unwrap_or(32768); + config.attention_bias = Some(config.attention_bias.unwrap_or(false)); + config.bos_token_id = Some( + config + .bos_token_id + .unwrap_or(super::TokenID(either::Either::Left(Some(128256)))), + ); + if config.quantization_config.is_some() { + config.quant = Some( + config + .quantization_config + .as_ref() + .unwrap() + .quant_method + .clone(), + ); + } else if isq.is_some() { + config.quant = Some(isq.unwrap().to_string()); } + Ok(config) } } @@ -198,34 +167,35 @@ impl SelfAttention { ) -> Result { let hidden_sz = cfg.hidden_size; let num_heads = cfg.num_attention_heads; - let num_kv_heads = cfg.num_key_value_heads; + let num_kv_heads = cfg.num_key_value_heads.unwrap(); + let attention_bias = cfg.attention_bias.unwrap_or(false); let head_dim = cfg.head_dim.unwrap_or(hidden_sz / num_heads); let q_proj = TensorParallelColumnLinear::load_with_hints( hidden_sz, num_heads * head_dim, - cfg.attention_bias, + attention_bias, vb.pp("q_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let k_proj = TensorParallelColumnLinear::load_with_hints( hidden_sz, num_kv_heads * head_dim, - cfg.attention_bias, + attention_bias, vb.pp("k_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let v_proj = TensorParallelColumnLinear::load_with_hints( hidden_sz, num_kv_heads * head_dim, - cfg.attention_bias, + attention_bias, vb.pp("v_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; @@ -235,18 +205,18 @@ impl SelfAttention { false, vb.pp("o_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; assert!(cfg.num_attention_heads >= comm.world_size()); assert!(cfg.num_attention_heads % comm.world_size() == 0); - assert!(cfg.num_key_value_heads >= comm.world_size()); - assert!(cfg.num_key_value_heads % comm.world_size() == 0); + assert!(cfg.num_key_value_heads.unwrap() >= comm.world_size()); + assert!(cfg.num_key_value_heads.unwrap() % comm.world_size() == 0); let attention_heads = cfg.num_attention_heads / comm.world_size(); - let kv_heads = cfg.num_key_value_heads / comm.world_size(); + let kv_heads = cfg.num_key_value_heads.unwrap() / comm.world_size(); Ok(Self { q_proj, k_proj, @@ -332,7 +302,7 @@ impl MLP { 2, vb.pp("gate_up_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; @@ -342,7 +312,7 @@ impl MLP { false, vb.pp("down_proj"), comm, - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; diff --git a/src/openai/models/llama.rs b/src/openai/models/llama.rs index 9bbc61b7..2085dd59 100644 --- a/src/openai/models/llama.rs +++ b/src/openai/models/llama.rs @@ -1,82 +1,46 @@ -use super::{Config, QuantConfig}; +use super::Config; +use crate::backend::progress::{ProgressLike, ProgressReporter}; use crate::openai::distributed::{ embedding, rms_norm, Comm, ReplicatedLinear, TensorParallelColumnLinear, TensorParallelRowLinear, VarBuilder, }; use crate::paged_attention::input_metadata::InputMetadata; use crate::paged_attention::PagedAttention; -use crate::SpecificConfig; use candle::{DType, Device, IndexOp, Result, Tensor}; use candle_core as candle; use candle_nn::{Embedding, Module, RmsNorm}; -pub const MAX_SEQ_LEN: usize = 4096; -use crate::backend::progress::{ProgressLike, ProgressReporter}; -use crate::openai::models::TokenID; use std::iter::zip; +use std::path::PathBuf; pub use std::rc::Rc; use std::sync::{Arc, RwLock}; -#[derive(Debug, Clone, serde::Deserialize)] -pub struct LlamaConfig { - pub hidden_size: usize, - pub intermediate_size: usize, - pub vocab_size: usize, - pub num_hidden_layers: usize, - pub num_attention_heads: usize, - pub num_key_value_heads: Option, - pub rms_norm_eps: f64, - #[serde(default = "default_rope")] - pub rope_theta: f32, - pub bos_token_id: TokenID, - pub eos_token_id: TokenID, - pub sliding_window: Option, - pub max_position_embeddings: Option, - pub quantization_config: Option, -} -fn default_rope() -> f32 { - 10_000.0 -} - -impl LlamaConfig { - pub fn into_config( - self, - use_flash_attn: bool, - kv_cache_dtype: DType, - scfg: &SpecificConfig, - ) -> Config { - Config { - hidden_size: self.hidden_size, - head_dim: Some(self.hidden_size / self.num_attention_heads), - intermediate_size: self.intermediate_size, - vocab_size: self.vocab_size, - num_hidden_layers: self.num_hidden_layers, - num_attention_heads: self.num_attention_heads, - num_key_value_heads: self.num_key_value_heads.unwrap_or(self.num_attention_heads), - rms_norm_eps: self.rms_norm_eps, - rope_theta: f64::from(self.rope_theta), - rope_local_base_freq: None, - use_flash_attn, - bos_token_id: self.bos_token_id, - eos_token_id: self.eos_token_id, - max_seq_len: self.max_position_embeddings.unwrap_or(MAX_SEQ_LEN), - sliding_window: self.sliding_window, - sliding_window_pattern: None, - hidden_act: None, - tie_word_embeddings: false, - rope_scaling: None, - original_max_position_embeddings: None, - attention_bias: false, - partial_rotary_factor: None, - qk_layer_rms_norm: None, - kv_cache_dtype, - use_qkv_bias: None, - custom_stop_tokens: None, - specific_config: scfg.clone(), - attn_logit_softcapping: None, - final_logit_softcapping: None, - quantization_config: self.quantization_config, - moe_config: None, +impl Llama { + pub fn load_config(filename: &PathBuf, isq: Option) -> Result { + let mut config = Config::load_config(filename.clone())?; + config.head_dim = Some( + config + .head_dim + .unwrap_or(config.hidden_size / config.num_attention_heads), + ); + config.num_key_value_heads = Some( + config + .num_key_value_heads + .unwrap_or(config.num_attention_heads), + ); + config.max_seq_len = config.max_position_embeddings.unwrap_or(config.max_seq_len); + if config.quantization_config.is_some() { + config.quant = Some( + config + .quantization_config + .as_ref() + .unwrap() + .quant_method + .clone(), + ); + } else if isq.is_some() { + config.quant = Some(isq.unwrap().to_string()); } + Ok(config) } } @@ -202,14 +166,15 @@ impl CausalSelfAttention { ) -> Result { let size_in = cfg.hidden_size; let size_q = (cfg.hidden_size / cfg.num_attention_heads) * cfg.num_attention_heads; - let size_kv = (cfg.hidden_size / cfg.num_attention_heads) * cfg.num_key_value_heads; + let size_kv = + (cfg.hidden_size / cfg.num_attention_heads) * cfg.num_key_value_heads.unwrap(); let q_proj = TensorParallelColumnLinear::load_with_hints( size_in, size_q, false, vb.pp("q_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let k_proj = TensorParallelColumnLinear::load_with_hints( @@ -218,7 +183,7 @@ impl CausalSelfAttention { false, vb.pp("k_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let v_proj = TensorParallelColumnLinear::load_with_hints( @@ -227,7 +192,7 @@ impl CausalSelfAttention { false, vb.pp("v_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; @@ -237,12 +202,12 @@ impl CausalSelfAttention { false, vb.pp("o_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let head_dim = cfg.hidden_size / cfg.num_attention_heads; let attention_heads = cfg.num_attention_heads / comm.world_size(); - let kv_heads = cfg.num_key_value_heads / comm.world_size(); + let kv_heads = cfg.num_key_value_heads.unwrap() / comm.world_size(); Ok(Self { q_proj, k_proj, @@ -286,7 +251,7 @@ impl Mlp { false, vb.pp("gate_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let c_fc2 = TensorParallelColumnLinear::load_with_hints( @@ -295,7 +260,7 @@ impl Mlp { false, vb.pp("up_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let c_proj = TensorParallelRowLinear::load_with_hints( @@ -304,7 +269,7 @@ impl Mlp { false, vb.pp("down_proj"), comm, - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; Ok(Self { diff --git a/src/openai/models/mistral.rs b/src/openai/models/mistral.rs index 8cd2c20b..cb5aeaf1 100644 --- a/src/openai/models/mistral.rs +++ b/src/openai/models/mistral.rs @@ -1,77 +1,45 @@ -use super::{Config, QuantConfig}; +use super::Config; use crate::backend::progress::{ProgressLike, ProgressReporter}; use crate::openai::distributed::{ embedding, rms_norm, Comm, ReplicatedLinear, TensorParallelColumnLinear, TensorParallelRowLinear, VarBuilder, }; -use crate::openai::models::TokenID; use crate::paged_attention::input_metadata::InputMetadata; use crate::paged_attention::PagedAttention; -use crate::SpecificConfig; use candle_core::{DType, Device, IndexOp, Module, Result, Tensor}; use candle_nn::{Activation, RmsNorm}; use std::iter::zip; +use std::path::PathBuf; use std::rc::Rc; use std::sync::{Arc, RwLock}; -#[derive(Debug, Clone, serde::Deserialize)] -pub struct MistralConfig { - pub vocab_size: usize, - pub hidden_size: usize, - pub intermediate_size: usize, - pub num_hidden_layers: usize, - pub num_attention_heads: usize, - pub num_key_value_heads: usize, - pub hidden_act: Activation, - pub max_position_embeddings: usize, - pub rms_norm_eps: f64, - pub rope_theta: f64, - pub sliding_window: Option, - pub bos_token_id: TokenID, - pub eos_token_id: TokenID, - pub tie_word_embeddings: Option, - pub quantization_config: Option, -} -impl MistralConfig { - pub fn into_config( - self, - use_flash_attn: bool, - kv_cache_dtype: DType, - scfg: &SpecificConfig, - ) -> Config { - Config { - hidden_size: self.hidden_size, - head_dim: Some(self.hidden_size / self.num_attention_heads), - intermediate_size: self.intermediate_size, - vocab_size: self.vocab_size, - num_hidden_layers: self.num_hidden_layers, - num_attention_heads: self.num_attention_heads, - num_key_value_heads: self.num_key_value_heads, - rms_norm_eps: self.rms_norm_eps, - rope_theta: self.rope_theta, - rope_local_base_freq: None, - use_flash_attn, - bos_token_id: self.bos_token_id, - eos_token_id: self.eos_token_id, - max_seq_len: self.max_position_embeddings, - sliding_window: self.sliding_window, - sliding_window_pattern: None, - hidden_act: Some(self.hidden_act), - tie_word_embeddings: self.tie_word_embeddings.unwrap_or(false), - rope_scaling: None, - original_max_position_embeddings: None, - attention_bias: false, - partial_rotary_factor: None, - qk_layer_rms_norm: None, - kv_cache_dtype, - use_qkv_bias: None, - custom_stop_tokens: None, - specific_config: scfg.clone(), - attn_logit_softcapping: None, - final_logit_softcapping: None, - quantization_config: self.quantization_config, - moe_config: None, +impl Mistral { + pub fn load_config(filename: &PathBuf, isq: Option) -> Result { + let mut config = Config::load_config(filename.clone())?; + config.head_dim = Some( + config + .head_dim + .unwrap_or(config.hidden_size / config.num_attention_heads), + ); + config.num_key_value_heads = Some( + config + .num_key_value_heads + .unwrap_or(config.num_attention_heads), + ); + config.max_seq_len = config.max_position_embeddings.unwrap_or(config.max_seq_len); + if config.quantization_config.is_some() { + config.quant = Some( + config + .quantization_config + .as_ref() + .unwrap() + .quant_method + .clone(), + ); + } else if isq.is_some() { + config.quant = Some(isq.unwrap().to_string()); } + Ok(config) } } @@ -143,7 +111,7 @@ impl Mlp { false, vb.pp("gate_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let up_proj = TensorParallelColumnLinear::load_with_hints( @@ -152,7 +120,7 @@ impl Mlp { false, vb.pp("up_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let down_proj = TensorParallelRowLinear::load_with_hints( @@ -161,7 +129,7 @@ impl Mlp { false, vb.pp("down_proj"), comm, - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; Ok(Self { @@ -202,7 +170,7 @@ impl Attention { ) -> Result { let hidden_sz = cfg.hidden_size; let num_heads = cfg.num_attention_heads; - let num_kv_heads = cfg.num_key_value_heads; + let num_kv_heads = cfg.num_key_value_heads.unwrap(); let head_dim = hidden_sz / num_heads; let q_proj = TensorParallelColumnLinear::load_with_hints( hidden_sz, @@ -210,7 +178,7 @@ impl Attention { false, vb.pp("q_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let k_proj = TensorParallelColumnLinear::load_with_hints( @@ -219,7 +187,7 @@ impl Attention { false, vb.pp("k_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let v_proj = TensorParallelColumnLinear::load_with_hints( @@ -228,7 +196,7 @@ impl Attention { false, vb.pp("v_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; @@ -238,11 +206,11 @@ impl Attention { false, vb.pp("o_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let attention_heads = cfg.num_attention_heads / comm.world_size(); - let kv_heads = cfg.num_key_value_heads / comm.world_size(); + let kv_heads = cfg.num_key_value_heads.unwrap() / comm.world_size(); Ok(Self { q_proj, k_proj, diff --git a/src/openai/models/mod.rs b/src/openai/models/mod.rs index a2f9b84b..6088bddf 100644 --- a/src/openai/models/mod.rs +++ b/src/openai/models/mod.rs @@ -17,15 +17,18 @@ pub mod yi; use crate::openai::distributed::Comm; use crate::paged_attention::input_metadata::InputMetadata; use crate::paged_attention::PagedAttention; -use crate::SpecificConfig; use candle_core::{DType, Device, Result, Tensor}; use either::Either; use serde::Deserialize; use std::cell::RefCell; use std::collections::HashMap; +use std::path::PathBuf; use std::rc::Rc; #[derive(Deserialize, Debug, Clone)] -pub struct RopeScaling(#[serde(with = "either::serde_untagged")] pub Either, String>); +pub struct ScalingValue(#[serde(with = "either::serde_untagged")] pub Either>); + +#[derive(Deserialize, Debug, Clone)] +pub struct RopeScaling(#[serde(with = "either::serde_untagged")] pub Either); #[derive(Deserialize, Debug, Clone)] pub struct TokenID( @@ -60,7 +63,7 @@ pub enum ScoringFunc { Sigmoid, } -#[derive(Debug, Clone)] +#[derive(Deserialize, Debug, Clone)] pub struct MoEConfig { pub num_experts_per_tok: Option, pub n_routed_experts: usize, @@ -117,42 +120,89 @@ pub enum DeepSeekRopeScaling { }, } -#[derive(Debug, Clone)] +#[doc(hidden)] +#[macro_export] +macro_rules! serde_default_cfg { + ($t:ty, $name:ident, $v:expr) => { + fn $name() -> $t { + $v + } + }; +} +serde_default_cfg!(usize, max_seq_len, 8192); +serde_default_cfg!(usize, original_max_position_embeddings, 8192); +serde_default_cfg!(bool, tie_word_embeddings, false); +serde_default_cfg!(f64, rope_theta, 10_000.0f64); +serde_default_cfg!(bool, qk_layernorm, false); + +#[derive(Deserialize, Debug, Clone)] pub struct Config { + pub architectures: Option>, pub hidden_size: usize, pub head_dim: Option, pub intermediate_size: usize, pub vocab_size: usize, pub num_hidden_layers: usize, pub num_attention_heads: usize, - pub num_key_value_heads: usize, - pub use_flash_attn: bool, + pub num_key_value_heads: Option, pub rms_norm_eps: f64, + #[serde(default = "rope_theta")] pub rope_theta: f64, pub rope_local_base_freq: Option, - pub bos_token_id: TokenID, + pub bos_token_id: Option, pub eos_token_id: TokenID, + #[serde(default = "max_seq_len")] pub max_seq_len: usize, + #[serde(default = "original_max_position_embeddings")] + pub original_max_position_embeddings: usize, pub sliding_window: Option, pub sliding_window_pattern: Option, pub hidden_act: Option, + pub hidden_activation: Option, + #[serde(default = "tie_word_embeddings")] pub tie_word_embeddings: bool, pub rope_scaling: Option>, - pub original_max_position_embeddings: Option, - pub attention_bias: bool, + pub max_position_embeddings: Option, + pub attention_bias: Option, pub partial_rotary_factor: Option, - pub qk_layer_rms_norm: Option, - pub kv_cache_dtype: DType, + #[serde(default = "qk_layernorm")] + pub qk_layernorm: bool, pub use_qkv_bias: Option, pub custom_stop_tokens: Option>, - pub specific_config: SpecificConfig, pub attn_logit_softcapping: Option, pub final_logit_softcapping: Option, pub quantization_config: Option, pub moe_config: Option, + pub quant: Option, } impl Config { + pub fn load_config(filename: PathBuf) -> Result { + match std::fs::read(filename) { + Ok(f) => { + let config: Config = + serde_json::from_slice(&f).map_err(candle_core::Error::wrap)?; + Ok(config) + } + Err(e) => panic!( + "Unable to load config file {:?}\n ***Tips: use `--f` to specify GGUF file path!", + e + ), + } + } + + pub fn get_model_arch(filename: &PathBuf) -> Result { + let config = Config::load_config(filename.clone())?; + if config.architectures.is_none() { + candle_core::bail!("Missing architectures in config file!"); + } + let architectures = config.architectures.unwrap(); + if architectures.is_empty() { + candle_core::bail!("No architectures defined in config file!"); + } + Ok(architectures[0].clone()) + } + pub fn get_head_size(&self) -> usize { self.head_dim .unwrap_or(self.hidden_size / self.num_attention_heads) @@ -261,7 +311,7 @@ pub struct NaiveAttention { impl NaiveAttention { pub fn new(cfg: &Config, sliding_window: Option) -> Self { let num_heads = cfg.num_attention_heads; - let num_kv_heads = cfg.num_key_value_heads; + let num_kv_heads = cfg.num_key_value_heads.unwrap(); let num_kv_groups = num_heads / num_kv_heads; let scale = 1f64 / f64::sqrt(cfg.head_dim.unwrap() as f64); @@ -345,7 +395,7 @@ impl AttentionSelect { } else { let head_dim = cfg.head_dim.unwrap(); let attention_heads = cfg.num_attention_heads / comm.world_size(); - let kv_heads = cfg.num_key_value_heads / comm.world_size(); + let kv_heads = cfg.num_key_value_heads.unwrap() / comm.world_size(); AttentionSelect::Paged( PagedAttention::new( attention_heads, diff --git a/src/openai/models/phi2.rs b/src/openai/models/phi2.rs index 861078c9..ef7371ba 100644 --- a/src/openai/models/phi2.rs +++ b/src/openai/models/phi2.rs @@ -1,81 +1,45 @@ -use super::{Config, QuantConfig}; +use super::Config; use crate::backend::progress::{ProgressLike, ProgressReporter}; use crate::openai::distributed::{ embedding, layer_norm, Comm, ReplicatedLinear, TensorParallelColumnLinear, TensorParallelRowLinear, VarBuilder, }; -use crate::openai::models::TokenID; use crate::paged_attention::input_metadata::InputMetadata; use crate::paged_attention::PagedAttention; -use crate::SpecificConfig; use candle_core::{DType, Device, IndexOp, Module, Result, Tensor, D}; use candle_nn::{Activation, Embedding, LayerNorm}; -use serde::Deserialize; use std::iter::zip; +use std::path::PathBuf; use std::rc::Rc; use std::sync::{Arc, RwLock}; -#[derive(Debug, Clone, Deserialize)] -pub struct Phi2Config { - pub vocab_size: usize, - pub hidden_size: usize, - pub intermediate_size: usize, - pub num_hidden_layers: usize, - pub num_attention_heads: usize, - pub num_key_value_heads: Option, - pub hidden_act: Activation, - pub max_position_embeddings: usize, - pub layer_norm_eps: f64, - pub tie_word_embeddings: bool, - pub rope_theta: f64, - pub partial_rotary_factor: f32, - pub qk_layernorm: bool, - pub bos_token_id: TokenID, - pub eos_token_id: TokenID, - pub sliding_window: Option, - pub original_max_position_embeddings: Option, - pub quantization_config: Option, -} -impl Phi2Config { - pub fn into_config( - self, - use_flash_attn: bool, - kv_cache_dtype: DType, - scfg: &SpecificConfig, - ) -> Config { - Config { - hidden_size: self.hidden_size, - head_dim: Some(self.hidden_size / self.num_attention_heads), - intermediate_size: self.intermediate_size, - vocab_size: self.vocab_size, - num_hidden_layers: self.num_hidden_layers, - num_attention_heads: self.num_attention_heads, - num_key_value_heads: self.num_key_value_heads.unwrap_or(self.num_attention_heads), - rms_norm_eps: self.layer_norm_eps, - rope_theta: self.rope_theta, - rope_local_base_freq: None, - bos_token_id: self.bos_token_id, - eos_token_id: self.eos_token_id, - max_seq_len: self.max_position_embeddings, - sliding_window: self.sliding_window, - sliding_window_pattern: None, - hidden_act: Some(self.hidden_act), - tie_word_embeddings: false, - rope_scaling: None, - use_flash_attn, - original_max_position_embeddings: self.original_max_position_embeddings, - attention_bias: false, - partial_rotary_factor: Some(self.partial_rotary_factor), - qk_layer_rms_norm: Some(self.qk_layernorm), - kv_cache_dtype, - use_qkv_bias: None, - custom_stop_tokens: None, - specific_config: scfg.clone(), - attn_logit_softcapping: None, - final_logit_softcapping: None, - quantization_config: self.quantization_config, - moe_config: None, +impl Phi2 { + pub fn load_config(filename: &PathBuf, isq: Option) -> Result { + let mut config = Config::load_config(filename.clone())?; + config.head_dim = Some( + config + .head_dim + .unwrap_or(config.hidden_size / config.num_attention_heads), + ); + config.num_key_value_heads = Some( + config + .num_key_value_heads + .unwrap_or(config.num_attention_heads), + ); + config.max_seq_len = config.max_position_embeddings.unwrap_or(config.max_seq_len); + if config.quantization_config.is_some() { + config.quant = Some( + config + .quantization_config + .as_ref() + .unwrap() + .quant_method + .clone(), + ); + } else if isq.is_some() { + config.quant = Some(isq.unwrap().to_string()); } + Ok(config) } } @@ -137,7 +101,7 @@ impl Mlp { false, vb.pp("gate_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let fc2 = TensorParallelRowLinear::load_with_hints( @@ -146,7 +110,7 @@ impl Mlp { false, vb.pp("down_proj"), comm, - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; Ok(Self { @@ -182,7 +146,7 @@ struct Attention { impl Attention { fn new(cfg: &Config, vb: VarBuilder, comm: Rc) -> Result { let num_heads = cfg.num_attention_heads; - let num_kv_heads = cfg.num_key_value_heads; + let num_kv_heads = cfg.num_key_value_heads.unwrap(); let head_dim = cfg.hidden_size / cfg.num_attention_heads; let q_proj = TensorParallelColumnLinear::load_with_hints( @@ -191,7 +155,7 @@ impl Attention { false, vb.pp("q_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let k_proj = TensorParallelColumnLinear::load_with_hints( @@ -200,7 +164,7 @@ impl Attention { false, vb.pp("k_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let v_proj = TensorParallelColumnLinear::load_with_hints( @@ -209,7 +173,7 @@ impl Attention { false, vb.pp("v_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; @@ -219,12 +183,12 @@ impl Attention { false, vb.pp("o_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; // Alternative rope scalings are not supported. let rotary_emb = RotaryEmbedding::new(cfg, vb.dtype(), vb.device())?; - let (q_layernorm, k_layernorm) = if cfg.qk_layer_rms_norm.unwrap() { + let (q_layernorm, k_layernorm) = if cfg.qk_layernorm { let q_layernorm = layer_norm(head_dim, cfg.rms_norm_eps, true, vb.pp("q_layernorm"))?; let k_layernorm = layer_norm(head_dim, cfg.rms_norm_eps, true, vb.pp("k_layernorm"))?; (Some(q_layernorm), Some(k_layernorm)) @@ -232,7 +196,7 @@ impl Attention { (None, None) }; let attention_heads = cfg.num_attention_heads / comm.world_size(); - let kv_heads = cfg.num_key_value_heads / comm.world_size(); + let kv_heads = cfg.num_key_value_heads.unwrap() / comm.world_size(); Ok(Self { q_proj, k_proj, diff --git a/src/openai/models/phi3.rs b/src/openai/models/phi3.rs index 1b5b657b..d116f441 100644 --- a/src/openai/models/phi3.rs +++ b/src/openai/models/phi3.rs @@ -1,83 +1,49 @@ // This implementation is based on: // https://huggingface.co/microsoft/Phi-3-mini-4k-instruct/blob/main/modeling_phi3.py -use super::{Config, QuantConfig, RopeScaling}; +use super::{Config, RopeScaling, ScalingValue}; use crate::backend::progress::{ProgressLike, ProgressReporter}; use crate::openai::distributed::{ embedding, rms_norm, Comm, ReplicatedLinear, TensorParallelColumnLinear, TensorParallelRowLinear, VarBuilder, }; -use crate::openai::models::TokenID; use crate::paged_attention::input_metadata::InputMetadata; use crate::paged_attention::PagedAttention; -use crate::SpecificConfig; use candle::{DType, Device, IndexOp, Module, Result, Tensor, D}; use candle_core as candle; use candle_nn::RmsNorm; use either::Either; -use std::collections::HashMap; use std::iter::zip; +use std::path::PathBuf; use std::rc::Rc; use std::sync::{Arc, RwLock}; -#[derive(Debug, Clone, serde::Deserialize)] -pub struct PhiConfig { - pub vocab_size: usize, - pub hidden_act: candle_nn::Activation, - pub hidden_size: usize, - pub intermediate_size: usize, - pub num_hidden_layers: usize, - pub num_attention_heads: usize, - pub num_key_value_heads: usize, - pub rms_norm_eps: f64, - pub rope_theta: f64, - pub bos_token_id: TokenID, - pub eos_token_id: TokenID, - pub rope_scaling: Option>, - pub max_position_embeddings: usize, - pub original_max_position_embeddings: Option, - pub sliding_window: Option, - pub quantization_config: Option, -} -impl PhiConfig { - pub fn into_config( - self, - use_flash_attn: bool, - kv_cache_dtype: DType, - scfg: &SpecificConfig, - ) -> Config { - Config { - hidden_size: self.hidden_size, - head_dim: Some(self.hidden_size / self.num_attention_heads), - intermediate_size: self.intermediate_size, - vocab_size: self.vocab_size, - num_hidden_layers: self.num_hidden_layers, - num_attention_heads: self.num_attention_heads, - num_key_value_heads: self.num_key_value_heads, - rms_norm_eps: self.rms_norm_eps, - rope_theta: self.rope_theta, - rope_local_base_freq: None, - use_flash_attn, - bos_token_id: self.bos_token_id, - eos_token_id: self.eos_token_id, - max_seq_len: self.max_position_embeddings, - sliding_window: self.sliding_window, - sliding_window_pattern: None, - hidden_act: Some(self.hidden_act), - tie_word_embeddings: false, - rope_scaling: self.rope_scaling, - original_max_position_embeddings: self.original_max_position_embeddings, - attention_bias: false, - partial_rotary_factor: None, - qk_layer_rms_norm: None, - kv_cache_dtype, - use_qkv_bias: None, - custom_stop_tokens: None, - specific_config: scfg.clone(), - attn_logit_softcapping: None, - final_logit_softcapping: None, - quantization_config: self.quantization_config, - moe_config: None, +impl Phi { + pub fn load_config(filename: &PathBuf, isq: Option) -> Result { + let mut config = Config::load_config(filename.clone())?; + config.head_dim = Some( + config + .head_dim + .unwrap_or(config.hidden_size / config.num_attention_heads), + ); + config.num_key_value_heads = Some( + config + .num_key_value_heads + .unwrap_or(config.num_attention_heads), + ); + config.max_seq_len = config.max_position_embeddings.unwrap_or(config.max_seq_len); + if config.quantization_config.is_some() { + config.quant = Some( + config + .quantization_config + .as_ref() + .unwrap() + .quant_method + .clone(), + ); + } else if isq.is_some() { + config.quant = Some(isq.unwrap().to_string()); } + Ok(config) } } @@ -94,9 +60,10 @@ impl RotaryEmbedding { fn new(_dtype: DType, cfg: &Config, dev: &Device) -> Result { let dim = cfg.hidden_size / cfg.num_attention_heads; let max_seq_len = cfg.max_seq_len; + let rope_theta = cfg.rope_theta; let inv_freq: Vec<_> = (0..dim) .step_by(2) - .map(|i| 1f32 / cfg.rope_theta.powf(i as f64 / dim as f64) as f32) + .map(|i| 1f32 / rope_theta.powf(i as f64 / dim as f64) as f32) .collect(); let inv_freq_len = inv_freq.len(); let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(DType::F32)?; @@ -112,19 +79,18 @@ impl RotaryEmbedding { &rope_scaling["type"], ) { ( - RopeScaling(Either::Left(short_factor)), - RopeScaling(Either::Left(long_factor)), + RopeScaling(Either::Left(ScalingValue(Either::Right(short_factor)))), + RopeScaling(Either::Left(ScalingValue(Either::Right(long_factor)))), RopeScaling(Either::Right(tp)), ) => { - let scale = cfg.max_seq_len as f64 - / cfg.original_max_position_embeddings.unwrap() as f64; + let scale = + cfg.max_seq_len as f64 / cfg.original_max_position_embeddings as f64; let scaling_factor = if scale <= 1.0 { 1.0 } else { match tp.as_str() { "su" | "longrope" => (1.0 - + scale.ln() - / (cfg.original_max_position_embeddings.unwrap() as f64).ln()) + + scale.ln() / (cfg.original_max_position_embeddings as f64).ln()) .sqrt(), "yarn" => 0.1 * scale.ln() + 1.0, _ => 1.0, @@ -135,7 +101,7 @@ impl RotaryEmbedding { .step_by(2) .enumerate() .map(|(k, i)| { - (1f64 / (long_factor[k] * cfg.rope_theta.powf(i as f64 / dim as f64))) + (1f64 / (long_factor[k] * rope_theta.powf(i as f64 / dim as f64))) as f32 }) .collect::>(); @@ -143,7 +109,7 @@ impl RotaryEmbedding { .step_by(2) .enumerate() .map(|(k, i)| { - (1f64 / (short_factor[k] * cfg.rope_theta.powf(i as f64 / dim as f64))) + (1f64 / (short_factor[k] * rope_theta.powf(i as f64 / dim as f64))) as f32 }) .collect::>(); @@ -172,7 +138,9 @@ impl RotaryEmbedding { cos: short_cos, sin_long: Some(long_sin), cos_long: Some(long_cos), - original_max_position_embeddings: cfg.original_max_position_embeddings, + original_max_position_embeddings: Some( + cfg.original_max_position_embeddings, + ), }); } _ => { @@ -204,7 +172,6 @@ impl RotaryEmbedding { for (b, seqlen_offset) in zip(0..b_size, input_positions) { let (cos, sin) = if self.sin_long.as_ref().is_some() && self.cos_long.as_ref().is_some() - && self.original_max_position_embeddings.is_some() && seqlen_offset[0] > self.original_max_position_embeddings.unwrap() { let cos = self @@ -252,21 +219,21 @@ impl Attention { comm: Rc, ) -> Result { let num_heads = cfg.num_attention_heads; - let num_kv_heads = cfg.num_key_value_heads; + let num_kv_heads = cfg.num_key_value_heads.unwrap(); let head_dim = cfg.hidden_size / cfg.num_attention_heads; let op_size = num_heads * head_dim + 2 * num_kv_heads * head_dim; let qkv_proj = ReplicatedLinear::load_no_bias( cfg.hidden_size, op_size, vb.pp("qkv_proj"), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let o_proj = ReplicatedLinear::load_no_bias( num_heads * head_dim, cfg.hidden_size, vb.pp("o_proj"), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; //Fix the attention parallel @@ -275,7 +242,7 @@ impl Attention { "Packed qkv_proj is not supported under multi-gpu setting!" ); let attention_heads = cfg.num_attention_heads / comm.world_size(); - let kv_heads = cfg.num_key_value_heads / comm.world_size(); + let kv_heads = cfg.num_key_value_heads.unwrap() / comm.world_size(); Ok(Self { qkv_proj, o_proj, @@ -381,7 +348,7 @@ impl Mlp { false, vb.pp("gate_up_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let down_proj = TensorParallelRowLinear::load_with_hints( @@ -390,7 +357,7 @@ impl Mlp { false, vb.pp("down_proj"), comm, - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; Ok(Self { diff --git a/src/openai/models/quantized_glm4.rs b/src/openai/models/quantized_glm4.rs index cea6c143..f65d5423 100644 --- a/src/openai/models/quantized_glm4.rs +++ b/src/openai/models/quantized_glm4.rs @@ -3,7 +3,6 @@ use crate::backend::progress::{ProgressLike, ProgressReporter}; use crate::openai::models::glm4::RotaryEmbedding; use crate::paged_attention::input_metadata::InputMetadata; use crate::paged_attention::PagedAttention; -use crate::SpecificConfig; use candle_core::quantized::{gguf_file, QMatMul}; use candle_core::{DType, Device, IndexOp, Result, Tensor}; use candle_nn::{Embedding, Module}; @@ -152,44 +151,43 @@ impl GGUFGLM4 { block_count: usize, head_count: usize, head_count_kv: usize, - rope_theta: f32, + rope_theta: f64, rms_eps: f64, max_seq_len: usize, - kv_cache_dtype: DType, - s_cfg: SpecificConfig, ) -> Config { Config { + architectures: Some(vec!["glm4".to_string()]), hidden_size: embedding_length, head_dim: Some(head_dim), intermediate_size: i_size, vocab_size: 0, num_hidden_layers: block_count, num_attention_heads: head_count, - num_key_value_heads: head_count_kv, + num_key_value_heads: Some(head_count_kv), rms_norm_eps: rms_eps, - rope_theta: rope_theta as f64, + rope_theta, rope_local_base_freq: None, - use_flash_attn: false, - bos_token_id: super::TokenID(Either::Left(None)), + bos_token_id: None, eos_token_id: super::TokenID(Either::Left(None)), max_seq_len, sliding_window: None, sliding_window_pattern: None, hidden_act: None, + hidden_activation: None, tie_word_embeddings: false, rope_scaling: None, - original_max_position_embeddings: Some(max_seq_len), - attention_bias: false, + max_position_embeddings: Some(max_seq_len), + original_max_position_embeddings: max_seq_len, + attention_bias: Some(false), partial_rotary_factor: Some(0.5), - qk_layer_rms_norm: None, - kv_cache_dtype, + qk_layernorm: false, use_qkv_bias: None, custom_stop_tokens: None, - specific_config: s_cfg, attn_logit_softcapping: None, final_logit_softcapping: None, quantization_config: None, moe_config: None, + quant: Some("gguf".to_string()), } } @@ -206,7 +204,6 @@ impl GGUFGLM4 { reader: &mut R, device: &Device, dtype: DType, - s_cfg: SpecificConfig, progress_reporter: Arc>, ) -> Result { let md_get = |s: &str| match ct.metadata.get(s) { @@ -254,11 +251,9 @@ impl GGUFGLM4 { block_count, head_count, head_count_kv, - rope_freq_base, + rope_freq_base as f64, rms_norm_eps, context_length, - dtype, - s_cfg, ); let rotary_emb = Arc::new(RotaryEmbedding::new(&cfg, dtype, device)?); diff --git a/src/openai/models/quantized_llama.rs b/src/openai/models/quantized_llama.rs index fe969985..3035ade6 100644 --- a/src/openai/models/quantized_llama.rs +++ b/src/openai/models/quantized_llama.rs @@ -2,7 +2,6 @@ use super::Config; use crate::backend::progress::{ProgressLike, ProgressReporter}; use crate::paged_attention::input_metadata::InputMetadata; use crate::paged_attention::PagedAttention; -use crate::SpecificConfig; use candle_core::quantized::{ggml_file, gguf_file, QMatMul}; use candle_core::{DType, Device, IndexOp, Result, Tensor}; use candle_nn::{Embedding, Module}; @@ -10,7 +9,6 @@ use candle_transformers::quantized_nn::RmsNorm; use either::Either; use std::iter::zip; use std::sync::{Arc, RwLock}; -pub const MAX_SEQ_LEN: usize = 4096; #[derive(Debug, Clone)] struct Mlp { @@ -247,53 +245,46 @@ impl GGUFLLaMa { head_count_kv: usize, rms_eps: f64, max_seq_len: usize, - kv_cache_dtype: DType, - s_cfg: SpecificConfig, ) -> Config { Config { + architectures: Some(vec!["llama".to_string()]), hidden_size: embedding_length, head_dim: Some(head_dim), intermediate_size: i_size, vocab_size: 0, num_hidden_layers: block_count, num_attention_heads: head_count, - num_key_value_heads: head_count_kv, + num_key_value_heads: Some(head_count_kv), rms_norm_eps: rms_eps, - rope_theta: 0., + rope_theta: 10_000.0f64, rope_local_base_freq: None, - use_flash_attn: false, - bos_token_id: super::TokenID(Either::Left(Some(128256))), + bos_token_id: Some(super::TokenID(Either::Left(Some(128256)))), eos_token_id: super::TokenID(Either::Left(Some(128257))), max_seq_len, sliding_window: None, sliding_window_pattern: None, hidden_act: None, + hidden_activation: None, tie_word_embeddings: false, rope_scaling: None, - original_max_position_embeddings: Some(max_seq_len), - attention_bias: false, + max_position_embeddings: Some(max_seq_len), + original_max_position_embeddings: max_seq_len, + attention_bias: Some(false), partial_rotary_factor: None, - qk_layer_rms_norm: None, - kv_cache_dtype, + qk_layernorm: false, use_qkv_bias: None, custom_stop_tokens: Some(vec!["<|end_of_text|>".to_string()]), - specific_config: s_cfg, attn_logit_softcapping: None, final_logit_softcapping: None, quantization_config: None, moe_config: None, + quant: Some("gguf".to_string()), } } - pub fn from_ggml( - mut ct: ggml_file::Content, - gqa: usize, - dtype: DType, - s_cfg: SpecificConfig, - ) -> Result { + pub fn from_ggml(mut ct: ggml_file::Content, gqa: usize, dtype: DType) -> Result { let head_dim = (ct.hparams.n_embd / ct.hparams.n_head) as usize; - let (cos, sin) = - precomput_freqs_cis(head_dim, 10000., MAX_SEQ_LEN, &ct.device, DType::F32)?; + let (cos, sin) = precomput_freqs_cis(head_dim, 10000., 8192, &ct.device, DType::F32)?; let tok_embeddings = ct.remove("tok_embeddings.weight")?; let tok_embeddings = tok_embeddings.dequantize(&ct.device)?; let norm = RmsNorm::from_qtensor(ct.remove("norm.weight")?, 1e-5)?; @@ -356,8 +347,6 @@ impl GGUFLLaMa { ct.hparams.n_head as usize / gqa, 1e-5, 0, - dtype, - s_cfg, ), dtype, device: ct.device.clone(), @@ -377,7 +366,6 @@ impl GGUFLLaMa { reader: &mut R, device: &Device, dtype: DType, - s_cfg: SpecificConfig, progress_reporter: Arc>, ) -> Result { let md_get = |s: &str| match ct.metadata.get(s) { @@ -399,7 +387,7 @@ impl GGUFLLaMa { let embedding_length = md_get("llama.embedding_length")?.to_u32()? as usize; // let rope_dim = md_get("llama.rope.dimension_count")?.to_u32()? as usize; let context_length = md_get("llama.context_length")?.to_u32(); - let context_length = context_length.map_or(MAX_SEQ_LEN, |v| v as usize); + let context_length = context_length.map_or(8192, |v| v as usize); let head_dim = md_get("llama.attention.key_length"); let head_dim = if head_dim.is_ok() { @@ -513,8 +501,6 @@ impl GGUFLLaMa { head_count_kv, rms_norm_eps, context_length, - dtype, - s_cfg, ), dtype, device: device.clone(), diff --git a/src/openai/models/quantized_phi3.rs b/src/openai/models/quantized_phi3.rs index c05f9808..0a26ae21 100644 --- a/src/openai/models/quantized_phi3.rs +++ b/src/openai/models/quantized_phi3.rs @@ -2,7 +2,6 @@ use super::Config; use crate::backend::progress::{ProgressLike, ProgressReporter}; use crate::paged_attention::input_metadata::InputMetadata; use crate::paged_attention::PagedAttention; -use crate::SpecificConfig; use candle_core::quantized::gguf_file; use candle_core::quantized::QTensor; use candle_core::{DType, Device, IndexOp, Module, Result, Tensor, D}; @@ -199,41 +198,40 @@ impl GGUFPhi3 { head_count_kv: usize, rms_eps: f64, max_seq_len: usize, - kv_cache_dtype: DType, - s_cfg: SpecificConfig, ) -> Config { Config { + architectures: Some(vec!["phi3".to_string()]), hidden_size: embedding_length, head_dim: Some(embedding_length / head_count), intermediate_size: i_size, vocab_size: 0, num_hidden_layers: block_count, num_attention_heads: head_count, - num_key_value_heads: head_count_kv, + num_key_value_heads: Some(head_count_kv), rms_norm_eps: rms_eps, - rope_theta: 0., + rope_theta: 10_000.0f64, rope_local_base_freq: None, - use_flash_attn: false, - bos_token_id: super::TokenID(Either::Left(Some(1))), + bos_token_id: Some(super::TokenID(Either::Left(Some(1)))), eos_token_id: super::TokenID(Either::Left(Some(2))), max_seq_len, sliding_window: None, sliding_window_pattern: None, hidden_act: None, + hidden_activation: None, tie_word_embeddings: false, rope_scaling: None, - original_max_position_embeddings: Some(max_seq_len), - attention_bias: false, + max_position_embeddings: Some(max_seq_len), + original_max_position_embeddings: max_seq_len, + attention_bias: Some(false), partial_rotary_factor: None, - qk_layer_rms_norm: None, - kv_cache_dtype, + qk_layernorm: false, use_qkv_bias: None, custom_stop_tokens: None, - specific_config: s_cfg, attn_logit_softcapping: None, final_logit_softcapping: None, quantization_config: None, moe_config: None, + quant: Some("gguf".to_string()), } } @@ -250,7 +248,6 @@ impl GGUFPhi3 { reader: &mut R, device: &Device, dtype: DType, - s_cfg: SpecificConfig, progress_reporter: Arc>, ) -> Result { let md_get = |s: &str| match ct.metadata.get(s) { @@ -331,8 +328,6 @@ impl GGUFPhi3 { head_count_kv, rms_eps, max_seq_len, - dtype, - s_cfg, ), dtype, device: device.clone(), diff --git a/src/openai/models/quantized_qwen.rs b/src/openai/models/quantized_qwen.rs index 46eb4ae9..6c29ef91 100644 --- a/src/openai/models/quantized_qwen.rs +++ b/src/openai/models/quantized_qwen.rs @@ -2,7 +2,6 @@ use super::Config; use crate::backend::progress::{ProgressLike, ProgressReporter}; use crate::paged_attention::input_metadata::InputMetadata; use crate::paged_attention::PagedAttention; -use crate::SpecificConfig; use candle_core::quantized::{gguf_file, QMatMul}; use candle_core::{DType, Device, IndexOp, Result, Tensor}; use candle_nn::{Embedding, Module}; @@ -206,41 +205,40 @@ impl GGUFQWen { head_count_kv: usize, rms_eps: f64, max_seq_len: usize, - kv_cache_dtype: DType, - s_cfg: SpecificConfig, ) -> Config { Config { + architectures: Some(vec!["qwen".to_string()]), hidden_size: embedding_length, head_dim: Some(head_dim), intermediate_size: i_size, vocab_size: 0, num_hidden_layers: block_count, num_attention_heads: head_count, - num_key_value_heads: head_count_kv, + num_key_value_heads: Some(head_count_kv), rms_norm_eps: rms_eps, - rope_theta: 0., + rope_theta: 10_000.0f64, rope_local_base_freq: None, - use_flash_attn: false, - bos_token_id: super::TokenID(Either::Left(Some(151644))), + bos_token_id: Some(super::TokenID(Either::Left(Some(151644)))), eos_token_id: super::TokenID(Either::Left(Some(151645))), max_seq_len, sliding_window: None, sliding_window_pattern: None, hidden_act: None, + hidden_activation: None, tie_word_embeddings: false, rope_scaling: None, - original_max_position_embeddings: Some(max_seq_len), - attention_bias: false, + max_position_embeddings: Some(max_seq_len), + original_max_position_embeddings: max_seq_len, + attention_bias: Some(false), partial_rotary_factor: None, - qk_layer_rms_norm: None, - kv_cache_dtype, + qk_layernorm: false, use_qkv_bias: None, custom_stop_tokens: None, - specific_config: s_cfg, attn_logit_softcapping: None, final_logit_softcapping: None, quantization_config: None, moe_config: None, + quant: Some("gguf".to_string()), } } @@ -261,7 +259,6 @@ impl GGUFQWen { reader: &mut R, device: &Device, dtype: DType, - s_cfg: SpecificConfig, progress_reporter: Arc>, ) -> Result { let md_get = |s: &str| match ct.metadata.get(s) { @@ -433,8 +430,6 @@ impl GGUFQWen { head_count_kv, rms_norm_eps, context_length, - dtype, - s_cfg, ), dtype, device: device.clone(), diff --git a/src/openai/models/qwen.rs b/src/openai/models/qwen.rs index 140844d4..2537e2a9 100644 --- a/src/openai/models/qwen.rs +++ b/src/openai/models/qwen.rs @@ -1,90 +1,47 @@ -use super::{Config, QuantConfig}; +use super::Config; use crate::backend::progress::{ProgressLike, ProgressReporter}; use crate::openai::distributed::{ embedding, rms_norm, Comm, ReplicatedLinear, TensorParallelColumnLinear, TensorParallelRowLinear, VarBuilder, }; -use crate::openai::models::TokenID; use crate::paged_attention::input_metadata::InputMetadata; use crate::paged_attention::PagedAttention; -use crate::SpecificConfig; use candle::{DType, Device, IndexOp, Module, Result, Tensor}; use candle_core as candle; use candle_nn::RmsNorm; use std::iter::zip; +use std::path::PathBuf; use std::rc::Rc; use std::sync::{Arc, RwLock}; -#[derive(Debug, Clone, serde::Deserialize)] -pub struct QwenConfig { - pub vocab_size: usize, - pub hidden_size: usize, - pub head_dim: Option, - pub intermediate_size: usize, - pub num_hidden_layers: usize, - pub num_attention_heads: usize, - pub num_key_value_heads: usize, - pub max_position_embeddings: usize, - pub sliding_window: Option, - pub max_window_layers: usize, - pub tie_word_embeddings: bool, //shared weights between input/output embeddings - pub attention_bias: Option, - pub rope_theta: f64, - pub rms_norm_eps: f64, - pub use_sliding_window: Option, - pub hidden_act: candle_nn::Activation, - pub bos_token_id: TokenID, - pub eos_token_id: TokenID, - pub quantization_config: Option, -} -impl QwenConfig { - pub fn into_config( - self, - use_flash_attn: bool, - kv_cache_dtype: DType, - scfg: &SpecificConfig, - ) -> Config { - let sliding_window = if self.use_sliding_window.unwrap_or(false) { - self.sliding_window - } else { - None - }; - Config { - hidden_size: self.hidden_size, - head_dim: Some( - self.head_dim - .unwrap_or(self.hidden_size / self.num_attention_heads), - ), - intermediate_size: self.intermediate_size, - vocab_size: self.vocab_size, - num_hidden_layers: self.num_hidden_layers, - num_attention_heads: self.num_attention_heads, - num_key_value_heads: self.num_key_value_heads, - rms_norm_eps: self.rms_norm_eps, - rope_theta: self.rope_theta, - rope_local_base_freq: None, - use_flash_attn, - bos_token_id: self.bos_token_id, - eos_token_id: self.eos_token_id, - max_seq_len: self.max_position_embeddings, - sliding_window, - sliding_window_pattern: None, - hidden_act: Some(self.hidden_act), - tie_word_embeddings: self.tie_word_embeddings, - rope_scaling: None, - original_max_position_embeddings: None, - attention_bias: self.attention_bias.unwrap_or(true), - partial_rotary_factor: None, - qk_layer_rms_norm: None, - kv_cache_dtype, - use_qkv_bias: None, - custom_stop_tokens: None, - specific_config: scfg.clone(), - attn_logit_softcapping: None, - final_logit_softcapping: None, - quantization_config: self.quantization_config, - moe_config: None, +impl Qwen { + pub fn load_config(filename: &PathBuf, isq: Option) -> Result { + let mut config = Config::load_config(filename.clone())?; + config.head_dim = Some( + config + .head_dim + .unwrap_or(config.hidden_size / config.num_attention_heads), + ); + config.num_key_value_heads = Some( + config + .num_key_value_heads + .unwrap_or(config.num_attention_heads), + ); + config.max_seq_len = config.max_position_embeddings.unwrap_or(config.max_seq_len); + config.attention_bias = Some(config.attention_bias.unwrap_or(true)); + if config.quantization_config.is_some() { + config.quant = Some( + config + .quantization_config + .as_ref() + .unwrap() + .quant_method + .clone(), + ); + } else if isq.is_some() { + config.quant = Some(isq.unwrap().to_string()); } + Ok(config) } } @@ -157,7 +114,7 @@ impl Mlp { false, vb.pp("gate_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let up_proj = TensorParallelColumnLinear::load_with_hints( @@ -166,7 +123,7 @@ impl Mlp { false, vb.pp("up_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let down_proj = TensorParallelRowLinear::load_with_hints( @@ -175,7 +132,7 @@ impl Mlp { false, vb.pp("down_proj"), comm, - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; Ok(Self { @@ -219,34 +176,35 @@ impl Attention { ) -> Result { let hidden_sz = cfg.hidden_size; let num_heads = cfg.num_attention_heads; - let num_kv_heads = cfg.num_key_value_heads; + let num_kv_heads = cfg.num_key_value_heads.unwrap(); let head_dim = cfg.head_dim.unwrap_or(hidden_sz / num_heads); + let attention_bias = cfg.attention_bias.unwrap_or(false); let q_proj = TensorParallelColumnLinear::load_with_hints( hidden_sz, num_heads * head_dim, - cfg.attention_bias, + attention_bias, vb.pp("q_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let k_proj = TensorParallelColumnLinear::load_with_hints( hidden_sz, num_kv_heads * head_dim, - cfg.attention_bias, + attention_bias, vb.pp("k_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let v_proj = TensorParallelColumnLinear::load_with_hints( hidden_sz, num_kv_heads * head_dim, - cfg.attention_bias, + attention_bias, vb.pp("v_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; @@ -256,7 +214,7 @@ impl Attention { false, vb.pp("o_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; @@ -272,7 +230,7 @@ impl Attention { }; let attention_heads = cfg.num_attention_heads / comm.world_size(); - let kv_heads = cfg.num_key_value_heads / comm.world_size(); + let kv_heads = cfg.num_key_value_heads.unwrap() / comm.world_size(); Ok(Self { q_proj, k_proj, diff --git a/src/openai/models/stable_lm.rs b/src/openai/models/stable_lm.rs index 996b505b..1d31aa15 100644 --- a/src/openai/models/stable_lm.rs +++ b/src/openai/models/stable_lm.rs @@ -1,85 +1,47 @@ -use super::{Config, QuantConfig}; +use super::Config; use crate::backend::progress::{ProgressLike, ProgressReporter}; use crate::openai::distributed::{ embedding, layer_norm, Comm, ReplicatedLinear, TensorParallelColumnLinear, TensorParallelRowLinear, VarBuilder, }; -use crate::openai::models::TokenID; use crate::paged_attention::input_metadata::InputMetadata; use crate::paged_attention::PagedAttention; -use crate::SpecificConfig; use candle_core::{DType, Device, IndexOp, Module, Result, Tensor, D}; use candle_nn::{Activation, LayerNorm}; use std::iter::zip; +use std::path::PathBuf; use std::rc::Rc; use std::sync::{Arc, RwLock}; -#[derive(Debug, Clone, serde::Deserialize)] -pub struct StableLMConfig { - pub vocab_size: usize, - pub intermediate_size: usize, - pub hidden_size: usize, - pub num_hidden_layers: usize, - pub num_attention_heads: usize, - pub num_key_value_heads: usize, - pub hidden_act: Activation, - pub rope_theta: f64, - pub max_position_embeddings: usize, - pub norm_eps: f64, - pub use_cache: bool, - pub use_qkv_bias: Option, // Used in StableLM-2 - pub partial_rotary_factor: Option, - pub rope_pct: Option, - pub tie_word_embeddings: Option, - pub bos_token_id: TokenID, - pub eos_token_id: TokenID, - pub sliding_window: Option, - pub quantization_config: Option, -} - -impl StableLMConfig { - pub fn into_config( - self, - use_flash_attn: bool, - kv_cache_dtype: DType, - scfg: &SpecificConfig, - ) -> Config { - Config { - hidden_size: self.hidden_size, - head_dim: Some(self.hidden_size / self.num_attention_heads), - intermediate_size: self.intermediate_size, - vocab_size: self.vocab_size, - num_hidden_layers: self.num_hidden_layers, - num_attention_heads: self.num_attention_heads, - num_key_value_heads: self.num_key_value_heads, - rms_norm_eps: self.norm_eps, - rope_theta: self.rope_theta, - rope_local_base_freq: None, - use_flash_attn, - bos_token_id: self.bos_token_id, - eos_token_id: self.eos_token_id, - max_seq_len: self.max_position_embeddings, - sliding_window: self.sliding_window, - sliding_window_pattern: None, - hidden_act: Some(self.hidden_act), - tie_word_embeddings: self.tie_word_embeddings.unwrap_or(false), - rope_scaling: None, - original_max_position_embeddings: None, - attention_bias: false, - partial_rotary_factor: Some( - self.partial_rotary_factor - .unwrap_or(self.rope_pct.unwrap_or(0.25)), - ), - qk_layer_rms_norm: None, - kv_cache_dtype, - use_qkv_bias: Some(self.use_qkv_bias.unwrap_or(false)), - custom_stop_tokens: None, - specific_config: scfg.clone(), - attn_logit_softcapping: None, - final_logit_softcapping: None, - quantization_config: self.quantization_config, - moe_config: None, +impl StableLM { + pub fn load_config(filename: &PathBuf, isq: Option) -> Result { + let mut config = Config::load_config(filename.clone())?; + config.head_dim = Some( + config + .head_dim + .unwrap_or(config.hidden_size / config.num_attention_heads), + ); + config.num_key_value_heads = Some( + config + .num_key_value_heads + .unwrap_or(config.num_attention_heads), + ); + config.max_seq_len = config.max_position_embeddings.unwrap_or(config.max_seq_len); + config.partial_rotary_factor = Some(config.partial_rotary_factor.unwrap_or(0.25)); + config.use_qkv_bias = Some(config.use_qkv_bias.unwrap_or(false)); + if config.quantization_config.is_some() { + config.quant = Some( + config + .quantization_config + .as_ref() + .unwrap() + .quant_method + .clone(), + ); + } else if isq.is_some() { + config.quant = Some(isq.unwrap().to_string()); } + Ok(config) } } @@ -149,7 +111,7 @@ impl MLP { false, vb.pp("gate_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let up_proj = TensorParallelColumnLinear::load_with_hints( @@ -158,7 +120,7 @@ impl MLP { false, vb.pp("up_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let down_proj = TensorParallelRowLinear::load_with_hints( @@ -167,7 +129,7 @@ impl MLP { false, vb.pp("down_proj"), comm, - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; Ok(Self { @@ -208,7 +170,7 @@ impl Attention { ) -> Result { let hidden_sz = cfg.hidden_size; let num_heads = cfg.num_attention_heads; - let num_kv_heads = cfg.num_key_value_heads; + let num_kv_heads = cfg.num_key_value_heads.unwrap(); let head_dim = hidden_sz / num_heads; let q_proj = TensorParallelColumnLinear::load_with_hints( @@ -217,7 +179,7 @@ impl Attention { cfg.use_qkv_bias.unwrap_or(false), vb.pp("q_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let k_proj = TensorParallelColumnLinear::load_with_hints( @@ -226,7 +188,7 @@ impl Attention { cfg.use_qkv_bias.unwrap_or(false), vb.pp("k_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let v_proj = TensorParallelColumnLinear::load_with_hints( @@ -235,7 +197,7 @@ impl Attention { cfg.use_qkv_bias.unwrap_or(false), vb.pp("v_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; @@ -245,11 +207,11 @@ impl Attention { false, vb.pp("o_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let attention_heads = cfg.num_attention_heads / comm.world_size(); - let kv_heads = cfg.num_key_value_heads / comm.world_size(); + let kv_heads = cfg.num_key_value_heads.unwrap() / comm.world_size(); Ok(Self { q_proj, k_proj, diff --git a/src/openai/models/yi.rs b/src/openai/models/yi.rs index a87975e1..cfae7b0f 100644 --- a/src/openai/models/yi.rs +++ b/src/openai/models/yi.rs @@ -1,78 +1,46 @@ -use super::{Config, QuantConfig}; +use super::Config; use crate::backend::progress::{ProgressLike, ProgressReporter}; use crate::openai::distributed::{ embedding, rms_norm, Comm, ReplicatedLinear, TensorParallelColumnLinear, TensorParallelRowLinear, VarBuilder, }; -use crate::openai::models::TokenID; use crate::paged_attention::input_metadata::InputMetadata; use crate::paged_attention::PagedAttention; -use crate::SpecificConfig; use candle_core::{DType, Device, IndexOp, Module, Result, Tensor}; use candle_nn::{Activation, RmsNorm}; use std::iter::zip; +use std::path::PathBuf; use std::rc::Rc; use std::sync::{Arc, RwLock}; -#[derive(Debug, Clone, serde::Deserialize)] -pub struct YiConfig { - pub vocab_size: usize, - pub hidden_size: usize, - pub intermediate_size: usize, - pub num_hidden_layers: usize, - pub num_attention_heads: usize, - pub num_key_value_heads: usize, - pub hidden_act: Activation, - pub max_position_embeddings: usize, - pub rms_norm_eps: f64, - pub rope_theta: f64, - pub sliding_window: Option, - pub tie_word_embeddings: Option, - pub bos_token_id: TokenID, - pub eos_token_id: TokenID, - pub quantization_config: Option, -} - -impl YiConfig { - pub fn into_config( - self, - use_flash_attn: bool, - kv_cache_dtype: DType, - scfg: &SpecificConfig, - ) -> Config { - Config { - hidden_size: self.hidden_size, - head_dim: Some(self.hidden_size / self.num_attention_heads), - intermediate_size: self.intermediate_size, - vocab_size: self.vocab_size, - num_hidden_layers: self.num_hidden_layers, - num_attention_heads: self.num_attention_heads, - num_key_value_heads: self.num_key_value_heads, - rms_norm_eps: self.rms_norm_eps, - rope_theta: self.rope_theta, - rope_local_base_freq: None, - use_flash_attn, - bos_token_id: self.bos_token_id, - eos_token_id: self.eos_token_id, - max_seq_len: self.max_position_embeddings, - sliding_window: self.sliding_window, - sliding_window_pattern: None, - hidden_act: Some(self.hidden_act), - tie_word_embeddings: self.tie_word_embeddings.unwrap_or(false), - rope_scaling: None, - original_max_position_embeddings: None, - attention_bias: false, - partial_rotary_factor: None, - qk_layer_rms_norm: None, - kv_cache_dtype, - use_qkv_bias: None, - custom_stop_tokens: Some(vec!["<|im_end|>".to_string()]), - specific_config: scfg.clone(), - attn_logit_softcapping: None, - final_logit_softcapping: None, - quantization_config: self.quantization_config, - moe_config: None, +impl Yi { + pub fn load_config(filename: &PathBuf, isq: Option) -> Result { + let mut config = Config::load_config(filename.clone())?; + config.head_dim = Some( + config + .head_dim + .unwrap_or(config.hidden_size / config.num_attention_heads), + ); + config.num_key_value_heads = Some( + config + .num_key_value_heads + .unwrap_or(config.num_attention_heads), + ); + config.max_seq_len = config.max_position_embeddings.unwrap_or(config.max_seq_len); + config.custom_stop_tokens = Some(vec!["<|im_end|>".to_string()]); + if config.quantization_config.is_some() { + config.quant = Some( + config + .quantization_config + .as_ref() + .unwrap() + .quant_method + .clone(), + ); + } else if isq.is_some() { + config.quant = Some(isq.unwrap().to_string()); } + Ok(config) } } @@ -143,7 +111,7 @@ impl Mlp { false, vb.pp("gate_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let up_proj = TensorParallelColumnLinear::load_with_hints( @@ -152,7 +120,7 @@ impl Mlp { false, vb.pp("up_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let down_proj = TensorParallelRowLinear::load_with_hints( @@ -161,7 +129,7 @@ impl Mlp { false, vb.pp("down_proj"), comm, - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; Ok(Self { @@ -202,7 +170,7 @@ impl Attention { ) -> Result { let hidden_sz = cfg.hidden_size; let num_heads = cfg.num_attention_heads; - let num_kv_heads = cfg.num_key_value_heads; + let num_kv_heads = cfg.num_key_value_heads.unwrap(); let head_dim = hidden_sz / num_heads; let q_proj = TensorParallelColumnLinear::load_with_hints( @@ -211,7 +179,7 @@ impl Attention { false, vb.pp("q_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let k_proj = TensorParallelColumnLinear::load_with_hints( @@ -220,7 +188,7 @@ impl Attention { false, vb.pp("k_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let v_proj = TensorParallelColumnLinear::load_with_hints( @@ -229,7 +197,7 @@ impl Attention { false, vb.pp("v_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; @@ -239,11 +207,11 @@ impl Attention { false, vb.pp("o_proj"), comm.clone(), - &cfg.specific_config.quant, + &cfg.quant, &cfg.quantization_config, )?; let attention_heads = cfg.num_attention_heads / comm.world_size(); - let kv_heads = cfg.num_key_value_heads / comm.world_size(); + let kv_heads = cfg.num_key_value_heads.unwrap() / comm.world_size(); Ok(Self { q_proj, k_proj, diff --git a/src/openai/openai_server.rs b/src/openai/openai_server.rs index 81539ac4..41f5de1e 100644 --- a/src/openai/openai_server.rs +++ b/src/openai/openai_server.rs @@ -55,35 +55,35 @@ async fn get_gen_prompt( } } - Ok(conversation.get_prompt( - request - .thinking - .unwrap_or(data.pipeline_config.thinking.unwrap_or(false)), - )) + Ok(conversation.get_prompt(request.thinking.unwrap_or(false))) } async fn check_length( request: &ChatCompletionRequest, prompt: String, data: &OpenAIServerData, -) -> Result { - let token_ids = { +) -> Result<(Encoding, usize), APIError> { + let (token_ids, available_kv_tokens) = { let model = data.model.read(); + let available_kv_tokens = model.get_available_kv_tokens(); let pipeline = model .get_pipeline(0) .ok_or(APIError::new("Missing pipeline".to_string()))?; - pipeline - .0 - .tokenizer() - .encode_fast(prompt, false) - .map_err(APIError::from)? + ( + pipeline + .0 + .tokenizer() + .encode_fast(prompt, false) + .map_err(APIError::from)?, + available_kv_tokens, + ) }; let max_gen_tokens = request .max_tokens .unwrap_or(data.pipeline_config.default_max_tokens); - if token_ids.len() > data.pipeline_config.max_model_len { + if token_ids.len() >= data.pipeline_config.max_model_len { Err(APIError::new(format!( "This model's maximum context length is {} tokens. \ However, you requested {} tokens ({} in the messages, \ @@ -94,8 +94,23 @@ async fn check_length( token_ids.len(), max_gen_tokens ))) + } else if token_ids.len() >= available_kv_tokens { + Err(APIError::new(format!( + "Requested prompt({} tokens) is \ + larger than available kvcache (maximum {} tokens).\n \ + You can increase kvcache by setting `--mem` to a larger value!", + token_ids.len(), + available_kv_tokens + ))) } else { - Ok(token_ids) + let max_valid_request_tokens = std::cmp::min( + available_kv_tokens, + data.pipeline_config.max_model_len - token_ids.len(), + ) - 10; + Ok(( + token_ids, + std::cmp::min(max_gen_tokens, max_valid_request_tokens), + )) } } @@ -132,39 +147,56 @@ pub async fn chat_completions( Err(e) => return ChatResponder::ValidationError(e), }; - let token_ids: Encoding = match check_length(&request, prompt.clone(), &data).await { - Ok(ids) => ids, - Err(e) => return ChatResponder::ValidationError(e), - }; + let (token_ids, available_tokens): (Encoding, usize) = + match check_length(&request, prompt.clone(), &data).await { + Ok(ids) => ids, + Err(e) => return ChatResponder::ValidationError(e), + }; debug!("\n\n\nPrompt {:?}", prompt); let request_id = format!("cmpl-{}", Uuid::new_v4()); + let mut max_request_tokens = request + .max_tokens + .unwrap_or(data.pipeline_config.default_max_tokens); + + if max_request_tokens + token_ids.len() > available_tokens { + tracing::warn!("Requested max tokens {} larger than available tokens {}, max_tokens changed to {} ({} tokens reserved for prompt)!", max_request_tokens, available_tokens, available_tokens - token_ids.len(), token_ids.len()); + max_request_tokens = if available_tokens > token_ids.len() { + available_tokens - token_ids.len() + } else { + return ChatResponder::ValidationError(APIError::new(format!( + "Requested prompt({} tokens) is \ + larger than available kvcache (maximum {} tokens).\n \ + You can increase kvcache by setting `--mem` to a larger value!", + token_ids.len(), + available_tokens + ))); + } + } + let sampling_params = match SamplingParams::new( request.n.unwrap_or(1), request.best_of, request.presence_penalty.unwrap_or(0.0), request.frequency_penalty.unwrap_or(0.0), - request - .repetition_penalty - .unwrap_or(data.pipeline_config.penalty), - request.temperature.or(data.pipeline_config.temperature), - request.top_p.or(data.pipeline_config.top_p), - request.top_k.or(data.pipeline_config.top_k), + request.repetition_penalty, + request.repeat_last_n, + request.temperature, + request.top_p, + request.top_k, request.use_beam_search.unwrap_or(false), 1.0, EarlyStoppingCondition::UnlikelyBetterCandidates, request.stop.clone(), request.stop_token_ids.clone().unwrap_or_default(), request.ignore_eos.unwrap_or(false), - request - .max_tokens - .unwrap_or(data.pipeline_config.default_max_tokens), + max_request_tokens, None, None, request.skip_special_tokens.unwrap_or(true), - request.thinking.or(data.pipeline_config.thinking), + request.thinking, ) { Ok(params) => params, Err(e) => return ChatResponder::ValidationError(e), diff --git a/src/openai/pipelines/llm_engine.rs b/src/openai/pipelines/llm_engine.rs index 05eb4077..dac6ed8c 100644 --- a/src/openai/pipelines/llm_engine.rs +++ b/src/openai/pipelines/llm_engine.rs @@ -10,8 +10,8 @@ use crate::{ openai::{ models::Config, responses::{ - APIError, ChatChoice, ChatChoiceData, ChatCompletionChunk, ChatCompletionUsageResponse, - Choice, ChoiceData, WrapperLogprobs, + ChatChoice, ChatChoiceData, ChatCompletionChunk, ChatCompletionUsageResponse, Choice, + ChoiceData, WrapperLogprobs, }, sampling_params::SamplingParams, utils::get_created_time_secs, @@ -22,9 +22,8 @@ use crate::{ sequence::{Sequence, SequenceGroup, _Sequence}, SchedulerConfig, SchedulerOutput, }, - try_api, }; -use candle_core::{Device, Tensor}; +use candle_core::{Device, Result, Tensor}; use either::Either; use flume::Sender; use parking_lot::RwLock; @@ -40,6 +39,7 @@ use std::{ }; use tokenizers::Encoding; use tokio::sync::Notify; +#[allow(unused_imports)] use tracing::{debug, info, warn}; #[allow(dead_code)] struct PreparedInputs { @@ -50,6 +50,7 @@ struct PreparedInputs { const _PAD_SLOT_ID: i64 = -1; +#[allow(unused)] pub struct LLMEngine { pipelines: HashMap, CacheEngine)>, pub scheduler: Scheduler, @@ -99,7 +100,7 @@ impl LLMEngine { num_shards: usize, multi_process: bool, #[cfg(feature = "nccl")] daemon_manager: Option, - ) -> Result>, APIError> { + ) -> Result>> { let num_threads: usize = pipelines.len(); let engine = Arc::new(RwLock::new(Self { pipelines, @@ -209,6 +210,7 @@ impl LLMEngine { Ok(engine_clone) } + #[allow(unused_mut, unused_variables)] pub fn sync_waiting_task_to_group(&mut self) -> bool { let mut continue_loop = false; #[cfg(feature = "nccl")] @@ -426,7 +428,7 @@ impl LLMEngine { engine: Arc>, rank: usize, multi_process: bool, - ) -> Result, ChatCompletionUsageResponse)>, APIError> { + ) -> Result, ChatCompletionUsageResponse)>> { let mut responses = HashMap::, ChatCompletionUsageResponse)>::new(); let mut prompt_finish_times = HashMap::::new(); @@ -808,16 +810,16 @@ impl LLMEngine { &mut self, scheduler_output: &SchedulerOutput, rank: usize, - ) -> Result<(), APIError> { + ) -> Result<()> { let cache_engine = Box::new(&mut self.get_mut_pipeline(rank).unwrap().1); if !scheduler_output.blocks_to_swap_in.is_empty() { - try_api!(cache_engine.swap_in(scheduler_output.blocks_to_swap_in.clone())); + cache_engine.swap_in(scheduler_output.blocks_to_swap_in.clone())?; } if !scheduler_output.blocks_to_swap_out.is_empty() { - try_api!(cache_engine.swap_out(scheduler_output.blocks_to_swap_out.clone())); + cache_engine.swap_out(scheduler_output.blocks_to_swap_out.clone())?; } if !scheduler_output.blocks_to_copy.is_empty() { - try_api!(cache_engine.copy(scheduler_output.blocks_to_copy.clone())); + cache_engine.copy(scheduler_output.blocks_to_copy.clone())?; } Ok(()) } @@ -826,7 +828,7 @@ impl LLMEngine { &self, groups: &VecDeque>, device: &Device, - ) -> Result { + ) -> Result { let mut prompt_lens = Vec::new(); let mut input_tokens = Vec::new(); let mut input_positions = Vec::new(); @@ -914,7 +916,6 @@ impl LLMEngine { context_lens: None, block_tables: None, is_prompt: true, - kv_cache_dtype: "auto".to_string(), // TODO(EricLBuehler): specialize for models }, }) } @@ -923,7 +924,7 @@ impl LLMEngine { &self, groups: &VecDeque>, device: &Device, - ) -> Result { + ) -> Result { let mut input_tokens = Vec::new(); let mut input_positions = Vec::new(); let mut context_lens = Vec::new(); @@ -991,11 +992,11 @@ impl LLMEngine { let slot_mapping = _make_tensor_with_pad(slot_mappings, 1, _PAD_SLOT_ID, device)?; let max_context_len = context_lens.iter().max().unwrap(); - let context_lens = try_api!(Tensor::from_vec( + let context_lens = Tensor::from_vec( context_lens.iter().map(|x| *x as u32).collect::>(), (context_lens.len(),), device, - )); + )?; let max_block_table_len = block_tables.iter().map(|x| x.len()).max().unwrap(); let block_tables = _make_tensor_with_pad( @@ -1007,7 +1008,7 @@ impl LLMEngine { 0, device, )?; - let block_tables = try_api!(block_tables.reshape(((), max_block_table_len))); + let block_tables = block_tables.reshape(((), max_block_table_len))?; Ok(PreparedInputs { tokens: input_tokens, positions: input_positions, @@ -1018,7 +1019,6 @@ impl LLMEngine { context_lens: Some(context_lens), block_tables: Some(block_tables), is_prompt: false, - kv_cache_dtype: "auto".to_string(), // TODO(EricLBuehler): specialize for models }, }) } @@ -1103,4 +1103,8 @@ impl LLMEngine { self.seq_id += 1; self.group_id += 1; } + + pub fn get_available_kv_tokens(&self) -> usize { + self.scheduler.get_available_kv_tokens() + } } diff --git a/src/openai/pipelines/mod.rs b/src/openai/pipelines/mod.rs index 7967cb4b..53ebb719 100644 --- a/src/openai/pipelines/mod.rs +++ b/src/openai/pipelines/mod.rs @@ -1,7 +1,5 @@ -use super::responses::APIError; use crate::openai::sampling_params::Logprobs; -use crate::try_api; -use candle_core::{Device, Tensor, WithDType}; +use candle_core::{Device, Result, Tensor, WithDType}; use dirs; use either::Either; use std::collections::HashMap; @@ -18,7 +16,7 @@ fn _make_tensor_with_pad( max_len: usize, pad: D, device: &Device, -) -> Result { +) -> Result { let mut padded_x = Vec::new(); for mut x_i in x { if x_i.len() < max_len { @@ -31,29 +29,30 @@ fn _make_tensor_with_pad( .flat_map(|slice| slice.iter()) .map(|&xx| xx) .collect(); - Tensor::from_vec(flattened, (padded_x.len(), max_len), device).map_err(APIError::from) + Tensor::from_vec(flattened, (padded_x.len(), max_len), device) } -pub(crate) fn get_token( - hf_token: Option, - hf_token_path: Option, -) -> Result { +pub(crate) fn get_token(hf_token: Option, hf_token_path: Option) -> Result { Ok(match (hf_token, hf_token_path) { - (Some(envvar), None) => try_api!(env::var(envvar)).trim().to_string(), - (None, Some(path)) => try_api!(fs::read_to_string(path)).trim().to_string(), - (None, None) => try_api!(fs::read_to_string(format!( + (Some(envvar), None) => env::var(envvar) + .map_err(candle_core::Error::wrap)? + .trim() + .to_string(), + (None, Some(path)) => fs::read_to_string(path) + .map_err(candle_core::Error::wrap)? + .trim() + .to_string(), + (None, None) => fs::read_to_string(format!( "{}/.cache/huggingface/token", - dirs::home_dir() - .ok_or(APIError::new_str("No home directory"))? - .display() - ))) + dirs::home_dir().unwrap().display() + )) + .map_err(candle_core::Error::wrap)? .trim() .to_string(), - _ => { - return Err(APIError::new_str( - "Do not specify `hf_token` and `hf_token_path` at the same time.", - )) - } + (Some(_), Some(path)) => fs::read_to_string(path) + .map_err(candle_core::Error::wrap)? + .trim() + .to_string(), }) } diff --git a/src/openai/pipelines/pipeline.rs b/src/openai/pipelines/pipeline.rs index bdf64664..e49cf804 100644 --- a/src/openai/pipelines/pipeline.rs +++ b/src/openai/pipelines/pipeline.rs @@ -1,4 +1,5 @@ use super::{get_token, TokenOrFinishReason}; +use crate::backend::gguf; use crate::backend::progress::{progress_worker, ProgressReporter}; use crate::openai::logits_processor::LogitsProcessor; use crate::openai::models::TokenID; @@ -15,31 +16,17 @@ use crate::{ Conversation, }, models::{ - deepseek::{DeepSeek, DeepSeekConfig}, - gemma::{Gemma, GemmaConfig}, - gemma3::{Gemma3, Gemma3Config}, - glm4::{GLMConfig, GLM4}, - llama::{Llama, LlamaConfig}, - mistral::{Mistral, MistralConfig}, - phi2::{Phi2, Phi2Config}, - phi3::{Phi, PhiConfig}, - quantized_glm4::GGUFGLM4, - quantized_llama::GGUFLLaMa, - quantized_phi3::GGUFPhi3, - quantized_qwen::GGUFQWen, - qwen::{Qwen, QwenConfig}, - stable_lm::{StableLM, StableLMConfig}, - yi::{Yi, YiConfig}, - Config, + deepseek::DeepSeek, gemma::Gemma, gemma3::Gemma3, glm4::GLM4, llama::Llama, + mistral::Mistral, phi2::Phi2, phi3::Phi, quantized_glm4::GGUFGLM4, + quantized_llama::GGUFLLaMa, quantized_phi3::GGUFPhi3, quantized_qwen::GGUFQWen, + qwen::Qwen, stable_lm::StableLM, yi::Yi, Config, }, - responses::APIError, PipelineConfig, }, paged_attention::input_metadata::InputMetadata, - try_api, SpecificConfig, }; use candle_core::quantized::gguf_file; -use candle_core::{DType, Device, Tensor}; +use candle_core::{DType, Device, Result, Tensor}; use either::Either; use either::Either::{Left, Right}; use hf_hub::{api::sync::ApiBuilder, Repo, RepoType}; @@ -76,7 +63,6 @@ enum LLMModel { /// top-p, multinomial, and argmax sampling are implemented. Beam search is not implemented. pub struct DefaultPipeline { model: LLMModel, - args: SpecificConfig, tokenizer: Tokenizer, logits_processor: LogitsProcessor, conversation: DefaultConversation, @@ -89,8 +75,9 @@ pub struct DefaultPipeline { } pub struct DefaultLoader { - config: SpecificConfig, - name: String, + model_id: Option, + weight_path: Option, + weight_file: Option, } #[derive(Debug, Clone)] @@ -117,43 +104,132 @@ impl DefaultModelPaths { } impl DefaultLoader { - pub fn new(config: SpecificConfig, name: String) -> Self { - Self { config, name } + pub fn new( + model_id: Option, + weight_path: Option, + weight_file: Option, + ) -> Self { + Self { + model_id, + weight_path, + weight_file, + } } } impl DefaultLoader { + pub fn prepare_model_weights( + &self, + hf_token: Option, + hf_token_path: Option, + ) -> Result<(DefaultModelPaths, bool)> { + let (paths, gguf): (DefaultModelPaths, bool) = match (&self.weight_path, &self.weight_file) + { + //model in a folder (safetensor format, huggingface folder structure) + (Some(path), None) => ( + DefaultModelPaths { + tokenizer_filename: Path::new(path).join("tokenizer.json"), + tokenizer_config_filename: Path::new(path).join("tokenizer_config.json"), + config_filename: Path::new(path).join("config.json"), + filenames: if Path::new(path) + .join("model.safetensors.index.json") + .exists() + { + crate::hub_load_local_safetensors(path, "model.safetensors.index.json")? + } else { + //a single weight file case + let mut safetensors_files = Vec::::new(); + safetensors_files.insert(0, Path::new(path).join("model.safetensors")); + safetensors_files + }, + }, + false, + ), + //model in a quantized file (gguf/ggml format) + (path, Some(file)) => ( + DefaultModelPaths { + tokenizer_filename: PathBuf::new(), + tokenizer_config_filename: PathBuf::new(), + config_filename: PathBuf::new(), + filenames: { + let path = path.clone().unwrap_or("".to_string()); + if Path::new(&path).join(file).exists() { + vec![Path::new(&path).join(file)] + } else { + panic!("Model file not found {file}"); + } + }, + }, + true, + ), + _ => { + if self.weight_file.is_some() { + return Ok((self.download_gguf_model(None)?, true)); + }; + //try download model anonymously + let loaded = self.download_model(None, hf_token.clone(), hf_token_path.clone()); + if loaded.is_ok() { + (loaded.unwrap(), false) + } else { + //if it's failed, try using huggingface token + info!("Try request model using cached huggingface token..."); + if hf_token.is_none() && hf_token_path.is_none() { + //no token provided + let token_path = format!( + "{}/.cache/huggingface/token", + dirs::home_dir().unwrap().display() + ); + if !Path::new(&token_path).exists() { + //also no token cache + use std::io::Write; + let mut input_token = String::new(); + warn!("Unable to request model, please provide your huggingface token to download model:\n"); + std::io::stdin() + .read_line(&mut input_token) + .expect("Failed to read token!"); + std::fs::create_dir_all(Path::new(&token_path).parent().unwrap()) + .unwrap(); + let mut output = std::fs::File::create(token_path).unwrap(); + write!(output, "{}", input_token.trim()) + .expect("Failed to save token!"); + } + } + ( + self.download_model(None, hf_token.clone(), hf_token_path.clone())?, + false, + ) + } + } + }; + + Ok((paths, gguf)) + } + pub fn download_model( &self, - model_id: String, - weight_file: Option, - quant: Option, revision: Option, hf_token: Option, hf_token_path: Option, - ) -> Result { - if quant.is_some() && quant.as_ref().unwrap() == "gguf" && weight_file.is_some() { - info!( - "Downloading GGUF file {} from repo {}", - weight_file.as_ref().unwrap(), - model_id, - ); - return self.download_gguf_model(model_id, None, weight_file.clone().unwrap()); - }; - let api = try_api!(ApiBuilder::new() + ) -> Result { + assert!(self.model_id.is_some(), "No model id provided!"); + + let api = ApiBuilder::new() .with_progress(true) .with_token(Some(get_token(hf_token, hf_token_path)?)) - .build()); + .build() + .map_err(candle_core::Error::wrap)?; let revision = revision.unwrap_or("main".to_string()); let api = api.repo(Repo::with_revision( - model_id, + self.model_id.clone().unwrap(), RepoType::Model, revision.clone(), )); - let tokenizer_filename = try_api!(api.get("tokenizer.json")); + let tokenizer_filename = api + .get("tokenizer.json") + .map_err(candle_core::Error::wrap)?; - let config_filename = try_api!(api.get("config.json")); + let config_filename = api.get("config.json").map_err(candle_core::Error::wrap)?; let tokenizer_config_filename = match api.get("tokenizer_config.json") { Ok(f) => f, @@ -161,13 +237,15 @@ impl DefaultLoader { }; let mut filenames = vec![]; - for rfilename in try_api!(api.info()) + for rfilename in api + .info() + .map_err(candle_core::Error::wrap)? .siblings .iter() .map(|x| x.rfilename.clone()) .filter(|x| x.ends_with(".safetensors")) { - let filename = try_api!(api.get(&rfilename)); + let filename = api.get(&rfilename).map_err(candle_core::Error::wrap)?; filenames.push(filename); } @@ -179,22 +257,25 @@ impl DefaultLoader { }) } - pub fn download_gguf_model( - &self, - model_id: String, - revision: Option, - filename: String, - ) -> Result { + pub fn download_gguf_model(&self, revision: Option) -> Result { + assert!(self.model_id.is_some(), "No model id provided!"); + info!( + "Downloading GGUF file {} from repo {}", + self.weight_file.as_ref().unwrap(), + self.model_id.as_ref().unwrap(), + ); + let filename = self.weight_file.clone().unwrap(); let api = hf_hub::api::sync::Api::new().unwrap(); let revision = revision.unwrap_or("main".to_string()); let mut filenames = vec![]; - let filename = try_api!(api + let filename = api .repo(hf_hub::Repo::with_revision( - model_id, + self.model_id.clone().unwrap(), hf_hub::RepoType::Model, revision.to_string(), )) - .get(filename.as_str())); + .get(filename.as_str()) + .map_err(candle_core::Error::wrap)?; filenames.push(filename); Ok(DefaultModelPaths { @@ -206,114 +287,124 @@ impl DefaultLoader { } //support loading in both multithreaded and multiprocess mode + #[allow(unused_variables)] pub async fn load_model( &self, paths: DefaultModelPaths, dtype: DType, - quant: &Option, + gguf: bool, + isq: Option, device_ids: Vec, //pass only 1 device_id in multiprocess mode, otherwise, multiple device_ids in multithread mode #[cfg(feature = "nccl")] comm_id: Option, //must pass comm id in multiprocess mode local_rank: Option, //must pass current rank in multiprocess mode local_world_size: Option, //must pass the number of local devices used in multiprocess mode #[cfg(feature = "nccl")] global_rank: Option, //must pass current global rank in multi-node mode #[cfg(feature = "nccl")] global_world_size: Option, //must pass total number of devices used in multi-node mode - ) -> Result<(Vec>, PipelineConfig), APIError> { - let specific_args = self.config.clone(); + ) -> Result<(Vec>, PipelineConfig)> { let reporter = Arc::new(RwLock::new(ProgressReporter::new(local_rank.unwrap_or(0)))); let num_subprogress = local_world_size.map_or(0, |n| n - 1); - let (models, devices, config, sep_style) = if quant.is_some() - && matches!(quant.as_ref().unwrap().as_str(), "ggml" | "gguf") - { + let (models, devices, config, sep_style) = if gguf { let device = crate::new_device(device_ids[0]).unwrap(); let path = paths.get_weight_filenames()[0].clone(); - info!( - "Loading quantized {} model from file {}", - self.name, - path.display() - ); - let s_cfg = specific_args.clone(); - let nlayers = { - let mut file = try_api!(std::fs::File::open(path.clone())); - let content = try_api!( - gguf_file::Content::read(&mut file).map_err(|e| e.with_path(path.clone())) - ); - let nlayers = match self.name.as_str() { - "llama" | "llama3" => GGUFLLaMa::get_num_of_layers(content), - "phi3" => GGUFPhi3::get_num_of_layers(content), - "qwen2" | "qwen3" => { - GGUFQWen::get_num_of_layers(self.name.as_str() == "qwen3", content) + info!("Loading quantized model from file {}", path.display()); + let (arch, nlayers) = { + let mut file = match std::fs::File::open(path.clone()) + .map_err(candle_core::Error::wrap) + { + Ok(file) => file, + Err(e) => { + tracing::error!("Failed to open gguf file {}: {}\n ***Tips: use `--w` to load safetensors models.", path.display(), e); + return Err(e); } - "glm4" => GGUFGLM4::get_num_of_layers(content), - _ => panic!("Model not supported!"), }; - nlayers.unwrap() + let content = match gguf_file::Content::read(&mut file) + .map_err(|e| e.with_path(path.clone())) + .map_err(candle_core::Error::wrap) + { + Ok(content) => content, + Err(e) => { + tracing::error!("Failed to open gguf file {}: {}\n ***Tips: use `--w` to load safetensors models.", path.display(), e); + return Err(e); + } + }; + let (arch, nlayers) = + gguf::get_arch_and_num_of_layers(content).map_err(candle_core::Error::wrap)?; + if !matches!( + arch.as_str(), + "llama" | "llama3" | "phi3" | "qwen2" | "qwen3" | "glm4" + ) { + panic!("Model arch {} not supported!", arch); + } else { + info!("Quantized {} model has {} layers.", arch, nlayers,); + } + (arch, nlayers) }; let handle = progress_worker(Some(num_subprogress), nlayers, Arc::clone(&reporter)).await; - let mut file = try_api!(std::fs::File::open(path.clone())); - let content = try_api!( - gguf_file::Content::read(&mut file).map_err(|e| e.with_path(path.clone())) - ); - let (model, config, sep_style) = match self.name.as_str() { + let mut file = std::fs::File::open(path.clone()).map_err(candle_core::Error::wrap)?; + let content = gguf_file::Content::read(&mut file) + .map_err(|e| e.with_path(path.clone())) + .map_err(candle_core::Error::wrap)?; + let (model, config, sep_style) = match arch.as_str() { "llama" => { - let model = try_api!(GGUFLLaMa::from_gguf( + let model = GGUFLLaMa::from_gguf( &content, &mut file, &device, dtype, - s_cfg, Arc::clone(&reporter), - )); + ) + .map_err(candle_core::Error::wrap)?; let cfg = model.get_config().clone(); (LLMModel::LlamaGGUF(model), cfg, SeparatorStyle::Llama) } "llama3" => { - let model = try_api!(GGUFLLaMa::from_gguf( + let model = GGUFLLaMa::from_gguf( &content, &mut file, &device, dtype, - s_cfg, Arc::clone(&reporter), - )); + ) + .map_err(candle_core::Error::wrap)?; let cfg = model.get_config().clone(); (LLMModel::LlamaGGUF(model), cfg, SeparatorStyle::Llama3) } "phi3" => { - let model = try_api!(GGUFPhi3::from_gguf( + let model = GGUFPhi3::from_gguf( &content, &mut file, &device, dtype, - s_cfg, Arc::clone(&reporter), - )); + ) + .map_err(candle_core::Error::wrap)?; let cfg = model.get_config().clone(); (LLMModel::Phi3GGUF(model), cfg, SeparatorStyle::Phi) } "qwen2" | "qwen3" => { - let model = try_api!(GGUFQWen::from_gguf( - self.name.as_str() == "qwen3", + let model = GGUFQWen::from_gguf( + arch.as_str() == "qwen3", &content, &mut file, &device, dtype, - s_cfg, Arc::clone(&reporter), - )); + ) + .map_err(candle_core::Error::wrap)?; let cfg = model.get_config().clone(); (LLMModel::QWenGGUF(model), cfg, SeparatorStyle::Qwen) } "glm4" => { - let model = try_api!(GGUFGLM4::from_gguf( + let model = GGUFGLM4::from_gguf( &content, &mut file, &device, dtype, - s_cfg, Arc::clone(&reporter), - )); + ) + .map_err(candle_core::Error::wrap)?; let cfg = model.get_config().clone(); (LLMModel::GLM4GGUF(model), cfg, SeparatorStyle::GLM) } @@ -322,80 +413,29 @@ impl DefaultLoader { handle.join().unwrap(); (vec![model], vec![device], config.to_owned(), sep_style) } else { - let config = match self.name.as_str() { - "llama" | "llama3" => { - let config: LlamaConfig = try_api!(serde_json::from_slice(&try_api!( - std::fs::read(paths.get_config_filename()) - ),)); - config.into_config(false, dtype, &specific_args) - } - "phi2" => { - let config: Phi2Config = try_api!(serde_json::from_slice(&try_api!( - std::fs::read(paths.get_config_filename()) - ),)); - //Phi2 use F32 type for kvcache - config.into_config(false, DType::F32, &specific_args) - } - "phi3" => { - let config: PhiConfig = try_api!(serde_json::from_slice(&try_api!( - std::fs::read(paths.get_config_filename()) - ),)); - config.into_config(false, dtype, &specific_args) - } - "qwen2" | "qwen3" => { - let config: QwenConfig = try_api!(serde_json::from_slice(&try_api!( - std::fs::read(paths.get_config_filename()) - ),)); - config.into_config(false, dtype, &specific_args) - } - "gemma" => { - let config: GemmaConfig = try_api!(serde_json::from_slice(&try_api!( - std::fs::read(paths.get_config_filename()) - ),)); - config.into_config(false, dtype, &specific_args) - } - "gemma3" => { - let config: Gemma3Config = try_api!(serde_json::from_slice(&try_api!( - std::fs::read(paths.get_config_filename()) - ),)); - config.into_config(false, dtype, &specific_args) - } - "mistral" => { - let config: MistralConfig = try_api!(serde_json::from_slice(&try_api!( - std::fs::read(paths.get_config_filename()) - ),)); - config.into_config(false, dtype, &specific_args) - } - "yi" => { - let config: YiConfig = try_api!(serde_json::from_slice(&try_api!( - std::fs::read(paths.get_config_filename()) - ),)); - config.into_config(false, dtype, &specific_args) - } - "stablelm" => { - let config: StableLMConfig = try_api!(serde_json::from_slice(&try_api!( - std::fs::read(paths.get_config_filename()) - ),)); - config.into_config(false, dtype, &specific_args) - } - "glm4" => { - let config: GLMConfig = try_api!(serde_json::from_slice(&try_api!( - std::fs::read(paths.get_config_filename()) - ),)); - config.into_config(false, dtype, &specific_args) - } - "deepseek" => { - let config: DeepSeekConfig = try_api!(serde_json::from_slice(&try_api!( - std::fs::read(paths.get_config_filename()) - ),)); - config.into_config(false, dtype, &specific_args) + let cfile = paths.get_config_filename(); + let arch = Config::get_model_arch(&cfile)?; + + let config = match arch.as_str() { + "LlamaForCausalLM" => Llama::load_config(&cfile, isq)?, + "PhiForCausalLM" => Phi2::load_config(&cfile, isq)?, + "Phi3ForCausalLM" => Phi::load_config(&cfile, isq)?, + "Qwen2ForCausalLM" | "Qwen3ForCausalLM" => Qwen::load_config(&cfile, isq)?, + "Gemma2ForCausalLM" => Gemma::load_config(&cfile, isq)?, + "Gemma3ForConditionalGeneration" => Gemma3::load_config(&cfile, isq)?, + "MistralForCausalLM" => Mistral::load_config(&cfile, isq)?, + "yi" => Yi::load_config(&cfile, isq)?, + "StableLmForCausalLM" => StableLM::load_config(&cfile, isq)?, + "Glm4ForCausalLM" => GLM4::load_config(&cfile, isq)?, + "DeepseekV2ForCausalLM" | "DeepseekV3ForCausalLM" => { + DeepSeek::load_config(&cfile, isq)? } _ => panic!("Model not supported!"), }; info!("Model {:?}", config); - info!("Loading {} model.", self.name); + info!("Loading {} model.", arch); let handle = progress_worker( Some(num_subprogress), config.num_hidden_layers, @@ -468,22 +508,8 @@ impl DefaultLoader { .unwrap() }; - let (model, sep) = match self.name.as_str() { - "llama" => ( - LLMModel::Llama( - Llama::load( - vb, - &config, - dtype, - &device, - comm, - Arc::clone(&reporter), - ) - .unwrap(), - ), - SeparatorStyle::Llama, - ), - "llama3" => ( + let (model, sep) = match arch.as_str() { + "LlamaForCausalLM" => ( LLMModel::Llama( Llama::load( vb, @@ -497,24 +523,24 @@ impl DefaultLoader { ), SeparatorStyle::Llama3, ), - "phi2" => ( + "PhiForCausalLM" => ( LLMModel::Phi2( Phi2::new(vb, &config, dtype, &device, comm, Arc::clone(&reporter)) .unwrap(), ), SeparatorStyle::Phi, ), - "phi3" => ( + "Phi3ForCausalLM" => ( LLMModel::Phi3( Phi::new(vb, &config, dtype, &device, comm, Arc::clone(&reporter)) .unwrap(), ), SeparatorStyle::Phi, ), - "qwen2" | "qwen3" => ( + "Qwen2ForCausalLM" | "Qwen3ForCausalLM" => ( LLMModel::Qwen( Qwen::new( - self.name.as_str() == "qwen3", + matches!(arch.as_str(), "qwen3" | "Qwen3ForCausalLM"), vb, &config, dtype, @@ -526,7 +552,7 @@ impl DefaultLoader { ), SeparatorStyle::Qwen, ), - "gemma" => ( + "Gemma2ForCausalLM" => ( LLMModel::Gemma( Gemma::new( vb, @@ -540,7 +566,7 @@ impl DefaultLoader { ), SeparatorStyle::Gemma, ), - "gemma3" => ( + "Gemma3ForConditionalGeneration" => ( LLMModel::Gemma3( Gemma3::new( vb, @@ -554,7 +580,7 @@ impl DefaultLoader { ), SeparatorStyle::Gemma, ), - "mistral" => ( + "MistralForCausalLM" => ( LLMModel::Mistral( Mistral::new( vb, @@ -575,7 +601,7 @@ impl DefaultLoader { ), SeparatorStyle::Yi, ), - "stablelm" => ( + "StableLmForCausalLM" => ( LLMModel::StableLM( StableLM::new( vb, @@ -589,14 +615,14 @@ impl DefaultLoader { ), SeparatorStyle::StableLM, ), - "glm4" => ( + "Glm4ForCausalLM" => ( LLMModel::GLM4( GLM4::new(vb, &config, dtype, &device, comm, Arc::clone(&reporter)) .unwrap(), ), SeparatorStyle::Llama, ), - "deepseek" => ( + "DeepseekV2ForCausalLM" | "DeepseekV3ForCausalLM" => ( LLMModel::DeepSeek( DeepSeek::load( vb, @@ -641,20 +667,11 @@ impl DefaultLoader { warn!("Done loading."); //max and min number of tokens generated per request - let default_max_tokens = specific_args - .max_gen_tokens - .unwrap_or(config.max_seq_len / 2) - .clamp(MIN_GEN_TOKENS, MAX_GEN_TOKENS); + let default_max_tokens = (config.max_seq_len / 5).clamp(MIN_GEN_TOKENS, MAX_GEN_TOKENS); let pipeline_config = PipelineConfig { max_model_len: config.max_seq_len, default_max_tokens, - penalty: specific_args.penalty.unwrap_or(1.), - repeat_last_n: specific_args.repeat_last_n.unwrap_or(64), - temperature: specific_args.temperature, - top_k: specific_args.top_k, - top_p: specific_args.top_p, - thinking: Some(specific_args.thinking), }; #[cfg(feature = "nccl")] @@ -669,9 +686,9 @@ impl DefaultLoader { let logits_processor = { LogitsProcessor::new( SAMPLING_SEED, - pipeline_config.temperature, - specific_args.top_k, - specific_args.top_p, + None, + None, + None, ) }; let tokenizer_file = paths.get_tokenizer_filename(); @@ -682,8 +699,7 @@ impl DefaultLoader { Option, ) = if tokenizer_file.display().to_string() != "" && Path::exists(&tokenizer_file) { let tokenizer = Tokenizer::from_file(tokenizer_file.clone()) - .map_err(|x| APIError::new(x.to_string())) - .unwrap(); + .map_err(candle_core::Error::wrap).unwrap(); let tokenizer_cfg_file = paths.get_tokenizer_config_filename(); let (chat_template, bos, eos) = if Path::exists(&tokenizer_cfg_file) { @@ -715,7 +731,7 @@ impl DefaultLoader { (None, None, None) }; (tokenizer, chat_template, bos, eos) - } else if quant.is_some() && matches!(quant.as_ref().unwrap().as_str(), "ggml" | "gguf") { + } else if gguf { use crate::backend::gguf::{get_gguf_info, Content, GGUFInfo}; let filename = paths.get_weight_filenames()[0].clone(); let mut reader = std::fs::File::open(filename).unwrap(); @@ -746,7 +762,6 @@ impl DefaultLoader { info!("Warning: Missing tokenizer_config.json \n Warning: Chat Template not found, use built-in template which may not correct!"); } info!("{:?}", pipeline_config); - info!("{:?}", specific_args); } let mut stop_token_ids = Vec::::new(); @@ -812,11 +827,10 @@ impl DefaultLoader { tracing::warn!("stop_token_ids {:?}", stop_token_ids); Box::new(DefaultPipeline { model, - args: specific_args.clone(), tokenizer, logits_processor, conversation: DefaultConversation::new( - self.name.to_string(), + config.architectures.as_ref().unwrap()[0].clone(), chat_template.clone(), Vec::default(), sep_style.clone(), @@ -828,7 +842,7 @@ impl DefaultLoader { sep2: Some(" ".to_string()), }, ), - name: self.name.clone(), + name: config.architectures.as_ref().unwrap()[0].clone(), dtype, device: devices[rank].clone(), stop_token_ids, @@ -849,7 +863,7 @@ impl DefaultPipeline { input_positions: &[Vec], kv_cache: Option<&Vec<(Tensor, Tensor)>>, input_metadata: &InputMetadata, - ) -> Result { + ) -> Result { let input_tokens = if input_tokens.shape().dims().len() < 2 { input_tokens .reshape((1, input_tokens.shape().dims()[0])) @@ -859,51 +873,51 @@ impl DefaultPipeline { }; match &self.model { - LLMModel::Llama(llama) => llama - .forward(&input_tokens, input_positions, kv_cache, input_metadata) - .map_err(APIError::from), - LLMModel::Phi2(phi) => phi - .forward(&input_tokens, input_positions, kv_cache, input_metadata) - .map_err(APIError::from), - LLMModel::Phi3(phi) => phi - .forward(&input_tokens, input_positions, kv_cache, input_metadata) - .map_err(APIError::from), - LLMModel::Qwen(qwen) => qwen - .forward(&input_tokens, input_positions, kv_cache, input_metadata) - .map_err(APIError::from), - LLMModel::Gemma(gemma) => gemma - .forward(&input_tokens, input_positions, kv_cache, input_metadata) - .map_err(APIError::from), - LLMModel::Gemma3(gemma3) => gemma3 - .forward(&input_tokens, input_positions, kv_cache, input_metadata) - .map_err(APIError::from), - LLMModel::Mistral(mistral) => mistral - .forward(&input_tokens, input_positions, kv_cache, input_metadata) - .map_err(APIError::from), - LLMModel::Yi(yi) => yi - .forward(&input_tokens, input_positions, kv_cache, input_metadata) - .map_err(APIError::from), - LLMModel::StableLM(stablelm) => stablelm - .forward(&input_tokens, input_positions, kv_cache, input_metadata) - .map_err(APIError::from), - LLMModel::GLM4(glm4) => glm4 - .forward(&input_tokens, input_positions, kv_cache, input_metadata) - .map_err(APIError::from), - LLMModel::DeepSeek(deepseek) => deepseek - .forward(&input_tokens, input_positions, kv_cache, input_metadata) - .map_err(APIError::from), - LLMModel::Phi3GGUF(phi3) => phi3 - .forward(&input_tokens, input_positions, kv_cache, input_metadata) - .map_err(APIError::from), - LLMModel::LlamaGGUF(llama) => llama - .forward(&input_tokens, input_positions, kv_cache, input_metadata) - .map_err(APIError::from), - LLMModel::QWenGGUF(qwen) => qwen - .forward(&input_tokens, input_positions, kv_cache, input_metadata) - .map_err(APIError::from), - LLMModel::GLM4GGUF(glm4) => glm4 - .forward(&input_tokens, input_positions, kv_cache, input_metadata) - .map_err(APIError::from), + LLMModel::Llama(llama) => { + llama.forward(&input_tokens, input_positions, kv_cache, input_metadata) + } + LLMModel::Phi2(phi) => { + phi.forward(&input_tokens, input_positions, kv_cache, input_metadata) + } + LLMModel::Phi3(phi) => { + phi.forward(&input_tokens, input_positions, kv_cache, input_metadata) + } + LLMModel::Qwen(qwen) => { + qwen.forward(&input_tokens, input_positions, kv_cache, input_metadata) + } + LLMModel::Gemma(gemma) => { + gemma.forward(&input_tokens, input_positions, kv_cache, input_metadata) + } + LLMModel::Gemma3(gemma3) => { + gemma3.forward(&input_tokens, input_positions, kv_cache, input_metadata) + } + LLMModel::Mistral(mistral) => { + mistral.forward(&input_tokens, input_positions, kv_cache, input_metadata) + } + LLMModel::Yi(yi) => { + yi.forward(&input_tokens, input_positions, kv_cache, input_metadata) + } + LLMModel::StableLM(stablelm) => { + stablelm.forward(&input_tokens, input_positions, kv_cache, input_metadata) + } + LLMModel::GLM4(glm4) => { + glm4.forward(&input_tokens, input_positions, kv_cache, input_metadata) + } + LLMModel::DeepSeek(deepseek) => { + deepseek.forward(&input_tokens, input_positions, kv_cache, input_metadata) + } + LLMModel::Phi3GGUF(phi3) => { + phi3.forward(&input_tokens, input_positions, kv_cache, input_metadata) + } + LLMModel::LlamaGGUF(llama) => { + llama.forward(&input_tokens, input_positions, kv_cache, input_metadata) + } + LLMModel::QWenGGUF(qwen) => { + qwen.forward(&input_tokens, input_positions, kv_cache, input_metadata) + } + LLMModel::GLM4GGUF(glm4) => { + glm4.forward(&input_tokens, input_positions, kv_cache, input_metadata) + } } } @@ -911,7 +925,7 @@ impl DefaultPipeline { &mut self, logits: &Tensor, groups: &VecDeque>, - ) -> Result, APIError> { + ) -> Result> { let (tokens_generated, custom_stop_tokens, panalties, reference_tokens): ( Vec, Vec>, @@ -938,12 +952,12 @@ impl DefaultPipeline { _ => vec![], }; - let ref_tokens = if sampling_params.repetition_penalty != 1. - && self.args.repeat_last_n.unwrap_or(64) < generated + let ref_tokens = if sampling_params.repetition_penalty.unwrap_or(1.) != 1. + && sampling_params.repeat_last_n.unwrap_or(64) < generated { let start_at = tokens .len() - .saturating_sub(self.args.repeat_last_n.unwrap_or(64)); + .saturating_sub(sampling_params.repeat_last_n.unwrap_or(64)); tokens[start_at..].to_vec() } else { vec![] @@ -955,7 +969,7 @@ impl DefaultPipeline { generated as i32 }, custom_stop_token, - sampling_params.repetition_penalty, + sampling_params.repetition_penalty.unwrap_or(1.0), ref_tokens, ) }) diff --git a/src/openai/requests.rs b/src/openai/requests.rs index fe74fdb5..91261279 100644 --- a/src/openai/requests.rs +++ b/src/openai/requests.rs @@ -31,13 +31,14 @@ pub struct ChatCompletionRequest { pub stream: Option, //false #[serde(default)] pub presence_penalty: Option, //0.0 + pub repeat_last_n: Option, //0.0 #[serde(default)] pub frequency_penalty: Option, //0.0 #[serde(default)] pub logit_bias: Option>, //None #[serde(default)] pub user: Option, //None - pub top_k: Option, //-1 + pub top_k: Option, //-1 #[serde(default)] pub best_of: Option, //None #[serde(default)] @@ -51,5 +52,5 @@ pub struct ChatCompletionRequest { #[serde(default)] pub logprobs: Option, //false pub repetition_penalty: Option, //1.1 - pub thinking: Option, //false + pub thinking: Option, //false } diff --git a/src/openai/responses.rs b/src/openai/responses.rs index cb9a7a6c..a7ca3b9e 100644 --- a/src/openai/responses.rs +++ b/src/openai/responses.rs @@ -43,7 +43,7 @@ macro_rules! try_api { match $candle_result { Ok(v) => v, Err(e) => { - return Err(APIError::from(e)); + return Err(crate::openai::responses::APIError::from(e)); } } }; diff --git a/src/openai/sampling_params.rs b/src/openai/sampling_params.rs index 85f8f8f4..070af486 100644 --- a/src/openai/sampling_params.rs +++ b/src/openai/sampling_params.rs @@ -52,7 +52,8 @@ pub struct SamplingParams { pub frequency_penalty: f32, /// Penalize new tokens based upon whether their frequency in the generated text so far, >1 encourage new, <1 encourage repeat /// rec. default = 1 - pub repetition_penalty: f32, + pub repetition_penalty: Option, + pub repeat_last_n: Option, /// Randomness of sampling. /// rec. default = 1 pub temperature: Option, @@ -101,7 +102,8 @@ impl SamplingParams { best_of: Option, presence_penalty: f32, frequency_penalty: f32, - repetition_penalty: f32, + repetition_penalty: Option, + repeat_last_n: Option, temperature: Option, top_p: Option, top_k: Option, @@ -123,6 +125,7 @@ impl SamplingParams { presence_penalty, frequency_penalty, repetition_penalty, + repeat_last_n, temperature, top_p, top_k, @@ -177,18 +180,22 @@ impl SamplingParams { self.frequency_penalty ))); } - if !(Range { - start: 0.0, - end: 2.0, - }) - .contains(&self.repetition_penalty) - || self.repetition_penalty == 0.0 - { - return Err(APIError::new(format!( - "repetition_penalty must be in (0, 2], got {}", - self.repetition_penalty - ))); + if self.repetition_penalty.is_some() { + let repetition_penalty = self.repetition_penalty.unwrap(); + if !(Range { + start: 0.0, + end: 2.0, + }) + .contains(&repetition_penalty) + || repetition_penalty == 0.0 + { + return Err(APIError::new(format!( + "repetition_penalty must be in (0, 2], got {}", + repetition_penalty + ))); + } } + if self.temperature.unwrap_or(0.0f32) < 0.0f32 { return Err(APIError::new(format!( "temperature must be non-negative, got {}", diff --git a/src/paged_attention/input_metadata.rs b/src/paged_attention/input_metadata.rs index 5daaba4e..8d53fbf5 100644 --- a/src/paged_attention/input_metadata.rs +++ b/src/paged_attention/input_metadata.rs @@ -7,7 +7,6 @@ pub struct InputMetadata { pub slot_mapping: Tensor, // pub attn_bias: Option, pub is_prompt: bool, - pub kv_cache_dtype: String, } impl InputMetadata { @@ -16,14 +15,12 @@ impl InputMetadata { /// context_lens: the length of attention context for each generation token. /// max_context_len: The maximum context length. /// block_tables: The block tables. (Seq id -> list of physical block) - /// kv_cache_dtype: KV cache datatype (auto or fp8_e5m2) pub fn new( prompt_lens: Vec, max_context_len: Option, block_tables: Option, context_lens: Option, slot_mapping: Tensor, - kv_cache_dtype: String, ) -> Self { let is_prompt = !prompt_lens.is_empty(); Self { @@ -34,7 +31,6 @@ impl InputMetadata { slot_mapping, // attn_bias: None, is_prompt, - kv_cache_dtype, } } } diff --git a/src/paged_attention/mod.rs b/src/paged_attention/mod.rs index 48bcf8c9..1c0edd7e 100644 --- a/src/paged_attention/mod.rs +++ b/src/paged_attention/mod.rs @@ -58,6 +58,7 @@ impl PagedAttention { /// value_cache: shape = [num_blocks, num_kv_heads, head_size, /// block_size] /// input_metadata: metadata for paged attention. + #[allow(dead_code)] pub fn forward( &self, query: &Tensor, diff --git a/src/scheduler/cache_engine.rs b/src/scheduler/cache_engine.rs index 6e76e979..17eb8824 100644 --- a/src/scheduler/cache_engine.rs +++ b/src/scheduler/cache_engine.rs @@ -1,15 +1,11 @@ +use crate::openai::models::Config; +use candle_core::{DType, Device, Result, Tensor}; use std::{ collections::HashMap, sync::{Arc, Mutex, MutexGuard}, }; -use candle_core::{DType, Device, Tensor}; - -use crate::{ - backend::{copy_blocks, swap_blocks}, - openai::{models::Config, responses::APIError}, - try_api, -}; +use crate::backend::{copy_blocks, swap_blocks}; #[derive(Clone, Debug)] pub struct CacheConfig { @@ -51,7 +47,7 @@ impl CacheEngine { dtype: DType, device: &Device, num_shards: usize, - ) -> Result { + ) -> Result { Ok(Self { gpu_cache: Arc::new(Mutex::new(Self::allocate_gpu_cache( model_config, @@ -85,7 +81,7 @@ impl CacheEngine { dtype: DType, device: &Device, num_shards: usize, - ) -> Result, APIError> { + ) -> Result> { assert!(cache_config.fully_init); let key_block_shape = Self::calculate_key_block_shape( @@ -98,7 +94,7 @@ impl CacheEngine { Self::calculate_value_block_shape(model_config, cache_config.block_size, num_shards); let mut gpu_cache = Vec::new(); for _ in 0..model_config.num_hidden_layers { - let key_blocks = try_api!(Tensor::zeros( + let key_blocks = Tensor::zeros( ( cache_config.num_gpu_blocks.unwrap(), key_block_shape.0, @@ -108,8 +104,8 @@ impl CacheEngine { ), dtype, device, - )); - let value_blocks = try_api!(Tensor::zeros( + )?; + let value_blocks = Tensor::zeros( ( cache_config.num_gpu_blocks.unwrap(), value_block_shape.0, @@ -118,7 +114,7 @@ impl CacheEngine { ), dtype, device, - )); + )?; gpu_cache.push((key_blocks, value_blocks)); } Ok(gpu_cache) @@ -130,7 +126,7 @@ impl CacheEngine { dtype: DType, device: &Device, num_shards: usize, - ) -> Result, APIError> { + ) -> Result> { assert!(cache_config.fully_init); let key_block_shape = Self::calculate_key_block_shape( @@ -143,7 +139,7 @@ impl CacheEngine { Self::calculate_value_block_shape(model_config, cache_config.block_size, num_shards); let mut cpu_cache = Vec::new(); for _ in 0..model_config.num_hidden_layers { - let key_blocks = try_api!(Tensor::zeros( + let key_blocks = Tensor::zeros( ( cache_config.num_cpu_blocks.unwrap(), key_block_shape.0, @@ -153,8 +149,8 @@ impl CacheEngine { ), dtype, device, - )); - let value_blocks = try_api!(Tensor::zeros( + )?; + let value_blocks = Tensor::zeros( ( cache_config.num_cpu_blocks.unwrap(), value_block_shape.0, @@ -163,7 +159,7 @@ impl CacheEngine { ), dtype, device, - )); + )?; cpu_cache.push((key_blocks, value_blocks)); } Ok(cpu_cache) @@ -180,7 +176,7 @@ impl CacheEngine { let element_size = dtype.size_in_bytes(); let x = 16 / element_size; ( - model_config.num_key_value_heads / num_shards, + model_config.num_key_value_heads.unwrap() / num_shards, model_config.k_head_dim() / x, block_size, x, @@ -193,7 +189,7 @@ impl CacheEngine { num_shards: usize, ) -> (usize, usize, usize) { ( - model_config.num_key_value_heads / num_shards, + model_config.num_key_value_heads.unwrap() / num_shards, model_config.v_head_dim(), block_size, ) @@ -201,28 +197,20 @@ impl CacheEngine { } impl CacheEngine { - pub fn swap_in(&self, src_to_dst: HashMap) -> Result<(), APIError> { + pub fn swap_in(&self, src_to_dst: HashMap) -> Result<()> { for i in 0..self.num_layers { let (src_key_cache, src_value_cache) = self.cpu_cache.get(i).unwrap(); let mut gpu_cache = self.get_kv_cache(); let (dst_key_cache, dst_value_cache) = gpu_cache.get_mut(i).unwrap(); // Swap (copy) key blocks - try_api!(swap_blocks( - src_key_cache.clone(), - dst_key_cache, - src_to_dst.clone() - )); + swap_blocks(src_key_cache.clone(), dst_key_cache, src_to_dst.clone())?; // Swap (copy) key blocks - try_api!(swap_blocks( - src_value_cache.clone(), - dst_value_cache, - src_to_dst.clone() - )); + swap_blocks(src_value_cache.clone(), dst_value_cache, src_to_dst.clone())?; } Ok(()) } - pub fn swap_out(&mut self, src_to_dst: HashMap) -> Result<(), APIError> { + pub fn swap_out(&mut self, src_to_dst: HashMap) -> Result<()> { for i in 0..self.num_layers { let gpu_cache = self.get_kv_cache(); let (src_key_cache, src_value_cache) = gpu_cache.get(i).unwrap().clone(); @@ -230,22 +218,14 @@ impl CacheEngine { let (dst_key_cache, dst_value_cache) = self.cpu_cache.get_mut(i).unwrap(); // Swap (copy) key blocks - try_api!(swap_blocks( - src_key_cache.clone(), - dst_key_cache, - src_to_dst.clone() - )); + swap_blocks(src_key_cache.clone(), dst_key_cache, src_to_dst.clone())?; // Swap (copy) key blocks - try_api!(swap_blocks( - src_value_cache.clone(), - dst_value_cache, - src_to_dst.clone() - )); + swap_blocks(src_value_cache.clone(), dst_value_cache, src_to_dst.clone())?; } Ok(()) } - - pub fn copy(&mut self, src_to_dst: HashMap>) -> Result<(), APIError> { + #[allow(unused_unsafe)] + pub fn copy(&mut self, src_to_dst: HashMap>) -> Result<()> { let mut gpu_cache = self.get_kv_cache(); #[allow(clippy::map_identity)] let caches: (Vec<&mut Tensor>, Vec<&mut Tensor>) = @@ -253,8 +233,9 @@ impl CacheEngine { let (key_caches, value_caches) = caches; // NOTE(EricLBuehler): This may synchronize the CPU and GPU - try_api!(unsafe { copy_blocks(key_caches, value_caches, src_to_dst) }); - + unsafe { + copy_blocks(key_caches, value_caches, src_to_dst)?; + } Ok(()) } } diff --git a/src/scheduler/mod.rs b/src/scheduler/mod.rs index 7a528008..1330fdf9 100644 --- a/src/scheduler/mod.rs +++ b/src/scheduler/mod.rs @@ -231,6 +231,11 @@ impl Scheduler { free_blocks * self.block_engine.get_block_size() ); } + + pub fn get_available_kv_tokens(&self) -> usize { + let free_blocks = self.block_engine.get_num_free_blocks(); + free_blocks * self.block_engine.get_block_size() + } } impl Scheduler {