Add decoder modeling

.github/workflows/build_pr_documentation.yml

@@ -51,7 +51,7 @@ jobs:
           git clone --depth 1 --branch v3.5 https://github.com/Xilinx/Vitis-AI.git && cd Vitis-AI/src/vai_quantizer/vai_q_onnx && sh build.sh && pip install pkgs/*.whl
 
           cd ../../../../optimum-amd
-          pip install .[brevitas,tests]
+          pip install .[brevitas]
 
           pip install onnxruntime==1.14.0
           cd ..

.github/workflows/test_ryzenai_nightly.yaml

@@ -27,13 +27,21 @@ jobs:
       timeout_minutes: 1200
     secrets:
       hf_hub_read_token: ${{ secrets.HF_READ_TOKEN }}
+  run_tests_brevitas_quantized_decoder_llms:
+    uses: huggingface/hf-workflows/.github/workflows/ryzenai_ci.yaml@main
+    with:
+      pytest_marker: "brevitas_quantized_decoder_llms_test"
+      test_file: "tests/ryzenai/test_modeling.py"
+      report_name: "tests_brevitas_quantized_decoder_llms"
+      slow_test: true
   send_results:
     name: Send results to webhook
     runs-on: ubuntu-22.04
     if: always()
     needs: [
       run_tests_prequantized_models,
       run_tests_quantization,
+      run_tests_brevitas_quantized_decoder_llms,
     ]
     steps:
       - uses: actions/checkout@v3

.gitignore

@@ -166,3 +166,5 @@ cython_debug/
 runs
 sweeps
 experiments
+ryzen_cache
+RyzenAI-SW

optimum/amd/ryzenai/__init__.py

@@ -16,6 +16,9 @@
         "RyzenAIModelForImageToImage",
         "RyzenAIModelForObjectDetection",
     ],
+    "modeling_decoder": [
+        "RyzenAIModelForCausalLM",
+    ],
     "quantization": ["RyzenAIOnnxQuantizer"],
     "pipelines": ["pipeline"],
     "utils": ["DEFAULT_VAIP_CONFIG"],
@@ -33,6 +36,7 @@
         RyzenAIModelForObjectDetection,
         RyzenAIModelForSemanticSegmentation,
     )
+    from .modeling_decoder import RyzenAIModelForCausalLM
     from .pipelines import pipeline
     from .quantization import RyzenAIOnnxQuantizer
     from .utils import DEFAULT_VAIP_CONFIG

optimum/amd/ryzenai/configs/vaip_config_transformers.json

@@ -0,0 +1,35 @@
+{
+    "passes": [
+      {
+        "name": "init",
+        "plugin": "vaip-pass_init"
+      },
+      {
+        "name": "fuse_GEMM",
+        "plugin": "vaip-pass_py_ext",
+        "disabled": false,
+        "pyExt": {
+          "moduleName": "voe.passes.fuse_GEMM",
+          "methodName": "rules"
+        }
+      },
+      {
+        "name": "fuse_MATMUL",
+        "plugin": "vaip-pass_py_ext",
+        "disabled": false,
+        "pyExt": {
+          "moduleName": "voe.passes.fuse_MATMUL",
+          "methodName": "rules"
+        }
+      },
+      {
+        "name": "fuse_MATMULINTEGER",
+        "plugin": "vaip-pass_py_ext",
+        "disabled": false,
+        "pyExt": {
+          "moduleName": "voe.passes.fuse_MATMULINTEGER",
+          "methodName": "rules"
+        }
+      }
+    ]
+  }

optimum/amd/ryzenai/modeling.py

