Skip to content

GitLab

Unverified Commit 91276c57 authored by Jee Jee Li's avatar Jee Jee Li Committed by GitHub
Browse files

[Model] Adding torch compile annotations to chatglm (#15624) 


Signed-off-by: default avatarJee Jee Li <pandaleefree@gmail.com>
parent 0b416752
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 0 deletions
vllm/model_executor/models/chatglm.py
View file @ 91276c57
...@@ -10,6 +10,7 @@ from torch import nn ...@@ -10,6 +10,7 @@ from torch import nn
from torch.nn import LayerNorm from torch.nn import LayerNorm
from vllm.attention import Attention from vllm.attention import Attention
from vllm.compilation.decorators import support_torch_compile
from vllm.config import CacheConfig, VllmConfig from vllm.config import CacheConfig, VllmConfig
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.activation import SiluAndMul
...@@ -293,6 +294,7 @@ class GLMTransformer(nn.Module): ...@@ -293,6 +294,7 @@ class GLMTransformer(nn.Module):
return hidden_states return hidden_states
@support_torch_compile
class ChatGLMModel(nn.Module): class ChatGLMModel(nn.Module):
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment