examples/llama2/notebook/llama2_multiep/notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optimize Llama2 on multiple EPs(CPU/CUDA)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Olive does not only support the optimization of the model on a single execution provider but also on multiple execution providers. In this notebook, we will show you how to optimize the ResNet model on multiple execution providers (CPU and CUDA) using Olive.\n",
    "\n",
    "When enable multiple execution providers, Olive will try to optimize the input model on each execution provider one by one. As different execution providers have different environments and requirements, Olive will manage the environment for each execution provider separately. The user can specify the device and execution provider by themselves or let Olive choose the best execution provider for them.\n",
    "\n",
    "In Olive, we manage the environment by installing the required packages from the requirements_file or Dockerfile in the environment. Only Python environment system, Docker system, AzureML system support managed system. The device and execution_providers for managed system is mandatory. Otherwise, Olive will raise an error."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this case we will use the Llama2 model and optimize it on multiple execution providers."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prerequisite\n",
    "\n",
    "Before running this script, you need to install the required python packages. You can install them using the following command:\n",
    "\n",
    "```bash\n",
    "pip install -r ../../requirements.txt\n",
    "```\n",
    "\n",
    "Also, please ensure you already installed olive-ai. Please refer to the [installation guide](https://github.com/microsoft/Olive?tab=readme-ov-file#installation) for more information."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multiple Execution Providers\n",
    "\n",
    "Now, let's see how to optimize the model on multiple execution providers in Olive. Instead of creating two separate environments for CPU and CUDA execution providers, we can use Olive to optimize the model on multiple execution providers.\n",
    "\n",
    "\n",
    "In Olive, we use python `venv` to manage the environment. The user can specify the requirements file for the environment. Olive will install the required packages from the requirements file in the environment. The user can also specify the device and execution providers for the environment. Olive will manage the environment for each execution provider separately.\n",
    "\n",
    "\"\"Note that\"\": Olive will leverage the system's site packages to reduce the environment size and creation time. That may require the user to try the multi-ep case on a clean environment(*virtual env like conda/venv may fail*)* to avoid any conflicts.\n",
    "\n",
    "```json\n",
    " \"systems\": {\n",
    "    \"python_system\": {\n",
    "        \"type\": \"PythonEnvironment\",\n",
    "        \"config\": {\n",
    "            \"accelerators\": [\n",
    "                {\n",
    "                    \"device\": \"GPU\",\n",
    "                    \"execution_providers\": [\n",
    "                        \"CPUExecutionProvider\",\n",
    "                        \"CUDAExecutionProvider\"\n",
    "                    ]\n",
    "                }\n",
    "            ],\n",
    "            \"olive_managed_env\": true,    // <---------- let olive to install dependencies automatically\n",
    "            \"requirements_file\": \"multiple_ep_requirements.txt\" // <---------- requirements file for the environment\n",
    "        }\n",
    "    }\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-04-17 15:04:11,408] [INFO] [run.py:261:run] Loading Olive module configuration from: test/lib/python3.8/site-packages/olive/olive_config.json\n",
      "[2024-04-17 15:04:11,409] [INFO] [run.py:267:run] Loading run configuration from: config_multi_ep.json\n",
      "[2024-04-17 15:04:11,472] [INFO] [accelerator.py:336:create_accelerators] Running workflow on accelerator specs: gpu-cpu,gpu-cuda\n",
      "[2024-04-17 15:04:11,472] [INFO] [engine.py:106:initialize] Using cache directory: cache_folder//cache\n",
      "[2024-04-17 15:04:11,472] [INFO] [engine.py:262:run] Running Olive on accelerator: gpu-cpu\n",
      "[2024-04-17 15:04:11,473] [INFO] [engine.py:1074:_create_system] Creating target system ...\n",
      "[2024-04-17 15:04:13,464] [INFO] [misc.py:68:create_managed_system] Virtual environment 'tmp/olive_python_env_653h69jt' created.\n",
      "[2024-04-17 15:04:21,545] [INFO] [engine.py:1077:_create_system] Target system created in 10.071915 seconds\n",
      "[2024-04-17 15:04:21,545] [INFO] [engine.py:1086:_create_system] Creating host system ...\n",
      "[2024-04-17 15:04:23,518] [INFO] [misc.py:68:create_managed_system] Virtual environment 'tmp/olive_python_env_3j9yd7lv' created.\n",
      "[2024-04-17 15:04:30,727] [INFO] [engine.py:1089:_create_system] Host system created in 9.181122 seconds\n",
      "[2024-04-17 15:05:33,364] [INFO] [engine.py:324:run_accelerator] Input model evaluation results: {\n",
      "  \"latency_prompt_processing-avg\": 52.9096,\n",
      "  \"latency_token_generation-avg\": 48.74422\n",
      "}\n",
      "[2024-04-17 15:05:33,365] [INFO] [engine.py:329:run_accelerator] Saved evaluation results of input model to cache_folder//models/llama2_gpu/meta-llama/Llama-2-7b-hf/gpu-cpu_input_model_metrics.json\n",
      "[2024-04-17 15:05:33,366] [INFO] [engine.py:864:_run_pass] Running pass conversion_merged:OnnxConversion\n",
      "[2024-04-17 15:05:33,369] [INFO] [engine.py:898:_run_pass] Loaded model from cache: 0_OnnxConversion-78f9207c2bfa458863121c42f50c49e1-dfaff1da61d127bb5e9dc2f31a708897 from cache_folder//cache/runs\n",
      "[2024-04-17 15:05:33,369] [INFO] [engine.py:864:_run_pass] Running pass transformers_optimization_fp32:OrtTransformersOptimization\n",
      "[2024-04-17 15:05:33,373] [INFO] [engine.py:898:_run_pass] Loaded model from cache: 7_OrtTransformersOptimization-0-7da73ab63dc53bb765de0b1bf0e06614-gpu-cpu from cache_folder//cache/runs\n",
      "[2024-04-17 15:05:33,374] [INFO] [engine.py:864:_run_pass] Running pass blockwise_quant_int4:OnnxMatMul4Quantizer\n",
      "[2024-04-17 15:05:33,377] [INFO] [engine.py:898:_run_pass] Loaded model from cache: 8_OnnxMatMul4Quantizer-7-fe98a8b48ad1e4a5fc22fd942d64d9e6 from cache_folder//cache/runs\n",
      "[2024-04-17 15:05:33,377] [INFO] [engine.py:842:_run_passes] Run model evaluation for the final model...\n",
      "[2024-04-17 15:07:32,955] [INFO] [engine.py:864:_run_pass] Running pass conversion_merged:OnnxConversion\n",
      "[2024-04-17 15:07:32,957] [INFO] [engine.py:898:_run_pass] Loaded model from cache: 0_OnnxConversion-78f9207c2bfa458863121c42f50c49e1-dfaff1da61d127bb5e9dc2f31a708897 from cache_folder//cache/runs\n",
      "[2024-04-17 15:07:32,957] [INFO] [engine.py:864:_run_pass] Running pass transformers_optimization_fp16:OrtTransformersOptimization\n",
      "[2024-04-17 15:07:32,958] [INFO] [transformer_optimization.py:154:validate_search_point] CPUExecutionProvider does not support float16 very well, please avoid to use float16.\n",
      "[2024-04-17 15:07:32,958] [WARNING] [engine.py:870:_run_pass] Invalid search point, prune\n",
      "[2024-04-17 15:07:32,958] [WARNING] [engine.py:847:_run_passes] Skipping evaluation as model was pruned\n",
      "[2024-04-17 15:07:32,958] [WARNING] [engine.py:434:run_no_search] Flow ['conversion_merged', 'transformers_optimization_fp16'] is pruned due to failed or invalid config for pass 'transformers_optimization_fp16'\n",
      "[2024-04-17 15:07:32,958] [INFO] [engine.py:361:run_accelerator] Save footprint to cache_folder//models/llama2_gpu/meta-llama/Llama-2-7b-hf/gpu-cpu_footprints.json.\n",
      "[2024-04-17 15:07:32,960] [INFO] [engine.py:1095:_create_system] Removing target system ...\n",
      "[2024-04-17 15:07:32,989] [INFO] [python_environment_system.py:188:remove] Virtual environment 'tmp/olive_python_env_653h69jt' removed.\n",
      "[2024-04-17 15:07:33,019] [INFO] [python_environment_system.py:195:remove] Temporary directory 'tmp' removed.\n",
      "[2024-04-17 15:07:33,019] [INFO] [engine.py:1099:_create_system] Removing host system ...\n",
      "[2024-04-17 15:07:33,019] [INFO] [engine.py:262:run] Running Olive on accelerator: gpu-cuda\n",
      "[2024-04-17 15:07:33,019] [INFO] [engine.py:1074:_create_system] Creating target system ...\n",
      "[2024-04-17 15:07:34,967] [INFO] [misc.py:68:create_managed_system] Virtual environment 'tmp/olive_python_env_buh7851z' created.\n",
      "[2024-04-17 15:07:40,486] [INFO] [engine.py:1077:_create_system] Target system created in 7.466205 seconds\n",
      "[2024-04-17 15:07:40,486] [INFO] [engine.py:1086:_create_system] Creating host system ...\n",
      "[2024-04-17 15:07:42,462] [INFO] [misc.py:68:create_managed_system] Virtual environment 'tmp/olive_python_env_4k_x_g6j' created.\n",
      "[2024-04-17 15:07:50,233] [INFO] [engine.py:1089:_create_system] Host system created in 9.747077 seconds\n",
      "[2024-04-17 15:07:50,363] [INFO] [engine.py:324:run_accelerator] Input model evaluation results: {\n",
      "  \"latency_prompt_processing-avg\": 53.67267,\n",
      "  \"latency_token_generation-avg\": 47.0224\n",
      "}\n",
      "[2024-04-17 15:07:50,364] [INFO] [engine.py:329:run_accelerator] Saved evaluation results of input model to cache_folder//models/llama2_gpu/meta-llama/Llama-2-7b-hf/gpu-cuda_input_model_metrics.json\n",
      "[2024-04-17 15:07:50,364] [INFO] [engine.py:864:_run_pass] Running pass conversion_merged:OnnxConversion\n",
      "[2024-04-17 15:07:50,365] [INFO] [engine.py:898:_run_pass] Loaded model from cache: 0_OnnxConversion-78f9207c2bfa458863121c42f50c49e1-dfaff1da61d127bb5e9dc2f31a708897 from cache_folder//cache/runs\n",
      "[2024-04-17 15:07:50,366] [INFO] [engine.py:864:_run_pass] Running pass transformers_optimization_fp32:OrtTransformersOptimization\n",
      "[2024-04-17 15:07:50,368] [INFO] [engine.py:898:_run_pass] Loaded model from cache: 3_OrtTransformersOptimization-0-8d800f56ee00b5607adb7b96f5076e59-gpu-cuda from cache_folder//cache/runs\n",
      "[2024-04-17 15:07:50,368] [INFO] [engine.py:864:_run_pass] Running pass blockwise_quant_int4:OnnxMatMul4Quantizer\n",
      "[2024-04-17 15:07:50,370] [INFO] [engine.py:898:_run_pass] Loaded model from cache: 4_OnnxMatMul4Quantizer-3-fe98a8b48ad1e4a5fc22fd942d64d9e6 from cache_folder//cache/runs\n",
      "[2024-04-17 15:07:50,370] [INFO] [engine.py:842:_run_passes] Run model evaluation for the final model...\n",
      "[2024-04-17 15:07:54,703] [INFO] [engine.py:864:_run_pass] Running pass conversion_merged:OnnxConversion\n",
      "[2024-04-17 15:07:54,704] [INFO] [engine.py:898:_run_pass] Loaded model from cache: 0_OnnxConversion-78f9207c2bfa458863121c42f50c49e1-dfaff1da61d127bb5e9dc2f31a708897 from cache_folder//cache/runs\n",
      "[2024-04-17 15:07:54,704] [INFO] [engine.py:864:_run_pass] Running pass transformers_optimization_fp16:OrtTransformersOptimization\n",
      "[2024-04-17 15:07:54,706] [INFO] [engine.py:898:_run_pass] Loaded model from cache: 5_OrtTransformersOptimization-0-9a2d6ea047237bbe7e923e3ede6e3570-gpu-cuda from cache_folder//cache/runs\n",
      "[2024-04-17 15:07:54,706] [INFO] [engine.py:842:_run_passes] Run model evaluation for the final model...\n",
      "[2024-04-17 15:08:06,459] [INFO] [engine.py:361:run_accelerator] Save footprint to cache_folder//models/llama2_gpu/meta-llama/Llama-2-7b-hf/gpu-cuda_footprints.json.\n",
      "[2024-04-17 15:08:06,461] [INFO] [engine.py:1095:_create_system] Removing target system ...\n",
      "[2024-04-17 15:08:06,475] [INFO] [python_environment_system.py:188:remove] Virtual environment 'tmp/olive_python_env_buh7851z' removed.\n",
      "[2024-04-17 15:08:06,504] [INFO] [python_environment_system.py:195:remove] Temporary directory 'tmp' removed.\n",
      "[2024-04-17 15:08:06,505] [INFO] [engine.py:1099:_create_system] Removing host system ...\n",
      "[2024-04-17 15:08:06,505] [INFO] [engine.py:279:run] Run history for gpu-cpu:\n",
      "[2024-04-17 15:08:06,514] [INFO] [engine.py:567:dump_run_history] run history:\n",
      "+------------------------------------------------------------------------------------+------------------------------------------------------------------------------------+-----------------------------+----------------+-----------------------------------------------+\n",
      "| model_id                                                                           | parent_model_id                                                                    | from_pass                   |   duration_sec | metrics                                       |\n",
      "+====================================================================================+====================================================================================+=============================+================+===============================================+\n",
      "| 78f9207c2bfa458863121c42f50c49e1                                                   |                                                                                    |                             |                | {                                             |\n",
      "|                                                                                    |                                                                                    |                             |                |   \"latency_prompt_processing-avg\": 52.9096,   |\n",
      "|                                                                                    |                                                                                    |                             |                |   \"latency_token_generation-avg\": 48.74422    |\n",
      "|                                                                                    |                                                                                    |                             |                | }                                             |\n",
      "+------------------------------------------------------------------------------------+------------------------------------------------------------------------------------+-----------------------------+----------------+-----------------------------------------------+\n",
      "| 0_OnnxConversion-78f9207c2bfa458863121c42f50c49e1-dfaff1da61d127bb5e9dc2f31a708897 | 78f9207c2bfa458863121c42f50c49e1                                                   | OnnxConversion              |       439.78   |                                               |\n",
      "+------------------------------------------------------------------------------------+------------------------------------------------------------------------------------+-----------------------------+----------------+-----------------------------------------------+\n",
      "| 7_OrtTransformersOptimization-0-7da73ab63dc53bb765de0b1bf0e06614-gpu-cpu           | 0_OnnxConversion-78f9207c2bfa458863121c42f50c49e1-dfaff1da61d127bb5e9dc2f31a708897 | OrtTransformersOptimization |       128.437  |                                               |\n",
      "+------------------------------------------------------------------------------------+------------------------------------------------------------------------------------+-----------------------------+----------------+-----------------------------------------------+\n",
      "| 8_OnnxMatMul4Quantizer-7-fe98a8b48ad1e4a5fc22fd942d64d9e6                          | 7_OrtTransformersOptimization-0-7da73ab63dc53bb765de0b1bf0e06614-gpu-cpu           | OnnxMatMul4Quantizer        |        71.4255 | {                                             |\n",
      "|                                                                                    |                                                                                    |                             |                |   \"latency_prompt_processing-avg\": 1772.6851, |\n",
      "|                                                                                    |                                                                                    |                             |                |   \"latency_token_generation-avg\": 1743.38576  |\n",
      "|                                                                                    |                                                                                    |                             |                | }                                             |\n",
      "+------------------------------------------------------------------------------------+------------------------------------------------------------------------------------+-----------------------------+----------------+-----------------------------------------------+\n",
      "[2024-04-17 15:08:06,514] [INFO] [engine.py:279:run] Run history for gpu-cuda:\n",
      "[2024-04-17 15:08:06,516] [INFO] [engine.py:567:dump_run_history] run history:\n",
      "+------------------------------------------------------------------------------------+------------------------------------------------------------------------------------+-----------------------------+----------------+-----------------------------------------------+\n",
      "| model_id                                                                           | parent_model_id                                                                    | from_pass                   |   duration_sec | metrics                                       |\n",
      "+====================================================================================+====================================================================================+=============================+================+===============================================+\n",
      "| 78f9207c2bfa458863121c42f50c49e1                                                   |                                                                                    |                             |                | {                                             |\n",
      "|                                                                                    |                                                                                    |                             |                |   \"latency_prompt_processing-avg\": 53.67267,  |\n",
      "|                                                                                    |                                                                                    |                             |                |   \"latency_token_generation-avg\": 47.0224     |\n",
      "|                                                                                    |                                                                                    |                             |                | }                                             |\n",
      "+------------------------------------------------------------------------------------+------------------------------------------------------------------------------------+-----------------------------+----------------+-----------------------------------------------+\n",
      "| 0_OnnxConversion-78f9207c2bfa458863121c42f50c49e1-dfaff1da61d127bb5e9dc2f31a708897 | 78f9207c2bfa458863121c42f50c49e1                                                   | OnnxConversion              |       439.78   |                                               |\n",
      "+------------------------------------------------------------------------------------+------------------------------------------------------------------------------------+-----------------------------+----------------+-----------------------------------------------+\n",
      "| 3_OrtTransformersOptimization-0-8d800f56ee00b5607adb7b96f5076e59-gpu-cuda          | 0_OnnxConversion-78f9207c2bfa458863121c42f50c49e1-dfaff1da61d127bb5e9dc2f31a708897 | OrtTransformersOptimization |       124.263  |                                               |\n",
      "+------------------------------------------------------------------------------------+------------------------------------------------------------------------------------+-----------------------------+----------------+-----------------------------------------------+\n",
      "| 4_OnnxMatMul4Quantizer-3-fe98a8b48ad1e4a5fc22fd942d64d9e6                          | 3_OrtTransformersOptimization-0-8d800f56ee00b5607adb7b96f5076e59-gpu-cuda          | OnnxMatMul4Quantizer        |        70.1778 | {                                             |\n",
      "|                                                                                    |                                                                                    |                             |                |   \"latency_prompt_processing-avg\": 207.52853, |\n",
      "|                                                                                    |                                                                                    |                             |                |   \"latency_token_generation-avg\": 191.34901   |\n",
      "|                                                                                    |                                                                                    |                             |                | }                                             |\n",
      "+------------------------------------------------------------------------------------+------------------------------------------------------------------------------------+-----------------------------+----------------+-----------------------------------------------+\n",
      "| 5_OrtTransformersOptimization-0-9a2d6ea047237bbe7e923e3ede6e3570-gpu-cuda          | 0_OnnxConversion-78f9207c2bfa458863121c42f50c49e1-dfaff1da61d127bb5e9dc2f31a708897 | OrtTransformersOptimization |       511.793  | {                                             |\n",
      "|                                                                                    |                                                                                    |                             |                |   \"latency_prompt_processing-avg\": 25.82543,  |\n",
      "|                                                                                    |                                                                                    |                             |                |   \"latency_token_generation-avg\": 25.14717    |\n",
      "|                                                                                    |                                                                                    |                             |                | }                                             |\n",
      "+------------------------------------------------------------------------------------+------------------------------------------------------------------------------------+-----------------------------+----------------+-----------------------------------------------+\n",
      "[2024-04-17 15:08:06,516] [INFO] [engine.py:285:run] Package top ranked 3 models as artifacts\n",
      "[2024-04-17 15:08:06,517] [INFO] [packaging_generator.py:67:_package_candidate_models] Packaging output models to PackagingType.Zipfile\n"
     ]
    }
   ],
   "source": [
    "! python3 -m olive.workflows.run --config config_multi_ep.json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the support of multiple execution providers, Olive will optimize the model on each execution provider one by one and package the optimized model into a single file under the `output_dir` set by users.\n",
    "\n",
    "You can check the `models_rank.json` under the zipped file to see the optimized model details:\n",
    "\n",
    "Here is the optimal model output for the Llama2 model on multiple execution providers:\n",
    "```json\n",
    "{\n",
    "    \"rank\": 1,\n",
    "    \"model_config\": {\n",
    "        \"type\": \"ONNXModel\",\n",
    "        \"config\": {\n",
    "            \"model_path\": \"CandidateModels/gpu-cuda/BestCandidateModel_1\",\n",
    "            \"onnx_file_name\": \"model.onnx\",\n",
    "            \"inference_settings\": null,\n",
    "            \"use_ort_extensions\": false,\n",
    "            \"external_initializers_file_name\": null,\n",
    "            \"constant_inputs_file_name\": null,\n",
    "            \"model_attributes\": {\n",
    "                \"vocab_size\": 32000,\n",
    "                \"max_position_embeddings\": 4096,\n",
    "                \"hidden_size\": 4096,\n",
    "                \"intermediate_size\": 11008,\n",
    "                \"num_hidden_layers\": 32,\n",
    "                \"num_attention_heads\": 32,\n",
    "                \"num_key_value_heads\": 32,\n",
    "                \"hidden_act\": \"silu\",\n",
    "                \"initializer_range\": 0.02,\n",
    "                \"rms_norm_eps\": 1e-05,\n",
    "                \"pretraining_tp\": 1,\n",
    "                \"use_cache\": true,\n",
    "                \"rope_theta\": 10000.0,\n",
    "                \"rope_scaling\": null,\n",
    "                \"attention_bias\": false,\n",
    "                \"attention_dropout\": 0.0,\n",
    "                \"return_dict\": true,\n",
    "                \"output_hidden_states\": false,\n",
    "                \"output_attentions\": false,\n",
    "                \"torchscript\": false,\n",
    "                \"torch_dtype\": \"float16\",\n",
    "                \"use_bfloat16\": false,\n",
    "                \"tf_legacy_loss\": false,\n",
    "                \"pruned_heads\": {},\n",
    "                \"tie_word_embeddings\": false,\n",
    "                \"chunk_size_feed_forward\": 0,\n",
    "                \"is_encoder_decoder\": false,\n",
    "                \"is_decoder\": false,\n",
    "                \"cross_attention_hidden_size\": null,\n",
    "                \"add_cross_attention\": false,\n",
    "                \"tie_encoder_decoder\": false,\n",
    "                \"max_length\": 20,\n",
    "                \"min_length\": 0,\n",
    "                \"do_sample\": false,\n",
    "                \"early_stopping\": false,\n",
    "                \"num_beams\": 1,\n",
    "                \"num_beam_groups\": 1,\n",
    "                \"diversity_penalty\": 0.0,\n",
    "                \"temperature\": 1.0,\n",
    "                \"top_k\": 50,\n",
    "                \"top_p\": 1.0,\n",
    "                \"typical_p\": 1.0,\n",
    "                \"repetition_penalty\": 1.0,\n",
    "                \"length_penalty\": 1.0,\n",
    "                \"no_repeat_ngram_size\": 0,\n",
    "                \"encoder_no_repeat_ngram_size\": 0,\n",
    "                \"bad_words_ids\": null,\n",
    "                \"num_return_sequences\": 1,\n",
    "                \"output_scores\": false,\n",
    "                \"return_dict_in_generate\": false,\n",
    "                \"forced_bos_token_id\": null,\n",
    "                \"forced_eos_token_id\": null,\n",
    "                \"remove_invalid_values\": false,\n",
    "                \"exponential_decay_length_penalty\": null,\n",
    "                \"suppress_tokens\": null,\n",
    "                \"begin_suppress_tokens\": null,\n",
    "                \"architectures\": [\n",
    "                    \"LlamaForCausalLM\"\n",
    "                ],\n",
    "                \"finetuning_task\": null,\n",
    "                \"id2label\": {\n",
    "                    \"0\": \"LABEL_0\",\n",
    "                    \"1\": \"LABEL_1\"\n",
    "                },\n",
    "                \"label2id\": {\n",
    "                    \"LABEL_0\": 0,\n",
    "                    \"LABEL_1\": 1\n",
    "                },\n",
    "                \"tokenizer_class\": null,\n",
    "                \"prefix\": null,\n",
    "                \"bos_token_id\": 1,\n",
    "                \"pad_token_id\": null,\n",
    "                \"eos_token_id\": 2,\n",
    "                \"sep_token_id\": null,\n",
    "                \"decoder_start_token_id\": null,\n",
    "                \"task_specific_params\": null,\n",
    "                \"problem_type\": null,\n",
    "                \"_name_or_path\": \"meta-llama/Llama-2-7b-hf\",\n",
    "                \"transformers_version\": \"4.39.3\",\n",
    "                \"model_type\": \"llama\"\n",
    "            }\n",
    "        }\n",
    "    },\n",
    "    \"metrics\": {\n",
    "        \"latency_prompt_processing-avg\": {\n",
    "            \"value\": 25.82543,\n",
    "            \"priority\": 1,\n",
    "            \"higher_is_better\": false\n",
    "        },\n",
    "        \"latency_token_generation-avg\": {\n",
    "            \"value\": 25.14717,\n",
    "            \"priority\": -1,\n",
    "            \"higher_is_better\": false\n",
    "        }\n",
    "    }\n",
    "}\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Legacy way to optimize the model on multiple execution providers\n",
    "\n",
    "Before we start, let's see how to optimize the model on multiple execution providers in the legacy way.\n",
    "Since CPU execution provider requires `onnxruntime`, but CUDA execution provider requires `onnxruntime-gpu`, and they cannot be installed in the same environment, we need to create two separate environments for CPU and CUDA execution providers then run Olive to optimize."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us try this kind of optimization in both CPU and GPU firstly to see the difference.\n",
    "\n",
    "Optimization stack: \n",
    "1. OnnxConversion: Convert the model to ONNX format.\n",
    "2. OrtTransformersOptimization: Optimize the model with transformers optimization.\n",
    "3. OnnxMatMul4Quantizer: Quantize the model with blockwise quantization."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### CPU System config\n",
    "\n",
    "For CPU system, we need to install `onnxruntime` in the environment firstly,\n",
    "then we can add the following system config to the Olive config file.\n",
    "```json\n",
    "\"systems\": {\n",
    "    \"local_system\": {\n",
    "        \"type\": \"LocalSystem\",\n",
    "        \"config\": {\n",
    "            \"accelerators\": [\n",
    "                {\n",
    "                    \"device\": \"CPU\",\n",
    "                    \"execution_providers\": [\n",
    "                        \"CPUExecutionProvider\"\n",
    "                    ]\n",
    "                }\n",
    "            ]\n",
    "        }\n",
    "    }\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "metadata": {}
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found existing installation: onnxruntime 1.17.3\n",
      "Uninstalling onnxruntime-1.17.3:\n",
      "  Successfully uninstalled onnxruntime-1.17.3\n",
      "\u001b[33mWARNING: Skipping onnxruntime-gpu as it is not installed.\u001b[0m\u001b[33m\n",
      "\u001b[0mDefaulting to user installation because normal site-packages is not writeable\n",
      "Collecting onnxruntime\n",
      "  Using cached onnxruntime-1.17.3-cp38-cp38-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (4.4 kB)\n",
      "Requirement already satisfied: coloredlogs in test/lib/python3.8/site-packages (from onnxruntime) (15.0.1)\n",
      "Requirement already satisfied: flatbuffers in test/lib/python3.8/site-packages (from onnxruntime) (23.5.26)\n",
      "Requirement already satisfied: numpy>=1.21.6 in test/lib/python3.8/site-packages (from onnxruntime) (1.24.4)\n",
      "Requirement already satisfied: packaging in test/lib/python3.8/site-packages (from onnxruntime) (23.2)\n",
      "Requirement already satisfied: protobuf in test/lib/python3.8/site-packages (from onnxruntime) (3.20.2)\n",
      "Requirement already satisfied: sympy in test/lib/python3.8/site-packages (from onnxruntime) (1.12)\n",
      "Requirement already satisfied: humanfriendly>=9.1 in test/lib/python3.8/site-packages (from coloredlogs->onnxruntime) (10.0)\n",
      "Requirement already satisfied: mpmath>=0.19 in test/lib/python3.8/site-packages (from sympy->onnxruntime) (1.3.0)\n",
      "Using cached onnxruntime-1.17.3-cp38-cp38-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (6.8 MB)\n",
      "Installing collected packages: onnxruntime\n",
      "Successfully installed onnxruntime-1.17.3\n"
     ]
    }
   ],
   "source": [
    "! python3 -m pip uninstall onnxruntime onnxruntime-gpu -y\n",
    "! python3 -m pip install onnxruntime -U"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "metadata": {}
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-04-17 14:57:23,246] [INFO] [run.py:261:run] Loading Olive module configuration from: test/lib/python3.8/site-packages/olive/olive_config.json\n",
      "[2024-04-17 14:57:23,247] [INFO] [run.py:267:run] Loading run configuration from: config_cpu.json\n",
      "[2024-04-17 14:57:23,305] [INFO] [accelerator.py:336:create_accelerators] Running workflow on accelerator specs: cpu-cpu\n",
      "[2024-04-17 14:57:23,306] [INFO] [engine.py:106:initialize] Using cache directory: cache_folder//cache\n",
      "[2024-04-17 14:57:23,306] [INFO] [engine.py:262:run] Running Olive on accelerator: cpu-cpu\n",
      "[2024-04-17 14:57:23,306] [INFO] [engine.py:1074:_create_system] Creating target system ...\n",
      "[2024-04-17 14:57:23,307] [INFO] [engine.py:1077:_create_system] Target system created in 0.000136 seconds\n",
      "[2024-04-17 14:57:23,307] [INFO] [engine.py:1086:_create_system] Creating host system ...\n",
      "[2024-04-17 14:57:23,307] [INFO] [engine.py:1089:_create_system] Host system created in 0.000122 seconds\n",
      "data loader kwargs: {'batch_size': 2, 'model_framework': <Framework.PYTORCH: 'PyTorch'>, 'model_id': 'meta-llama/Llama-2-7b-hf', 'past_seq_length': 0, 'seq_length': 8, 'max_seq_length': 2048}\n",
      "[2024-04-17 14:57:24,611] [INFO] [hf_config.py:68:load_hf_model] Loading Huggingface model from meta-llama/Llama-2-7b-hf\n",
      "Loading checkpoint shards: 100%|██████████████████| 2/2 [00:02<00:00,  1.20s/it]\n",
      "data loader kwargs: {'batch_size': 2, 'model_framework': <Framework.PYTORCH: 'PyTorch'>, 'model_id': 'meta-llama/Llama-2-7b-hf', 'past_seq_length': 8, 'seq_length': 1, 'max_seq_length': 2048}\n",
      "[2024-04-17 14:58:05,610] [INFO] [engine.py:324:run_accelerator] Input model evaluation results: {\n",
      "  \"latency_prompt_processing-avg\": 590.32982,\n",
      "  \"latency_token_generation-avg\": 501.17834\n",
      "}\n",
      "[2024-04-17 14:58:05,611] [INFO] [engine.py:329:run_accelerator] Saved evaluation results of input model to cache_folder//models/llama2_gpu/meta-llama/Llama-2-7b-hf/cpu-cpu_input_model_metrics.json\n",
      "[2024-04-17 14:58:05,611] [INFO] [engine.py:864:_run_pass] Running pass conversion_merged:OnnxConversion\n",
      "[2024-04-17 14:58:05,612] [INFO] [engine.py:898:_run_pass] Loaded model from cache: 0_OnnxConversion-78f9207c2bfa458863121c42f50c49e1-dfaff1da61d127bb5e9dc2f31a708897 from cache_folder//cache/runs\n",
      "[2024-04-17 14:58:05,612] [INFO] [engine.py:864:_run_pass] Running pass transformers_optimization_fp32:OrtTransformersOptimization\n",
      "[2024-04-17 14:58:05,614] [INFO] [engine.py:898:_run_pass] Loaded model from cache: 1_OrtTransformersOptimization-0-7da73ab63dc53bb765de0b1bf0e06614-cpu-cpu from cache_folder//cache/runs\n",
      "[2024-04-17 14:58:05,614] [INFO] [engine.py:864:_run_pass] Running pass blockwise_quant_int4:OnnxMatMul4Quantizer\n",
      "[2024-04-17 14:58:05,616] [INFO] [engine.py:898:_run_pass] Loaded model from cache: 2_OnnxMatMul4Quantizer-1-fe98a8b48ad1e4a5fc22fd942d64d9e6 from cache_folder//cache/runs\n",
      "[2024-04-17 14:58:05,616] [INFO] [engine.py:842:_run_passes] Run model evaluation for the final model...\n",
      "data loader kwargs: {'batch_size': 2, 'model_framework': <Framework.ONNX: 'ONNX'>, 'model_id': 'meta-llama/Llama-2-7b-hf', 'past_seq_length': 0, 'seq_length': 8, 'max_seq_length': 2048}\n",
      "data loader kwargs: {'batch_size': 2, 'model_framework': <Framework.ONNX: 'ONNX'>, 'model_id': 'meta-llama/Llama-2-7b-hf', 'past_seq_length': 8, 'seq_length': 1, 'max_seq_length': 2048}\n",
      "[2024-04-17 15:00:53,176] [INFO] [engine.py:864:_run_pass] Running pass conversion_merged:OnnxConversion\n",
      "[2024-04-17 15:00:53,178] [INFO] [engine.py:898:_run_pass] Loaded model from cache: 0_OnnxConversion-78f9207c2bfa458863121c42f50c49e1-dfaff1da61d127bb5e9dc2f31a708897 from cache_folder//cache/runs\n",
      "[2024-04-17 15:00:53,179] [INFO] [engine.py:864:_run_pass] Running pass transformers_optimization_fp16:OrtTransformersOptimization\n",
      "[2024-04-17 15:00:53,180] [INFO] [transformer_optimization.py:154:validate_search_point] CPUExecutionProvider does not support float16 very well, please avoid to use float16.\n",
      "[2024-04-17 15:00:53,180] [WARNING] [engine.py:870:_run_pass] Invalid search point, prune\n",
      "[2024-04-17 15:00:53,180] [WARNING] [engine.py:847:_run_passes] Skipping evaluation as model was pruned\n",
      "[2024-04-17 15:00:53,180] [WARNING] [engine.py:434:run_no_search] Flow ['conversion_merged', 'transformers_optimization_fp16'] is pruned due to failed or invalid config for pass 'transformers_optimization_fp16'\n",
      "[2024-04-17 15:00:53,181] [INFO] [engine.py:361:run_accelerator] Save footprint to cache_folder//models/llama2_gpu/meta-llama/Llama-2-7b-hf/cpu-cpu_footprints.json.\n",
      "[2024-04-17 15:00:53,184] [INFO] [engine.py:279:run] Run history for cpu-cpu:\n",
      "[2024-04-17 15:00:53,195] [INFO] [engine.py:567:dump_run_history] run history:\n",
      "+------------------------------------------------------------------------------------+------------------------------------------------------------------------------------+-----------------------------+----------------+-----------------------------------------------+\n",
      "| model_id                                                                           | parent_model_id                                                                    | from_pass                   |   duration_sec | metrics                                       |\n",
      "+====================================================================================+====================================================================================+=============================+================+===============================================+\n",
      "| 78f9207c2bfa458863121c42f50c49e1                                                   |                                                                                    |                             |                | {                                             |\n",
      "|                                                                                    |                                                                                    |                             |                |   \"latency_prompt_processing-avg\": 590.32982, |\n",
      "|                                                                                    |                                                                                    |                             |                |   \"latency_token_generation-avg\": 501.17834   |\n",
      "|                                                                                    |                                                                                    |                             |                | }                                             |\n",
      "+------------------------------------------------------------------------------------+------------------------------------------------------------------------------------+-----------------------------+----------------+-----------------------------------------------+\n",
      "| 0_OnnxConversion-78f9207c2bfa458863121c42f50c49e1-dfaff1da61d127bb5e9dc2f31a708897 | 78f9207c2bfa458863121c42f50c49e1                                                   | OnnxConversion              |       439.78   |                                               |\n",
      "+------------------------------------------------------------------------------------+------------------------------------------------------------------------------------+-----------------------------+----------------+-----------------------------------------------+\n",
      "| 1_OrtTransformersOptimization-0-7da73ab63dc53bb765de0b1bf0e06614-cpu-cpu           | 0_OnnxConversion-78f9207c2bfa458863121c42f50c49e1-dfaff1da61d127bb5e9dc2f31a708897 | OrtTransformersOptimization |       110.275  |                                               |\n",
      "+------------------------------------------------------------------------------------+------------------------------------------------------------------------------------+-----------------------------+----------------+-----------------------------------------------+\n",
      "| 2_OnnxMatMul4Quantizer-1-fe98a8b48ad1e4a5fc22fd942d64d9e6                          | 1_OrtTransformersOptimization-0-7da73ab63dc53bb765de0b1bf0e06614-cpu-cpu           | OnnxMatMul4Quantizer        |        65.0451 | {                                             |\n",
      "|                                                                                    |                                                                                    |                             |                |   \"latency_prompt_processing-avg\": 764.51641, |\n",
      "|                                                                                    |                                                                                    |                             |                |   \"latency_token_generation-avg\": 192.5402    |\n",
      "|                                                                                    |                                                                                    |                             |                | }                                             |\n",
      "+------------------------------------------------------------------------------------+------------------------------------------------------------------------------------+-----------------------------+----------------+-----------------------------------------------+\n",
      "[2024-04-17 15:00:53,195] [INFO] [engine.py:294:run] No packaging config provided, skip packaging artifacts\n"
     ]
    }
   ],
   "source": [
    "# Then run the workflow:\n",
    "! python3 -m olive.workflows.run --config config_cpu.json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we can get the:\n",
    "\n",
    "- `latency_prompt_processing`: The time taken to process the prompt where the KV_cache is not used.\n",
    "- `latency_token_generation`: The time taken to generate the tokens where the KV_cache is used.\n",
    "\n",
    "\n",
    "| Model Type | latency_prompt_processing(ms) | latency_token_generation(ms) |\n",
    "|------------|---------------------------|-------------------------|\n",
    "| Torch Model | 590.32982 | 501.17834 |\n",
    "| Olive Optimized Model | 764.51641 | 192.5402 |\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### GPU System config\n",
    "\n",
    "For GPU system, we need to install `onnxruntime-gpu` in the environment firstly,\n",
    "then we can add the following system config to the Olive config file.\n",
    "```json\n",
    "\"systems\": {\n",
    "    \"local_system\": {\n",
    "        \"type\": \"LocalSystem\",\n",
    "        \"config\": {\n",
    "            \"accelerators\": [\n",
    "                {\n",
    "                    \"device\": \"gpu\",\n",
    "                    \"execution_providers\": [\n",
    "                        \"GPUExecutionProvider\"\n",
    "                    ]\n",
    "                }\n",
    "            ]\n",
    "        }\n",
    "    }\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "metadata": {}
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found existing installation: onnxruntime 1.17.3\n",
      "Uninstalling onnxruntime-1.17.3:\n",
      "  Successfully uninstalled onnxruntime-1.17.3\n",
      "\u001b[33mWARNING: Skipping onnxruntime-gpu as it is not installed.\u001b[0m\u001b[33m\n",
      "\u001b[0mDefaulting to user installation because normal site-packages is not writeable\n",
      "Collecting onnxruntime-gpu\n",
      "  Using cached onnxruntime_gpu-1.17.1-cp38-cp38-manylinux_2_28_x86_64.whl.metadata (4.3 kB)\n",
      "Requirement already satisfied: coloredlogs in test/lib/python3.8/site-packages (from onnxruntime-gpu) (15.0.1)\n",
      "Requirement already satisfied: flatbuffers in test/lib/python3.8/site-packages (from onnxruntime-gpu) (23.5.26)\n",
      "Requirement already satisfied: numpy>=1.21.6 in test/lib/python3.8/site-packages (from onnxruntime-gpu) (1.24.4)\n",
      "Requirement already satisfied: packaging in test/lib/python3.8/site-packages (from onnxruntime-gpu) (23.2)\n",
      "Requirement already satisfied: protobuf in test/lib/python3.8/site-packages (from onnxruntime-gpu) (3.20.2)\n",
      "Requirement already satisfied: sympy in test/lib/python3.8/site-packages (from onnxruntime-gpu) (1.12)\n",
      "Requirement already satisfied: humanfriendly>=9.1 in test/lib/python3.8/site-packages (from coloredlogs->onnxruntime-gpu) (10.0)\n",
      "Requirement already satisfied: mpmath>=0.19 in test/lib/python3.8/site-packages (from sympy->onnxruntime-gpu) (1.3.0)\n",
      "Using cached onnxruntime_gpu-1.17.1-cp38-cp38-manylinux_2_28_x86_64.whl (192.1 MB)\n",
      "Installing collected packages: onnxruntime-gpu\n",
      "Successfully installed onnxruntime-gpu-1.17.1\n"
     ]
    }
   ],
   "source": [
    "! python3 -m pip uninstall onnxruntime onnxruntime-gpu -y\n",
    "! python3 -m pip install onnxruntime-gpu -U"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "metadata": {}
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-04-17 15:02:03,396] [INFO] [run.py:261:run] Loading Olive module configuration from: test/lib/python3.8/site-packages/olive/olive_config.json\n",
      "[2024-04-17 15:02:03,397] [INFO] [run.py:267:run] Loading run configuration from: config_gpu.json\n",
      "[2024-04-17 15:02:03,464] [INFO] [accelerator.py:336:create_accelerators] Running workflow on accelerator specs: gpu-cuda\n",
      "[2024-04-17 15:02:03,464] [INFO] [engine.py:106:initialize] Using cache directory: cache_folder//cache\n",
      "[2024-04-17 15:02:03,465] [INFO] [engine.py:262:run] Running Olive on accelerator: gpu-cuda\n",
      "[2024-04-17 15:02:03,465] [INFO] [engine.py:1074:_create_system] Creating target system ...\n",
      "[2024-04-17 15:02:03,466] [INFO] [engine.py:1077:_create_system] Target system created in 0.000136 seconds\n",
      "[2024-04-17 15:02:03,466] [INFO] [engine.py:1086:_create_system] Creating host system ...\n",
      "[2024-04-17 15:02:03,466] [INFO] [engine.py:1089:_create_system] Host system created in 0.000114 seconds\n",
      "data loader kwargs: {'batch_size': 2, 'model_framework': <Framework.PYTORCH: 'PyTorch'>, 'model_id': 'meta-llama/Llama-2-7b-hf', 'past_seq_length': 0, 'seq_length': 8, 'max_seq_length': 2048}\n",
      "[2024-04-17 15:02:04,837] [INFO] [hf_config.py:68:load_hf_model] Loading Huggingface model from meta-llama/Llama-2-7b-hf\n",
      "Loading checkpoint shards: 100%|██████████████████| 2/2 [00:01<00:00,  1.15it/s]\n",
      "data loader kwargs: {'batch_size': 2, 'model_framework': <Framework.PYTORCH: 'PyTorch'>, 'model_id': 'meta-llama/Llama-2-7b-hf', 'past_seq_length': 8, 'seq_length': 1, 'max_seq_length': 2048}\n",
      "[2024-04-17 15:02:57,153] [INFO] [engine.py:324:run_accelerator] Input model evaluation results: {\n",
      "  \"latency_prompt_processing-avg\": 53.67267,\n",
      "  \"latency_token_generation-avg\": 47.0224\n",
      "}\n",
      "[2024-04-17 15:02:57,154] [INFO] [engine.py:329:run_accelerator] Saved evaluation results of input model to cache_folder//models/llama2_gpu/meta-llama/Llama-2-7b-hf/gpu-cuda_input_model_metrics.json\n",
      "[2024-04-17 15:02:57,154] [INFO] [engine.py:864:_run_pass] Running pass conversion_merged:OnnxConversion\n",
      "[2024-04-17 15:02:57,155] [INFO] [engine.py:898:_run_pass] Loaded model from cache: 0_OnnxConversion-78f9207c2bfa458863121c42f50c49e1-dfaff1da61d127bb5e9dc2f31a708897 from cache_folder//cache/runs\n",
      "[2024-04-17 15:02:57,155] [INFO] [engine.py:864:_run_pass] Running pass transformers_optimization_fp32:OrtTransformersOptimization\n",
      "[2024-04-17 15:02:57,157] [INFO] [engine.py:898:_run_pass] Loaded model from cache: 3_OrtTransformersOptimization-0-8d800f56ee00b5607adb7b96f5076e59-gpu-cuda from cache_folder//cache/runs\n",
      "[2024-04-17 15:02:57,157] [INFO] [engine.py:864:_run_pass] Running pass blockwise_quant_int4:OnnxMatMul4Quantizer\n",
      "[2024-04-17 15:02:57,159] [INFO] [engine.py:898:_run_pass] Loaded model from cache: 4_OnnxMatMul4Quantizer-3-fe98a8b48ad1e4a5fc22fd942d64d9e6 from cache_folder//cache/runs\n",
      "[2024-04-17 15:02:57,159] [INFO] [engine.py:842:_run_passes] Run model evaluation for the final model...\n",
      "data loader kwargs: {'batch_size': 2, 'model_framework': <Framework.ONNX: 'ONNX'>, 'model_id': 'meta-llama/Llama-2-7b-hf', 'past_seq_length': 0, 'seq_length': 8, 'max_seq_length': 2048}\n",
      "data loader kwargs: {'batch_size': 2, 'model_framework': <Framework.ONNX: 'ONNX'>, 'model_id': 'meta-llama/Llama-2-7b-hf', 'past_seq_length': 8, 'seq_length': 1, 'max_seq_length': 2048}\n",
      "[2024-04-17 15:03:18,943] [INFO] [engine.py:864:_run_pass] Running pass conversion_merged:OnnxConversion\n",
      "[2024-04-17 15:03:18,944] [INFO] [engine.py:898:_run_pass] Loaded model from cache: 0_OnnxConversion-78f9207c2bfa458863121c42f50c49e1-dfaff1da61d127bb5e9dc2f31a708897 from cache_folder//cache/runs\n",
      "[2024-04-17 15:03:18,945] [INFO] [engine.py:864:_run_pass] Running pass transformers_optimization_fp16:OrtTransformersOptimization\n",
      "[2024-04-17 15:03:18,946] [INFO] [engine.py:898:_run_pass] Loaded model from cache: 5_OrtTransformersOptimization-0-9a2d6ea047237bbe7e923e3ede6e3570-gpu-cuda from cache_folder//cache/runs\n",
      "[2024-04-17 15:03:18,947] [INFO] [engine.py:842:_run_passes] Run model evaluation for the final model...\n",
      "data loader kwargs: {'batch_size': 2, 'model_framework': <Framework.ONNX: 'ONNX'>, 'model_id': 'meta-llama/Llama-2-7b-hf', 'past_seq_length': 0, 'seq_length': 8, 'max_seq_length': 2048}\n",
      "data loader kwargs: {'batch_size': 2, 'model_framework': <Framework.ONNX: 'ONNX'>, 'model_id': 'meta-llama/Llama-2-7b-hf', 'past_seq_length': 8, 'seq_length': 1, 'max_seq_length': 2048}\n",
      "[2024-04-17 15:03:41,343] [INFO] [engine.py:361:run_accelerator] Save footprint to cache_folder//models/llama2_gpu/meta-llama/Llama-2-7b-hf/gpu-cuda_footprints.json.\n",
      "[2024-04-17 15:03:41,345] [INFO] [engine.py:279:run] Run history for gpu-cuda:\n",
      "[2024-04-17 15:03:41,352] [INFO] [engine.py:567:dump_run_history] run history:\n",
      "+------------------------------------------------------------------------------------+------------------------------------------------------------------------------------+-----------------------------+----------------+-----------------------------------------------+\n",
      "| model_id                                                                           | parent_model_id                                                                    | from_pass                   |   duration_sec | metrics                                       |\n",
      "+====================================================================================+====================================================================================+=============================+================+===============================================+\n",
      "| 78f9207c2bfa458863121c42f50c49e1                                                   |                                                                                    |                             |                | {                                             |\n",
      "|                                                                                    |                                                                                    |                             |                |   \"latency_prompt_processing-avg\": 53.67267,  |\n",
      "|                                                                                    |                                                                                    |                             |                |   \"latency_token_generation-avg\": 47.0224     |\n",
      "|                                                                                    |                                                                                    |                             |                | }                                             |\n",
      "+------------------------------------------------------------------------------------+------------------------------------------------------------------------------------+-----------------------------+----------------+-----------------------------------------------+\n",
      "| 0_OnnxConversion-78f9207c2bfa458863121c42f50c49e1-dfaff1da61d127bb5e9dc2f31a708897 | 78f9207c2bfa458863121c42f50c49e1                                                   | OnnxConversion              |       439.78   |                                               |\n",
      "+------------------------------------------------------------------------------------+------------------------------------------------------------------------------------+-----------------------------+----------------+-----------------------------------------------+\n",
      "| 3_OrtTransformersOptimization-0-8d800f56ee00b5607adb7b96f5076e59-gpu-cuda          | 0_OnnxConversion-78f9207c2bfa458863121c42f50c49e1-dfaff1da61d127bb5e9dc2f31a708897 | OrtTransformersOptimization |       124.263  |                                               |\n",
      "+------------------------------------------------------------------------------------+------------------------------------------------------------------------------------+-----------------------------+----------------+-----------------------------------------------+\n",
      "| 4_OnnxMatMul4Quantizer-3-fe98a8b48ad1e4a5fc22fd942d64d9e6                          | 3_OrtTransformersOptimization-0-8d800f56ee00b5607adb7b96f5076e59-gpu-cuda          | OnnxMatMul4Quantizer        |        70.1778 | {                                             |\n",
      "|                                                                                    |                                                                                    |                             |                |   \"latency_prompt_processing-avg\": 207.52853, |\n",
      "|                                                                                    |                                                                                    |                             |                |   \"latency_token_generation-avg\": 191.34901   |\n",
      "|                                                                                    |                                                                                    |                             |                | }                                             |\n",
      "+------------------------------------------------------------------------------------+------------------------------------------------------------------------------------+-----------------------------+----------------+-----------------------------------------------+\n",
      "| 5_OrtTransformersOptimization-0-9a2d6ea047237bbe7e923e3ede6e3570-gpu-cuda          | 0_OnnxConversion-78f9207c2bfa458863121c42f50c49e1-dfaff1da61d127bb5e9dc2f31a708897 | OrtTransformersOptimization |       511.793  | {                                             |\n",
      "|                                                                                    |                                                                                    |                             |                |   \"latency_prompt_processing-avg\": 25.82543,  |\n",
      "|                                                                                    |                                                                                    |                             |                |   \"latency_token_generation-avg\": 25.14717    |\n",
      "|                                                                                    |                                                                                    |                             |                | }                                             |\n",
      "+------------------------------------------------------------------------------------+------------------------------------------------------------------------------------+-----------------------------+----------------+-----------------------------------------------+\n",
      "[2024-04-17 15:03:41,352] [INFO] [engine.py:294:run] No packaging config provided, skip packaging artifacts\n"
     ]
    }
   ],
   "source": [
    "! python3 -m olive.workflows.run --config config_gpu.json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we can get the:\n",
    "\n",
    "- `latency_prompt_processing`: The time taken to process the prompt where the KV_cache is not used.\n",
    "- `latency_token_generation`: The time taken to generate the tokens where the KV_cache is used.\n",
    "\n",
    "\n",
    "| Model Type | latency_prompt_processing(ms) | latency_token_generation(ms) |\n",
    "|------------|---------------------------|-------------------------|\n",
    "| Torch Model | 53.67267 | 47.0224 |\n",
    "| Olive Optimized Model | 25.82543 | 25.14717 |\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}