Fix (brevitas_examples/llm): fix transformers tests (#1446)

Author: Giuseppe5

requirements/requirements-llm.txt

@@ -9,4 +9,4 @@ pandas
 pydantic
 torch>=2.4
 tqdm
-transformers[sentencepiece]
+transformers[sentencepiece]<5.0

src/brevitas_examples/llm/llm_quant/data_utils.py

@@ -37,7 +37,6 @@
 import numpy as np
 from optimum.utils.normalized_config import NormalizedConfigManager
 import torch
-from torch.utils.data import DataLoader
 from transformers import AutoConfig
 
 from brevitas_examples.llm.llm_quant.data import get_clm_dataset

src/brevitas_examples/llm/llm_quant/rotation_optimization.py

@@ -16,7 +16,12 @@
 import transformers
 from transformers import Trainer
 from transformers.data.data_collator import InputDataClass
-from transformers.tokenization_utils import PreTrainedTokenizerBase
+
+try:
+    from transformers.tokenization_utils import PreTrainedTokenizerBase
+except:
+    # This has changed in transformers v5
+    from transformers.tokenization_utils_base import PreTrainedTokenizerBase
 
 from brevitas.graph.calibrate import quantization_status_manager
 from brevitas.optim.cailey_sgd import CaileySGD

src/brevitas_examples/llm/llm_quant/run_utils.py

@@ -20,35 +20,16 @@
 """
 
 from contextlib import contextmanager
-import inspect
 
 from optimum.utils.normalized_config import NormalizedConfigManager
 import torch
 from torch.utils._python_dispatch import TorchDispatchMode
 from torch.utils._pytree import tree_map
 from transformers import AutoConfig
-from transformers.utils.fx import symbolic_trace
 
 from brevitas.fx.value_tracer import ValueProxy
 
 
-def get_fx(model, is_export=True):
-    forward_signature = inspect.signature(model.forward).parameters
-    if all(input_name in forward_signature
-           for input_name in ["input_ids", "attention_mask", "past_key_values"]):
-        input_names = ["input_ids", "attention_mask", "past_key_values"]
-        if not is_export:
-            input_names.remove('past_key_values')
-    else:
-        raise ValueError(
-            f"Quantization with an FX graph is currently only supported for models taking `input_ids`, `attention_mask` and `past_key_values` as inputs. The model only has the following inputs: {forward_signature}"
-        )
-
-    with torch.no_grad():
-        model = symbolic_trace(model, input_names)
-    return model
-
-
 def modify_dataloader(model_name_or_path, data, dtype):
     config = AutoConfig.from_pretrained(model_name_or_path)
 

tests/brevitas_examples/test_llm_cases.py

@@ -116,7 +116,7 @@ class LLMPerplexityCases:
                 "input_scale_type": "dynamic",
                 "input_quant_type": "sym",
                 "float_ppl": 32428.475,
-                "quant_ppl": 32428.383},
+                "quant_ppl": 32447.685546875},
             {
                 "model": "hf-internal-testing/tiny-random-MistralForCausalLM",
                 "act_equalization": "layerwise",