@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 189 lines 5.26 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import copy
5 from enum import Enum
6 from itertools import product
7
8 import torch
9 from tqdm import tqdm
10
11 from vllm.config import VllmConfig, set_current_vllm_config
12 from vllm.model_executor.layers.fused_moe.config import FUSED_MOE_UNQUANTIZED_CONFIG
13 from vllm.utils.torch_utils import set_random_seed
14
15 from .common import (
16 Config,
17 RankTensors,
18 WeightTensors,
19 reference_moe_impl,
20 run_modular_kernel,
21 )
22 from .mk_objects import (
23 MK_FUSED_EXPERT_TYPES,
24 MK_MULTI_GPU_PREPARE_FINALIZE_TYPES,
25 MK_QUANT_CONFIGS,
26 )
27 from .parallel_utils import ProcessGroupInfo, parallel_launch_with_config
28
29
30 class Result(Enum):
31 PASS = 1
32 FAIL = 2
33 SKIP = 3
34
35
36 def rank_worker(
37 pgi: ProcessGroupInfo,
38 vllm_config: VllmConfig,
39 cpu_group,
40 config: Config,
41 weights: WeightTensors,
42 ):
43 set_random_seed(pgi.rank)
44
45 # get weights to this device
46 weights.to_current_device()
47
48 Ms = config.Ms
49 assert isinstance(Ms, list)
50 TOPKs = config.topks
51 assert isinstance(TOPKs, list)
52
53 for m, topk in product(Ms, TOPKs):
54 print(f"Running m={m}, topk={topk} ...")
55 # override m and topk
56 cfgx = copy.deepcopy(config)
57 cfgx.Ms = m
58 cfgx.topks = topk
59
60 # inputs for rank
61 rank_tensors = RankTensors.make(cfgx, pgi)
62
63 # modular kernel out
64 mk_out = run_modular_kernel(pgi, vllm_config, cfgx, weights, rank_tensors)
65
66 with set_current_vllm_config(vllm_config):
67 ref_out = reference_moe_impl(cfgx, weights, rank_tensors)
68
69 torch.testing.assert_close(ref_out, mk_out, atol=3e-2, rtol=3e-2)
70
71
72 def make_feature_matrix(csv_file_path: str):
73 from dataclasses import asdict
74
75 import pandas as pd
76
77 def add_to_results(
78 config: Config, success: Result, results_df: pd.DataFrame | None = None
79 ):
80 config_dict = asdict(config)
81 config_dict["prepare_finalize_type"] = config_dict[
82 "prepare_finalize_type"
83 ].__name__
84 config_dict["fused_experts_type"] = config_dict["fused_experts_type"].__name__
85 config_dict["per_tensor_act_quant"] = config.is_per_tensor_act_quant
86 quant_config_dict = config_dict["quant_config"]
87 del config_dict["quant_config"]
88 if quant_config_dict is None:
89 quant_config = FUSED_MOE_UNQUANTIZED_CONFIG
90 quant_config_dict = asdict(quant_config)
91
92 config_dict |= quant_config_dict
93 result_dict = config_dict | {"success": success.name}
94
95 result_df = pd.DataFrame([result_dict])
96 if results_df is None:
97 results_df = result_df
98 else:
99 results_df = pd.concat([results_df, result_df], ignore_index=True)
100
101 return results_df
102
103 Ms = [64]
104 Ks = [7168] # hidden sizes
105 Ns = [2048]
106 TOPKs = [[4, 1]]
107 Es = [32]
108 DTYPEs = [torch.bfloat16]
109 PF_TYPES = MK_MULTI_GPU_PREPARE_FINALIZE_TYPES
110 FE_TYPES = MK_FUSED_EXPERT_TYPES
111 Q_TYPES = MK_QUANT_CONFIGS
112
113 combinations = list(
114 product(Ms, Ks, Ns, Es, TOPKs, DTYPEs, PF_TYPES, FE_TYPES, Q_TYPES)
115 )
116
117 results_df: pd.DataFrame | None = None
118 for m, k, n, e, topks, dtype, pf_type, experts_type, quant_config in tqdm(
119 combinations
120 ):
121 config = Config(
122 Ms=[m],
123 K=k,
124 N=n,
125 E=e,
126 topks=topks,
127 dtype=dtype,
128 prepare_finalize_type=pf_type,
129 fused_experts_type=experts_type,
130 quant_config=quant_config,
131 world_size=2,
132 )
133
134 success = None
135 if config.is_valid()[0]:
136 print(f"Running config : {config.describe()} ...")
137 try:
138 weights: WeightTensors = WeightTensors.make(config)
139 vllm_config, env_dict = config.make_env_data()
140 parallel_launch_with_config(
141 config.world_size,
142 rank_worker,
143 vllm_config,
144 env_dict,
145 config,
146 weights,
147 )
148 success = Result.PASS
149 except Exception as _:
150 success = Result.FAIL
151 else:
152 success = Result.SKIP
153
154 results_df = add_to_results(config, success, results_df)
155
156 if results_df is not None:
157 results_df.to_csv(f"{csv_file_path}")
158
159
160 if __name__ == "__main__":
161 import argparse
162 from pathlib import Path
163
164 parser = argparse.ArgumentParser(
165 description=(
166 "Make ModularKernel feature matrix \n"
167 "Example : python3 -m tests.kernels.moe.modular_kernel_tools.make_feature_matrix " # noqa: E501
168 "-f ./feature_matrices/feature_matrix.csv"
169 )
170 )
171
172 parser.add_argument(
173 "-f",
174 "--feature-matrix-csv-file-path",
175 type=str,
176 required=True,
177 help="File name to Generate a .csv file",
178 )
179 args = parser.parse_args()
180
181 csv_path = args.feature_matrix_csv_file_path
182 assert csv_path.endswith("csv"), (
183 f"Need a file path ending with .csv, got {csv_path}"
184 )
185 assert Path(csv_path).parent.is_dir(), (
186 f"Cannot find parent directory for {Path(csv_path).parent}"
187 )
188
189 make_feature_matrix(args.feature_matrix_csv_file_path)