From 6ffdebbf16f78324c6b7d4e6a936dd53ebf13b93 Mon Sep 17 00:00:00 2001 From: zyw_hw Date: Mon, 23 Jun 2025 20:48:22 +0800 Subject: [PATCH] fix pickle security question --- mindformers/core/callback/callback.py | 2 +- .../dataset/blended_datasets/blended_dataset.py | 8 ++++---- mindformers/dataset/blended_datasets/gpt_dataset.py | 12 ++++++------ mindformers/models/clip/convert_weight.py | 2 +- mindformers/models/llama/convert_weight.py | 2 +- mindformers/tools/moe_token_distribution_tools.py | 2 +- research/llm_boost/convert_weight.py | 2 +- research/telechat2/convert_weight.py | 2 +- .../test_parallel_random/run_parallel_random.py | 2 +- .../test_multi_latent_attention.py | 4 ++-- 10 files changed, 19 insertions(+), 19 deletions(-) diff --git a/mindformers/core/callback/callback.py b/mindformers/core/callback/callback.py index 298363522..712afdbbb 100644 --- a/mindformers/core/callback/callback.py +++ b/mindformers/core/callback/callback.py @@ -933,7 +933,7 @@ class TrainingStateMonitor(Callback): dump_id, prefix, param_name = parsed_name if not begin_id < dump_id < end_id: continue - data = np.load(os.path.join(self.dump_path, f), allow_pickle=True) + data = np.load(os.path.join(self.dump_path, f), allow_pickle=False) if prefix == 'device_local_norm': self._output(f'device_local_norm', data, self.dump_step, self.device_local_norm_format) elif prefix == 'local_loss': diff --git a/mindformers/dataset/blended_datasets/blended_dataset.py b/mindformers/dataset/blended_datasets/blended_dataset.py index 6cee62c01..85e251b62 100644 --- a/mindformers/dataset/blended_datasets/blended_dataset.py +++ b/mindformers/dataset/blended_datasets/blended_dataset.py @@ -157,8 +157,8 @@ class BlendedDataset(): with open(path_to_description, "wt") as writer: writer.write(self.unique_description) # Save the indexes - numpy.save(path_to_dataset_index, dataset_index, allow_pickle=True) - numpy.save(path_to_dataset_sample_index, dataset_sample_index, allow_pickle=True) + numpy.save(path_to_dataset_index, dataset_index, allow_pickle=False) + numpy.save(path_to_dataset_sample_index, dataset_sample_index, allow_pickle=False) else: logger.info(f"Unable to save the {type(self).__name__} indexes because path_to_cache is None") @@ -171,14 +171,14 @@ class BlendedDataset(): logger.info(f"\tLoad the dataset index from {path_to_dataset_index}") t_beg = time.time() - dataset_index = numpy.load(path_to_dataset_index, allow_pickle=True, mmap_mode='r') + dataset_index = numpy.load(path_to_dataset_index, allow_pickle=False, mmap_mode='r') t_end = time.time() logger.info(f"\t> time elapsed: {t_end - t_beg:4f} seconds") logger.info(f"\tLoad the dataset sample index from {path_to_dataset_sample_index}") t_beg = time.time() dataset_sample_index = numpy.load( - path_to_dataset_sample_index, allow_pickle=True, mmap_mode='r' + path_to_dataset_sample_index, allow_pickle=False, mmap_mode='r' ) t_end = time.time() logger.info(f"\t> time elapsed: {t_end - t_beg:4f} seconds") diff --git a/mindformers/dataset/blended_datasets/gpt_dataset.py b/mindformers/dataset/blended_datasets/gpt_dataset.py index 7e6517ac3..3e4483ddb 100644 --- a/mindformers/dataset/blended_datasets/gpt_dataset.py +++ b/mindformers/dataset/blended_datasets/gpt_dataset.py @@ -465,9 +465,9 @@ class GPTDataset(MegatronDataset): # Write the description with open(path_to_description, "wt") as writer: writer.write(self.unique_description) - numpy.save(path_to_document_index, document_index, allow_pickle=True) - numpy.save(path_to_sample_index, sample_index, allow_pickle=True) - numpy.save(path_to_shuffle_index, shuffle_index, allow_pickle=True) + numpy.save(path_to_document_index, document_index, allow_pickle=False) + numpy.save(path_to_sample_index, sample_index, allow_pickle=False) + numpy.save(path_to_shuffle_index, shuffle_index, allow_pickle=False) else: logger.warning( f"Unable to save the {type(self).__name__} indexes because path_to_cache is None", @@ -489,7 +489,7 @@ class GPTDataset(MegatronDataset): f"\tLoad the document index from {os.path.basename(path_to_document_index)}", ) t_beg = time.time() - document_index = numpy.load(path_to_document_index, allow_pickle=True, mmap_mode='r') + document_index = numpy.load(path_to_document_index, allow_pickle=False, mmap_mode='r') t_end = time.time() logger.debug(f"\t> time elapsed: {t_end - t_beg:4f} seconds") @@ -497,7 +497,7 @@ class GPTDataset(MegatronDataset): f"\tLoad the sample index from {os.path.basename(path_to_sample_index)}" ) t_beg = time.time() - sample_index = numpy.load(path_to_sample_index, allow_pickle=True, mmap_mode='r') + sample_index = numpy.load(path_to_sample_index, allow_pickle=False, mmap_mode='r') t_end = time.time() logger.debug(f"\t> time elapsed: {t_end - t_beg:4f} seconds") @@ -505,7 +505,7 @@ class GPTDataset(MegatronDataset): f"\tLoad the shuffle index from {os.path.basename(path_to_shuffle_index)}", ) t_beg = time.time() - shuffle_index = numpy.load(path_to_shuffle_index, allow_pickle=True, mmap_mode='r') + shuffle_index = numpy.load(path_to_shuffle_index, allow_pickle=False, mmap_mode='r') t_end = time.time() logger.debug(f"\t> time elapsed: {t_end - t_beg:4f} seconds") diff --git a/mindformers/models/clip/convert_weight.py b/mindformers/models/clip/convert_weight.py index 43167b558..bc50f068e 100644 --- a/mindformers/models/clip/convert_weight.py +++ b/mindformers/models/clip/convert_weight.py @@ -32,7 +32,7 @@ def convert_weight(torch_path="ViT-B-32.pt", mindspore_path： The save path for clip_vit_b_32.ckpt. """ - param_dict = torch.load(torch_path, map_location=torch.device('cpu')) + param_dict = torch.load(torch_path, map_location=torch.device('cpu'), weights_only=True) new_dict = [] for name, param in param_dict.items(): diff --git a/mindformers/models/llama/convert_weight.py b/mindformers/models/llama/convert_weight.py index 27b8ab46f..3f974c2f5 100644 --- a/mindformers/models/llama/convert_weight.py +++ b/mindformers/models/llama/convert_weight.py @@ -148,7 +148,7 @@ def convert_megatron_to_ms(input_path, output_path, dtype=None, **kwargs): except Exception as e: raise RuntimeError("Unexpected error occurred when loading PyTorch checkpoint.") from e try: - megatron_ckpt = torch.load(input_path, map_location='cpu') + megatron_ckpt = torch.load(input_path, map_location='cpu', weights_only=True) # pylint: disable=W0703 except Exception as e: print(f"Fail to load meagtron checkpoint '{input_path}', Error {e.message}.", flush=True) diff --git a/mindformers/tools/moe_token_distribution_tools.py b/mindformers/tools/moe_token_distribution_tools.py index 7267de5f0..23248a705 100644 --- a/mindformers/tools/moe_token_distribution_tools.py +++ b/mindformers/tools/moe_token_distribution_tools.py @@ -58,7 +58,7 @@ def merge_npy_files(str_layer, path_npy_list): sub_path_list.sort(key=lambda x: int(x.split(str_layer)[1].split('.')[0])) temp = [] for path in sub_path_list: - real_data = np.load(path, allow_pickle=True) + real_data = np.load(path, allow_pickle=False) temp.append(real_data) return temp diff --git a/research/llm_boost/convert_weight.py b/research/llm_boost/convert_weight.py index b73da684e..921c6d2c4 100644 --- a/research/llm_boost/convert_weight.py +++ b/research/llm_boost/convert_weight.py @@ -148,7 +148,7 @@ def convert_megatron_to_ms(input_path, output_path, dtype=None, **kwargs): except Exception as e: raise RuntimeError("Unexpected error occurred when loading PyTorch checkpoint.") from e try: - megatron_ckpt = torch.load(input_path, map_location='cpu') + megatron_ckpt = torch.load(input_path, map_location='cpu', weights_only=True) # pylint: disable=W0703 except Exception as e: print(f"Fail to load megatron checkpoint '{input_path}', Error {e.message}.", flush=True) diff --git a/research/telechat2/convert_weight.py b/research/telechat2/convert_weight.py index 8b46f8fa7..6f744fb12 100644 --- a/research/telechat2/convert_weight.py +++ b/research/telechat2/convert_weight.py @@ -97,7 +97,7 @@ def convert_pt_to_ms(input_path, output_path, dtype=None, **kwargs): if convert_safetensors: pt_states = load_file(per_file) else: - pt_states = torch.load(per_file, map_location='cpu') + pt_states = torch.load(per_file, map_location='cpu', weights_only=True) pt_states_list.append(pt_states) ckpt_list = [] diff --git a/tests/st/test_pynative/test_parallel_random/run_parallel_random.py b/tests/st/test_pynative/test_parallel_random/run_parallel_random.py index d1817f42b..b85e1b541 100644 --- a/tests/st/test_pynative/test_parallel_random/run_parallel_random.py +++ b/tests/st/test_pynative/test_parallel_random/run_parallel_random.py @@ -93,7 +93,7 @@ def run_random_tracer_parallel(): for mode in range(3): candidate = [[] for _ in range(world_size)] for rank_id in range(world_size): - candidate[rank_id] = np.load(f"./result{mode}-rank{rank_id}.npy", allow_pickle=True) + candidate[rank_id] = np.load(f"./result{mode}-rank{rank_id}.npy", allow_pickle=False) if mode == 0: # in raw mode, the value should be the same in each rank assert np.allclose(candidate[0], candidate[1], rtol=1e-4, atol=1e-4) assert np.allclose(candidate[1], candidate[2], rtol=1e-4, atol=1e-4) diff --git a/tests/st/test_ut/test_parallel_core/test_training_graph/test_transformer/test_multi_latent_attention/test_multi_latent_attention.py b/tests/st/test_ut/test_parallel_core/test_training_graph/test_transformer/test_multi_latent_attention/test_multi_latent_attention.py index dd714f6f1..0d7f4d313 100644 --- a/tests/st/test_ut/test_parallel_core/test_training_graph/test_transformer/test_multi_latent_attention/test_multi_latent_attention.py +++ b/tests/st/test_ut/test_parallel_core/test_training_graph/test_transformer/test_multi_latent_attention/test_multi_latent_attention.py @@ -228,8 +228,8 @@ class TestMultiLatentAttention: prefix, mode_num = tmp_path.name.split('mode') mega_output_dir = prefix + 'mode' + str(int(mode_num) - 1) mega_output_file = tmp_path.parent / mega_output_dir / "output_mla_ms_megatron_multi.npz" - output_ms_dict_mega = np.load(mega_output_file, allow_pickle=True)['output'] - output_ms_dict_mind = np.load(output_file_path, allow_pickle=True)['output'] + output_ms_dict_mega = np.load(mega_output_file, allow_pickle=False)['output'] + output_ms_dict_mind = np.load(output_file_path, allow_pickle=False)['output'] assert np.allclose(output_ms_dict_mind, output_ms_dict_mega) class TestMultiLatentAttentionSingleCard(TestMultiLatentAttention): -- Gitee