wehub-resource-sync
561 lines
20 KiB
Python
561 lines
20 KiB
Python
"""Debug compiled models with TVM instrument"""
|
|
|
|
import json
|
|
import random
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional, Tuple, Union # noqa: UP035
|
|
|
|
import numpy as np
|
|
import tvm
|
|
import tvm_ffi
|
|
from tvm import DataType, relax
|
|
from tvm.contrib import tvmjs
|
|
from tvm.runtime import Device, Module, Object
|
|
from tvm.runtime.vm import VirtualMachine
|
|
from tvm_ffi import Shape
|
|
|
|
from mlc_llm.conversation_template import ConvTemplateRegistry
|
|
from mlc_llm.interface.help import HELP
|
|
from mlc_llm.protocol.mlc_chat_config import MLCChatConfig
|
|
from mlc_llm.serve import data, engine_utils
|
|
from mlc_llm.support.argparse import ArgumentParser
|
|
from mlc_llm.support.auto_device import detect_device
|
|
from mlc_llm.support.style import green, red
|
|
from mlc_llm.tokenizers import Tokenizer
|
|
|
|
|
|
def _extract_metadata(mod: Module):
|
|
return json.loads(VirtualMachine(mod, tvm.runtime.device("cpu"))["_metadata"]())
|
|
|
|
|
|
def _load_params(
|
|
model_weight_path: str,
|
|
device: Device,
|
|
model_metadata: Dict[str, Any], # noqa: UP006
|
|
) -> List[tvm.runtime.Tensor]: # noqa: UP006
|
|
params, meta = tvmjs.load_tensor_cache(model_weight_path, device)
|
|
param_names = [param["name"] for param in model_metadata["params"]]
|
|
assert len(param_names) == meta["ParamSize"]
|
|
|
|
plist = []
|
|
for param_name in param_names:
|
|
plist.append(params[param_name])
|
|
return plist
|
|
|
|
|
|
def _get_tvm_module(
|
|
model_weight_path: str,
|
|
lib_path: str,
|
|
device: Device,
|
|
instrument: Union[tvm_ffi.Function, None],
|
|
):
|
|
ex = tvm.runtime.load_module(lib_path)
|
|
vm = relax.VirtualMachine(ex, device)
|
|
if instrument is not None:
|
|
vm.set_instrument(instrument)
|
|
metadata = _extract_metadata(ex)
|
|
params = _load_params(model_weight_path, device, metadata)
|
|
return vm.module, params, metadata
|
|
|
|
|
|
class DefaultDebugInstrument:
|
|
"""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.
|
|
"""
|
|
|
|
def __init__(self, debug_out: Path):
|
|
"""Constructor
|
|
|
|
Parameters
|
|
----------
|
|
debug_out : Path
|
|
the directory to dump the .npz files
|
|
"""
|
|
self.counter = 0
|
|
self.first_nan_occurred = False
|
|
self.first_inf_occurred = False
|
|
self.debug_out = debug_out
|
|
debug_out.mkdir(exist_ok=True, parents=True)
|
|
|
|
def reset(self, debug_out: Path):
|
|
"""Reset the state of the Instrument class
|
|
|
|
Parameters
|
|
----------
|
|
debug_out : Path
|
|
the directory to dump the .npz files
|
|
"""
|
|
self.counter = 0
|
|
self.first_nan_occurred = False
|
|
self.first_inf_occurred = False
|
|
self.debug_out = debug_out
|
|
debug_out.mkdir(exist_ok=True, parents=True)
|
|
|
|
def __call__(self, func, name, before_run, ret_val, *args):
|
|
# Determine what functions to look at
|
|
if before_run: # Whether before the function is called or after
|
|
return
|
|
if self.first_nan_occurred:
|
|
return
|
|
if self.first_inf_occurred:
|
|
return
|
|
if (
|
|
name.startswith("vm.builtin.")
|
|
and "call_tir_dyn" not in name
|
|
and "attention_with_fused_qkv" not in name
|
|
and "self_attention" not in name
|
|
and "cross_attention" not in name
|
|
):
|
|
return
|
|
|
|
# Decide what to print or save about the function's arguments (where args[-1] is the
|
|
# buffer we write the result to)
|
|
func_name = f"f{self.counter}_{name}"
|
|
|
|
# Write your own behavior below. For example, we can count the number of INF/NaN in args[-1]
|
|
def _check_nan_inf(npy):
|
|
num_nans = np.sum(np.isnan(npy))
|
|
num_infs = np.sum(np.isinf(npy))
|
|
if num_nans > 0:
|
|
print(f"{red(f'{func_name} has NaN')}: {num_nans}")
|
|
self.first_nan_occurred = True
|
|
if num_infs > 0:
|
|
print(f"{red(f'{func_name} has INF')}: {num_infs}")
|
|
self.first_inf_occurred = True
|
|
|
|
# Save the arguments to npz
|
|
arg_dict = {}
|
|
for i, arg in enumerate(args):
|
|
if isinstance(arg, tvm.runtime.Tensor):
|
|
if np.prod(arg.shape) * (DataType(arg.dtype).bits // 8) > 2147483648:
|
|
# We skip dump large tensors
|
|
arg_dict[f"arg_{i}"] = np.zeros(())
|
|
elif arg.dtype in ["bfloat16", "float8_e4m3fn"]:
|
|
arg_dict[f"arg_{i}"] = arg.numpy().astype(np.float32)
|
|
else:
|
|
arg_dict[f"arg_{i}"] = arg.numpy()
|
|
_check_nan_inf(arg.numpy())
|
|
np.savez(self.debug_out / f"{func_name}.npz", **arg_dict)
|
|
|
|
self.counter += 1
|
|
|
|
|
|
class DebugChat:
|
|
"""A chat interface used only for debugging purpose.
|
|
|
|
It debugs auto-regressive decoding fully in Python via the prefill and
|
|
decode interface. It supports debugging instrument (either default or
|
|
customized) to dump intermediate values for each VM function call.
|
|
|
|
Given a prompt, it also prints out the parsed prompt, input tokens, output
|
|
tokens and output text.
|
|
|
|
Sample usage:
|
|
|
|
dc = DebugChat(
|
|
model="./dist/Llama-2-7b-chat-hf-q4f16_1-MLC",
|
|
debug_dir=Path("./debug-llama-2"),
|
|
model_lib="./dist/llama-2-7b-chat-q4f16_1-metal.so",
|
|
)
|
|
dc.generate("hello world", 3)
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
model: str,
|
|
model_lib: str,
|
|
debug_dir: Path,
|
|
device: Optional[str] = "auto",
|
|
debug_instrument: Optional[Any] = None,
|
|
is_image_model: Optional[bool] = False,
|
|
disable_instrument: Optional[bool] = False,
|
|
):
|
|
"""_summary_
|
|
|
|
Parameters
|
|
----------
|
|
model: str
|
|
The model folder after compiling with MLC-LLM build process. The parameter
|
|
can either be the model name with its quantization scheme
|
|
(e.g. ``Llama-2-7b-chat-hf-q4f16_1``), or a full path to the model
|
|
folder. In the former case, we will use the provided name to search
|
|
for the model folder over possible paths.
|
|
|
|
model_lib : str
|
|
The full path to the model library file to use (e.g. a ``.so`` file).
|
|
|
|
debug_dir: Path
|
|
The output folder to store the dumped debug files.
|
|
|
|
device : Optional[str]
|
|
The description of the device to run on. User should provide a string in the
|
|
form of 'device_name:device_id' or 'device_name', where 'device_name' is one of
|
|
'cuda', 'metal', 'vulkan', 'rocm', 'opencl', 'auto' (automatically detect the
|
|
local device), and 'device_id' is the device id to run on. If no 'device_id'
|
|
is provided, it will be set to 0 by default.
|
|
|
|
chat_config : Optional[ChatConfig]
|
|
A ``ChatConfig`` instance partially filled. Will be used to override the
|
|
``mlc-chat-config.json``.
|
|
|
|
debug_instrument : Optional[Any]
|
|
An instrument function that will be called before/after each Call instruction.
|
|
The function have the following signature:
|
|
|
|
.. code:: python
|
|
|
|
def instrument(
|
|
func: Union[VMClosure, Function],
|
|
func_symbol: str,
|
|
before_run: bool,
|
|
ret_value: any,
|
|
*args) -> bool:
|
|
pass
|
|
|
|
The instrument takes the following parameters:
|
|
- func: function object to be called.
|
|
- func_symbol: the symbol name of the function.
|
|
- before_run: whether it is before or after call.
|
|
- ret_value: the return value of the call, only valid after run.
|
|
- args: the arguments being passed to call.
|
|
|
|
is_image_model: Optional[bool]
|
|
Whether the model support image input. If so, will look for image embedding method.
|
|
Default to False.
|
|
|
|
disable_instrument: Optional[bool]
|
|
If true, will not use debug instrument for faster generation. Default to False.
|
|
"""
|
|
self.debug_dir = debug_dir
|
|
self.device = detect_device(device)
|
|
if disable_instrument:
|
|
self.instrument = None
|
|
else:
|
|
self.instrument = (
|
|
debug_instrument
|
|
if debug_instrument
|
|
else DefaultDebugInstrument(debug_dir / "prefill")
|
|
)
|
|
self.mod, self.params, self.metadata = _get_tvm_module(
|
|
model, model_lib, self.device, self.instrument
|
|
)
|
|
self.model_path = Path(model)
|
|
self.config_file_path = self.model_path / "mlc-chat-config.json"
|
|
with open(self.config_file_path, encoding="utf-8") as file:
|
|
self.chat_config = MLCChatConfig.model_validate_json(file.read())
|
|
|
|
conv_template = self.chat_config.conv_template
|
|
|
|
self.conversation = (
|
|
ConvTemplateRegistry.get_conv_template(conv_template)
|
|
if isinstance(conv_template, str)
|
|
else conv_template
|
|
)
|
|
self.tokenizer = Tokenizer(str(self.model_path))
|
|
|
|
self.add_sequence_func = tvm.get_global_func("vm.builtin.kv_state_add_sequence")
|
|
self.begin_forward_func = tvm.get_global_func("vm.builtin.kv_state_begin_forward")
|
|
self.end_forward_func = tvm.get_global_func("vm.builtin.kv_state_end_forward")
|
|
self.nd_view_func = tvm.get_global_func("vm.builtin.reshape")
|
|
self.sample_topp_from_prob_func = tvm.get_global_func("vm.builtin.sample_top_p_from_prob")
|
|
|
|
try:
|
|
self.embed_func = self.mod["embed"]
|
|
except AttributeError as exc:
|
|
raise RuntimeError("DebugChat only supports separate embedding layer") from exc
|
|
|
|
if is_image_model:
|
|
try:
|
|
self.embed_image_func = self.mod["image_embed"]
|
|
except AttributeError as exc:
|
|
raise RuntimeError(
|
|
"Expect the model to be an image model, but cannot find `image_embed`."
|
|
) from exc
|
|
|
|
self.prefill_func = self.mod["prefill"]
|
|
self.decode_func = self.mod["decode"]
|
|
self.create_kv_cache_func = None
|
|
if self.mod.implements_function("create_flashinfer_paged_kv_cache"):
|
|
self.create_kv_cache_func = self.mod["create_flashinfer_paged_kv_cache"]
|
|
elif self.mod.implements_function("create_tir_paged_kv_cache"):
|
|
self.create_kv_cache_func = self.mod["create_tir_paged_kv_cache"]
|
|
else:
|
|
# TODO: Support RNN KVState
|
|
raise RuntimeError("DebugChat cannot find create KV cache function")
|
|
|
|
self.appeared_token_freq: Dict[int, int] = {} # noqa: UP006
|
|
|
|
def _preprocess_prompts(
|
|
self, prompt: str, image_url: Optional[str] = None
|
|
) -> List[Union[List[int], data.ImageData]]: # noqa: UP006
|
|
print("======================= Starts Tokenization & Embedding =======================")
|
|
# Step 0. Generate prompt string using conversation template
|
|
if image_url is None:
|
|
self.conversation.messages.append(("user", prompt))
|
|
else:
|
|
self.conversation.messages.append(
|
|
(
|
|
"user",
|
|
[
|
|
{"type": "image_url", "image_url": image_url},
|
|
{"type": "text", "text": prompt},
|
|
],
|
|
)
|
|
)
|
|
self.conversation.messages.append(("assistant", None))
|
|
|
|
with open(self.config_file_path, encoding="utf-8") as file:
|
|
config = json.load(file)
|
|
parsed_prompt = self.conversation.as_prompt(config)
|
|
print(
|
|
"Parsed prompt using conversation template "
|
|
f"{green(self.conversation.name)}: {parsed_prompt}"
|
|
)
|
|
tokens = engine_utils.process_prompts(parsed_prompt, self.tokenizer.encode)
|
|
|
|
if self.conversation.system_prefix_token_ids is not None:
|
|
tokens[0] = self.conversation.system_prefix_token_ids + tokens[0]
|
|
|
|
return tokens
|
|
|
|
def _embed(
|
|
self,
|
|
data_inputs: List[Union[List[int], data.ImageData]], # noqa: UP006
|
|
) -> Tuple[tvm.runtime.Tensor, int]: # noqa: UP006
|
|
# We currently convert to numpy after embedded, concat in numpy, then convert back to
|
|
# tvm tensor; could be more optimized; but may suffice for debug purposes.
|
|
embeddings = []
|
|
for data_input in data_inputs:
|
|
if isinstance(data_input, data.ImageData):
|
|
# Process image data
|
|
# print(f"data_input.get_embed_size(): {data_input.embed_size}")
|
|
image_input = data_input.image
|
|
if data_input.image.device != self.device:
|
|
image_input = data_input.image.copyto(self.device)
|
|
embeddings.append(self.embed_image_func(image_input, self.params).asnumpy())
|
|
else:
|
|
# Process token data
|
|
data_input = tvm.runtime.tensor(
|
|
np.array(data_input).astype("int32"), device=self.device
|
|
)
|
|
embeddings.append(self.embed_func(data_input, self.params).asnumpy())
|
|
# for embedding in embeddings:
|
|
# print(f"embedding.shape: {embedding.shape}")
|
|
|
|
# Concatenate
|
|
concat_embeddings = tvm.runtime.tensor(
|
|
np.concatenate(embeddings, axis=0), device=self.device
|
|
)
|
|
concat_embeddings = self.nd_view_func(
|
|
concat_embeddings,
|
|
Shape([1, concat_embeddings.shape[0], concat_embeddings.shape[1]]),
|
|
)
|
|
input_len = concat_embeddings.shape[1]
|
|
|
|
return concat_embeddings, input_len
|
|
|
|
def _prefill(self, embedding: tvm.runtime.Tensor, input_len: int):
|
|
print("======================= Starts Prefill =======================")
|
|
seq_len_shape = Shape([input_len])
|
|
max_num_sequence = 1
|
|
page_size = 16
|
|
sliding_window_size = (
|
|
self.chat_config.sliding_window_size
|
|
if self.chat_config.sliding_window_size
|
|
else self.metadata["sliding_window_size"]
|
|
)
|
|
context_window_size = (
|
|
self.chat_config.context_window_size
|
|
if self.chat_config.context_window_size
|
|
else self.metadata["context_window_size"]
|
|
)
|
|
prefill_chunk_size = (
|
|
self.chat_config.prefill_chunk_size
|
|
if self.chat_config.prefill_chunk_size
|
|
else self.metadata["prefill_chunk_size"]
|
|
)
|
|
max_total_sequence_length = (
|
|
sliding_window_size if context_window_size == -1 else context_window_size
|
|
)
|
|
support_sliding_window = int(sliding_window_size != -1)
|
|
|
|
kv_caches = self.create_kv_cache_func(
|
|
Shape([max_num_sequence]),
|
|
Shape([max_total_sequence_length]),
|
|
Shape([prefill_chunk_size]),
|
|
Shape([page_size]),
|
|
Shape([support_sliding_window]),
|
|
)
|
|
self.add_sequence_func(kv_caches, 0)
|
|
self.begin_forward_func(kv_caches, Shape([0]), seq_len_shape)
|
|
logits, kv_caches = self.prefill_func(embedding, kv_caches, self.params)
|
|
self.end_forward_func(kv_caches)
|
|
return logits, kv_caches
|
|
|
|
def _decode(self, token: int, kv_caches: Object):
|
|
embedding, _ = self._embed([[token]])
|
|
self.begin_forward_func(kv_caches, Shape([0]), Shape([1]))
|
|
logits, kv_caches = self.decode_func(embedding, kv_caches, self.params)
|
|
self.end_forward_func(kv_caches)
|
|
return logits
|
|
|
|
def _softmax_with_temperature(self, logits: np.ndarray, temperature: float):
|
|
# Adjust logits based on the temperature
|
|
logits = np.array(logits) / temperature
|
|
logits -= np.max(logits, axis=-1, keepdims=True)
|
|
|
|
exp_logits = np.exp(logits, logits)
|
|
exp_logits /= np.sum(exp_logits, axis=-1, keepdims=True)
|
|
return exp_logits
|
|
|
|
def _apply_presence_and_freq_penalty(
|
|
self, logits: np.ndarray, presence_penalty: float, freq_penalty: float
|
|
):
|
|
for token_id, freq in self.appeared_token_freq.items():
|
|
logits[:, :, token_id] -= freq * freq_penalty + presence_penalty
|
|
|
|
def _sample_token_from_logits(
|
|
self,
|
|
logits: tvm.runtime.Tensor,
|
|
*,
|
|
temperature=1.0,
|
|
top_p=1.0,
|
|
presence_penalty=0.0,
|
|
frequency_penalty=0.0,
|
|
):
|
|
logits_np = logits.numpy()
|
|
|
|
if presence_penalty != 0.0 or frequency_penalty != 0.0:
|
|
self._apply_presence_and_freq_penalty(logits_np, presence_penalty, frequency_penalty)
|
|
|
|
logits_np = self._softmax_with_temperature(logits_np, temperature)
|
|
if self.instrument is not None:
|
|
np.savez(self.instrument.debug_out / "logits.npz", logits_np)
|
|
|
|
logits = logits.copyfrom(logits_np)
|
|
next_token = self.sample_topp_from_prob_func(logits, top_p, random.random())
|
|
return next_token
|
|
|
|
def generate(
|
|
self,
|
|
prompt: str,
|
|
generate_length: int,
|
|
image_url: Optional[str] = None,
|
|
):
|
|
"""Generates the response from the model given a user prompt. User will need to
|
|
specify the generation length for debugging purpose. For example, a generation
|
|
length of 3 will include 1 prefill step and 2 decode steps.
|
|
|
|
Parameters
|
|
----------
|
|
prompt : str
|
|
The user input prompt.
|
|
|
|
generate_length : int
|
|
How many tokens to generate.
|
|
"""
|
|
out_tokens = []
|
|
|
|
data_inputs = self._preprocess_prompts(prompt, image_url)
|
|
print(f"{green('Data inputs: ')}: {data_inputs}")
|
|
embedding, input_len = self._embed(data_inputs)
|
|
logits, kv_caches = self._prefill(embedding, input_len)
|
|
next_token = self._sample_token_from_logits(logits)
|
|
out_tokens.append(next_token)
|
|
if self.instrument is not None:
|
|
path_str = (self.debug_dir / "prefill").as_posix()
|
|
print(f"Debug instrument output dumped to {green(path_str)}")
|
|
|
|
print("======================= Starts Decode =======================")
|
|
for i in range(generate_length - 1):
|
|
if self.instrument is not None:
|
|
self.instrument.reset(self.debug_dir / f"decode_{i}")
|
|
logits = self._decode(next_token, kv_caches)
|
|
next_token = self._sample_token_from_logits(logits)
|
|
out_tokens.append(next_token)
|
|
if self.instrument is not None:
|
|
path_str = (self.debug_dir / f"decode_{i}").as_posix()
|
|
print(f"Debug instrument output dumped to {green(path_str)}")
|
|
|
|
if next_token in self.conversation.stop_token_ids:
|
|
break
|
|
|
|
print(f"{green('Generated output tokens')}: {np.array(out_tokens)}")
|
|
|
|
out_text = self.tokenizer.decode(out_tokens)
|
|
print(f"{green('Generated output text')}: {out_text}")
|
|
|
|
|
|
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(
|
|
"--image-url",
|
|
type=str,
|
|
required=False,
|
|
help="Image to prefill into the model, can only be set for image models",
|
|
)
|
|
parser.add_argument(
|
|
"--disable-instrument",
|
|
action="store_true",
|
|
help=(
|
|
"Disable dumping customizable detailed information of kernel input "
|
|
+ "and output, hence making generation faster."
|
|
),
|
|
)
|
|
parsed = parser.parse_args()
|
|
dc = DebugChat(
|
|
model=parsed.model,
|
|
model_lib=parsed.model_lib,
|
|
debug_dir=Path(parsed.debug_dir),
|
|
device=parsed.device,
|
|
is_image_model=parsed.image_url is not None,
|
|
disable_instrument=parsed.disable_instrument,
|
|
)
|
|
|
|
dc.generate(parsed.prompt, parsed.generate_len, parsed.image_url)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|