chore: import upstream snapshot with attribution

python/mlc_llm/cli/model_metadata.py

@@ -0,0 +1,194 @@
+"""A tool that inspects the metadata of a model lib."""
+
+import json
+import math
+from dataclasses import asdict
+from pathlib import Path
+from typing import Any, Dict, List, Union  # noqa: UP035
+
+from tvm.runtime import DataType
+
+from mlc_llm.support import logging
+from mlc_llm.support.argparse import ArgumentParser
+from mlc_llm.support.config import ConfigBase
+from mlc_llm.support.style import green, red
+
+logger = logging.getLogger(__name__)
+
+
+def _extract_metadata(model_lib: Path) -> Dict[str, Any]:  # noqa: UP006
+    from tvm.runtime import device, load_module
+    from tvm.runtime.vm import VirtualMachine
+
+    return json.loads(VirtualMachine(load_module(model_lib), device("cpu"))["_metadata"]())
+
+
+def _report_all(metadata: Dict[str, Any]) -> None:  # noqa: UP006
+    # Print JSON with aesthetic values that packs each parameter into one line,
+    # while keeping the rest indented.
+    indent = 2
+    indents = " " * indent
+    params = metadata.pop("params")
+    params = indents * 2 + (",\n" + indents * 2).join(json.dumps(p) for p in params)
+    lines = json.dumps(
+        metadata,
+        sort_keys=True,
+        indent=indent,
+    ).splitlines()
+    lines.insert(1, indents + '"params": [\n' + params + "\n" + indents + "],")
+    beautified_json = "\n".join(lines)
+    print(beautified_json)
+
+
+def _read_dynamic_shape(shape: List[Union[int, str]], config: Union[Dict, ConfigBase]) -> List[int]:  # noqa: UP006
+    if isinstance(config, ConfigBase):
+        config = asdict(config)
+    param_shape = []
+    for s in shape:
+        if isinstance(s, int):
+            param_shape.append(s)
+        else:
+            if config is None:
+                logger.error(
+                    "%s: Encountered dynamic shape %s, need to specify `--mlc-chat-config` for "
+                    + "memory usage calculation.",
+                    red("FAILED"),
+                    red(s),
+                )
+                raise AttributeError
+            if s not in config:
+                logger.error(
+                    "%s to retrieve concrete %s for dynamic shape from %s.",
+                    red("FAILED"),
+                    red(s),
+                    config,
+                )
+                raise KeyError
+            param_shape.append(config[s])
+    return param_shape
+
+
+def _compute_memory_usage(metadata: Dict[str, Any], config: Union[Dict, ConfigBase]):  # noqa: UP006
+    params_bytes = 0.0
+    for param in metadata["params"]:
+        if all(isinstance(v, int) for v in param["shape"]):
+            assert all(v > 0 for v in param["shape"]), "All shapes should be strictly positive."
+            param_shape = param["shape"]
+        else:
+            # Contains dynamic shape; use config to look up concrete values
+            param_shape = _read_dynamic_shape(param["shape"], config)
+        params_bytes += math.prod(param_shape) * DataType(param["dtype"]).itemsize
+    temp_func_bytes = 0.0
+    for _func_name, func_bytes in metadata["memory_usage"].items():
+        temp_func_bytes = max(temp_func_bytes, func_bytes)
+
+    return params_bytes, temp_func_bytes
+
+
+def _report_memory_usage(metadata: Dict[str, Any], config: Union[Dict, ConfigBase]) -> None:  # noqa: UP006
+    params_bytes, temp_func_bytes = _compute_memory_usage(metadata, config)
+    total_size = params_bytes + temp_func_bytes
+    logger.info(
+        "%s: %.2f MB (Parameters: %.2f MB. Temporary buffer: %.2f MB)",
+        green("Total memory usage without KV cache"),
+        total_size / 1024 / 1024,
+        params_bytes / 1024 / 1024,
+        temp_func_bytes / 1024 / 1024,
+    )
+
+    # Compute KV cache size per token of context window.
+    if isinstance(config, ConfigBase):
+        config = asdict(config)
+    if (
+        "head_dim" in config
+        and "num_hidden_layers" in config
+        and "num_key_value_heads" in config
+        and "quantization" in metadata
+    ):
+        quantization_type = metadata["quantization"]
+        dtype_bytes = None
+        if "f32" in quantization_type:
+            dtype_bytes = 4
+        elif "bf16" in quantization_type:
+            dtype_bytes = 2
+        elif "f16" in quantization_type:
+            dtype_bytes = 2
+        # TODO: If support quantized KV in future, need to change this
+        if dtype_bytes is not None:
+            bytes_per_token = (
+                config["head_dim"]
+                * config["num_hidden_layers"]
+                * config["num_key_value_heads"]
+                * dtype_bytes
+                * 2  # 2 for key and value
+            )
+            logger.info(
+                "%s: %.2f MB per token in the context window",
+                green("KV cache size"),
+                bytes_per_token / 1024 / 1024,
+            )
+            logger.info(
+                "%s: %.2f MB",
+                green("Total memory usage with a 4K KV cache"),
+                (total_size + bytes_per_token * 4096) / 1024 / 1024,
+            )
+
+    logger.info(
+        "To reduce memory usage, "
+        "tweak `prefill_chunk_size`, `context_window_size` and `sliding_window_size`"
+    )
+
+
+def main():
+    """Entry point for the model metadata tool."""
+    parser = ArgumentParser(description="A tool that inspects the metadata of a model lib.")
+    parser.add_argument(
+        "model_lib",
+        type=Path,
+        help="""The compiled model library. In MLC LLM, an LLM is compiled to a shared or static
+        library (.so or .a), which contains GPU computation to efficiently run the LLM. MLC Chat,
+        as the runtime of MLC LLM, depends on the compiled model library to generate tokens.
+        """,
+    )
+    parser.add_argument(
+        "--mlc-chat-config",
+        type=Path,
+        help="""The `mlc-chat-config.json` file specific to a model variant. This is only required
+        when `memory-only` is true and `model_lib` contains a dynamic parameter shape (i.e. using
+        a variable to represent the shape). For instance, `model.embed_tokens.q_weight` can have
+        shape `["vocab_size", 512]`. In these cases, we look up the concrete value in
+        `mlc-chat-config.json`.
+        """,
+    )
+    parser.add_argument(
+        "--memory-only",
+        action="store_true",
+        help="""If set, only inspect the metadata in memory usage and print richer analysis.
+        Otherwise, the tool will load all the metadata from the model library file but only print
+        the basic information in JSON.
+        """,
+    )
+    parsed = parser.parse_args()
+    # Load metadata from model lib
+    try:
+        metadata = _extract_metadata(parsed.model_lib)
+    except Exception:
+        logger.exception("%s to read metadata section in legacy model lib.", red("FAILED"))
+        return
+    # Load mlc_chat_config if provided
+    cfg = None
+    if parsed.mlc_chat_config:
+        mlc_chat_config_path = Path(parsed.mlc_chat_config)
+        if not mlc_chat_config_path.exists():
+            raise ValueError(f"{mlc_chat_config_path} does not exist.")
+        with open(mlc_chat_config_path, encoding="utf-8") as config_file:
+            cfg = json.load(config_file)
+    # Main body
+    if parsed.memory_only:
+        _report_memory_usage(metadata, cfg)
+    else:
+        _report_all(metadata)
+
+
+if __name__ == "__main__":
+    main()