update pt_weights_iterator

vllm/model_executor/model_loader/weight_utils.py

@@ -693,26 +693,39 @@ def pt_weights_iterator(
     pt_load_map_location: Union[str, dict[str, str]] = "cpu",
 ) -> Generator[tuple[str, torch.Tensor], None, None]:
     """Iterate over the weights in the model bin/pt files."""
-    total_count = 0
-    for bin_file in hf_weights_files:
-        state = torch.load(bin_file, map_location=pt_load_map_location, weights_only=True)
-        total_count += len(state)
-        del state
-        
-    current_count = 0
-    for bin_file in tqdm(
-            hf_weights_files,
-            desc="Loading pt checkpoint shards",
-            disable=not enable_tqdm(use_tqdm_on_load),
-            bar_format=_BAR_FORMAT,
-    ):
-        state = torch.load(bin_file, map_location=pt_load_map_location, weights_only=True)
-        for name, param in state.items():
-            current_count += 1
-            param.current_count = current_count
-            param.total_count = total_count
-            yield name, param
-        del state
+    if os.environ.get('LLAMA_NN') == '1':
+        total_count = 0
+        for bin_file in hf_weights_files:
+            state = torch.load(bin_file, map_location=pt_load_map_location, weights_only=True)
+            total_count += len(state)
+            del state
+            
+        current_count = 0
+        for bin_file in tqdm(
+                hf_weights_files,
+                desc="Loading pt checkpoint shards",
+                disable=not enable_tqdm(use_tqdm_on_load),
+                bar_format=_BAR_FORMAT,
+        ):
+            state = torch.load(bin_file, map_location=pt_load_map_location, weights_only=True)
+            for name, param in state.items():
+                current_count += 1
+                param.current_count = current_count
+                param.total_count = total_count
+                yield name, param
+            del state
+    else:
+        for bin_file in tqdm(
+                hf_weights_files,
+                desc="Loading pt checkpoint shards",
+                disable=not enable_tqdm(use_tqdm_on_load),
+                bar_format=_BAR_FORMAT,
+        ):
+            state = torch.load(bin_file,
+                            map_location=pt_load_map_location,
+                            weights_only=True)
+            yield from state.items()
+            del state
 
 
 def multi_thread_pt_weights_iterator(