Enable quantizing local checkpoints in model release script

.github/scripts/torchao_model_releases/quantize_and_upload.py

@@ -568,7 +568,7 @@ def _untie_weights_and_save_locally(model_id):
 """
 
 
-def quantize_and_upload(model_id, quant):
+def quantize_and_upload(model_id, quant, push_to_hub):
     _int8_int4_linear_config = Int8DynamicActivationIntxWeightConfig(
         weight_dtype=torch.int4,
         weight_granularity=PerGroup(32),
@@ -657,9 +657,13 @@ def quantize_and_upload(model_id, quant):
     card = ModelCard(content)
 
     # Push to hub
-    quantized_model.push_to_hub(quantized_model_id, safe_serialization=False)
-    tokenizer.push_to_hub(quantized_model_id)
-    card.push_to_hub(quantized_model_id)
+    if push_to_hub:
+        quantized_model.push_to_hub(quantized_model_id, safe_serialization=False)
+        tokenizer.push_to_hub(quantized_model_id)
+        card.push_to_hub(quantized_model_id)
+    else:
+        quantized_model.save_pretrained(quantized_model_id, safe_serialization=False)
+        tokenizer.save_pretrained(quantized_model_id)
 
     # Manual Testing
     prompt = "Hey, are you conscious? Can you talk to me?"
@@ -700,5 +704,11 @@ def quantize_and_upload(model_id, quant):
         type=str,
         help="Quantization method. Options are FP8, INT4, INT8_INT4, AWQ-INT4",
     )
+    parser.add_argument(
+        "--push_to_hub",
+        action="store_true",
+        default=False,
+        help="Flag to indicate whether push to huggingface hub or not",
+    )
     args = parser.parse_args()
-    quantize_and_upload(args.model_id, args.quant)
+    quantize_and_upload(args.model_id, args.quant, args.push_to_hub)

.github/scripts/torchao_model_releases/release.sh

@@ -14,6 +14,7 @@
 
 # Default quantization options
 default_quants=("FP8" "INT4" "INT8-INT4")
+push_to_hub=""
 # Parse arguments
 while [[ $# -gt 0 ]]; do
   case "$1" in
@@ -29,6 +30,10 @@ while [[ $# -gt 0 ]]; do
         shift
       done
       ;;
+     --push_to_hub)
+      push_to_hub="--push_to_hub"
+      shift
+      ;;
     *)
       echo "Unknown option: $1"
       exit 1
@@ -38,14 +43,14 @@ done
 # Use default quants if none specified
 if [[ -z "$model_id" ]]; then
   echo "Error: --model_id is required"
-  echo "Usage: $0 --model_id <model_id> [--quants <quant1> [quant2 ...]]"
+  echo "Usage: $0 --model_id <model_id> [--quants <quant1> [quant2 ...]] [--push_to_hub]"
   exit 1
 fi
 if [[ ${#quants[@]} -eq 0 ]]; then
   quants=("${default_quants[@]}")
 fi
 # Run the python command for each quantization option
 for quant in "${quants[@]}"; do
-  echo "Running: python quantize_and_upload.py --model_id $model_id --quant $quant"
-  python quantize_and_upload.py --model_id "$model_id" --quant "$quant"
+  echo "Running: python quantize_and_upload.py --model_id $model_id --quant $quant $push_to_hub"
+  python quantize_and_upload.py --model_id "$model_id" --quant "$quant" $push_to_hub
 done

.github/workflows/release_model.yml

@@ -43,4 +43,4 @@ jobs:
         pip install .
         HF_MODEL_ID=${{ github.event.inputs.hf_model_id }}
         cd .github/scripts/torchao_model_releases
-        ./release.sh --model_id $HF_MODEL_ID
+        ./release.sh --model_id $HF_MODEL_ID --push_to_hub