diff --git a/tests/models/safetensors_loader_test.py b/tests/models/safetensors_loader_test.py index e35f8cfd0..769ef455e 100644 --- a/tests/models/safetensors_loader_test.py +++ b/tests/models/safetensors_loader_test.py @@ -168,6 +168,45 @@ def test_load_and_create_model_from_gcs(self): loaded_state, ) + def test_load_and_create_model_raises_on_duplicate_key(self): + # Two distinct source keys that both map to the same jax key. In the + # 'original' loader these are processed on separate threads, so without a + # lock around the duplicate-key check and write the second write can + # silently overwrite the first (issue #1259). With the lock the duplicate + # is always detected. + try: + st_dir = self.create_tempdir().full_path + except Exception: # pylint: disable=broad-except + st_dir = tempfile.TemporaryDirectory().name + os.makedirs(st_dir, exist_ok=True) + + tensors = { + 'lm_head.kernel': np.zeros((2, 2), dtype=np.float32), + 'lm_head.weight': np.zeros((2, 2), dtype=np.float32), + } + filename = os.path.join(st_dir, 'model.safetensors') + stnp.save_file(tensors, filename) + + def duplicate_key_mapping(config): + del config + return { + r'^lm_head\.kernel$': ('lm_head.kernel', None), + r'^lm_head\.weight$': ('lm_head.kernel', None), + } + + # The duplicate ValueError is re-raised wrapped as RuntimeError by the + # loader, so match on RuntimeError. mode='original' is required because the + # default optimized path is single-threaded and does not run this check. + with self.assertRaisesRegex(RuntimeError, 'Duplicate key'): + safetensors_loader.load_and_create_model( + st_dir, + test_common.ToyTransformer, + self.model.config, + duplicate_key_mapping, + dtype=jnp.float32, + mode='original', + ) + if __name__ == '__main__': absltest.main() diff --git a/tunix/models/safetensors_loader.py b/tunix/models/safetensors_loader.py index 0f429623a..08a84cc70 100644 --- a/tunix/models/safetensors_loader.py +++ b/tunix/models/safetensors_loader.py @@ -194,6 +194,7 @@ def load_and_create_model_orig( key_map = key_mapping(config) file_lock = threading.Lock() + dict_lock = threading.Lock() # guards the duplicate-key check and write # Load tensors from all files skipped_keys = [] @@ -225,11 +226,12 @@ def process_key(k_name, f, sf_file, file_loaded_tensors): if dtype and current_arr.dtype != dtype: current_arr = current_arr.astype(dtype) - if jax_key_mapped in file_loaded_tensors: - raise ValueError( - f'Duplicate key {jax_key_mapped} found within file {f.name}.' - ) - file_loaded_tensors[jax_key_mapped] = current_arr + with dict_lock: + if jax_key_mapped in file_loaded_tensors: + raise ValueError( + f'Duplicate key {jax_key_mapped} found within file {f.name}.' + ) + file_loaded_tensors[jax_key_mapped] = current_arr except Exception as e: raise RuntimeError(