github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-06-20 12:43:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Infer/Fix] Fix Dependency in test - RMSNorm kernel (#5399)

Browse Source 
fix dependency in pytest
This commit is contained in:
Yuanheng Zhao 2024-02-26 16:17:47 +08:00 committed by GitHub
parent bc1da87366
commit 19061188c3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 6 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
tests/test_infer/test_ops/triton/test_rmsnorm_triton.py
View File

@ -3,13 +3,12 @@ import torch
import triton import triton
from packaging import version from packaging import version
from transformers.models.llama.modeling_llama import LlamaRMSNorm from transformers.models.llama.modeling_llama import LlamaRMSNorm
from vllm.model_executor.layers.layernorm import RMSNorm
from colossalai.kernel.triton import rms_layernorm from colossalai.kernel.triton import rms_layernorm
from colossalai.testing.utils import parameterize from colossalai.testing.utils import parameterize
try: try:
pass import triton # noqa
HAS_TRITON = True HAS_TRITON = True
except ImportError: except ImportError:
@ -85,6 +84,11 @@ def benchmark_rms_layernorm(
SEQUENCE_TOTAL: int, SEQUENCE_TOTAL: int,
HIDDEN_SIZE: int, HIDDEN_SIZE: int,
): ):
try:
from vllm.model_executor.layers.layernorm import RMSNorm
except ImportError:
raise ImportError("Please install vllm from https://github.com/vllm-project/vllm")
warmup = 10 warmup = 10
rep = 1000 rep = 1000