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

Small fix for the order of apply_torchao_config (#2495)

Browse Source
This commit is contained in:
Lianmin Zheng
2024-12-16 19:21:11 -08:00
committed by GitHub
parent 56198b45d9
commit bd6196163e
1 changed files with 5 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -157,6 +157,11 @@ class ModelRunner:
self.sampler = Sampler() self.sampler = Sampler()
self.load_model() self.load_model()
# Apply torchao quantization
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 +170,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()