Guard duplicate-key check and write in safetensors loader with a lock

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()

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(