Support uploaded quantised models for SCC25 (#612)

* Add support for using existing quantized llama2 70b models

* use r2-downloader tool for model download

* Fix path issue for r2-downloader

* copy llama2 checkpoint to scratch space

* update preprocessed openorca dataset path

* make sure target data path is created before softlink

* Initialize git lfs when building TRT-LLM

* Update download URL from staging to production

* Update meta.yaml

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Arjun Suresh <[email protected]>

Author: 3 people

script/app-mlperf-inference-nvidia/customize.py

@@ -318,18 +318,6 @@ def preprocess(i):
         model_path = fp8_model_path
 
     elif "llama2" in env["MLC_MODEL"]:
-        # path to which the data file is present
-        target_data_path = os.path.join(
-            env['MLPERF_SCRATCH_PATH'],
-            'data',
-            'llama2-70b')
-        # path to the dataset file
-        target_data_file_path = os.path.join(
-            env['MLPERF_SCRATCH_PATH'],
-            'data',
-            'llama2-70b',
-            'open_orca_gpt4_tokenized_llama.sampled_24576.pkl')
-
         preprocessed_data_for_accuracy_checker = os.path.join(
             env['MLPERF_SCRATCH_PATH'],
             'preprocessed_data',
@@ -345,41 +333,13 @@ def preprocess(i):
 
         tmp_tp_size = env['MLC_NVIDIA_TP_SIZE']
         tmp_pp_size = env['MLC_NVIDIA_PP_SIZE']
-        if tmp_tp_size == "1":
-            fp8_model_path = os.path.join(
-                env['MLPERF_SCRATCH_PATH'],
-                'models',
-                'Llama2',
-                'fp8-quantized-ammo',
-                f'llama2-70b-chat-hf-tp{tmp_tp_size}pp1-fp8-02072024')
-        else:
-            fp8_model_path = os.path.join(
-                env['MLPERF_SCRATCH_PATH'],
-                'models',
-                'Llama2',
-                'fp8-quantized-ammo',
-                f'llama2-70b-chat-hf-tp{tmp_tp_size}pp{tmp_pp_size}-fp8')
-
-        # check the presence of validation dataset
-        if not os.path.exists(target_data_file_path):
-            if env.get('MLC_DATASET_OPENORCA_PREPROCESSED_PATH', '') == '':
-                return {
-                    'return': 1, 'error': 'Llama2 70B validation dataset not present.'}
-            if not os.path.exists(target_data_path):
-                cmds.append(f"mkdir -p {target_data_path}")
-            cmds.append(
-                f"ln -sf {env['MLC_DATASET_OPENORCA_PREPROCESSED_PATH']} {target_data_file_path}")
 
-        # check the presence of calibration dataset
-        if not env.get('LLAMA2_PRE_QUANTIZED_CHECKPOINT_PATH'):
-            if not os.path.exists(target_calibration_data_file_path):
-                if env.get('MLC_DATASET_OPENORCA_CALIBRATION_PATH', '') == '':
-                    return {
-                        'return': 1, 'error': 'Llama2 70B calibration dataset not present.'}
-                if not os.path.exists(target_data_path):
-                    cmds.append(f"mkdir -p {target_data_path}")
-                cmds.append(
-                    f"ln -sf {env['MLC_DATASET_OPENORCA_CALIBRATION_PATH']} {target_calibration_data_file_path}")
+        fp8_model_path = os.path.join(
+            env['MLPERF_SCRATCH_PATH'],
+            'models',
+            'Llama2',
+            'fp8-quantized-ammo',
+            f'llama-2-70b-chat-hf-tp{tmp_tp_size}pp{tmp_pp_size}-fp8')
 
         if not os.path.exists(preprocessed_data_for_accuracy_checker):
             if not os.path.exists(preprocessed_data_for_accuracy_checker):
@@ -476,7 +436,15 @@ def preprocess(i):
                 'open_orca',
                 'input_ids_padded.npy')
             if not os.path.exists(target_preprocessed_data_path):
-                cmds.append(f"make preprocess_data BENCHMARKS='{model_name}'")
+                cmds.append(
+                    f"mkdir -p {os.path.dirname(target_preprocessed_data_path)}")
+                if env.get('MLC_DATASET_OPENORCA_PREPROCESSED_PATH'):
+                    cmds.append(
+                        f"ln -sf {env['MLC_DATASET_OPENORCA_PREPROCESSED_PATH']} {os.path.join(env['MLPERF_SCRATCH_PATH'], "preprocessed_data", "open_orca")}"
+                    )
+                else:
+                    cmds.append(
+                        f"make preprocess_data BENCHMARKS='{model_name}'")
         else:
             cmds.append(f"make preprocess_data BENCHMARKS='{model_name}'")
 

script/build-mlperf-inference-server-nvidia/run.sh

@@ -3,6 +3,10 @@ CUR=$PWD
 
 cd ${MLC_MLPERF_INFERENCE_NVIDIA_CODE_PATH}
 
+if [[ ${BUILD_TRTLLM} == "1" ]]; then
+  git lfs install
+fi
+
 if [[ ${MLC_MAKE_CLEAN} == "yes" ]]; then
   make clean
 fi

script/get-ml-model-llama2/customize.py

@@ -8,7 +8,7 @@ def preprocess(i):
     os_info = i['os_info']
     env = i['env']
 
-    if env.get('MLC_TMP_ML_MODEL_PROVIDER', '') == 'nvidia':
+    if env.get('MLC_TMP_ML_MODEL_QUANTIZE_LOCALLY', '') == 'nvidia':
         if is_true(env.get('MLC_ML_MODEL_QUANTIZE_LOCALLY')):
             i['run_script_input']['script_name'] = 'run-nvidia'
             gpu_arch = int(
@@ -18,7 +18,28 @@ def preprocess(i):
             env['MLC_GPU_ARCH'] = gpu_arch
             env['MLC_TMP_REQUIRE_DOWNLOAD'] = 'no'
         else:
-            run_cmd = f"ln -sf {env['LLAMA2_CHECKPOINT_PATH']} {env['MLC_NVIDIA_MLPERF_SCRATCH_PATH']}/models/Llama2/fp8-quantized-ammo/llama-2-70b-chat-hf-tp{env['MLC_NVIDIA_TP_SIZE']}pp{env['MLC_NVIDIA_PP_SIZE']}-{env['MLC_ML_MODEL_PRECISION']}"
+            target_quantized_model_dir = os.path.join(
+                env['MLC_NVIDIA_MLPERF_SCRATCH_PATH'],
+                "models",
+                "Llama2",
+                "fp8-quantized-ammo",
+                f"llama-2-70b-chat-hf-tp{env['MLC_NVIDIA_TP_SIZE']}pp{env['MLC_NVIDIA_PP_SIZE']}-{env['MLC_ML_MODEL_PRECISION']}"
+            )
+
+            target_model_dir = os.path.join(
+                env['MLC_NVIDIA_MLPERF_SCRATCH_PATH'],
+                "models",
+                "Llama2",
+                "Llama-2-70b-chat-hf"
+            )
+
+            # Ensure target directory exists
+            os.makedirs(target_quantized_model_dir, exist_ok=True)
+            os.makedirs(target_model_dir, exist_ok=True)
+
+            run_cmd = f"cp -r {env['LLAMA2_QUANTIZED_CHECKPOINT_PATH']}/* {env['MLC_NVIDIA_MLPERF_SCRATCH_PATH']}/models/Llama2/fp8-quantized-ammo/llama-2-70b-chat-hf-tp{env['MLC_NVIDIA_TP_SIZE']}pp{env['MLC_NVIDIA_PP_SIZE']}-{env['MLC_ML_MODEL_PRECISION']}"
+            run_cmd += f" && cp -r {env['LLAMA2_CHECKPOINT_PATH']}/* {env['MLC_NVIDIA_MLPERF_SCRATCH_PATH']}/models/Llama2/Llama-2-70b-chat-hf"
+
             env['MLC_RUN_CMD'] = run_cmd
     else:
         path = env.get('LLAMA2_CHECKPOINT_PATH', '').strip()

script/get-ml-model-llama2/meta.yaml

@@ -51,13 +51,33 @@ tests:
 variations:
   pre-quantized:
     group: quantization
+    deps:
+      - enable_if_env:
+          MLC_TMP_ML_MODEL_PRE_QUANTIZED:
+          - 'yes'
+        skip_if_env:
+          LLAMA2_QUANTIZED_CHECKPOINT_PATH:
+          - 'yes'
+        env:
+          MLC_DOWNLOAD_FINAL_ENV_NAME: LLAMA2_QUANTIZED_CHECKPOINT_PATH
+          MLC_EXTRACT_FINAL_ENV_NAME: LLAMA2_QUANTIZED_CHECKPOINT_PATH
+        force_cache: true
+        extra_cache_tags: llama2,llama2-model,llama2-checkpoint,llama2-70b
+        names:
+          - dae-quantized-models
+        tags: download-and-extract
+        force_env_keys:
+          - MLC_OUTDIRNAME  
+        update_tags_from_env_with_prefix:
+          _url.:
+            - MLC_DOWNLOAD_URL
     env:
-      MLC_ML_MODEL_PRE_QUANTIZED: 'yes'
+      MLC_TMP_ML_MODEL_PRE_QUANTIZED: 'yes'
   quantize-locally:
     default: true
     group: quantization
     env:
-      MLC_ML_MODEL_QUANTIZE_LOCALLY: 'yes'
+      MLC_TMP_ML_MODEL_QUANTIZE_LOCALLY: 'yes'
   amd:
     default_env:
       MLC_LLAMA2_QUANTIZATION_DEVICE: ''
@@ -157,8 +177,11 @@ variations:
     env:
       MLC_DOWNLOAD_URL: https://llama2.mlcommons-storage.org/metadata/llama-2-7b-chat-hf.uri
   mlc,r2-downloader,70b,pre-quantized,fp8:
+    add_deps_recursive:
+      dae-quantized-models:
+        tags: _r2-downloader
     env:
-      MLC_DOWNLOAD_URL: https://llama2.mlcommons-storage.org/metadata/llama-2-70b-chat-hf-tp<<MLC_NVIDIA_TP_SIZE>>pp<<MLC_NVIDIA_PP_SIZE>>-<<<MLC_ML_MODEL_PRECISION>>>.uri
+      MLC_DOWNLOAD_URL: https://llama2.mlcommons-storage.org/metadata/llama2-70b-chat-hf-tp<<<MLC_NVIDIA_TP_SIZE>>>pp<<<MLC_NVIDIA_PP_SIZE>>>-<<<MLC_ML_MODEL_PRECISION>>>.uri
   hf:
     group: download-source
     env:
@@ -212,6 +235,10 @@ variations:
       - tags: get,nvidia,scratch,space
         names:
         - mlperf-inference-nvidia-scratch-space
+      - env: {}
+        force_new_env_keys:
+        - LLAMA2_CHECKPOINT_PATH
+        tags: get,ml-model,llama2-70b,_fp32,_pytorch
     group: model-provider
   pytorch:
     default: true
@@ -242,7 +269,7 @@ variations:
     - tags: get,generic-python-lib,_package.compressed_tensors
   pytorch,fp32:
     env: {}
-  pytorch,nvidia,v5.0:
+  pytorch,nvidia,v5.0,quantize-locally:
     deps:
       - env:
           MLC_GIT_CHECKOUT_PATH_ENV_NAME: MLC_TENSORRT_LLM_CHECKOUT_PATH
@@ -258,10 +285,6 @@ variations:
       - cuda
       tags: get,cuda
     - tags: get,cuda-devices,_with-pycuda
-    - env: {}
-      force_new_env_keys:
-      - LLAMA2_CHECKPOINT_PATH
-      tags: get,ml-model,llama2-70b,_fp32,_pytorch
     - names:
       - nvidia-inference-common-code
       tags: get,nvidia,inference,common-code

script/get-ml-model-llama2/run-nvidia.sh

@@ -15,13 +15,9 @@ cd ${MLC_TENSORRT_LLM_CHECKOUT_PATH}
 make -C docker build
 test $? -eq 0 || exit $?
 
-if [ "${MLC_NVIDIA_TP_SIZE}" -eq 1 ]; then
-  RUN_CMD="bash -c 'git lfs install && git lfs pull && python3 scripts/build_wheel.py -a=${MLC_GPU_ARCH} --clean --install --use_ccache --benchmarks --trt_root /usr/local/tensorrt/ && python examples/quantization/quantize.py --dtype=float16  --output_dir=/mnt/models/Llama2/fp8-quantized-ammo/llama2-70b-chat-hf-tp${MLC_NVIDIA_TP_SIZE}pp${MLC_NVIDIA_PP_SIZE}-fp8-02072024 --model_dir=/mnt/models/Llama2/Llama-2-70b-chat-hf --qformat=fp8 --kv_cache_dtype=fp8 --tp_size ${MLC_NVIDIA_TP_SIZE} --pp_size ${MLC_NVIDIA_PP_SIZE} --calib_dataset=/calib_dataset'"
-  echo "$RUN_CMD"
-else
-  RUN_CMD="bash -c 'git lfs install && git lfs pull && python3 scripts/build_wheel.py -a=${MLC_GPU_ARCH} --clean --install --use_ccache --benchmarks --trt_root /usr/local/tensorrt/ && python examples/quantization/quantize.py --dtype=float16  --output_dir=/mnt/models/Llama2/fp8-quantized-ammo/llama2-70b-chat-hf-tp${MLC_NVIDIA_TP_SIZE}pp${MLC_NVIDIA_PP_SIZE}-fp8 --model_dir=/mnt/models/Llama2/Llama-2-70b-chat-hf --qformat=fp8 --kv_cache_dtype=fp8 --tp_size ${MLC_NVIDIA_TP_SIZE} --pp_size ${MLC_NVIDIA_PP_SIZE} --calib_dataset=/calib_dataset'"
-  echo "$RUN_CMD"
-fi
+RUN_CMD="bash -c 'git lfs install && git lfs pull && python3 scripts/build_wheel.py -a=${MLC_GPU_ARCH} --clean --install --use_ccache --benchmarks --trt_root /usr/local/tensorrt/ && python examples/quantization/quantize.py --dtype=float16  --output_dir=/mnt/models/Llama2/fp8-quantized-ammo/llama-2-70b-chat-hf-tp${MLC_NVIDIA_TP_SIZE}pp${MLC_NVIDIA_PP_SIZE}-fp8 --model_dir=/mnt/models/Llama2/Llama-2-70b-chat-hf --qformat=fp8 --kv_cache_dtype=fp8 --tp_size ${MLC_NVIDIA_TP_SIZE} --pp_size ${MLC_NVIDIA_PP_SIZE} --calib_dataset=/calib_dataset'"
+echo "$RUN_CMD"
+
 # TODO: check whether --device nvidia.com/gpu=all would work for docker
 DOCKER_RUN_ARGS=" -v ${MLC_NVIDIA_MLPERF_SCRATCH_PATH}:/mnt -v ${MLC_NVIDIA_PREPROCESSED_CALIBRATION_DATASET_PATH}:/calib_dataset -u $(id -u):$(id -g) --userns=keep-id --device nvidia.com/gpu=all -e NVIDIA_VISIBLE_DEVICES=all"
 export DOCKER_RUN_ARGS="$DOCKER_RUN_ARGS"

script/get-preprocessed-dataset-openorca/customize.py

@@ -51,22 +51,32 @@ def postprocess(i):
         env['PREPROCESSED_DATA_DIR'] = os.path.dirname(
             env['MLC_OPENORCA_PREPROCESSED_ROOT'])
         if is_true(env.get('MLC_DATASET_CALIBRATION', '')):
-            env['MLC_DATASET_CALIBRATION_PATH'] = os.path.join(
-                env['MLC_OPENORCA_PREPROCESSED_ROOT'],
-                "open_orca_gpt4_tokenized_llama.calibration_1000.pkl")
             if env.get('MLC_TMP_DATASET_PREPROCESS_STEP_PROVIDER',
                        '') == "nvidia":
                 env['MLC_NVIDIA_PREPROCESSED_CALIBRATION_DATASET_PATH'] = os.path.join(
                     env['MLC_OPENORCA_PREPROCESSED_ROOT'],
                     "preprocessed_data",
                     "mlperf_llama2_openorca_calibration_1k")
-            env['MLC_GET_DEPENDENT_CACHED_PATH'] = env['MLC_DATASET_CALIBRATION_PATH']
-            env['MLC_DATASET_OPENORCA_CALIBRATION_PATH'] = env['MLC_DATASET_CALIBRATION_PATH']
+            else:
+                env['MLC_DATASET_CALIBRATION_PATH'] = os.path.join(
+                    env['MLC_OPENORCA_PREPROCESSED_ROOT'],
+                    "open_orca_gpt4_tokenized_llama.calibration_1000.pkl")
+                env['MLC_GET_DEPENDENT_CACHED_PATH'] = env['MLC_DATASET_CALIBRATION_PATH']
+                env['MLC_DATASET_OPENORCA_CALIBRATION_PATH'] = env['MLC_DATASET_CALIBRATION_PATH']
         else:
             env['MLC_DATASET_PREPROCESSED_PATH'] = os.path.join(
                 env['MLC_OPENORCA_PREPROCESSED_ROOT'],
                 "open_orca_gpt4_tokenized_llama.sampled_24576.pkl")
-            env['MLC_GET_DEPENDENT_CACHED_PATH'] = env['MLC_DATASET_PREPROCESSED_PATH']
+            if env.get('MLC_TMP_DATASET_PREPROCESS_STEP_PROVIDER',
+                       '') == "nvidia":
+                env['MLC_DATASET_OPENORCA_NVIDIA_PREPROCESSED_PATH'] = os.path.join(
+                    env['MLC_OPENORCA_PREPROCESSED_ROOT'],
+                    "preprocessed")
+                env['MLC_GET_DEPENDENT_CACHED_PATH'] = env['MLC_DATASET_OPENORCA_NVIDIA_PREPROCESSED_PATH']
+            else:
+                env['MLC_GET_DEPENDENT_CACHED_PATH'] = env['MLC_DATASET_PREPROCESSED_PATH']
+            # The openorca mlc preprocessed dataset is used in nvidia
+            # implementation for checking accuracy
             env['MLC_DATASET_OPENORCA_PREPROCESSED_PATH'] = env['MLC_DATASET_PREPROCESSED_PATH']
 
     return {'return': 0}

script/get-preprocessed-dataset-openorca/meta.yaml

@@ -144,3 +144,6 @@ variations:
   nvidia,calibration:
     new_env_keys:
     - MLC_NVIDIA_PREPROCESSED_CALIBRATION_DATASET_PATH
+  nvidia,validation:
+    new_env_keys:
+    - MLC_DATASET_OPENORCA_NVIDIA_PREPROCESSED_PATH