fix insecure deserialization when calling torch.load() (#4202)

Author: lvhan028

lmdeploy/lite/apis/auto_awq.py

@@ -98,7 +98,7 @@ def auto_awq(model: str,
     layer_type = LAYER_TYPE_MAP[type(model).__name__]
     fc2fcs = FC_FCS_MAP[layer_type]
     norm2fcs = NORM_FCS_MAP[layer_type]
-    input_stats = torch.load(osp.join(work_dir, 'inputs_stats.pth'))
+    input_stats = torch.load(osp.join(work_dir, 'inputs_stats.pth'), weights_only=True)
     layers = collect_target_modules(model, layer_type)
     fcs = {}
     for l_name, layer in layers.items():

lmdeploy/lite/apis/get_small_sharded_hf.py

@@ -38,7 +38,7 @@ def main():
 
     checkpoints = set(index['weight_map'].values())
     for ckpt in checkpoints:
-        state_dict = torch.load(os.path.join(args.src_dir, ckpt), map_location='cuda')
+        state_dict = torch.load(os.path.join(args.src_dir, ckpt), map_location='cuda', weights_only=True)
         keys = sorted(list(state_dict.keys()))
         for k in keys:
             new_state_dict_name = 'pytorch_model-{:05d}-of-{:05d}.bin'.format(cnt, n_shard)

lmdeploy/lite/apis/kv_qparams.py

@@ -58,7 +58,7 @@ def smooth_quant(model: str,
 
     # calibrate function exports the calibration statistics
     # (inputs, outputs, keys and values) to `work_dir`.
-    inp_stats = torch.load(work_dir / 'inputs_stats.pth')
+    inp_stats = torch.load(work_dir / 'inputs_stats.pth', weights_only=True)
     act_scales = inp_stats['absmax']
 
     model_type = type(model).__name__

lmdeploy/lite/apis/smooth_quant.py

@@ -119,7 +119,7 @@ def items(self):
         params = defaultdict(dict)
         for shard in self.shards:
             misc = {}
-            tmp = torch.load(shard, map_location='cpu')
+            tmp = torch.load(shard, map_location='cpu', weights_only=True)
             for k, v in tmp.items():
                 match = re.findall(self.pattern, k)
                 if not match:

lmdeploy/turbomind/deploy/loader.py

@@ -19,7 +19,7 @@ def load_weight_ckpt(ckpt: str) -> Dict[str, torch.Tensor]:
     if ckpt.endswith('.safetensors'):
         return load_file(ckpt)
     else:
-        return torch.load(ckpt)
+        return torch.load(ckpt, weights_only=True)
 
 
 def get_used_weight_files(folder: str, state_dict: Dict[str, torch.Tensor]) -> List[str]: