make --device fast the default

.ci/scripts/validate.sh

@@ -25,6 +25,7 @@ function generate_compiled_model_output() {
     local MODEL_DIR="${CHECKPOINT_PATH%/*}"
     local MODEL_NAME=$(basename "$CHECKPOINT_PATH" | sed 's/\.[^.]*$//')
 
+
     if [[ $CHECKPOINT_PATH != *"stories"* && $TARGET_DEVICE == "cuda" ]]; then
         DTYPES="bfloat16"
         EXCLUDE_INT8_QUANT=true
@@ -74,7 +75,7 @@ function generate_compiled_model_output() {
         python3 -W ignore generate.py --dtype ${DTYPE} --compile --quant '{"embedding" : {"bitwidth": 4, "groupsize": 8, "packed": "True"}}' --checkpoint-path "$CHECKPOINT_PATH" --temperature 0 --device "$TARGET_DEVICE" > "$MODEL_DIR/output_compiled" || exit 1
         cat "$MODEL_DIR/output_compiled"
 
-        if [ "$EXCLUDE_INT8_QUANT" = false ]; then
+        if [ "${EXCLUDE_INT8_QUANT:-false}" == false ]; then
             echo "******************************************"
             echo "******* INT8 channel-wise quantized ******"
             echo "******************************************"
@@ -109,17 +110,24 @@ function generate_compiled_model_output() {
 function generate_aoti_model_output() {
     local CHECKPOINT_PATH="$1"
     local TARGET_DEVICE="${2:-cpu}"
+    local DTYPES="${3:-default}"
     local MODEL_DIR="${CHECKPOINT_PATH%/*}"
     local MODEL_NAME=$(basename "$CHECKPOINT_PATH" | sed 's/\.[^.]*$//')
 
-    if [[ $CHECKPOINT_PATH != *"stories"* && $TARGET_DEVICE == "cuda" ]]; then
-        DTYPES="bfloat16"
-        EXCLUDE_INT8_QUANT=true
-    else
-        DTYPES="float32 bfloat16 float16"
-        EXCLUDE_INT8_QUANT=false
+    echo "Local DTYPES=$DTYPES"
+
+    if [[ $DTYPES == "default" ]]; then
+        if [[ $CHECKPOINT_PATH != *"stories"* && $TARGET_DEVICE == "cuda" ]]; then
+            DTYPES="bfloat16"
+            EXCLUDE_INT8_QUANT=true
+        else
+            DTYPES="float32 bfloat16 float16"
+            EXCLUDE_INT8_QUANT=false
+        fi
     fi
 
+    echo "Local after default DTYPES=$DTYPES"
+
     for DTYPE in $DTYPES; do
         echo ""############### Run inference with AOT Inductor  for dtype $DTYPE "###############"
         echo ""
@@ -158,7 +166,7 @@ function generate_aoti_model_output() {
         python3 -W ignore generate.py --dtype ${DTYPE} --checkpoint-path "$CHECKPOINT_PATH" --temperature 0 --dso-path ${MODEL_DIR}/${MODEL_NAME}.so --device "$TARGET_DEVICE" > "$MODEL_DIR/output_aoti" || exit 1
         cat "$MODEL_DIR/output_aoti"
 
-        if [ "$EXCLUDE_INT8_QUANT" = false ]; then
+        if [ "${EXCLUDE_INT8_QUANT:-false}" == false ]; then
             echo "******************************************"
             echo "******* INT8 channel-wise quantized ******"
             echo "******************************************"
@@ -247,10 +255,11 @@ function eval_model() {
 function eval_model_sanity_check() {
     local CHECKPOINT_PATH="$1"
     local TARGET_DEVICE="${2:-cpu}"
+    local DTYPES="$3"
     local MODEL_DIR="${CHECKPOINT_PATH%/*}"
     local MODEL_NAME=$(basename "$CHECKPOINT_PATH" | sed 's/\.[^.]*$//')
 
-    for DTYPE in float32 bfloat16 float16; do
+    for DTYPE in $DTYPES; do
         echo ""############### Run eval with torch.compile for dtype $DTYPE "###############"
         echo ""
         echo "******************************************"
@@ -295,11 +304,12 @@ function run_compile() {
 }
 
 function run_aoti() {
-    generate_aoti_model_output "$CHECKPOINT_PATH" "$TARGET_DEVICE" || exit 1
+    echo "Passing DTYPES=$DTYPES"
+    generate_aoti_model_output "$CHECKPOINT_PATH" "$TARGET_DEVICE" "$DTYPES" || exit 1
 }
 
 function run_executorch() {
-    if [ "$TARGET_DEVICE" = "cpu" ]; then
+    if [ "$TARGET_DEVICE" == "cpu" ]; then
         generate_executorch_model_output "$CHECKPOINT_PATH" "$TARGET_DEVICE" || exit 1
     else
         echo "Skipped: Executorch doesn't run on ${TARGET_DEVICE}"
@@ -311,31 +321,68 @@ function run_eval(){
 }
 
 function run_eval_sanity_check(){
-    eval_model_sanity_check "$CHECKPOINT_PATH" "$TARGET_DEVICE" || exit 1
+    echo "Passing DTYPES=$DTYPES"
+    eval_model_sanity_check "$CHECKPOINT_PATH" "$TARGET_DEVICE" "$DTYPES" || exit 1
 }
 
 CHECKPOINT_PATH="$1"
 TARGET_DEVICE="${2:-cpu}"
 PROMPT="Hello, my name is"
 
-
 if [ "$#" -gt 2 ]; then
     # Additional arguments provided
     for arg in "${@:3}"; do
         case "$arg" in
             "compile")
+                echo "arg:$arg"
                 run_compile || exit 1
                 ;;
             "aoti")
+                echo "arg:$arg"
+                DTYPES="default"
+                run_aoti || exit 1
+                ;;
+            "aoti-bfloat16")
+                echo "arg:$arg"
+                DTYPES="bfloat16"
+                run_aoti || exit 1
+                ;;
+            "aoti-float16")
+                echo "arg:$arg"
+                DTYPES="float16"
+                run_aoti || exit 1
+                ;;
+            "aoti-float32")
+                echo "arg:$arg"
+                DTYPES="float32"
                 run_aoti || exit 1
                 ;;
             "executorch")
+                echo "arg:$arg"
                 run_executorch || exit 1
                 ;;
             "eval")
+                echo "arg:$arg"
                 run_eval || exit 1
                 ;;
             "eval_sanity_check")
