diff --git a/superbench/benchmarks/micro_benchmarks/_export_torch_to_onnx.py b/superbench/benchmarks/micro_benchmarks/_export_torch_to_onnx.py index dc2f7f585..00bd58fe9 100644 --- a/superbench/benchmarks/micro_benchmarks/_export_torch_to_onnx.py +++ b/superbench/benchmarks/micro_benchmarks/_export_torch_to_onnx.py @@ -3,20 +3,26 @@ """Export PyTorch models to ONNX format.""" +import sys from pathlib import Path from packaging import version import torch.hub import torch.onnx import torchvision.models -from transformers import BertConfig, GPT2Config, LlamaConfig, MixtralConfig +from transformers import BertConfig, GPT2Config, LlamaConfig from superbench.benchmarks.model_benchmarks.pytorch_bert import BertBenchmarkModel from superbench.benchmarks.model_benchmarks.pytorch_gpt2 import GPT2BenchmarkModel from superbench.benchmarks.model_benchmarks.pytorch_lstm import LSTMBenchmarkModel from superbench.benchmarks.model_benchmarks.pytorch_llama import LlamaBenchmarkModel -from superbench.benchmarks.model_benchmarks.pytorch_mixtral import MixtralBenchmarkModel +# Check Python version and skip Mixtral if Python is 3.7 or lower +if sys.version_info <= (3, 7): + MixtralBenchmarkModel = None +else: + from transformers import MixtralConfig + from superbench.benchmarks.model_benchmarks.pytorch_mixtral import MixtralBenchmarkModel class torch2onnxExporter(): """PyTorch model to ONNX exporter.""" @@ -122,33 +128,37 @@ def __init__(self): ), self.num_classes, ), - 'mixtral-8x7b': - lambda: MixtralBenchmarkModel( - MixtralConfig( - hidden_size=4096, - num_hidden_layers=32, - num_attention_heads=32, - num_key_value_heads=8, - intermediate_size=14336, - max_position_embeddings=32768, - router_aux_loss_coef=0.02, + } + + # Only include Mixtral models if MixtralBenchmarkModel is available + if MixtralBenchmarkModel is not None: + self.benchmark_models.update({ + 'mixtral-8x7b': lambda: MixtralBenchmarkModel( + MixtralConfig( + hidden_size=4096, + num_hidden_layers=32, + num_attention_heads=32, + num_key_value_heads=8, + intermediate_size=14336, + max_position_embeddings=32768, + router_aux_loss_coef=0.02, + ), + self.num_classes, ), - self.num_classes, - ), - 'mixtral-8x22b': - lambda: MixtralBenchmarkModel( - MixtralConfig( - hidden_size=6144, - num_hidden_layers=56, - num_attention_heads=48, - num_key_value_heads=8, - intermediate_size=16384, - max_position_embeddings=65536, - router_aux_loss_coef=0.001, + 'mixtral-8x22b': lambda: MixtralBenchmarkModel( + MixtralConfig( + hidden_size=6144, + num_hidden_layers=56, + num_attention_heads=48, + num_key_value_heads=8, + intermediate_size=16384, + max_position_embeddings=65536, + router_aux_loss_coef=0.001, + ), + self.num_classes, ), - self.num_classes, - ), - } + }) + self._onnx_model_path = Path(torch.hub.get_dir()) / 'onnx' self._onnx_model_path.mkdir(parents=True, exist_ok=True) diff --git a/superbench/benchmarks/model_benchmarks/__init__.py b/superbench/benchmarks/model_benchmarks/__init__.py index e62214e47..0997287f4 100644 --- a/superbench/benchmarks/model_benchmarks/__init__.py +++ b/superbench/benchmarks/model_benchmarks/__init__.py @@ -3,6 +3,8 @@ """A module containing all the e2e model related benchmarks.""" +import sys + from superbench.benchmarks.model_benchmarks.model_base import ModelBenchmark from superbench.benchmarks.model_benchmarks.pytorch_bert import PytorchBERT from superbench.benchmarks.model_benchmarks.pytorch_gpt2 import PytorchGPT2 @@ -10,9 +12,13 @@ from superbench.benchmarks.model_benchmarks.pytorch_lstm import PytorchLSTM from superbench.benchmarks.model_benchmarks.megatron_gpt3 import MegatronGPT from superbench.benchmarks.model_benchmarks.pytorch_llama import PytorchLlama -from superbench.benchmarks.model_benchmarks.pytorch_mixtral import PytorchMixtral __all__ = [ 'ModelBenchmark', 'PytorchBERT', 'PytorchGPT2', 'PytorchCNN', 'PytorchLSTM', 'MegatronGPT', 'PytorchLlama', 'PytorchMixtral' ] + +# Check for Python version > 3.7 and conditionally import PytorchMixtral +if sys.version_info > (3, 7): + from superbench.benchmarks.model_benchmarks.pytorch_mixtral import PytorchMixtral + __all__.append('PytorchMixtral') diff --git a/superbench/benchmarks/model_benchmarks/pytorch_mixtral.py b/superbench/benchmarks/model_benchmarks/pytorch_mixtral.py index fd1a3e83f..6a3d49995 100644 --- a/superbench/benchmarks/model_benchmarks/pytorch_mixtral.py +++ b/superbench/benchmarks/model_benchmarks/pytorch_mixtral.py @@ -255,6 +255,7 @@ def _inference_step(self, precision): # Register Mixtral benchmark with 8x7b parameters. +# Ref: https://huggingface.co/mistralai/Mixtral-8x7B-v0.1/blob/main/config.json BenchmarkRegistry.register_benchmark( 'pytorch-mixtral-8x7b', PytorchMixtral, @@ -263,6 +264,7 @@ def _inference_step(self, precision): ) # Register Mixtral benchmark with 8x22b parameters. +# Ref: https://huggingface.co/mistralai/Mixtral-8x22B-v0.1/blob/main/config.json BenchmarkRegistry.register_benchmark( 'pytorch-mixtral-8x22b', PytorchMixtral, diff --git a/tests/benchmarks/model_benchmarks/test_pytorch_mixtral.py b/tests/benchmarks/model_benchmarks/test_pytorch_mixtral.py index 8c9ad2f35..6eb41d313 100644 --- a/tests/benchmarks/model_benchmarks/test_pytorch_mixtral.py +++ b/tests/benchmarks/model_benchmarks/test_pytorch_mixtral.py @@ -3,9 +3,14 @@ """Tests for mixtral model benchmarks.""" +import sys + from tests.helper import decorator from superbench.benchmarks import BenchmarkRegistry, Platform, Framework, BenchmarkType, ReturnCode -from superbench.benchmarks.model_benchmarks.pytorch_mixtral import PytorchMixtral + +# Check for Python version 3.8 or greater and conditionally import PytorchMixtral +if sys.version_info > (3, 7): + from superbench.benchmarks.model_benchmarks.pytorch_mixtral import PytorchMixtral @decorator.cuda_test