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

Fix llama for classification (#855)

Browse Source
This commit is contained in:
Ying Sheng
2024-07-31 15:48:31 -07:00
committed by GitHub
parent bc3eaac2b8
commit 5e7dd984fe
1 changed files with 5 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
python/sglang/srt/models/llama2.py
View File

@@ -26,6 +26,11 @@ from vllm.config import CacheConfig
from vllm.distributed import get_tensor_model_parallel_world_size from vllm.distributed import get_tensor_model_parallel_world_size
from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.activation import SiluAndMul
from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (
MergedColumnParallelLinear,
QKVParallelLinear,
RowParallelLinear,
)
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.rotary_embedding import get_rope
from vllm.model_executor.layers.vocab_parallel_embedding import ( from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -38,10 +43,6 @@ from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.model_executor.model_runner import InputMetadata from sglang.srt.model_executor.model_runner import InputMetadata
MergedColumnParallelLinear = None
QKVParallelLinear = None
RowParallelLinear = None
class LlamaMLP(nn.Module): class LlamaMLP(nn.Module):
def __init__( def __init__(
@@ -295,23 +296,6 @@ class LlamaForCausalLM(nn.Module):
cache_config: Optional[CacheConfig] = None, cache_config: Optional[CacheConfig] = None,
efficient_weight_load=False, efficient_weight_load=False,
) -> None: ) -> None:
global MergedColumnParallelLinear
global QKVParallelLinear
global RowParallelLinear
if efficient_weight_load:
from sglang.srt.layers.linear import (
MergedColumnParallelLinear,
QKVParallelLinear,
RowParallelLinear,
)
else:
from vllm.model_executor.layers.linear import (
MergedColumnParallelLinear,
QKVParallelLinear,
RowParallelLinear,
)
super().__init__() super().__init__()
self.config = config self.config = config
self.quant_config = quant_config self.quant_config = quant_config