From bd6196163ec3293b5254ecb5c6f14c16cb3577b6 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Mon, 16 Dec 2024 19:21:11 -0800 Subject: [PATCH] Small fix for the order of apply_torchao_config (#2495) --- python/sglang/srt/model_executor/model_runner.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index db024c5c7..2612f8840 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -157,6 +157,11 @@ class ModelRunner: self.sampler = Sampler() 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 supports_torch_tp = getattr(self.model, "supports_torch_tp", False) if self.tp_size > 1 and supports_torch_tp: @@ -165,10 +170,6 @@ class ModelRunner: else: self.torch_tp_applied = False - apply_torchao_config_to_model( - self.model, global_server_args_dict["torchao_config"] - ) - # Init memory pool and attention backends if server_args.lora_paths is not None: self.init_lora_manager()