minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Remove the torchao integration (--torchao-config) (#34304)

Browse Source 
Co-authored-by: Brayden Zhong <brayden@radixark.ai>
This commit is contained in:
Brayden Zhong
2026-08-14 21:49:11 +08:00
committed by GitHub
co-authored by Brayden Zhong
parent fe0c18effd
commit 5e65dd01a7
23 changed files with 6 additions and 309 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/pyproject.toml
-1
View File
@@ -79,7 +79,6 @@ dependencies = [
"tokenspeed_mla==0.1.8",
"torch==2.13.0",
"torch_memory_saver>=0.0.9.post1",
"torchao==0.17.0",
"torchaudio==2.11.0",
"torchcodec==0.15.0 ; sys_platform != 'linux' or (sys_platform == 'linux' and platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')", # Not available on Linux ARM.
"torchvision",
python/pyproject_cpu.toml
-1
View File
@@ -61,7 +61,6 @@ dependencies = [
"tiktoken",
"timm==1.0.16",
"torch==2.12.0",
"torchao==0.17.0",
"torchaudio==2.11.0",
"torchvision==0.27.0",
"tqdm",
python/pyproject_npu.toml
-1
View File
@@ -60,7 +60,6 @@ dependencies = [
"soundfile==0.13.1",
"tiktoken",
"timm==1.0.16",
"torchao==0.9.0",
"tqdm",
"transformers==5.12.1",
"uvicorn",
python/pyproject_other.toml
-2
View File
@@ -73,7 +73,6 @@ runtime_common = [
"compressed-tensors",
"outlines==0.1.11",
"timm==1.0.16",
"torchao==0.9.0",
"xgrammar==0.2.1",
]
@@ -159,7 +158,6 @@ srt_mps = [
"mlx-lm",
"sglang[runtime_common]",
"torch==2.11.0",
"torchao==0.9.0",
"torchaudio==2.11.0",
"torchvision",
]
python/pyproject_xpu.toml
-1
View File
@@ -60,7 +60,6 @@ dependencies = [
"tiktoken",
"timm==1.0.16",
"torch==2.12.0+xpu",
"torchao==0.17.0+xpu",
"torchaudio==2.11.0+xpu",
"torchcodec==0.12.0 ; sys_platform != 'linux' or (sys_platform == 'linux' and platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')", # torch==2.12.0 on XPU uses torchcodec 0.12.0
"torchvision==0.27.0+xpu",
python/sglang/check_env.py
-1
View File
@@ -40,7 +40,6 @@ PACKAGE_LIST = [
"pydantic",
"python-multipart",
"pyzmq",
"torchao",
"uvicorn",
"uvloop",
"vllm",
python/sglang/srt/layers/torchao_utils.py
-95
View File
@@ -1,95 +0,0 @@
"""
Common utilities for torchao.
"""
import logging
from typing import Callable, Optional
import torch
logger = logging.getLogger(__name__)
def proj_filter(
module: torch.nn.Module,
fqn: str,
):
"""Filter function for quantizing projection layers."""
return "proj" in fqn
# TODO: implement a more general filter function
def proj_filter_conv3d(
module: torch.nn.Module,
fqn: str,
):
if isinstance(module, torch.nn.Conv3d):
logger.warning(f"Quantize: skipping {fqn} because it's a Conv3d")
return False
return "proj" in fqn
def apply_torchao_config_to_model(
model: torch.nn.Module,
torchao_config: str,
filter_fn: Optional[Callable] = proj_filter,
):
"""Quantize a modelwith torchao quantization specified by torchao_config
Args:
`model`: a model to be quantized based on torchao_config
`torchao_config` (str): type of quantization and their arguments we want to use to
quantize the model, e.g. int4wo-128 means int4 weight only quantization with group_size
128
"""
if torchao_config == "" or torchao_config is None:
return model
# Lazy import to suppress some warnings
from torchao.quantization import (
float8_dynamic_activation_float8_weight,
float8_weight_only,
int4_weight_only,
int8_dynamic_activation_int8_weight,
int8_weight_only,
quantize_,
)
from torchao.quantization.observer import PerRow, PerTensor
if "int8wo" in torchao_config:
quantize_(model, int8_weight_only(), filter_fn=proj_filter_conv3d)
elif "int8dq" in torchao_config:
quantize_(model, int8_dynamic_activation_int8_weight(), filter_fn=filter_fn)
elif "int4wo" in torchao_config:
group_size = int(torchao_config.split("-")[-1])
assert group_size in [
32,
64,
128,
256,
], f"int4wo groupsize needs to be one of [32, 64, 128, 256] but got {group_size}"
quantize_(model, int4_weight_only(group_size=group_size), filter_fn=filter_fn)
elif "fp8wo" in torchao_config:
# this requires newer hardware
# [rank0]: AssertionError: fp8e4nv data type is not supported on CUDA arch < 89
quantize_(model, float8_weight_only(), filter_fn=proj_filter_conv3d)
elif "fp8dq" in torchao_config:
granularity = torchao_config.split("-")[-1]
GRANULARITY_MAP = {
"per_row": PerRow(),
"per_tensor": PerTensor(),
}
assert (
granularity in GRANULARITY_MAP
), f"Supported granularity are: {GRANULARITY_MAP.keys()}, got {granularity}"
quantize_(
model,
float8_dynamic_activation_float8_weight(
granularity=GRANULARITY_MAP[granularity]
),
filter_fn=proj_filter_conv3d,
)
else:
raise ValueError(f"Unexpected config: {torchao_config}")
return model
python/sglang/srt/model_executor/model_runner.py
-5
View File
@@ -80,7 +80,6 @@ from sglang.srt.layers.cp.utils import (
)
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.layers.sampler import create_sampler
from sglang.srt.layers.torchao_utils import apply_torchao_config_to_model
from sglang.srt.layers.utils.cp_utils import is_mla_prefill_cp_enabled
from sglang.srt.lora.lora_manager import LoRAManager, init_lora_cuda_graph_moe_buffers
from sglang.srt.lora.lora_registry import LoRARef
@@ -754,10 +753,6 @@ class ModelRunner:
)
def maybe_apply_post_load_model_transforms(self):
# In layered loading, torchao may have been applied
torchao_applied = getattr(self.model, "torchao_applied", False)
if not torchao_applied:
apply_torchao_config_to_model(self.model, get_exec().graph.torchao_config)
supports_torch_tp = getattr(self.model, "supports_torch_tp", False)
if self.ps.tp_size > 1 and supports_torch_tp:
self.apply_torch_tp()
python/sglang/srt/model_loader/loader.py
-11
View File
@@ -1079,9 +1079,6 @@ class LayeredModelLoader(DefaultModelLoader):
model_config: ModelConfig,
device_config: DeviceConfig,
) -> nn.Module:
from sglang.srt.layers.torchao_utils import apply_torchao_config_to_model
torchao_config = get_exec().graph.torchao_config
target_device = torch.device(device_config.device)
quant_config = _get_quantization_config(model_config, self.load_config)
@@ -1122,18 +1119,10 @@ class LayeredModelLoader(DefaultModelLoader):
fqn_path,
weights,
)
# Quantize weights if applicable
if torchao_config and "proj" in fqn_path:
# Note: `None` here is needed to indicate no filter, see
# `apply_torchao_config_to_model` for details.
apply_torchao_config_to_model(module, torchao_config, None)
# Start calling on root module
fill_module(model, [], weights)
if torchao_config:
model.torchao_applied = True
return model.eval()
python/sglang/srt/models/paddleocr_vl.py
+1 -1
View File
@@ -290,7 +290,7 @@ class SiglipMLP(nn.Module):
self.config = config
self.activation_fn = get_act_fn(config.hidden_act)
if quant_config and quant_config.get_name() in ["bitsandbytes", "torchao"]:
if quant_config and quant_config.get_name() == "bitsandbytes":
quantizable = True
else:
quantizable = (
python/sglang/srt/server_args.py
+1 -7
View File
@@ -2025,7 +2025,7 @@ class ServerArgs:
] = False
# -------------------------------------------------------------------------
# Torch compile and torchao
# Torch compile
# -------------------------------------------------------------------------
enable_torch_compile: A[
bool,
@@ -2038,12 +2038,6 @@ class ServerArgs:
torch_compile_max_bs: A[
int, "Set the maximum batch size when using torch compile.", NS("exec.graph")
] = 32
torchao_config: A[
str,
"Optimize the model with torchao. Experimental feature. Current choices are: int8dq, int8wo, int4wo-<group_size>, fp8wo, fp8dq-per_tensor, fp8dq-per_row",
NS("exec.graph"),
] = ""
# -------------------------------------------------------------------------
# Speculative decoding
# -------------------------------------------------------------------------
python/sglang/test/runners.py
-2
View File
@@ -577,7 +577,6 @@ class SRTRunner:
speculative_num_draft_tokens: Optional[int] = None,
disable_overlap_schedule: bool = False,
disable_custom_all_reduce: bool = False,
torchao_config: Optional[str] = None,
cuda_graph_max_bs_decode: int = 4,
sleep_on_idle=False,
max_lora_rank: Optional[int] = None,
@@ -615,7 +614,6 @@ class SRTRunner:
dtype=get_dtype_str(torch_dtype),
port=port,
model_impl=model_impl,
torchao_config=torchao_config,
mem_fraction_static=mem_fraction_static,
trust_remote_code=trust_remote_code,
is_embedding=not self.is_generation,