fix bug for support deepspeed tp load int4 checkpoint (#3396)

* fix bug.

* flake format.

Author: RanTao123

examples/cpu/llm/inference/distributed/run_accuracy_with_deepspeed.py

@@ -404,13 +404,13 @@ def write_checkpoints_json():
         self.model = self.model.module
         import pathlib
 
+        low_precision_checkpoint = None
         if args.low_precision_checkpoint != "":
             pathname = args.low_precision_checkpoint
             assert os.path.exists(
                 pathname
             ), f"Checkpoint file does not exist: {pathname}"
             if os.path.isfile(pathname):
-                low_precision_checkpoint = None
                 if pathname.endswith(".pt") or pathname.endswith(".pth"):
                     low_precision_checkpoint = torch.load(pathname, weights_only=True)
                 elif pathname.endswith(".safetensors"):
@@ -625,9 +625,12 @@ def write_checkpoints_json():
                             low_precision_checkpoint_dict[key] = data[
                                 :, q_head_start * dim : q_head_end * dim
                             ]
-                    low_precision_dict = (low_precision_checkpoint_dict, quant_method)
+                    low_precision_checkpoint = (
+                        low_precision_checkpoint_dict,
+                        quant_method,
+                    )
                 else:
-                    low_precision_dict = None
+                    low_precision_checkpoint = None
 
             self.model = ipex.llm.optimize(
                 self.model.eval(),
@@ -636,7 +639,7 @@ def write_checkpoints_json():
                 inplace=True,
                 deployment_mode=False,
                 cache_weight_for_large_batch=args.cache_weight_for_large_batch,
-                low_precision_checkpoint=low_precision_dict,
+                low_precision_checkpoint=low_precision_checkpoint,
             )
 
         self.base_model = self.model

examples/cpu/llm/inference/distributed/run_generation_with_deepspeed.py

@@ -489,14 +489,14 @@ def write_checkpoints_json():
 # to ipex
 if use_ipex:
     ipex_woq_enabled = args.ipex_weight_only_quantization
+    low_precision_checkpoint = None
     if ipex_woq_enabled:
         if args.low_precision_checkpoint != "":
             pathname = args.low_precision_checkpoint
             assert os.path.exists(
                 pathname
             ), f"Checkpoint file does not exist: {pathname}"
             if os.path.isfile(pathname):
-                low_precision_checkpoint = None
                 if pathname.endswith(".pt") or pathname.endswith(".pth"):
                     low_precision_checkpoint = torch.load(pathname, weights_only=True)
                 elif pathname.endswith(".safetensors"):
@@ -703,10 +703,10 @@ def write_checkpoints_json():
                     low_precision_checkpoint_dict[key] = data[
                         :, q_head_start * dim : q_head_end * dim
                     ]
-            low_precision_dict = (low_precision_checkpoint_dict, quant_method)
+            low_precision_checkpoint = (low_precision_checkpoint_dict, quant_method)
 
         else:
-            low_precision_dict = None
+            low_precision_checkpoint = None
 
     model = ipex.llm.optimize(
         model.eval(),
@@ -715,7 +715,7 @@ def write_checkpoints_json():
         inplace=True,
         deployment_mode=args.deployment_mode,
         cache_weight_for_large_batch=args.cache_weight_for_large_batch,
-        low_precision_checkpoint=low_precision_dict,
+        low_precision_checkpoint=low_precision_checkpoint,
     )
 # Generate
 print_rank0(f"*** Starting to generate {num_tokens} tokens with bs={args.batch_size}")