Repository: huggingface/text-generation-inference

Issue:
fix: Add V100 (older) GPU Support for Mistral 7b Models
# What does this PR do?

## Introduction
This PR introduces changes that significantly improve support for running Mistral 7b models on V100 GPU architectures. By update the latest `transformers` package, we ensure both compatibility and performance for users with V100 GPUs.

## Changes
- Altered the model instantiation logic in the scenario where the model type is 'mistral'.
- In environments where `FLASH_ATTENTION` is not available, the system defaults to using `CausalLM`. This ensures that Mistral models remain operational even without Flash Attention support.

## Impact
These updates are particularly beneficial for users with V100 GPU architectures, which previously faced compatibility issues due to `flash-attention` package limitations. With these changes, we can expand our hardware support, reduce barriers to entry, and streamline the user experience for a significant segment of our user base.

## Additional Notes
The `transformers` package has been updated to version 4.35.2, which now includes support for Mistral models. This update is essential since the `flash-attention` package does not support V100 GPUs, as discussed in [Dao-AILab/flash-attention#148](https://github.com/Dao-AILab/flash-attention/issues/148).

## Testing
- Comprehensive tests have been conducted to ensure that Mistral models perform as expected on V100 GPUs.
- Additional benchmarks have been added to compare the performance on V100 GPUs with previous versions.

I invite reviewers to pull this branch, test rigorously with V100 GPUs, and provide feedback on any aspects of this enhancement.

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
#1253
#1208
#319
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

@abhijithnair1



Discussion:

Resolving patch:
diff --git a/server/text_generation_server/models/__init__.py b/server/text_generation_server/models/__init__.py
--- a/server/text_generation_server/models/__init__.py
+++ b/server/text_generation_server/models/__init__.py
@@ -55,10 +55,14 @@
         FlashSantacoderSharded,
     )
     from text_generation_server.models.idefics import IDEFICSSharded
+    from text_generation_server.models.flash_mistral import FlashMistral
+    from text_generation_server.models.flash_mixtral import FlashMixtral
+    from text_generation_server.utils.flash_attn import HAS_FLASH_ATTN_V2_CUDA
 
 except ImportError as e:
     logger.warning(f"Could not import Flash Attention enabled models: {e}")
     FLASH_ATTENTION = False
+    HAS_FLASH_ATTN_V2_CUDA = False
 
 if FLASH_ATTENTION:
     __all__.append(FlashNeoXSharded)
@@ -66,25 +70,7 @@
     __all__.append(FlashSantacoderSharded)
     __all__.append(FlashLlama)
     __all__.append(IDEFICSSharded)
-
-MISTRAL = True
-try:
-    from text_generation_server.models.flash_mistral import FlashMistral
-except ImportError as e:
-    logger.warning(f"Could not import Mistral model: {e}")
-    MISTRAL = False
-
-if MISTRAL:
     __all__.append(FlashMistral)
-
-MIXTRAL = True
-try:
-    from text_generation_server.models.flash_mixtral import FlashMixtral
-except ImportError as e:
-    logger.warning(f"Could not import Mixtral model: {e}")
-    MIXTRAL = False
-
-if MIXTRAL:
     __all__.append(FlashMixtral)
 
 
@@ -295,7 +281,9 @@ def get_model(
                 )
 
     if model_type == "mistral":
