Revert changes to weight_utils.py (#18759)

python/sglang/srt/model_loader/weight_utils.py

@@ -751,12 +751,8 @@ def safetensors_weights_iterator(
         if disable_mmap:
             with open(st_file, "rb") as f:
                 result = safetensors.torch.load(f.read())
-                # NOTE(xiuyu): safetensors.torch.load() may return keys in a
+                for name, param in result.items():
-                # different order per TP-shard file. Sort to ensure that cross-rank
+                    yield name, param
-                # collectives (e.g. all_gather in sync_quantize_weight) stay
-                # aligned across TP ranks.
-                for name in sorted(result.keys()):
-                    yield name, result[name]
         else:
             with safetensors.safe_open(st_file, framework="pt", device="cpu") as f:
                 for name in f.keys():
@@ -850,31 +846,23 @@ def multi_thread_safetensors_weights_iterator(
         return safetensors.torch.load_file(st_file, device="cpu")
 
     with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
-        # Yield in deterministic order to keep cross-rank collectives aligned.
+        futures = [executor.submit(_load_file, st_file) for st_file in hf_weights_files]
-        future_to_idx = {
-            executor.submit(_load_file, st_file): idx
-            for idx, st_file in enumerate(hf_weights_files)
-        }
-        results_by_idx: Dict[int, dict] = {}
 
         if enable_tqdm:
             futures_iter = tqdm(
-                concurrent.futures.as_completed(future_to_idx),
+                concurrent.futures.as_completed(futures),
                 total=len(hf_weights_files),
                 desc="Multi-thread loading shards",
                 disable=not enable_tqdm,
                 bar_format=BAR_FORMAT,
             )
         else:
-            futures_iter = concurrent.futures.as_completed(future_to_idx)
+            futures_iter = concurrent.futures.as_completed(futures)
 
         for future in futures_iter:
-            results_by_idx[future_to_idx[future]] = future.result()
+            state_dict = future.result()
-
+            for name, param in state_dict.items():
-        for idx in range(len(hf_weights_files)):
+                yield name, param
-            state_dict = results_by_idx[idx]
-            for name in sorted(state_dict.keys()):
-                yield name, state_dict[name]
 
 
 def _load_pt_file(bin_file: str) -> dict: