wehub-resource-sync
228 lines
8.2 KiB
Python
228 lines
8.2 KiB
Python
"""Flags for overriding model config."""
|
|
|
|
import dataclasses
|
|
import enum
|
|
from io import StringIO
|
|
from typing import Optional
|
|
|
|
from mlc_llm.support import argparse, logging
|
|
from mlc_llm.support.config import ConfigOverrideBase
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class IPCAllReduceStrategyType(enum.IntEnum):
|
|
"""The all-reduce strategy."""
|
|
|
|
NONE = 0
|
|
ONESHOT = 1
|
|
TWOSHOT = 2
|
|
AUTO = 3
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class OptimizationFlags:
|
|
"""Optimization flags"""
|
|
|
|
flashinfer: bool = False
|
|
cublas_gemm: bool = False
|
|
faster_transformer: bool = False
|
|
cudagraph: bool = False
|
|
cutlass: bool = False
|
|
ipc_allreduce_strategy: IPCAllReduceStrategyType = IPCAllReduceStrategyType.NONE
|
|
|
|
def __repr__(self) -> str:
|
|
out = StringIO()
|
|
print(f"flashinfer={int(self.flashinfer)}", file=out, end="")
|
|
print(f";cublas_gemm={int(self.cublas_gemm)}", file=out, end="")
|
|
print(f";faster_transformer={int(self.faster_transformer)}", file=out, end="")
|
|
print(f";cudagraph={int(self.cudagraph)}", file=out, end="")
|
|
print(f";cutlass={int(self.cutlass)}", file=out, end="")
|
|
print(
|
|
f";ipc_allreduce_strategy={self.ipc_allreduce_strategy.name}",
|
|
file=out,
|
|
end="",
|
|
)
|
|
return out.getvalue().rstrip()
|
|
|
|
@staticmethod
|
|
def from_str(source: str) -> "OptimizationFlags":
|
|
"""Parse optimization flags from a string."""
|
|
|
|
if source in OPT_FLAG_PRESET:
|
|
return OPT_FLAG_PRESET[source]
|
|
|
|
def boolean(value: str) -> bool:
|
|
if value == "0":
|
|
return False
|
|
if value == "1":
|
|
return True
|
|
raise ValueError(f"Invalid boolean value: {value}")
|
|
|
|
parser = argparse.ArgumentParser(description="optimization flags")
|
|
parser.add_argument("--flashinfer", type=boolean, default=True)
|
|
parser.add_argument("--cublas_gemm", type=boolean, default=False)
|
|
parser.add_argument("--faster_transformer", type=boolean, default=False)
|
|
parser.add_argument("--cudagraph", type=boolean, default=False)
|
|
parser.add_argument("--cutlass", type=boolean, default=False)
|
|
parser.add_argument(
|
|
"--ipc_allreduce_strategy",
|
|
type=str,
|
|
choices=["NONE", "ONESHOT", "TWOSHOT", "AUTO"],
|
|
default="NONE",
|
|
)
|
|
results = parser.parse_args([f"--{i}" for i in source.split(";") if i])
|
|
return OptimizationFlags(
|
|
flashinfer=results.flashinfer,
|
|
cublas_gemm=results.cublas_gemm,
|
|
faster_transformer=results.faster_transformer,
|
|
cudagraph=results.cudagraph,
|
|
cutlass=results.cutlass,
|
|
ipc_allreduce_strategy=IPCAllReduceStrategyType[results.ipc_allreduce_strategy],
|
|
)
|
|
|
|
def update(self, target, quantization) -> None:
|
|
"""Update optimization flags based on additional information."""
|
|
|
|
def _flashinfer(target) -> bool:
|
|
from mlc_llm.support.auto_target import (
|
|
detect_cuda_arch_list,
|
|
)
|
|
|
|
if not self.flashinfer:
|
|
return False
|
|
if target.kind.name != "cuda":
|
|
return False
|
|
arch_list = detect_cuda_arch_list(target)
|
|
for arch in arch_list:
|
|
if arch < 80:
|
|
logger.warning("flashinfer is not supported on CUDA arch < 80")
|
|
return False
|
|
return True
|
|
|
|
def _cublas_gemm(target, quantization) -> bool:
|
|
"""correct cublas_gemm flag"""
|
|
if target.kind.name not in ["cuda", "rocm"]:
|
|
return False
|
|
if not (
|
|
quantization.name in ["q0f16", "q0bf16", "q0f32"]
|
|
or "e4m3" in quantization.name
|
|
or "e5m2" in quantization.name
|
|
):
|
|
return False
|
|
return self.cublas_gemm
|
|
|
|
def _faster_transformer(target) -> bool:
|
|
"""correct faster_transformer flag"""
|
|
if not target.kind.name == "cuda":
|
|
return False
|
|
return self.faster_transformer
|
|
|
|
def _cutlass(target) -> bool:
|
|
"""correct cutlass flag"""
|
|
if not target.kind.name == "cuda":
|
|
return False
|
|
return self.cutlass
|
|
|
|
def _cudagraph(target) -> bool:
|
|
"""correct cudagraph flag"""
|
|
if not target.kind.name == "cuda":
|
|
return False
|
|
return self.cudagraph
|
|
|
|
self.flashinfer = _flashinfer(target)
|
|
self.cublas_gemm = _cublas_gemm(target, quantization)
|
|
self.faster_transformer = _faster_transformer(target)
|
|
self.cutlass = _cutlass(target)
|
|
self.cudagraph = _cudagraph(target)
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class ModelConfigOverride(ConfigOverrideBase):
|
|
"""Flags for overriding model config."""
|
|
|
|
context_window_size: Optional[int] = None
|
|
sliding_window_size: Optional[int] = None
|
|
prefill_chunk_size: Optional[int] = None
|
|
attention_sink_size: Optional[int] = None
|
|
max_batch_size: Optional[int] = None
|
|
tensor_parallel_shards: Optional[int] = None
|
|
pipeline_parallel_stages: Optional[int] = None
|
|
disaggregation: Optional[bool] = None
|
|
|
|
def __repr__(self) -> str:
|
|
out = StringIO()
|
|
print(f"context_window_size={self.context_window_size}", file=out, end="")
|
|
print(f";sliding_window_size={self.sliding_window_size}", file=out, end="")
|
|
print(f";prefill_chunk_size={self.prefill_chunk_size}", file=out, end="")
|
|
print(f";attention_sink_size={self.attention_sink_size}", file=out, end="")
|
|
print(f";max_batch_size={self.max_batch_size}", file=out, end="")
|
|
print(f";tensor_parallel_shards={self.tensor_parallel_shards}", file=out, end="")
|
|
print(
|
|
f";pipeline_parallel_stages={self.pipeline_parallel_stages}",
|
|
file=out,
|
|
end="",
|
|
)
|
|
print(f";disaggregation={self.disaggregation}", file=out, end="")
|
|
return out.getvalue().rstrip()
|
|
|
|
@staticmethod
|
|
def from_str(source: str) -> "ModelConfigOverride":
|
|
"""Parse model config override values from a string."""
|
|
parser = argparse.ArgumentParser(description="model config override values")
|
|
parser.add_argument("--context_window_size", type=int, default=None)
|
|
parser.add_argument("--sliding_window_size", type=int, default=None)
|
|
parser.add_argument("--prefill_chunk_size", type=int, default=None)
|
|
parser.add_argument("--attention_sink_size", type=int, default=None)
|
|
parser.add_argument("--max_batch_size", type=int, default=None)
|
|
parser.add_argument("--tensor_parallel_shards", type=int, default=None)
|
|
parser.add_argument("--pipeline_parallel_stages", type=int, default=None)
|
|
parser.add_argument(
|
|
"--disaggregation",
|
|
type=lambda x: str(x).lower() in ["true", "1", "yes", "True"],
|
|
default=None,
|
|
)
|
|
results = parser.parse_args([f"--{i}" for i in source.split(";") if i])
|
|
return ModelConfigOverride(
|
|
context_window_size=results.context_window_size,
|
|
sliding_window_size=results.sliding_window_size,
|
|
prefill_chunk_size=results.prefill_chunk_size,
|
|
attention_sink_size=results.attention_sink_size,
|
|
max_batch_size=results.max_batch_size,
|
|
tensor_parallel_shards=results.tensor_parallel_shards,
|
|
pipeline_parallel_stages=results.pipeline_parallel_stages,
|
|
disaggregation=results.disaggregation,
|
|
)
|
|
|
|
|
|
OPT_FLAG_PRESET = {
|
|
"O0": OptimizationFlags(
|
|
flashinfer=False,
|
|
cublas_gemm=False,
|
|
cudagraph=False,
|
|
),
|
|
"O1": OptimizationFlags(
|
|
flashinfer=False,
|
|
cublas_gemm=True,
|
|
faster_transformer=True,
|
|
cudagraph=False,
|
|
cutlass=True,
|
|
),
|
|
"O2": OptimizationFlags(
|
|
flashinfer=True,
|
|
cublas_gemm=True,
|
|
faster_transformer=False,
|
|
cudagraph=True,
|
|
cutlass=True,
|
|
ipc_allreduce_strategy=IPCAllReduceStrategyType.NONE,
|
|
),
|
|
"O3": OptimizationFlags(
|
|
flashinfer=True,
|
|
cublas_gemm=True,
|
|
faster_transformer=True,
|
|
cudagraph=True,
|
|
cutlass=True,
|
|
ipc_allreduce_strategy=IPCAllReduceStrategyType.AUTO,
|
|
),
|
|
}
|