Implementation:Mlc ai Mlc llm Auto Weight

Overview

Auto Weight is a support module in MLC LLM that provides helper functions for automatically detecting weight file paths and weight storage formats. It is located at python/mlc_llm/support/auto_weight.py (178 lines).

When converting or loading model weights, MLC LLM needs to determine both where the weight files reside and what format they are stored in (PyTorch, SafeTensors, AWQ, etc.). This module centralizes that detection logic, offering both explicit format specification and automatic guessing.

Purpose

The module serves two primary purposes:

1. Weight path resolution -- Given a configuration JSON path, locate the directory containing model weight files. The path can be explicitly provided, extracted from a config.json field, or defaulted to the directory containing the config file.
2. Weight format detection -- Determine whether weights are stored as HuggingFace PyTorch checkpoints, HuggingFace SafeTensors, or AWQ quantized weights, either by explicit specification or by probing the filesystem.

Architecture

The module follows a strategy pattern where each supported weight format has a corresponding checker function registered in the CHECK_FORMAT_METHODS dictionary. The main entry point delegates to these checkers for both validation and auto-detection.

Key Constants

CHECK_FORMAT_METHODS = {
    "huggingface-torch": _check_pytorch,
    "huggingface-safetensor": _check_safetensor,
}

# "ggml", "gguf" are not supported yet.
AVAILABLE_WEIGHT_FORMAT = ["huggingface-torch", "huggingface-safetensor", "awq"]

The FOUND and NOT_FOUND constants are styled (green/red) log markers used for user-facing output:

FOUND = green("Found")
NOT_FOUND = red("Not found")

Core Functions

detect_weight

def detect_weight(
    weight_path: Path,
    config_json_path: Path,
    weight_format: str = "auto",
) -> Tuple[Path, str]:

This is the primary entry point. It accepts a weight path, a config JSON path, and an optional format hint. The function proceeds as follows:

1. If weight_path is None, it attempts to find the path from the "weight_path" key inside the config JSON, falling back to the config file's parent directory.
2. If weight_path is provided, it validates that the path exists.
3. If weight_format is "auto", it delegates to _guess_weight_format.
4. Otherwise, it validates the explicitly specified format using the corresponding checker function from CHECK_FORMAT_METHODS.

Returns: A tuple of (weight_config_path, weight_format) where weight_config_path points to the weight index file (e.g., pytorch_model.bin.index.json) or the weight directory.

_guess_weight_format

def _guess_weight_format(weight_path: Path) -> Tuple[Path, str]:

Iterates through all registered format checkers in CHECK_FORMAT_METHODS and collects all formats whose files are present in the weight directory. If no format matches, a ValueError is raised advising the user to specify --source-format explicitly. When multiple formats match, the first one found is selected.

_check_pytorch

def _check_pytorch(weight_path: Path) -> Optional[Path]:

Probes for HuggingFace PyTorch weight files by checking for:

• pytorch_model.bin.index.json (sharded weights index)
• pytorch_model.bin (single-file weights)

Returns the path to the found file, or None if neither exists.

_check_safetensor

def _check_safetensor(weight_path: Path) -> Optional[Path]:

Probes for HuggingFace SafeTensors weight files by checking for:

• model.safetensors.index.json (sharded weights index)
• model.safetensors (single-file weights)

When a single model.safetensors file is found without an accompanying index JSON, the function generates the index file on the fly. It loads the SafeTensors file, extracts the key names, and writes a model.safetensors.index.json containing the weight map:

from safetensors.torch import load_file
weights = load_file(safetensor_file_path, device="cpu")
weight_map = {key: "model.safetensors" for key in weights}
with open(safetensor_json_path, "w", encoding="utf-8") as file:
    json.dump({"weight_map": weight_map}, file, indent=2)

Supported Weight Formats

Format	Identifier	Detection Method
HuggingFace PyTorch	huggingface-torch	Checks for pytorch_model.bin.index.json or pytorch_model.bin
HuggingFace SafeTensors	huggingface-safetensor	Checks for model.safetensors.index.json or model.safetensors
AWQ	awq	Listed as available but has no dedicated checker; validated by path existence only

Dependencies

• json, pathlib.Path -- Standard library for JSON parsing and path manipulation
• mlc_llm.support.logging -- Internal logging module
• mlc_llm.support.style -- Terminal styling helpers (bold, green, red)
• safetensors.torch -- Conditionally imported when generating SafeTensors index files

File Location

python/mlc_llm/support/auto_weight.py