mlc-ai--mlc-llm/python/mlc_llm/cli/model_metadata.py

"""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()