chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,265 @@
|
||||
"""Python entrypoint of compilation."""
|
||||
|
||||
import dataclasses
|
||||
from io import StringIO
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple # noqa: UP035
|
||||
|
||||
from tvm import IRModule, relax, tirx
|
||||
from tvm.ir.transform import Pass, PassContext
|
||||
from tvm.relax.frontend import nn
|
||||
from tvm.target import Target
|
||||
|
||||
from mlc_llm import compiler_pass as _ # noqa: F401
|
||||
from mlc_llm import op as op_ext
|
||||
from mlc_llm.cli.model_metadata import _report_memory_usage
|
||||
from mlc_llm.model import Model
|
||||
from mlc_llm.quantization import Quantization
|
||||
from mlc_llm.support import logging
|
||||
from mlc_llm.support.config import ConfigBase
|
||||
from mlc_llm.support.style import bold
|
||||
|
||||
from .compiler_flags import ModelConfigOverride, OptimizationFlags
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class CompileArgs:
|
||||
"""Arguments to MLC LLM's compiler."""
|
||||
|
||||
config: Path
|
||||
quantization: Quantization
|
||||
model: Model
|
||||
target: Target
|
||||
opt: OptimizationFlags
|
||||
build_func: Callable[[IRModule, "CompileArgs", Pass], None]
|
||||
system_lib_prefix: str
|
||||
output: Path
|
||||
overrides: ModelConfigOverride
|
||||
debug_dump: Optional[Path]
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self.opt.update(self.target, self.quantization)
|
||||
|
||||
def display(self) -> None:
|
||||
"""Display the arguments to stdout."""
|
||||
out = StringIO()
|
||||
print(f"{bold('Compiling with arguments:')}", file=out)
|
||||
print(f" {bold('--config'):<25} {self.config}", file=out)
|
||||
print(f" {bold('--quantization'):<25} {self.quantization}", file=out)
|
||||
print(f" {bold('--model-type'):<25} {self.model.name}", file=out)
|
||||
print(f" {bold('--target'):<25} {self.target.export()}", file=out)
|
||||
print(f" {bold('--opt'):<25} {self.opt}", file=out)
|
||||
print(f' {bold("--system-lib-prefix"):<25} "{self.system_lib_prefix}"', file=out)
|
||||
print(f" {bold('--output'):<25} {self.output}", file=out)
|
||||
print(f" {bold('--overrides'):<25} {self.overrides}", file=out)
|
||||
# As it's debug only, no need to display
|
||||
# print(f" {bold('--debug-dump'):<25} {self.debug_dump}", file=out)
|
||||
print(out.getvalue().rstrip())
|
||||
|
||||
|
||||
def _apply_preproc_to_params_and_check_pipeline(
|
||||
named_params: List[Tuple[str, nn.Parameter]], # noqa: UP006
|
||||
model_config,
|
||||
) -> Dict[str, tirx.PrimFunc]: # noqa: UP006
|
||||
extra_tirs: Dict[str, tirx.PrimFunc] = {} # noqa: UP006
|
||||
for name, param in named_params:
|
||||
preprocs = param.attrs.get("preprocs", [])
|
||||
shard_strategy = param.attrs.get("shard_strategy", None)
|
||||
if shard_strategy is not None and model_config.tensor_parallel_shards > 1:
|
||||
preprocs.append(
|
||||
shard_strategy.gen_shard_info(
|
||||
shards=model_config.tensor_parallel_shards,
|
||||
weight=param,
|
||||
)
|
||||
)
|
||||
if shard_strategy.name not in extra_tirs:
|
||||
extra_tirs[shard_strategy.name] = shard_strategy.gen_tir(
|
||||
shards=model_config.tensor_parallel_shards,
|
||||
weight=param,
|
||||
)
|
||||
param.attrs["preprocs"] = preprocs
|
||||
|
||||
pipeline_parallel_stages = getattr(model_config, "pipeline_parallel_stages", 1)
|
||||
if pipeline_parallel_stages != 1:
|
||||
assert "pipeline_stages" in param.attrs, (
|
||||
f'The pipeline stage is undefined for parameter "{name}" when the number '
|
||||
f"of pipeline parallel stages is {pipeline_parallel_stages}"
|
||||
)
|
||||
param.attrs["pipeline_stages"] = (
|
||||
[0]
|
||||
if "pipeline_stages" not in param.attrs
|
||||
else list(set(param.attrs["pipeline_stages"]))
|
||||
)
|
||||
return extra_tirs
|
||||
|
||||
|
||||
def _infer_kv_state_kind(model_type) -> str:
|
||||
if "rwkv" in model_type:
|
||||
return "rnn_state"
|
||||
if "medusa" in model_type:
|
||||
return "none"
|
||||
if "qwen3_5" in model_type:
|
||||
return "hybrid"
|
||||
return "kv_cache"
|
||||
|
||||
|
||||
def _compile(args: CompileArgs, model_config: ConfigBase):
|
||||
def _get_variable_bounds(model_config) -> Dict[str, int]: # noqa: UP006
|
||||
if hasattr(model_config, "sliding_window_size"):
|
||||
return {
|
||||
"rolling_cache_len": model_config.sliding_window_size,
|
||||
"kv_seq_len": model_config.sliding_window_size + model_config.prefill_chunk_size,
|
||||
"seq_len": model_config.prefill_chunk_size,
|
||||
"batch_size": getattr(model_config, "max_batch_size", 1),
|
||||
}
|
||||
return {
|
||||
"total_seq_len": model_config.context_window_size,
|
||||
"seq_len": model_config.prefill_chunk_size,
|
||||
"batch_size": getattr(model_config, "max_batch_size", 1),
|
||||
}
|
||||
|
||||
def _get_param_metadata(name: str, param: nn.Parameter) -> Dict[str, Any]: # noqa: UP006
|
||||
return {
|
||||
"name": name,
|
||||
# Record dynamic shape as -1 (e.g. vocab_size)
|
||||
"shape": [s if isinstance(s, int) else s.name for s in param.shape],
|
||||
"dtype": str(param.dtype),
|
||||
"preprocs": param.attrs["preprocs"],
|
||||
"pipeline_stages": param.attrs.get("pipeline_stages", [0]),
|
||||
}
|
||||
|
||||
logger.info("TOP LEVEL MODEL CONFIG BEFORE OVERRIDES: %s", str(model_config))
|
||||
_kwargs = getattr(model_config, "kwargs", {})
|
||||
model_config = args.overrides.apply(model_config)
|
||||
with args.target:
|
||||
op_ext.enable(
|
||||
target=args.target,
|
||||
flashinfer=args.opt.flashinfer,
|
||||
faster_transformer=args.opt.faster_transformer,
|
||||
cutlass=args.opt.cutlass,
|
||||
)
|
||||
# Step 1. Create the quantized model
|
||||
logger.info("Creating model from: %s", model_config)
|
||||
if (
|
||||
args.quantization.kind == "ft-quant"
|
||||
and hasattr(model_config, "tensor_parallel_shards")
|
||||
and model_config.tensor_parallel_shards > 1
|
||||
):
|
||||
raise NotImplementedError
|
||||
if (
|
||||
hasattr(args.quantization, "linear_weight_layout")
|
||||
and args.quantization.linear_weight_layout == "KN"
|
||||
and hasattr(model_config, "tensor_parallel_shards")
|
||||
and model_config.tensor_parallel_shards > 1
|
||||
):
|
||||
raise NotImplementedError(
|
||||
"KN layout (q3f16_0 and q4f16_0) is not supported for tensor parallelism"
|
||||
)
|
||||
model, _ = args.model.quantize[args.quantization.kind](model_config, args.quantization)
|
||||
# Step 2. Exporting the model to TVM
|
||||
logger.info("Exporting the model to TVM compiler")
|
||||
mod, named_params, ext_mods = model.export_tvm(
|
||||
spec=model.get_default_spec(),
|
||||
allow_extern=True,
|
||||
)
|
||||
# Step 3. Running relax compilation pipeline
|
||||
logger.info("Running optimizations using TVM")
|
||||
additional_tirs = _apply_preproc_to_params_and_check_pipeline(named_params, model_config)
|
||||
variable_bounds = _get_variable_bounds(model_config)
|
||||
cuda_graph_symbolic_capture_hints = {
|
||||
"batch_decode": ["batch_size"],
|
||||
"batch_decode_to_last_hidden_states": ["batch_size"],
|
||||
"batch_verify": ["batch_size", "seq_len"],
|
||||
"batch_verify_to_last_hidden_states": ["batch_size", "seq_len"],
|
||||
}
|
||||
avs = _kwargs.get("active_vocab_size", None)
|
||||
if avs is not None and avs <= 0:
|
||||
avs = None
|
||||
metadata = {
|
||||
"model_type": args.model.name,
|
||||
"quantization": args.quantization.name,
|
||||
"context_window_size": getattr(model_config, "context_window_size", -1),
|
||||
"sliding_window_size": getattr(model_config, "sliding_window_size", -1),
|
||||
"attention_sink_size": getattr(model_config, "attention_sink_size", -1),
|
||||
"prefill_chunk_size": model_config.prefill_chunk_size,
|
||||
"tensor_parallel_shards": model_config.tensor_parallel_shards,
|
||||
"pipeline_parallel_stages": getattr(model_config, "pipeline_parallel_stages", 1),
|
||||
"disaggregation": getattr(model_config, "disaggregation", False),
|
||||
"kv_state_kind": _infer_kv_state_kind(args.model.name),
|
||||
"max_batch_size": getattr(model_config, "max_batch_size", 1),
|
||||
"active_vocab_size": avs,
|
||||
"model_task": args.model.model_task,
|
||||
}
|
||||
if args.model.embedding_metadata:
|
||||
metadata["embedding_metadata"] = dataclasses.asdict(args.model.embedding_metadata)
|
||||
logger.info("Registering metadata: %s", metadata)
|
||||
metadata["params"] = [_get_param_metadata(name, param) for name, param in named_params]
|
||||
pass_config = {"relax.backend.use_cuda_graph": args.opt.cudagraph}
|
||||
# TODO: Remove this workaround when the TVM CSE regression is fixed.
|
||||
# Temporary workaround for TVM CSE regression that can produce
|
||||
# dangling `cse_v*` vars during host codegen.
|
||||
pass_config["tirx.disable_cse_tir"] = True
|
||||
|
||||
with PassContext(config=pass_config):
|
||||
args.build_func(
|
||||
mod,
|
||||
args,
|
||||
pipeline=relax.get_pipeline(
|
||||
"mlc_llm",
|
||||
target=args.target,
|
||||
flashinfer=args.opt.flashinfer,
|
||||
cublas_gemm=args.opt.cublas_gemm,
|
||||
faster_transformer=args.opt.faster_transformer,
|
||||
allreduce_strategy=args.opt.ipc_allreduce_strategy,
|
||||
variable_bounds=variable_bounds,
|
||||
cuda_graph_symbolic_capture_hints=cuda_graph_symbolic_capture_hints,
|
||||
additional_tirs=additional_tirs,
|
||||
ext_mods=ext_mods,
|
||||
metadata=metadata,
|
||||
debug_dump=args.debug_dump,
|
||||
),
|
||||
)
|
||||
_report_memory_usage(metadata=metadata, config=model_config)
|
||||
logger.info("Generated: %s", bold(str(args.output)))
|
||||
|
||||
|
||||
def compile(
|
||||
config: Dict[str, Any], # noqa: UP006
|
||||
quantization: Quantization,
|
||||
model_type: Model,
|
||||
target: Target,
|
||||
opt: OptimizationFlags,
|
||||
build_func: Callable[[IRModule, CompileArgs, Pass], None],
|
||||
system_lib_prefix: str,
|
||||
output: Path,
|
||||
overrides: ModelConfigOverride,
|
||||
debug_dump: Optional[Path] = None,
|
||||
):
|
||||
"""Compile a model given its configuration and quantization format to a specific target."""
|
||||
avs = None
|
||||
if "active_vocab_size" in config:
|
||||
avs = config.pop("active_vocab_size")
|
||||
logger.info("Active vocab size from input config: %s", str(avs))
|
||||
if "model_config" in config:
|
||||
model_config = config.pop("model_config")
|
||||
model_config.update(config)
|
||||
model_config = model_type.config.from_dict(model_config)
|
||||
else:
|
||||
model_config = model_type.config.from_dict(config)
|
||||
model_config.kwargs = {"active_vocab_size": avs} if avs is not None else {}
|
||||
args = CompileArgs(
|
||||
model_config,
|
||||
quantization,
|
||||
model_type,
|
||||
target,
|
||||
opt,
|
||||
build_func,
|
||||
system_lib_prefix,
|
||||
output,
|
||||
overrides,
|
||||
debug_dump,
|
||||
)
|
||||
args.display()
|
||||
_compile(args, model_config)
|
||||
Reference in New Issue
Block a user