@@ -111,7 +111,7 @@ def shared_attributes_init(
     def __init__(
         self,
         model: ort.InferenceSession,
-        config: PretrainedConfig,
+        config: Optional[PretrainedConfig] = None,
         vaip_config: Union[str, Path] = None,
         model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
         preprocessors: Optional[List] = None,
@@ -175,11 +175,12 @@ def load_model(
         else:
             providers_options = None
 
-        is_dynamic = RyzenAIModel._check_uses_static_shape(path)
-        if is_dynamic and provider == "VitisAIExecutionProvider":
-            raise ValueError(
-                "The model provided has dynamic axes in input/output. Please provide model with static shapes for inference with RyzenAI."
-            )
+        from .utils import matmul_group_onnx
+
+        # from pdb import set_trace; set_trace()
+
+        path = matmul_group_onnx(path)
+        # from pdb import set_trace; set_trace()
 
         return ort.InferenceSession(
             path,
@@ -225,7 +226,7 @@ def _generate_regular_names_for_filename(filename: str):
         return [filename, f"{name}_quantized.{extension}", f"{name}_optimized.{extension}"]
 
     @classmethod
-    def _from_pretrained(
+    def _load_model_and_processors(
         cls,
         model_id: Union[str, Path],
         config: PretrainedConfig,
@@ -298,6 +299,8 @@ def _from_pretrained(
 
         preprocessors = None
         if model_path.is_dir():
+            cls.validate_static_shape_compatibility(model_path / file_name, provider)
+
             model = RyzenAIModel.load_model(
                 model_path / file_name,
                 provider=provider,
@@ -334,6 +337,8 @@ def _from_pretrained(
                 # model doesn't use external data
                 pass
 
+            cls.validate_static_shape_compatibility(model_cache_path, provider)
+
             model = RyzenAIModel.load_model(
                 model_cache_path,
                 provider=provider,
@@ -348,6 +353,49 @@ def _from_pretrained(
         if model_save_dir is None:
             model_save_dir = new_model_save_dir
 
+        return model, vaip_config, model_save_dir, preprocessors
+
+    @staticmethod
+    def validate_static_shape_compatibility(path: Union[str, Path], provider: str):
+        return True
+
+    @classmethod
+    def _from_pretrained(
+        cls,
+        model_id: Union[str, Path],
+        config: PretrainedConfig,
+        vaip_config: Optional[str] = None,
+        use_auth_token: Optional[Union[bool, str]] = None,
+        revision: Optional[str] = None,
+        force_download: bool = False,
+        cache_dir: Optional[str] = None,
+        file_name: Optional[str] = None,
+        subfolder: str = "",
+        local_files_only: bool = False,
+        provider: str = "VitisAIExecutionProvider",
+        session_options: Optional[ort.SessionOptions] = None,
+        provider_options: Optional[Dict[str, Any]] = None,
+        model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
+        **kwargs,
+    ) -> "RyzenAIModel":
+        model, vaip_config, model_save_dir, preprocessors = cls._load_model_and_processors(
+            model_id,
+            config,
+            vaip_config,
+            use_auth_token,
+            revision,
+            force_download,
+            cache_dir,
+            file_name,
+            subfolder,
+            local_files_only,
+            provider,
+            session_options,
+            provider_options,
+            model_save_dir,
+            **kwargs,
+        )
+
         return cls(
             model=model,
             config=config,
@@ -400,21 +448,41 @@ def from_pretrained(
         **kwargs,
     ):
         """
-        provider (`str`, defaults to `"VitisAIExecutionProvider"`):
-            ONNX Runtime provider to use for loading the model. See https://onnxruntime.ai/docs/execution-providers/ for
-            possible providers.
-        session_options (`Optional[onnxruntime.SessionOptions]`, defaults to `None`),:
-            ONNX Runtime session options to use for loading the model.
-        provider_options (`Optional[Dict[str, Any]]`, defaults to `None`):
-            Provider option dictionaries corresponding to the provider used. See available options
-            for each provider: https://onnxruntime.ai/docs/api/c/group___global.html .
-        kwargs (`Dict[str, Any]`):
-            Will be passed to the underlying model loading methods.
-
-        > Parameters for decoder models (RyzenAIForSpeechSeq2Seq)
-
-        use_cache (`Optional[bool]`, defaults to `True`):
-            Whether or not past key/values cache should be used. Defaults to `True`.
+        Instantiate a RyzenAIModel model from a model identifier.
+
+        Args:
+            model_id (`Union[str, Path]`):
+                The model identifier to instantiate the model from.
+            vaip_config (`str`, defaults to `None`):
+                The path to the Vitis AI config file.
+            export (`bool`, defaults to `False`):
+                Whether to export the model to ONNX before loading it.
+            force_download (`bool`, defaults to `False`):
+                Whether to force the download of the model files.
+            use_auth_token (`Optional[str]`, defaults to `None`):
+                The authorization token to use for downloading the model.
+            cache_dir (`Optional[str]`, defaults to `None`):
+                The directory to cache the model files.
+            subfolder (`str`, defaults to `""`):
+                The subfolder to look for the model files.
+            config (`Optional[PretrainedConfig]`, defaults to `None`):
+                The configuration to use for the model.
+            local_files_only (`bool`, defaults to `False`):
+                Whether to only look for the model files locally.
+            provider (`str`, defaults to `"VitisAIExecutionProvider"`):
+                ONNX Runtime provider to use for loading the model.
+            session_options (`Optional[onnxruntime.SessionOptions]`, defaults to `None`):
+                ONNX Runtime session options to use for loading the model.
+            provider_options (`Optional[Dict[str, Any]]`, defaults to `None`):
+                Provider option dictionaries corresponding to the provider used.
+            trust_remote_code (`bool`, defaults to `False`):
+                Whether to trust the remote code when exporting the model.
+            revision (`Optional[str]`, defaults to `None`):
+                The revision of the model to load.
+            library_name (`Optional[Dict[str, Any]]`, defaults to `None`):
+                The library name to use for the model.
+            kwargs (`Dict[str, Any]`):
+                Will be passed to the underlying model loading methods.
 
         Returns:
             `RyzenAIModel`: The loaded RyzenAIModel model.
@@ -564,6 +632,12 @@ def reshape(
 
         return model_path
 
+    def _convert_to_numpy(self, value, use_torch):
+        return value.cpu().detach().numpy() if use_torch else value
+
+    def _convert_to_tensor(self, value, use_torch):
+        return torch.from_numpy(value) if use_torch else torch.from_numpy(value)
+
 
 class RyzenAIModelForCustomTasks(RyzenAIModel):
     def forward(self, **kwargs):
@@ -672,6 +746,14 @@ def _export(
             **kwargs,
         )
 
+    @staticmethod
+    def validate_static_shape_compatibility(path: Union[str, Path], provider: str):
+        is_dynamic = RyzenAIModel._check_uses_static_shape(path)
+        if is_dynamic and provider == "VitisAIExecutionProvider":
+            raise ValueError(
+                "The model provided has dynamic axes in input/output. Please provide model with static shapes for inference with RyzenAI."
+            )
+
 
 class RyzenAIModelForObjectDetection(RyzenAIModelForCustomTasks):
     def forward(self, pixel_values):