move apply_torchao_config_ to model_runner (#2342)

This commit is contained in:
Jerry Zhang
2024-12-04 17:26:42 -08:00
committed by GitHub
parent d693ec0427
commit 9cc733b38c
8 changed files with 25 additions and 71 deletions

View File

@@ -35,12 +35,10 @@ from sglang.srt.layers.linear import (
from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.layers.torchao_utils import apply_torchao_config_
from sglang.srt.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
)
from sglang.srt.managers.schedule_batch import global_server_args_dict
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_loader.loader import DefaultModelLoader
from sglang.srt.model_loader.weight_utils import default_weight_loader
@@ -290,7 +288,6 @@ class Grok1ForCausalLM(nn.Module):
super().__init__()
self.config = config
self.quant_config = quant_config
self.torchao_config = global_server_args_dict["torchao_config"]
self.model = Grok1Model(config, quant_config=quant_config)
self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size)
self.logits_processor = LogitsProcessor(config)
@@ -374,8 +371,6 @@ class Grok1ForCausalLM(nn.Module):
)
weight_loader(param, loaded_weight)
apply_torchao_config_(self, params_dict, set(["proj.weight"]))
class Grok1ModelForCausalLM(Grok1ForCausalLM):
"""An alias for backward-compatbility."""