-        if MISTRAL:
+        if (config_dict["sliding_window"] is None and FLASH_ATTENTION) or (
+            config_dict["sliding_window"] > 0 and HAS_FLASH_ATTN_V2_CUDA
+        ):
             return FlashMistral(
                 model_id,
                 revision,
@@ -303,10 +291,11 @@ def get_model(
                 dtype=dtype,
                 trust_remote_code=trust_remote_code,
             )
-        raise NotImplementedError("Mistral models requires flash attention v2")
 
     if model_type == "mixtral":
-        if MIXTRAL:
+        if (config_dict["sliding_window"] is None and FLASH_ATTENTION) or (
+            config_dict["sliding_window"] > 0 and HAS_FLASH_ATTN_V2_CUDA
+        ):
             return FlashMixtral(
                 model_id,
                 revision,
@@ -314,9 +303,6 @@ def get_model(
                 dtype=dtype,
                 trust_remote_code=trust_remote_code,
             )
-        raise NotImplementedError(
-            "Mixtral models requires flash attention v2, stk and megablocks"
-        )
 
     if model_type == "opt":
         return OPTSharded(
@@ -348,17 +334,17 @@ def get_model(
             raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Idefics"))
 
     if sharded:
-        raise ValueError("sharded is not supported for AutoModel")
+        raise NotImplementedError("sharded is not supported for AutoModel")
     if quantize == "gptq":
-        raise ValueError(
+        raise NotImplementedError(
             "gptq quantization is not supported for AutoModel, you can try to quantize it with `text-generation-server quantize ORIGINAL_MODEL_ID NEW_MODEL_ID`"
         )
     if quantize == "awq":
-        raise ValueError("awq quantization is not supported for AutoModel")
+        raise NotImplementedError("awq quantization is not supported for AutoModel")
     elif (quantize == "bitsandbytes-fp4") or (quantize == "bitsandbytes-nf4"):
-        raise ValueError("4bit quantization is not supported for AutoModel")
+        raise NotImplementedError("4bit quantization is not supported for AutoModel")
     elif quantize == "eetq":
-        raise ValueError("Eetq quantization is not supported for AutoModel")
+        raise NotImplementedError("Eetq quantization is not supported for AutoModel")
     if model_type in modeling_auto.MODEL_FOR_CAUSAL_LM_MAPPING_NAMES:
         return CausalLM(
             model_id,
diff --git a/server/text_generation_server/models/custom_modeling/flash_mistral_modeling.py b/server/text_generation_server/models/custom_modeling/flash_mistral_modeling.py
--- a/server/text_generation_server/models/custom_modeling/flash_mistral_modeling.py
+++ b/server/text_generation_server/models/custom_modeling/flash_mistral_modeling.py
@@ -27,11 +27,6 @@
 from typing import Optional, List, Tuple
 
 from text_generation_server.utils import paged_attention, flash_attn
-from text_generation_server.utils.flash_attn import (
-    attention,
-    HAS_FLASH_ATTN_V2_ROCM,
-    HAS_FLASH_ATTN_V2_CUDA,
-)
 from text_generation_server.utils.layers import (
     TensorParallelRowLinear,
     TensorParallelColumnLinear,
@@ -43,10 +38,6 @@
 )
 
 
-if not HAS_FLASH_ATTN_V2_CUDA and not HAS_FLASH_ATTN_V2_ROCM:
-    raise ImportError("Mistral model requires flash attn v2")
-
-
 class MistralConfig(PretrainedConfig):
     model_type = "mistral"
 
diff --git a/server/text_generation_server/models/custom_modeling/flash_mixtral_modeling.py b/server/text_generation_server/models/custom_modeling/flash_mixtral_modeling.py
--- a/server/text_generation_server/models/custom_modeling/flash_mixtral_modeling.py
+++ b/server/text_generation_server/models/custom_modeling/flash_mixtral_modeling.py
@@ -27,12 +27,9 @@
 from transformers.activations import ACT2FN
 from transformers.configuration_utils import PretrainedConfig
 from typing import Optional, List, Tuple
+from loguru import logger
 
 from text_generation_server.utils import paged_attention, flash_attn
-from text_generation_server.utils.flash_attn import (
-    HAS_FLASH_ATTN_V2_ROCM,
-    HAS_FLASH_ATTN_V2_CUDA,
-)
 from text_generation_server.utils.layers import (
     FastLinear,
     FastRMSNorm,
@@ -44,18 +41,13 @@
     get_linear,
 )
 
-if not HAS_FLASH_ATTN_V2_CUDA and not HAS_FLASH_ATTN_V2_ROCM:
-    raise ImportError("Mixtral model requires flash attn v2")
-
-try:
-    import megablocks.ops as ops
-except ImportError:
-    raise ImportError("Mixtral model requires megablocks to be installed")
-
+HAS_MEGABLOCKS = True
 try:
     import stk
+    import megablocks.ops as ops
 except ImportError:
-    raise ImportError("Mixtral model requires stk to be installed")
+    logger.warning("Mixtral: megablocks is not installed")
+    HAS_MEGABLOCKS = False
 
 
 class MixtralConfig(PretrainedConfig):
@@ -590,7 +582,7 @@ def dense_forward(self, x: torch.Tensor) -> torch.Tensor:
         return out
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        if len(x) > 256:
+        if len(x) > 256 and HAS_MEGABLOCKS:
             return self.sparse_forward(x)
         # This is faster when there is not a lot of tokens
         return self.dense_forward(x)


Test patch: