Update inference/bf16_cast_block_int8.py

inference/bf16_cast_block_int8.py

@@ -3,6 +3,7 @@ import json
 from argparse import ArgumentParser
 from glob import glob
 from tqdm import tqdm
+import json
 
 import torch
 from safetensors.torch import load_file, save_file
@@ -14,15 +15,40 @@ def main(bf16_path, int8_path, model_name="deepseek-ai/DeepSeek-R1"):
     torch.set_default_dtype(torch.bfloat16)
     os.makedirs(int8_path, exist_ok=True)
     model_index_file = os.path.join(int8_path, "model.safetensors.index.json")
-    
-    if not os.path.exists(model_index_file):
+    config_file = os.path.join(int8_path, "config.json")
+     
+    if not os.path.exists(model_index_file) or not os.path.exists(config_file):
         snapshot_download(
             repo_id=model_name,
-            allow_patterns=["model.safetensors.index.json"],
+            ignore_patterns=["*.safetensors"],
             local_dir=int8_path,
             local_dir_use_symlinks=False
         )
-        print(f"model index file downloaded to {model_index_file}")
+        print(f"model index file and config file downloaded to {int8_path}")
+
+        # modify config.json and save it
+        config = json.load(open(config_file))
+        if "quantization_config" in config:
+            quant_config = config["quantization_config"]
+            quant_config.pop("fmt", None)
+            quant_config["quant_method"] = "blockwise_int8"
+            quant_config["weight_block_size"] = [
+                128,
+                128
+            ]
+            quant_config["activation_scheme"] = "dynamic"
+        else:
+            config["quantization_config"] = {
+                "activation_scheme": "dynamic",
+                "quant_method": "blockwise_int8",
+                "weight_block_size": [
+                    128,
+                    128
+                ]
+            }
+        with open(config_file, "w", encoding="utf-8") as f:
+            json.dump(config, f, indent=2, ensure_ascii=False, sort_keys=True)
+        print(f"config.json modified and saved to {config_file}")
 
     with open(model_index_file, "r") as f:
         model_index = json.load(f)
@@ -60,4 +86,3 @@ if __name__ == "__main__":
     args = parser.parse_args()
     main(args.input_bf16_hf_path, args.output_int8_hf_path, args.model_name)
     print("done")
-