Adding "torch compile" annotations to moe models (#9758)

This commit is contained in:
Yongzao
2024-10-29 04:49:56 +08:00
committed by GitHub
parent 5f8d8075f9
commit aa0addb397
4 changed files with 8 additions and 0 deletions

View File

@@ -5,6 +5,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,
@@ -360,6 +361,7 @@ class ArcticDecoderLayer(nn.Module):
return hidden_states
@support_torch_compile
class ArcticModel(nn.Module):
def __init__(