@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 94 lines 2.11 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 WEIGHT_SHAPES = {
5 "ideal": [[4 * 256 * 32, 256 * 32]],
6 "mistralai/Mistral-7B-v0.1/TP1": [
7 [4096, 6144],
8 [4096, 4096],
9 [4096, 28672],
10 [14336, 4096],
11 ],
12 "mistralai/Mistral-7B-v0.1/TP2": [
13 [4096, 3072],
14 [2048, 4096],
15 [4096, 14336],
16 [7168, 4096],
17 ],
18 "mistralai/Mistral-7B-v0.1/TP4": [
19 [4096, 1536],
20 [1024, 4096],
21 [4096, 7168],
22 [3584, 4096],
23 ],
24 "meta-llama/Llama-2-7b-hf/TP1": [
25 [4096, 12288],
26 [4096, 4096],
27 [4096, 22016],
28 [11008, 4096],
29 ],
30 "meta-llama/Llama-2-7b-hf/TP2": [
31 [4096, 6144],
32 [2048, 4096],
33 [4096, 11008],
34 [5504, 4096],
35 ],
36 "meta-llama/Llama-2-7b-hf/TP4": [
37 [4096, 3072],
38 [1024, 4096],
39 [4096, 5504],
40 [2752, 4096],
41 ],
42 "meta-llama/Llama-2-13b-hf/TP1": [
43 [5120, 15360],
44 [5120, 5120],
45 [5120, 27648],
46 [13824, 5120],
47 ],
48 "meta-llama/Llama-2-13b-hf/TP2": [
49 [5120, 7680],
50 [2560, 5120],
51 [5120, 13824],
52 [6912, 5120],
53 ],
54 "meta-llama/Llama-2-13b-hf/TP4": [
55 [5120, 3840],
56 [1280, 5120],
57 [5120, 6912],
58 [3456, 5120],
59 ],
60 "meta-llama/Llama-2-70b-hf/TP1": [
61 [8192, 10240],
62 [8192, 8192],
63 [8192, 57344],
64 [28672, 8192],
65 ],
66 "meta-llama/Llama-2-70b-hf/TP2": [
67 [8192, 5120],
68 [4096, 8192],
69 [8192, 28672],
70 [14336, 8192],
71 ],
72 "meta-llama/Llama-2-70b-hf/TP4": [
73 [8192, 2560],
74 [2048, 8192],
75 [8192, 14336],
76 [7168, 8192],
77 ],
78 }
79
80 WEIGHT_SHAPES_MOE = {
81 "mistralai/Mixtral-8x7B-Instruct-v0.1": [
82 [8, 2, 4096, 28672],
83 [8, 2, 14336, 4096],
84 ],
85 "deepseek-ai/DeepSeek-V2-Lite": [
86 [64, 6, 2048, 1408],
87 ],
88 "ibm-granite/granite-3.0-1b-a400m": [
89 [32, 8, 1024, 1024],
90 ],
91 "ibm-granite/granite-3.0-3b-a800m": [
92 [40, 8, 1024, 1536],
93 ],
94 }