Add Mistral fp16 config

microsoft · Mar 1, 2024 · 04d895e · 04d895e
1 parent 2587f17
commit 04d895e
Show file tree

Hide file tree

Showing 2 changed files with 114 additions and 0 deletions.
diff --git a/examples/mistral/mistral_fp16_optimize.json b/examples/mistral/mistral_fp16_optimize.json
@@ -0,0 +1,109 @@
+{
+    "input_model": {
+        "type": "PyTorchModel",
+        "config": {
+            "hf_config": {
+                "model_name": "mistralai/Mistral-7B-v0.1",
+                "model_class": "MistralForCausalLM"
+            }
+        }
+    },
+    "systems": {
+        "local_system": {
+            "type": "LocalSystem",
+            "config": {
+                "accelerators": [
+                    "gpu"
+                ]
+            }
+        }
+    },
+    "evaluators": {
+        "common_evaluator": {
+            "metrics": [
+                {
+                    "name": "latency",
+                    "type": "latency",
+                    "sub_types": [
+                        {
+                            "name": "avg",
+                            "priority": 1
+                        }
+                    ],
+                    "user_config": {
+                        "user_script": "user_script.py",
+                        "dataloader_func": "create_dataloader",
+                        "batch_size": 1,
+                        "inference_settings": {
+                            "onnx": {
+                                "session_options": {
+                                    "enable_profiling": false
+                                }
+                            }
+                        }
+                    }
+                }
+            ]
+        }
+    },
+    "passes": {
+        "convert": {
+            "type": "OptimumConversion",
+            "config": {
+                "target_opset": 14,
+                "extra_args": {
+                    "legacy": false,
+                    "no_post_process": false
+                }
+            }
+        },
+        "optimize": {
+            "type": "OrtTransformersOptimization",
+            "config": {
+                "model_type": "gpt2",
+                "use_gpu": true,
+                "keep_io_types": false,
+                "num_heads": 32,
+                "hidden_size": 4096,
+                "opt_level": 0,
+                "optimization_options": {
+                    "use_multi_head_attention": false
+                },
+                "save_as_external_data": true,
+                "all_tensors_to_one_file": true,
+                "float16": true,
+                "use_gqa": true
+            }
+        },
+        "perf_tuning": {
+            "type": "OrtPerfTuning",
+            "config": {
+                "user_script": "user_script.py",
+                "dataloader_func": "create_dataloader",
+                "batch_size": 1,
+                "enable_profiling": false
+            }
+        }
+    },
+    "pass_flows": [
+        [
+            "convert",
+            "optimize",
+            "perf_tuning"
+        ]
+    ],
+    "engine": {
+        "evaluate_input_model": false,
+        "evaluator": "common_evaluator",
+        "target": "local_system",
+        "cache_dir": "cache",
+        "output_name": "mistral",
+        "output_dir": "mistral_gptq",
+        "execution_providers": [
+            "CUDAExecutionProvider"
+        ],
+        "clean_cache": false,
+        "log_severity_level": 0,
+        "log_to_file": true
+    }
+}
diff --git a/examples/mistral/readme.md b/examples/mistral/readme.md
@@ -28,10 +28,15 @@ git config --system core.longpaths true
 ```
 
 ## Usage
+CPU:
 ```bash
 python -m olive.workflows.run --config mistral_optimize.json
 ```
 
+GPU:
+```bash
+python -m olive.workflows.run --config mistral_fp16_optimize.json
+```
 ### Local model
 if the input model is saved locally, you can specify the configuration like the following:
 ```json