try except flash-attn

config.json

@@ -24,6 +24,7 @@
     "tf_legacy_loss": false,
     "pruned_heads": {},
     "tie_word_embeddings": false,
+    "chunk_size_feed_forward": 0,
     "is_encoder_decoder": false,
     "is_decoder": false,
     "cross_attention_hidden_size": null,
@@ -46,7 +47,6 @@
     "encoder_no_repeat_ngram_size": 0,
     "bad_words_ids": null,
     "num_return_sequences": 1,
-    "chunk_size_feed_forward": 0,
     "output_scores": false,
     "return_dict_in_generate": false,
     "forced_bos_token_id": null,
@@ -77,7 +77,7 @@
     "task_specific_params": null,
     "problem_type": null,
     "_name_or_path": "",
-    "transformers_version": "4.36.2",
+    "transformers_version": "4.37.1",
     "aqlm": {
         "nbits_per_codebook": 16,
         "num_codebooks": 1,

modeling_llama_aqlm.py

@@ -25,6 +25,7 @@ from typing import List, Optional, Tuple, Union
 import torch
 import torch.nn.functional as F
 import torch.utils.checkpoint
+from aqlm import QuantizedLinear
 from torch import nn
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 from transformers.activations import ACT2FN
@@ -53,11 +54,13 @@ from transformers.utils import (
 from transformers.utils.import_utils import is_torch_fx_available
 
 from .configuration_llama_aqlm import LlamaConfig
-from aqlm import QuantizedLinear
 
 if is_flash_attn_2_available():
-    from flash_attn import flash_attn_func, flash_attn_varlen_func
-    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
+    try:
+        from flash_attn import flash_attn_func, flash_attn_varlen_func
+        from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
+    except:
+        pass
 
 
 # This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.