chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,452 @@
|
||||
"""Continuous model delivery for MLC LLM models."""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Tuple, Type, TypeVar, Union # noqa: UP035
|
||||
|
||||
from huggingface_hub import HfApi, snapshot_download
|
||||
from huggingface_hub.utils import HfHubHTTPError
|
||||
from pydantic import BaseModel, Field, ValidationError
|
||||
|
||||
from mlc_llm.support import logging
|
||||
from mlc_llm.support.argparse import ArgumentParser
|
||||
from mlc_llm.support.style import bold, green, red
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
GEN_CONFIG_OPTIONAL_ARGS = [
|
||||
"context_window_size",
|
||||
"sliding_window_size",
|
||||
"prefill_chunk_size",
|
||||
"attention_sink_size",
|
||||
"tensor_parallel_shards",
|
||||
"pipeline_parallel_stages",
|
||||
]
|
||||
|
||||
T = TypeVar("T", bound="BaseModel")
|
||||
|
||||
|
||||
class OverrideConfigs(BaseModel):
|
||||
"""
|
||||
The class that specifies the override configurations.
|
||||
"""
|
||||
|
||||
context_window_size: Optional[int] = None
|
||||
sliding_window_size: Optional[int] = None
|
||||
prefill_chunk_size: Optional[int] = None
|
||||
attention_sink_size: Optional[int] = None
|
||||
tensor_parallel_shards: Optional[int] = None
|
||||
pipeline_parallel_stages: Optional[int] = None
|
||||
|
||||
|
||||
class ModelDeliveryTask(BaseModel):
|
||||
"""
|
||||
Example:
|
||||
{
|
||||
"model_id": "Phi-3-mini-128k-instruct",
|
||||
"model": "HF://microsoft/Phi-3-mini-128k-instruct",
|
||||
"conv_template": "phi-3",
|
||||
"quantization": ["q3f16_1"],
|
||||
"overrides": {
|
||||
"q3f16_1": {
|
||||
"context_window_size": 512
|
||||
}
|
||||
}
|
||||
}
|
||||
"""
|
||||
|
||||
model_id: str
|
||||
model: str
|
||||
conv_template: str
|
||||
quantization: Union[List[str], str] = Field(default_factory=list) # noqa: UP006
|
||||
overrides: Dict[str, OverrideConfigs] = Field(default_factory=dict) # noqa: UP006
|
||||
destination: Optional[str] = None
|
||||
gen_config_only: Optional[bool] = False
|
||||
|
||||
|
||||
class ModelDeliveryList(BaseModel):
|
||||
"""
|
||||
The class that specifies the model delivery list.
|
||||
"""
|
||||
|
||||
tasks: List[ModelDeliveryTask] # noqa: UP006
|
||||
# For delivered log, the default destination and quantization fields are optional
|
||||
default_destination: Optional[str] = None
|
||||
default_quantization: List[str] = Field(default_factory=list) # noqa: UP006
|
||||
default_overrides: Dict[str, OverrideConfigs] = Field(default_factory=dict) # noqa: UP006
|
||||
|
||||
@classmethod
|
||||
def from_json(cls: Type[T], json_dict: Dict[str, Any]) -> T: # noqa: UP006
|
||||
"""
|
||||
Convert from a json dictionary.
|
||||
"""
|
||||
try:
|
||||
return ModelDeliveryList.model_validate(json_dict)
|
||||
except ValidationError as e:
|
||||
logger.error("Error validating ModelDeliveryList: %s", e)
|
||||
raise e
|
||||
|
||||
def to_json(self) -> Dict[str, Any]: # noqa: UP006
|
||||
"""
|
||||
Convert to a json dictionary.
|
||||
"""
|
||||
return self.model_dump(exclude_none=True)
|
||||
|
||||
|
||||
def _clone_repo(model: Union[str, Path], hf_local_dir: Optional[str]) -> str:
|
||||
if isinstance(model, Path):
|
||||
if not model.exists():
|
||||
raise ValueError(f"Invalid model source: {model}")
|
||||
return str(model)
|
||||
prefixes, mlc_prefix = ["HF://", "https://huggingface.co/"], ""
|
||||
mlc_prefix = next(p for p in prefixes if model.startswith(p))
|
||||
if mlc_prefix:
|
||||
repo_name = model[len(mlc_prefix) :]
|
||||
model_name = repo_name.split("/")[-1]
|
||||
if hf_local_dir:
|
||||
hf_local_dir = os.path.join(hf_local_dir, model_name)
|
||||
logger.info("[HF] Downloading model to %s", hf_local_dir)
|
||||
return snapshot_download(repo_id=repo_name, local_dir=hf_local_dir)
|
||||
result = Path(model)
|
||||
if result.exists():
|
||||
return model
|
||||
raise ValueError(f"Invalid model source: {model}")
|
||||
|
||||
|
||||
def _run_quantization(
|
||||
model_info: ModelDeliveryTask,
|
||||
repo: str,
|
||||
api: HfApi,
|
||||
output_dir: str,
|
||||
) -> bool:
|
||||
logger.info("[HF] Creating repo https://huggingface.co/%s", repo)
|
||||
try:
|
||||
api.create_repo(repo_id=repo, private=False)
|
||||
except HfHubHTTPError as error:
|
||||
if error.response.status_code != 409:
|
||||
raise
|
||||
logger.info("[HF] Repo already exists. Skipping creation.")
|
||||
succeeded = True
|
||||
log_path = Path(output_dir) / "logs.txt"
|
||||
with log_path.open("a", encoding="utf-8") as log_file:
|
||||
assert isinstance(model_info.quantization, str)
|
||||
logger.info("[MLC] Processing in directory: %s", output_dir)
|
||||
# Required arguments
|
||||
cmd = [
|
||||
sys.executable,
|
||||
"-m",
|
||||
"mlc_llm",
|
||||
"gen_config",
|
||||
model_info.model,
|
||||
"--quantization",
|
||||
model_info.quantization,
|
||||
"--conv-template",
|
||||
model_info.conv_template,
|
||||
"--output",
|
||||
output_dir,
|
||||
]
|
||||
# Optional arguments
|
||||
for optional_arg in GEN_CONFIG_OPTIONAL_ARGS:
|
||||
optional_arg_val = getattr(model_info, optional_arg, None)
|
||||
if optional_arg_val is not None:
|
||||
# e.g. --context-window-size 4096
|
||||
cmd += ["--" + optional_arg.replace("_", "-"), str(optional_arg_val)]
|
||||
|
||||
print(" ".join(cmd), file=log_file, flush=True)
|
||||
subprocess.run(cmd, check=True, stdout=log_file, stderr=subprocess.STDOUT, env=os.environ)
|
||||
if not model_info.gen_config_only:
|
||||
cmd = [
|
||||
sys.executable,
|
||||
"-m",
|
||||
"mlc_llm",
|
||||
"convert_weight",
|
||||
str(model_info.model),
|
||||
"--quantization",
|
||||
model_info.quantization,
|
||||
"--output",
|
||||
output_dir,
|
||||
]
|
||||
print(" ".join(cmd), file=log_file, flush=True)
|
||||
subprocess.run(
|
||||
cmd,
|
||||
check=False,
|
||||
stdout=log_file,
|
||||
stderr=subprocess.STDOUT,
|
||||
env=os.environ,
|
||||
)
|
||||
logger.info("[MLC] Complete!")
|
||||
if not (Path(output_dir) / "tensor-cache.json").exists() and not model_info.gen_config_only:
|
||||
logger.error(
|
||||
"[%s] Model %s. Quantization %s. No weights metadata found.",
|
||||
red("FAILED"),
|
||||
model_info.model_id,
|
||||
model_info.quantization,
|
||||
)
|
||||
succeeded = False
|
||||
logger.info("[HF] Uploading to: https://huggingface.co/%s", repo)
|
||||
for _retry in range(10):
|
||||
try:
|
||||
api.upload_folder(
|
||||
folder_path=output_dir,
|
||||
repo_id=repo,
|
||||
ignore_patterns=["logs.txt"],
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error("[%s] %s. Retrying...", red("FAILED"), exc)
|
||||
else:
|
||||
break
|
||||
else:
|
||||
raise RuntimeError("Failed to upload to HuggingFace Hub with 10 retries")
|
||||
return succeeded
|
||||
|
||||
|
||||
def _get_current_log(log: str) -> ModelDeliveryList:
|
||||
log_path = Path(log)
|
||||
if not log_path.exists():
|
||||
with log_path.open("w", encoding="utf-8") as o_f:
|
||||
current_log = ModelDeliveryList(tasks=[])
|
||||
json.dump(current_log.to_json(), o_f, indent=4)
|
||||
else:
|
||||
with log_path.open("r", encoding="utf-8") as i_f:
|
||||
current_log = ModelDeliveryList.from_json(json.load(i_f))
|
||||
return current_log
|
||||
|
||||
|
||||
def _generate_model_delivery_diff(
|
||||
spec: ModelDeliveryList, log: ModelDeliveryList
|
||||
) -> ModelDeliveryList:
|
||||
diff_tasks = []
|
||||
default_quantization = spec.default_quantization
|
||||
default_overrides = spec.default_overrides
|
||||
|
||||
for task in spec.tasks:
|
||||
model_id = task.model_id
|
||||
conv_template = task.conv_template
|
||||
quantization = task.quantization
|
||||
overrides = {**default_overrides, **task.overrides}
|
||||
|
||||
logger.info(
|
||||
"Checking task: %s %s %s %s",
|
||||
model_id,
|
||||
conv_template,
|
||||
quantization,
|
||||
overrides,
|
||||
)
|
||||
log_tasks = [t for t in log.tasks if t.model_id == model_id]
|
||||
delivered_quantizations = set()
|
||||
gen_config_only = set()
|
||||
|
||||
for log_task in log_tasks:
|
||||
log_quantization = log_task.quantization
|
||||
assert isinstance(log_quantization, str)
|
||||
log_override = log_task.overrides.get(log_quantization, OverrideConfigs())
|
||||
override = overrides.get(log_quantization, OverrideConfigs())
|
||||
if log_override == override:
|
||||
if log_task.conv_template == conv_template:
|
||||
delivered_quantizations.add(log_quantization)
|
||||
else:
|
||||
gen_config_only.add(log_quantization)
|
||||
|
||||
all_quantizations = set(default_quantization) | set(quantization)
|
||||
quantization_diff = all_quantizations - set(delivered_quantizations)
|
||||
|
||||
if quantization_diff:
|
||||
for q in quantization_diff:
|
||||
logger.info("Adding task %s %s %s to the diff.", model_id, conv_template, q)
|
||||
task_copy = task.model_copy()
|
||||
task_copy.quantization = [q]
|
||||
task_copy.overrides = {q: overrides.get(q, OverrideConfigs())}
|
||||
task_copy.gen_config_only = task_copy.gen_config_only or q in gen_config_only
|
||||
diff_tasks.append(task_copy)
|
||||
else:
|
||||
logger.info("Task %s %s %s is up-to-date.", model_id, conv_template, quantization)
|
||||
|
||||
diff_config = spec.model_copy()
|
||||
diff_config.default_quantization = []
|
||||
diff_config.default_overrides = {}
|
||||
diff_config.tasks = diff_tasks
|
||||
|
||||
logger.info(
|
||||
"Model delivery diff: %s",
|
||||
diff_config.model_dump_json(indent=4, exclude_none=True),
|
||||
)
|
||||
|
||||
return diff_config
|
||||
|
||||
|
||||
def _main(
|
||||
username: str,
|
||||
api: HfApi,
|
||||
spec: ModelDeliveryList,
|
||||
log: str,
|
||||
hf_local_dir: Optional[str],
|
||||
output: str,
|
||||
dry_run: bool,
|
||||
):
|
||||
delivery_diff = _generate_model_delivery_diff(spec, _get_current_log(log))
|
||||
if dry_run:
|
||||
logger.info("Dry run. No actual delivery.")
|
||||
return
|
||||
|
||||
failed_cases: List[Tuple[str, str]] = [] # noqa: UP006
|
||||
delivered_log = _get_current_log(log)
|
||||
for task_index, task in enumerate(delivery_diff.tasks, 1):
|
||||
logger.info(
|
||||
bold("[{task_index}/{total_tasks}] Processing model: ").format(
|
||||
task_index=task_index,
|
||||
total_tasks=len(delivery_diff.tasks),
|
||||
)
|
||||
+ green(task.model_id)
|
||||
)
|
||||
model = _clone_repo(task.model, hf_local_dir)
|
||||
|
||||
quantizations = []
|
||||
|
||||
if delivery_diff.default_quantization:
|
||||
quantizations += delivery_diff.default_quantization
|
||||
|
||||
if task.quantization:
|
||||
if isinstance(task.quantization, str):
|
||||
quantizations.append(task.quantization)
|
||||
else:
|
||||
quantizations += task.quantization
|
||||
|
||||
default_destination = (
|
||||
delivery_diff.default_destination or "{username}/{model_id}-{quantization}-MLC"
|
||||
)
|
||||
for quantization in quantizations:
|
||||
repo = default_destination.format(
|
||||
username=username,
|
||||
model_id=task.model_id,
|
||||
quantization=quantization,
|
||||
)
|
||||
model_info = ModelDeliveryTask(
|
||||
model=model,
|
||||
quantization=quantization,
|
||||
destination=repo,
|
||||
**task.model_dump(exclude_none=True, exclude={"model", "quantization"}),
|
||||
)
|
||||
logger.info("Model info: %s", model_info.model_dump_json(indent=4))
|
||||
output_dir = os.path.join(
|
||||
output, f"{model_info.model_id}-{model_info.quantization}-MLC"
|
||||
)
|
||||
if not os.path.exists(output_dir):
|
||||
os.makedirs(output_dir)
|
||||
|
||||
result = _run_quantization(
|
||||
model_info=model_info,
|
||||
repo=repo,
|
||||
api=api,
|
||||
output_dir=output_dir,
|
||||
)
|
||||
if not result:
|
||||
failed_cases.append(
|
||||
(task.model_id, quantization),
|
||||
)
|
||||
else:
|
||||
delivered_log.tasks = [
|
||||
task
|
||||
for task in delivered_log.tasks
|
||||
if task.model_id != model_info.model_id
|
||||
or task.quantization != model_info.quantization
|
||||
]
|
||||
delivered_log.tasks.append(model_info)
|
||||
if failed_cases:
|
||||
logger.info("Total %s %s:", len(failed_cases), red("failures"))
|
||||
for model_id, quantization in failed_cases:
|
||||
logger.info(" Model %s. Quantization %s.", model_id, quantization)
|
||||
|
||||
delivered_log.tasks.sort(key=lambda task: task.model_id)
|
||||
logger.info("Writing log to %s", log)
|
||||
with open(log, "w", encoding="utf-8") as o_f:
|
||||
json.dump(delivered_log.to_json(), o_f, indent=4)
|
||||
|
||||
|
||||
def main():
|
||||
"""Entry point."""
|
||||
|
||||
def _load_spec(path_spec: str) -> ModelDeliveryList:
|
||||
path = Path(path_spec)
|
||||
if not path.exists():
|
||||
raise argparse.ArgumentTypeError(f"Spec file does not exist: {path}")
|
||||
with path.open("r", encoding="utf-8") as i_f:
|
||||
return ModelDeliveryList.from_json(json.load(i_f))
|
||||
|
||||
def _get_default_hf_token() -> str:
|
||||
# Try to get the token from the environment variable
|
||||
hf_token = os.getenv("HF_TOKEN")
|
||||
if hf_token:
|
||||
logger.info("HF token found in environment variable HF_TOKEN")
|
||||
return hf_token
|
||||
|
||||
# If not found, look for the token in the default cache folder
|
||||
token_file_path = os.path.expanduser("~/.cache/huggingface/token")
|
||||
if os.path.exists(token_file_path):
|
||||
with open(token_file_path, encoding="utf-8") as token_file:
|
||||
hf_token = token_file.read().strip()
|
||||
if hf_token:
|
||||
logger.info("HF token found in ~/.cache/huggingface/token")
|
||||
return hf_token
|
||||
|
||||
raise OSError("HF token not found")
|
||||
|
||||
parser = ArgumentParser("MLC LLM continuous model delivery")
|
||||
parser.add_argument(
|
||||
"--username",
|
||||
type=str,
|
||||
required=True,
|
||||
help="HuggingFace username",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--token",
|
||||
type=str,
|
||||
default=_get_default_hf_token(),
|
||||
help="HuggingFace access token, obtained under https://huggingface.co/settings/tokens",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--spec",
|
||||
type=_load_spec,
|
||||
default="model-delivery-config.json",
|
||||
help="Path to the model delivery file" + ' (default: "%(default)s")',
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log",
|
||||
type=str,
|
||||
default="model-delivered-log.json",
|
||||
help="Path to the output log file" + ' (default: "%(default)s")',
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Directory to store the output MLC models",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hf-local-dir",
|
||||
type=str,
|
||||
required=False,
|
||||
help="Local directory to store the downloaded HuggingFace model",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dry-run",
|
||||
action="store_true",
|
||||
help="Dry run without uploading to HuggingFace Hub",
|
||||
)
|
||||
parsed = parser.parse_args()
|
||||
_main(
|
||||
parsed.username,
|
||||
spec=parsed.spec,
|
||||
log=parsed.log,
|
||||
api=HfApi(token=parsed.token),
|
||||
hf_local_dir=parsed.hf_local_dir,
|
||||
output=parsed.output,
|
||||
dry_run=parsed.dry_run,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user