chore: support quantized loading

tincans-ai · Mar 5, 2024 · 6e24ea3 · 6e24ea3
1 parent 7600d8e
commit 6e24ea3
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/gazelle/modeling_gazelle.py b/gazelle/modeling_gazelle.py
@@ -278,9 +278,9 @@ def forward(self, audio_features: torch.Tensor) -> torch.Tensor:
 )
 class GazellePreTrainedModel(PreTrainedModel):
     config_class = GazelleConfig
-    base_model_prefix = "model"
+    base_model_prefix = "gazelle"
     supports_gradient_checkpointing = True
-    _no_split_modules = ["GazelleAudioAttention"]
+    _no_split_modules = ["GazelleAudioAttention", "Wav2Vec2Model"]
     _skip_keys_device_placement = "past_key_values"
     _supports_flash_attn_2 = True
 

diff --git a/pyproject.toml b/pyproject.toml
@@ -11,8 +11,8 @@ classifiers = ["License :: OSI Approved :: Apache Software License"]
 dynamic = ["version", "description"]
 dependencies = [
     "transformers>=4.37.0",
-    "torch>=2.2",
-    "torchaudio>=2.2",
+    "torch>=2.1",
+    "torchaudio>=2.1",
     "sentencepiece",
 ]