[torch.compile] Adding torch compile annotations to some models (#9639)

Signed-off-by: youkaichao <youkaichao@gmail.com>
Co-authored-by: youkaichao <youkaichao@gmail.com>

docs/source/models/supported_models.rst

@@ -144,7 +144,7 @@ Text Generation
     - ✅︎
   * - :code:`JAISLMHeadModel`
     - Jais
-    - :code:`core42/jais-13b`, :code:`core42/jais-13b-chat`, :code:`core42/jais-30b-v3`, :code:`core42/jais-30b-chat-v3`, etc.
+    - :code:`inceptionai/jais-13b`, :code:`inceptionai/jais-13b-chat`, :code:`inceptionai/jais-30b-v3`, :code:`inceptionai/jais-30b-chat-v3`, etc.
     -
     - ✅︎
   * - :code:`JambaForCausalLM`

tests/distributed/test_pipeline_parallel.py

@@ -145,7 +145,7 @@ TEXT_GENERATION_MODELS = {
     # Uses Llama
     # "internlm/internlm-chat-7b": PPTestSettings.fast(),
     "internlm/internlm2-chat-7b": PPTestSettings.fast(trust_remote_code=True),
-    "core42/jais-13b-chat": PPTestSettings.fast(),
+    "inceptionai/jais-13b-chat": PPTestSettings.fast(),
     # TODO: Implement PP
     # "ai21labs/AI21-Jamba-1.5-Mini": PPTestSettings.fast(),
     "meta-llama/Meta-Llama-3-8B": PPTestSettings.detailed(),

vllm/model_executor/models/jais.py

@@ -1,6 +1,6 @@
 # coding=utf-8
 # Adapted from
-# https://huggingface.co/core42/jais-30b-chat-v3/blob/main/modeling_jais.py
+# https://huggingface.co/inceptionai/jais-30b-chat-v3/blob/main/modeling_jais.py
 # Copyright 2023 The vLLM team.
 # Copyright 2023 the Jais authors and HuggingFace Inc. team.  All rights
 # reserved.
@@ -26,6 +26,7 @@ import torch
 from torch import nn
 
 from vllm.attention import Attention, AttentionMetadata
+from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig
 from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,
                               get_tensor_model_parallel_world_size)
@@ -212,6 +213,7 @@ class JAISBlock(nn.Module):
         return hidden_states
 
 
+@support_torch_compile
 class JAISModel(nn.Module):
 
     def __init__(

vllm/model_executor/models/minicpm.py

@@ -29,6 +29,7 @@ from torch import nn
 from transformers import PretrainedConfig
 
 from vllm.attention import Attention, AttentionMetadata
+from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, LoRAConfig
 from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,
                               get_tensor_model_parallel_world_size,
@@ -348,6 +349,7 @@ class MiniCPMDecoderLayer(nn.Module):
         return hidden_states, None
 
 
+@support_torch_compile
 class MiniCPMModel(nn.Module):
 
     def __init__(

vllm/model_executor/models/mpt.py

@@ -7,6 +7,7 @@ import torch
 import torch.nn as nn
 
 from vllm.attention import Attention, AttentionMetadata
+from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig
 from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,
                               get_tensor_model_parallel_world_size)
@@ -204,6 +205,7 @@ class MPTBlock(nn.Module):
         return hidden_states
 
 
+@support_torch_compile
 class MPTModel(nn.Module):
 
     def __init__(

vllm/model_executor/models/nemotron.py

@@ -27,6 +27,7 @@ import torch
 from torch import nn
 
 from vllm.attention import Attention, AttentionMetadata
+from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, LoRAConfig
 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
 from vllm.model_executor.layers.activation import get_act_fn
@@ -290,6 +291,7 @@ class NemotronDecoderLayer(nn.Module):
         return hidden_states, residual
 
 
+@support_torch_compile
 class NemotronModel(nn.Module):
 
     def __init__(

vllm/model_executor/models/olmo.py

@@ -28,6 +28,7 @@ from torch import nn
 from transformers import OlmoConfig
 
 from vllm.attention import Attention, AttentionMetadata
+from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig
 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
 from vllm.model_executor.layers.activation import SiluAndMul
@@ -221,6 +222,7 @@ class OlmoDecoderLayer(nn.Module):
         return hidden_states
 
 
+@support_torch_compile
 class OlmoModel(nn.Module):
 
     def __init__(self,