+                echo "arg:$arg"
+                DTYPES="bfloat16 float16 float32"
+                run_eval_sanity_check || exit 1
+                ;;
+            "eval_sanity_check-bfloat16")
+                echo "arg:$arg"
+                DTYPES="bfloat16"
+                run_eval_sanity_check || exit 1
+                ;;
+            "eval_sanity_check-float16")
+                echo "arg:$arg"
+                DTYPES="float16"
+                run_eval_sanity_check || exit 1
+                ;;
+            "eval_sanity_check-float32")
+                echo "arg:$arg"
+                DTYPES="float32"
                 run_eval_sanity_check || exit 1
                 ;;
             *)

.github/workflows/hqq-dtype.yml

@@ -27,9 +27,9 @@ jobs:
 
         echo "::group::Download checkpoints"
         # Install requirements
-        pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu121
-        pip install -r requirements.txt
-        pip list
+        pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu121
+        pip3 install -r requirements.txt
+        pip3 list
         python3 -c 'import torch;print(f"torch: {torch.__version__, torch.version.git_version}")'
         echo "::endgroup::"
 

.github/workflows/periodic.yml

@@ -47,9 +47,9 @@ jobs:
           echo "$(uname -a)"
       - name: Install dependencies
         run: |
-          pip install --pre torch  --index-url https://download.pytorch.org/whl/nightly/cpu
-          pip install -r requirements.txt
-          pip list
+          pip3 install --pre torch  --index-url https://download.pytorch.org/whl/nightly/cpu
+          pip3 install -r requirements.txt
+          pip3 list
           python3 -c 'import torch;print(f"torch: {torch.__version__, torch.version.git_version}")'
       - name: Download checkpoints
         run: |
@@ -81,9 +81,9 @@ jobs:
           echo "$(uname -a)"
       - name: Install dependencies
         run: |
-          pip install --pre torch  --index-url https://download.pytorch.org/whl/nightly/cpu
-          pip install -r requirements.txt
-          pip list
+          pip3 install --pre torch  --index-url https://download.pytorch.org/whl/nightly/cpu
+          pip3 install -r requirements.txt
+          pip3 list
           python3 -c 'import torch;print(f"torch: {torch.__version__, torch.version.git_version}")'
       - name: Download checkpoints
         run: |
@@ -128,9 +128,9 @@ jobs:
         echo "::endgroup::"
 
         echo "::group::Install required packages"
-        pip install --progress-bar off --pre torch --index-url https://download.pytorch.org/whl/nightly/cu121
-        pip install -r ./requirements.txt
-        pip list
+        pip3 install --progress-bar off --pre torch --index-url https://download.pytorch.org/whl/nightly/cu121
+        pip3 install -r ./requirements.txt
+        pip3 list
         python3 -c 'import torch;print(f"torch: {torch.__version__, torch.version.git_version}")'
         echo "::endgroup::"