EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Small fixes for torchao quant (#2476)

Browse Source
This commit is contained in:
Jerry Zhang
2024-12-16 14:08:12 -08:00
committed by GitHub
parent 7154b4b1df
commit 82699474fd
2 changed files with 6 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/sglang/srt/layers/torchao_utils.py
View File

@@ -26,11 +26,12 @@ def apply_torchao_config_to_model(
quantize_, quantize_,
) )
from torchao.quantization.observer import PerRow, PerTensor from torchao.quantization.observer import PerRow, PerTensor
from torchao.quantization.quant_api import _is_linear
if filter_fn is None: if filter_fn is None:
def filter_fn(module, fqn): def filter_fn(module, fqn):
return "proj" in fqn return _is_linear(module) and "proj" in fqn
if torchao_config == "" or torchao_config is None: if torchao_config == "" or torchao_config is None:
return model return model

8
python/sglang/srt/model_executor/model_runner.py
View File

@@ -157,6 +157,10 @@ class ModelRunner:
self.sampler = Sampler() self.sampler = Sampler()
self.load_model() self.load_model()
apply_torchao_config_to_model(
self.model, global_server_args_dict["torchao_config"]
)
# Apply torch TP if the model supports it # Apply torch TP if the model supports it
supports_torch_tp = getattr(self.model, "supports_torch_tp", False) supports_torch_tp = getattr(self.model, "supports_torch_tp", False)
if self.tp_size > 1 and supports_torch_tp: if self.tp_size > 1 and supports_torch_tp:
@@ -165,10 +169,6 @@ class ModelRunner:
else: else:
self.torch_tp_applied = False self.torch_tp_applied = False
apply_torchao_config_to_model(
self.model, global_server_args_dict["torchao_config"]
)
# Init memory pool and attention backends # Init memory pool and attention backends
if server_args.lora_paths is not None: if server_args.lora_paths is not None:
self.init_lora_manager() self.init_lora_manager()