Add support for adding r2-downoader for GPT-j (#680)

* initial commit for supporting r2 downloader in gptj

* fix download-tool group issue

Author: anandhu-eng

script/get-dataset-cnndm/meta.yaml

@@ -35,6 +35,10 @@ deps:
   skip_if_env:
     MLC_TMP_ML_MODEL: 
       - llama3_1-8b
+- tags: get,generic-python-lib,_package.transformers
+  skip_if_env:
+    MLC_TMP_ML_MODEL: 
+      - llama3_1-8b
 - tags: get,generic-python-lib,_numpy
   skip_if_env:
     MLC_TMP_ML_MODEL: 

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

@@ -67,39 +67,42 @@ def postprocess(i):
 
     env = i['env']
 
-    if os.path.exists(os.path.join(
-            env['GPTJ_CHECKPOINT_PATH'], "checkpoint-final")):
-        env['GPTJ_CHECKPOINT_PATH'] = os.path.join(
-            env['GPTJ_CHECKPOINT_PATH'], "checkpoint-final")
-
-    is_saxml = env.get('MLC_TMP_MODEL_SAXML', '')
-    if is_saxml == "fp32":
-        if os.path.exists("pax_gptj_checkpoint"):
-            env['GPTJ_SAXML_CHECKPOINT_PATH'] = os.path.join(
-                os.getcwd(), "pax_gptj_checkpoint")
-            env['MLC_ML_MODEL_FILE_WITH_PATH'] = env['GPTJ_SAXML_CHECKPOINT_PATH']
-        else:
-            return {'return': 1, 'error': 'pax_gptj_checkpoint generation failed'}
+    if not env.get('MLC_DOWNLOAD_MODE', '') == "dry":
+        if os.path.exists(os.path.join(
+                env['GPTJ_CHECKPOINT_PATH'], "checkpoint-final")):
+            env['GPTJ_CHECKPOINT_PATH'] = os.path.join(
+                env['GPTJ_CHECKPOINT_PATH'], "checkpoint-final")
+
+        is_saxml = env.get('MLC_TMP_MODEL_SAXML', '')
+        if is_saxml == "fp32":
+            if os.path.exists("pax_gptj_checkpoint"):
+                env['GPTJ_SAXML_CHECKPOINT_PATH'] = os.path.join(
+                    os.getcwd(), "pax_gptj_checkpoint")
+                env['MLC_ML_MODEL_FILE_WITH_PATH'] = env['GPTJ_SAXML_CHECKPOINT_PATH']
+            else:
+                return {'return': 1,
+                        'error': 'pax_gptj_checkpoint generation failed'}
 
-    elif is_saxml == "int8":
-        if os.path.exists("int8_ckpt"):
-            env['GPTJ_SAXML_INT8_CHECKPOINT_PATH'] = os.path.join(
-                os.getcwd(), "int8_ckpt")
-            env['MLC_ML_MODEL_FILE_WITH_PATH'] = env['GPTJ_SAXML_INT8_CHECKPOINT_PATH']
+        elif is_saxml == "int8":
+            if os.path.exists("int8_ckpt"):
+                env['GPTJ_SAXML_INT8_CHECKPOINT_PATH'] = os.path.join(
+                    os.getcwd(), "int8_ckpt")
+                env['MLC_ML_MODEL_FILE_WITH_PATH'] = env['GPTJ_SAXML_INT8_CHECKPOINT_PATH']
+            else:
+                return {'return': 1,
+                        'error': 'pax_gptj_checkpoint generation failed'}
+        elif env.get('MLC_TMP_ML_MODEL_PROVIDER', '') == 'nvidia':
+            env['MLC_ML_MODEL_FILE_WITH_PATH'] = os.path.join(
+                env['MLC_NVIDIA_MLPERF_SCRATCH_PATH'],
+                'models',
+                'GPTJ-6B',
+                'fp8-quantized-ammo',
+                'GPTJ-FP8-quantized')
         else:
-            return {'return': 1, 'error': 'pax_gptj_checkpoint generation failed'}
-    elif env.get('MLC_TMP_ML_MODEL_PROVIDER', '') == 'nvidia':
-        env['MLC_ML_MODEL_FILE_WITH_PATH'] = os.path.join(
-            env['MLC_NVIDIA_MLPERF_SCRATCH_PATH'],
-            'models',
-            'GPTJ-6B',
-            'fp8-quantized-ammo',
-            'GPTJ-FP8-quantized')
-    else:
-        env['MLC_ML_MODEL_FILE_WITH_PATH'] = env['GPTJ_CHECKPOINT_PATH']
+            env['MLC_ML_MODEL_FILE_WITH_PATH'] = env['GPTJ_CHECKPOINT_PATH']
 
-    env['MLC_ML_MODEL_FILE'] = os.path.basename(
-        env['MLC_ML_MODEL_FILE_WITH_PATH'])
-    env['MLC_GET_DEPENDENT_CACHED_PATH'] = env['MLC_ML_MODEL_FILE_WITH_PATH']
+        env['MLC_ML_MODEL_FILE'] = os.path.basename(
+            env['MLC_ML_MODEL_FILE_WITH_PATH'])
+        env['MLC_GET_DEPENDENT_CACHED_PATH'] = env['MLC_ML_MODEL_FILE_WITH_PATH']
 
     return {'return': 0}

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

@@ -37,7 +37,8 @@ prehook_deps:
 tests:
   run_inputs:
     - variations_list:
-      - fp32
+      - fp32,pytorch,r2-downloader,dry-run
+      - fp32,pytorch,rclone,dry-run
 
 
 print_env_at_the_end:
@@ -102,10 +103,15 @@ variations:
     env:
       MLC_DOWNLOAD_CHECKSUM_NOT_USED: e677e28aaf03da84584bb3073b7ee315
       MLC_PACKAGE_URL: https://cloud.mlcommons.org/index.php/s/QAZ2oM94MkFtbQx/download
-      MLC_RCLONE_CONFIG_NAME: mlc-inference
-      MLC_RCLONE_URL: mlc-inference:mlcommons-inference-wg-public/gpt-j
       MLC_UNZIP: 'yes'
     required_disk_space: 22700
+  pytorch,fp32,rclone:
+    env:
+      MLC_PACKAGE_URL: mlc-inference:mlcommons-inference-wg-public/gpt-j
+      MLC_RCLONE_CONFIG_NAME: mlc-inference
+  pytorch,fp32,r2-downloader:
+    env:
+      MLC_DOWNLOAD_URL: https://inference.mlcommons-storage.org/metadata/gpt-j-model-checkpoint.uri
   pytorch,fp32,wget:
     add_deps_recursive:
       dae:
@@ -183,11 +189,28 @@ variations:
     add_deps_recursive:
       dae:
         tags: _rclone
-    default: true
     env:
       MLC_DOWNLOAD_FILENAME: checkpoint
-      MLC_DOWNLOAD_URL: <<<MLC_RCLONE_URL>>>
+      MLC_DOWNLOAD_URL: <<<MLC_PACKAGE_URL>>>
     group: download-tool
+  r2-downloader:
+    group: download-tool
+    default: true
+    add_deps_recursive:
+      dae:
+        tags: _r2-downloader
+    env:
+      MLC_DOWNLOAD_FILENAME: checkpoint
+  dry-run:
+    group: run-mode
+    env:
+      MLC_DOWNLOAD_MODE: dry
+  dry-run,r2-downloader:
+    env:
+      MLC_DOWNLOAD_EXTRA_OPTIONS: -x
+  dry-run,rclone:
+    env:
+      MLC_DOWNLOAD_EXTRA_OPTIONS: --dry-run
   saxml:
     group: framework
   saxml,fp32: