chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,560 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user