wehub-resource-sync
257 lines
7.7 KiB
Python
257 lines
7.7 KiB
Python
"""Debug compiled models with TVM instrument"""
|
|
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Dict, List, Set, Tuple # noqa: UP035
|
|
|
|
import tvm
|
|
from tvm import rpc, runtime
|
|
from tvm.relax.testing.lib_comparator import LibCompareVMInstrument
|
|
|
|
from mlc_llm.interface.help import HELP
|
|
from mlc_llm.support.argparse import ArgumentParser
|
|
from mlc_llm.testing.debug_chat import DebugChat
|
|
|
|
|
|
def _print_as_table(sorted_list):
|
|
print("=" * 100)
|
|
print(
|
|
"Name".ljust(50)
|
|
+ "Time (ms)".ljust(12)
|
|
+ "Count".ljust(8)
|
|
+ "Total time (ms)".ljust(18)
|
|
+ "Percentage (%)"
|
|
)
|
|
total_time = sum(record[1][0] * record[1][1] for record in sorted_list) * 1000
|
|
for record in sorted_list:
|
|
time = record[1][0] * 1000
|
|
weighted_time = time * record[1][1]
|
|
percentage = weighted_time / total_time * 100
|
|
print(
|
|
record[0].ljust(50)
|
|
+ f"{time:.4f}".ljust(12)
|
|
+ str(record[1][1]).ljust(8)
|
|
+ f"{weighted_time:.4f}".ljust(18)
|
|
+ f"{percentage:.2f}"
|
|
)
|
|
print(f"Total time: {total_time:.4f} ms")
|
|
|
|
|
|
class LibCompare(LibCompareVMInstrument):
|
|
"""The default debug instrument to use if users don't specify
|
|
a customized one.
|
|
|
|
This debug instrument will dump the arguments and output of each
|
|
VM Call instruction into a .npz file. It will also alert the user
|
|
if any function outputs are NaN or INF.
|
|
|
|
Parameters
|
|
----------
|
|
mod: runtime.Module
|
|
The module of interest to be validated.
|
|
|
|
device: runtime.Device
|
|
The device to run the target module on.
|
|
|
|
time_eval: bool
|
|
Whether to time evaluate the functions.
|
|
|
|
rtol: float
|
|
rtol used in validation
|
|
|
|
atol: float
|
|
atol used in validation
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
mod: runtime.Module,
|
|
device: runtime.Device,
|
|
debug_out: Path,
|
|
time_eval: bool = True,
|
|
rtol: float = 1e-2,
|
|
atol: float = 1,
|
|
skip_rounds: int = 0,
|
|
):
|
|
super().__init__(mod, device, True, rtol, atol)
|
|
self.debug_out = debug_out
|
|
self.time_eval = time_eval
|
|
self.time_eval_results: Dict[str, Tuple[float, int]] = {} # noqa: UP006
|
|
self.visited: Set[str] = set([]) # noqa: UP006
|
|
self.skip_rounds = skip_rounds
|
|
self.counter = 0
|
|
debug_out.mkdir(exist_ok=True, parents=True)
|
|
|
|
def reset(self, debug_out: Path):
|
|
"""Reset the state of the Instrument class
|
|
|
|
Note
|
|
----
|
|
`debug_out` is not used in this class.
|
|
|
|
Parameters
|
|
----------
|
|
debug_out : Path
|
|
the directory to dump the .npz files
|
|
"""
|
|
self.debug_out = debug_out
|
|
_print_as_table(
|
|
sorted(
|
|
self.time_eval_results.items(),
|
|
key=lambda x: -(x[1][0] * x[1][1]),
|
|
)
|
|
)
|
|
self.time_eval_results = {}
|
|
self.visited = set([])
|
|
self.counter = 0
|
|
debug_out.mkdir(exist_ok=True, parents=True)
|
|
|
|
def skip_instrument(self, func, name, before_run, ret_val, *args):
|
|
if name.startswith("shape_func"):
|
|
return True
|
|
if self.counter < self.skip_rounds:
|
|
self.counter += 1
|
|
print(f"[{self.counter}] Skip validating {name}..")
|
|
return True
|
|
if name in self.visited:
|
|
if self.time_eval and name in self.time_eval_results:
|
|
record = self.time_eval_results[name]
|
|
self.time_eval_results[name] = (record[0], record[1] + 1)
|
|
return True
|
|
self.visited.add(name)
|
|
return False
|
|
|
|
def compare(
|
|
self,
|
|
name: str,
|
|
ref_args: List[tvm.runtime.Tensor], # noqa: UP006
|
|
new_args: List[tvm.runtime.Tensor], # noqa: UP006
|
|
ret_indices: List[int], # noqa: UP006
|
|
):
|
|
super().compare(name, ref_args, new_args, ret_indices)
|
|
|
|
if self.time_eval and name not in self.time_eval_results:
|
|
res = self.mod.time_evaluator(
|
|
name,
|
|
self.device,
|
|
number=20,
|
|
repeat=3,
|
|
min_repeat_ms=100,
|
|
# cache_flush_bytes=256 * 10**6
|
|
)(*new_args)
|
|
self.time_eval_results[name] = (res.mean, 1)
|
|
print(f"Time-eval result {name} on {self.device}:\n {res}")
|
|
|
|
|
|
def get_instrument(args):
|
|
"""Get the debug instrument from the CLI arguments"""
|
|
if args.cmp_device is None:
|
|
assert args.cmp_lib_path is None, "cmp_lib_path must be None if cmp_device is None"
|
|
args.cmp_device = args.device
|
|
args.cmp_lib_path = args.model_lib
|
|
|
|
if args.cmp_device == "iphone":
|
|
assert args.cmp_lib_path.endswith(".dylib"), "Require a dylib file for iPhone"
|
|
proxy_host = os.environ.get("TVM_RPC_PROXY_HOST", "127.0.0.1")
|
|
proxy_port = int(os.environ.get("TVM_RPC_PROXY_PORT", "9090"))
|
|
sess = rpc.connect(proxy_host, proxy_port, "iphone")
|
|
sess.upload(args.cmp_lib_path)
|
|
lib = sess.load_module(os.path.basename(args.cmp_lib_path))
|
|
cmp_device = sess.metal()
|
|
elif args.cmp_device == "android":
|
|
assert args.cmp_lib_path.endswith(".so"), "Require a so file for Android"
|
|
tracker_host = os.environ.get("TVM_TRACKER_HOST", "0.0.0.0")
|
|
tracker_port = int(os.environ.get("TVM_TRACKER_PORT", "9190"))
|
|
tracker = rpc.connect_tracker(tracker_host, tracker_port)
|
|
sess = tracker.request("android")
|
|
sess.upload(args.cmp_lib_path)
|
|
lib = sess.load_module(os.path.basename(args.cmp_lib_path))
|
|
cmp_device = sess.cl(0)
|
|
else:
|
|
lib = tvm.runtime.load_module(args.cmp_lib_path)
|
|
cmp_device = tvm.device(args.cmp_device)
|
|
|
|
return LibCompare(
|
|
lib,
|
|
cmp_device,
|
|
time_eval=args.time_eval,
|
|
debug_out=Path(args.debug_dir),
|
|
)
|
|
|
|
|
|
def main():
|
|
"""The main function to start a DebugChat CLI"""
|
|
|
|
parser = ArgumentParser("MLC LLM Chat Debug Tool")
|
|
parser.add_argument(
|
|
"prompt",
|
|
type=str,
|
|
help="The user input prompt.",
|
|
)
|
|
parser.add_argument(
|
|
"--generate-len",
|
|
type=int,
|
|
help="Number of output tokens to generate.",
|
|
required=True,
|
|
)
|
|
parser.add_argument(
|
|
"--model",
|
|
type=str,
|
|
help="An MLC model directory that contains `mlc-chat-config.json`",
|
|
required=True,
|
|
)
|
|
parser.add_argument(
|
|
"--model-lib",
|
|
type=str,
|
|
help="The full path to the model library file to use (e.g. a ``.so`` file).",
|
|
required=True,
|
|
)
|
|
parser.add_argument(
|
|
"--debug-dir",
|
|
type=str,
|
|
help="The output folder to store the dumped debug files.",
|
|
required=True,
|
|
)
|
|
parser.add_argument(
|
|
"--device",
|
|
type=str,
|
|
default="auto",
|
|
help=HELP["device_compile"] + ' (default: "%(default)s")',
|
|
)
|
|
parser.add_argument(
|
|
"--cmp-device",
|
|
type=str,
|
|
default="none",
|
|
)
|
|
parser.add_argument(
|
|
"--cmp-lib-path",
|
|
type=str,
|
|
default="none",
|
|
)
|
|
parser.add_argument(
|
|
"--time-eval",
|
|
action="store_true",
|
|
help="Whether to time evaluate the functions.",
|
|
)
|
|
parsed = parser.parse_args()
|
|
instrument = get_instrument(parsed)
|
|
debug_chat = DebugChat(
|
|
model=parsed.model,
|
|
model_lib=parsed.model_lib,
|
|
debug_dir=Path(parsed.debug_dir),
|
|
device=parsed.device,
|
|
debug_instrument=instrument,
|
|
)
|
|
debug_chat.generate(parsed.prompt, parsed.generate_len)
|
|
# Only print decode for now
|
|
_print_as_table(
|
|
sorted(
|
|
instrument.time_eval_results.items(),
|
|
key=lambda x: -(x[1][0] * x[1][1]),
|
|
)
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|