Fix run_tutorials code (#1552)

* Fix run_tutorials code

Summary:
Last script actually has some errors but didn't error out, this PR
added the logic for the CI job to show error when some job fails
and also fixed remaining code

Test Plan:
CI

Reviewers:

Subscribers:

Tasks:

Tags:

* checking status code

* script

* more logs

* tensor paralell check

* change tp file check condition

* deps

* testing failing

* update loop

* try again

* try again

* done

* restore

* remove extra pint

Author: jerryzh168

.github/workflows/run_tutorials.yml

@@ -30,4 +30,4 @@ jobs:
           ${CONDA_RUN} pip install -r dev-requirements.txt
           ${CONDA_RUN} pip install .
           cd tutorials
-          ${CONDA_RUN} sh run_all.sh
+          ${CONDA_RUN} bash run_all.sh

dev-requirements.txt

@@ -21,6 +21,7 @@ lm_eval
 diskcache
 pycocotools
 tqdm
+importlib_metadata
 
 # Custom CUDA Extensions
 ninja

tutorials/developer_api_guide/my_dtype_tensor_subclass.py

@@ -16,7 +16,7 @@
     Layout,
     PlainLayout,
 )
-from torchao.quantization.quant_primitives import (
+from torchao.quantization import (
     MappingType,
     choose_qparams_affine,
     dequantize_affine,

tutorials/developer_api_guide/my_trainable_tensor_subclass.py

@@ -15,7 +15,11 @@
 from torch.utils._python_dispatch import return_and_correct_aliasing
 
 from torchao.dtypes.utils import Layout, PlainLayout
-from torchao.quantization.quant_primitives import MappingType, choose_qparams_affine
+from torchao.quantization import (
+    MappingType,
+    choose_qparams_affine,
+    quantize_affine,
+)
 
 aten = torch.ops.aten
 
@@ -40,10 +44,12 @@ def _quantize(
         Convert from a floating point tensor (fp32/fp16/bf16) to the desired dtype.
         """
         mapping_type = MappingType.SYMMETRIC
-        block_size = input_float.shape
-        dtype = torch.int16
-        scale, _ = choose_qparams_affine(input_float, mapping_type, block_size, dtype)
-        int_data = (input_float / scale).to(torch.int8)
+        block_size = (1, input_float.shape[-1])
+        dtype = torch.int8
+        scale, zero_point = choose_qparams_affine(
+            input_float, mapping_type, block_size, dtype
+        )
+        int_data = quantize_affine(input_float, block_size, scale, zero_point, dtype)
         tensor_impl_ctr = cls.get_tensor_impl_constructor(type(_layout))
         return tensor_impl_ctr(int_data, scale, _layout)
 

tutorials/run_all.sh

@@ -1,19 +1,37 @@
 #!/bin/bash
-find . -type d | while read dir; do
+FAILED=0
+for dir in $(find . -type d); do
   if [ -f "$dir/run.sh" ]; then
     echo "Running: $dir/run.sh"
-    pushd "$dir"
+    CURRENT_DIR=$(pwd)
+    cd "$dir"
     bash run.sh
-    popd
+    cd "$CURRENT_DIR"
   else
-    find "$dir" -maxdepth 1 -name "*.py" | while read file; do
-      if [[ "$file" == *"tensor_parallel"* ]]; then
+    for file in $(find "$dir" -maxdepth 1 -name "*.py"); do
+      filename=$(basename "$file")
+      if echo "$filename" | grep -q "tensor_parallel"; then
         echo "Running: torchrun --standalone --nnodes=1 --nproc-per-node=1 $file"
         torchrun --standalone --nnodes=1 --nproc-per-node=4 "$file"
+        STATUS=$?
       else
         echo "Running: python $file"
         python "$file"
+        STATUS=$?
+      fi
+
+      if [ $STATUS -ne 0 ]; then
+        FAILED=1
+        echo "Test failed: $file"
       fi
     done
   fi
 done
+
+if [ "$FAILED" -eq 1 ]; then
+  echo "One or more tests failed"
+  exit 1
+else
+  echo "All tests passed"
+  exit 0
+fi