[Model Loader] Apply EP weight filtering in InstantTensor (#61167)

Signed-off-by: Itay Etelis <92247226+Etelis@users.noreply.github.com>
Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Co-authored-by: Misha Goin <mgoin64@gmail.com>

requirements/cuda.txt

@@ -18,7 +18,7 @@ nvidia-cudnn-frontend>=1.19.1
 nvtx==0.2.15
 # Required for faster safetensors model loading
 fastsafetensors >= 0.3.3
-instanttensor >= 0.1.9
+instanttensor >= 0.2.2
 
 # QuACK and Cutlass DSL for FA4 (cute-DSL implementation)
 nvidia-cutlass-dsl[cu13]==4.8.0

requirements/test/cpu.txt

@@ -406,7 +406,7 @@ inflect==5.6.2
     # via datamodel-code-generator
 iniconfig==2.0.0
     # via pytest
-instanttensor==0.1.9
+instanttensor==0.2.2
     # via -r requirements/test/cuda.in
 interegular==0.3.3
     # via lm-format-enforcer

requirements/test/cuda.in

@@ -57,7 +57,7 @@ arctic-inference == 0.1.1; platform_machine == "x86_64" # Required for suffix de
 numba == 0.65.0 # Required for N-gram speculative decoding
 runai-model-streamer[s3,gcs,azure]==0.15.7
 fastsafetensors>=0.3.3
-instanttensor>=0.1.9; platform_machine == "x86_64"
+instanttensor>=0.2.2; platform_machine == "x86_64"
 decord==0.6.0; platform_machine == "x86_64"
 # terratorch is temporarily disabled while PyPI has the `lightning` package
 # in `quarantined` status (every published terratorch version transitively

requirements/test/cuda.txt

@@ -426,7 +426,7 @@ inflect==5.6.2
     # via datamodel-code-generator
 iniconfig==2.0.0
     # via pytest
-instanttensor==0.1.9
+instanttensor==0.2.2
     # via
     #   -c requirements/cuda.txt
     #   -r requirements/test/cuda.in

requirements/test/nightly-torch.txt

@@ -42,5 +42,5 @@ numba == 0.65.0 # Required for N-gram speculative decoding
 numpy
 runai-model-streamer[s3,gcs,azure]==0.15.7
 fastsafetensors>=0.3.3
-instanttensor>=0.1.9
+instanttensor>=0.2.2
 pydantic>=2.12 # 2.11 leads to error on python 3.13

setup.py

@@ -1692,7 +1692,7 @@ setup(
         "bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"],
         "tensorizer": ["tensorizer==2.10.1"],
         "fastsafetensors": ["fastsafetensors >= 0.3.3"],
-        "instanttensor": ["instanttensor >= 0.1.9"],
+        "instanttensor": ["instanttensor >= 0.2.2"],
         "runai": ["runai-model-streamer[s3,gcs,azure] >= 0.15.7"],
         "audio": [
             "av",

tests/model_executor/model_loader/instanttensor_loader/test_weight_utils.py

@@ -2,9 +2,13 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import glob
+import sys
+from types import SimpleNamespace
+from unittest.mock import MagicMock
 
 import pytest
 import torch
+from safetensors.torch import save_file
 
 from vllm.model_executor.model_loader.weight_utils import (
     download_weights_from_hf,
@@ -14,6 +18,60 @@ from vllm.model_executor.model_loader.weight_utils import (
 from vllm.platforms import current_platform
 
 
+@pytest.mark.parametrize(
+    ("weight_name", "local_ids", "collective"),
+    [
+        ("mlp.experts.gate_proj.weight", {0}, True),
+        ("mlp.experts.0.gate_proj.weight", {0}, False),
+        ("mlp.experts.0.gate_proj.weight", set(), False),
+    ],
+)
+def test_instanttensor_ep_collective_loading(
+    monkeypatch, tmp_path, weight_name, local_ids, collective
+):
+    """Keep fused checkpoints collective and expert decisions rank-independent."""
+    from vllm.model_executor.model_loader import weight_utils
+
+    path = tmp_path / "model.safetensors"
+    save_file({weight_name: torch.ones(1)}, str(path))
+    native_open = MagicMock()
+    reader = native_open.return_value.__enter__.return_value
+    reader.total_tensor_size = 0
+    reader.tensors.return_value = ()
+    monkeypatch.setitem(
+        sys.modules, "instanttensor", SimpleNamespace(safe_open=native_open)
+    )
+    group = SimpleNamespace(world_size=2, device_group=object())
+    monkeypatch.setattr(weight_utils, "get_world_group", lambda: group)
+    monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
+    monkeypatch.setattr(current_platform, "current_device", lambda: 0)
+    list(instanttensor_weights_iterator([str(path)], False, local_expert_ids=local_ids))
+    kwargs = native_open.call_args.kwargs
+    assert kwargs["process_group"] is (group.device_group if collective else None)
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda(),
+    reason="InstantTensor requires NVIDIA GPUs",
+)
+def test_instanttensor_ep_filter_cuda(tmp_path):
+    weights = {
+        "mlp.experts.0.gate_proj.weight": torch.ones(2),
+        "mlp.experts.1.gate_proj.weight": torch.ones(2),
+        "mlp.experts.1.gate_proj.weight_scale": torch.ones(1),
+        "model.layers[0].attn.weight": torch.ones(2),
+    }
+    path = tmp_path / "model.safetensors"
+    save_file(weights, str(path))
+    loaded = dict(
+        instanttensor_weights_iterator([str(path)], False, local_expert_ids={0})
+    )
+    del weights["mlp.experts.1.gate_proj.weight"]
+    assert loaded.keys() == weights.keys()
+    for name, tensor in loaded.items():
+        torch.testing.assert_close(tensor.cpu(), weights[name], rtol=0, atol=0)
+
+
 @pytest.mark.skipif(
     not current_platform.is_cuda(),
     reason="InstantTensor requires NVIDIA GPUs",

vllm/model_executor/model_loader/default_loader.py

@@ -318,6 +318,7 @@ class DefaultModelLoader(BaseModelLoader):
                 weights_iterator = instanttensor_weights_iterator(
                     hf_weights_files,
                     self.load_config.use_tqdm_on_load,
+                    local_expert_ids=self.local_expert_ids,
                 )
             else:
                 if extra_config.get("enable_multithread_load"):

vllm/model_executor/model_loader/weight_utils.py

@@ -1356,6 +1356,7 @@ def fastsafetensors_weights_iterator(
 def instanttensor_weights_iterator(
     hf_weights_files: list[str],
     use_tqdm_on_load: bool,
+    local_expert_ids: set[int] | None = None,
 ) -> Generator[tuple[str, torch.Tensor], None, None]:
     """Iterate over the weights in the model safetensor files
     using instanttensor library."""
@@ -1379,6 +1380,23 @@ def instanttensor_weights_iterator(
 
     device = current_platform.current_device()
 
+    extra_kwargs: dict[str, Any] = {}
+    if local_expert_ids is not None:
+        checkpoint_names: list[str] = []
+        for path in hf_weights_files:
+            with safe_open(path, framework="pt") as metadata:
+                checkpoint_names.extend(metadata.keys())  # noqa: SIM118
+        # Every rank must make the same decision about collective loading.
+        empty_expert_ids: set[int] = set()
+        if any(should_skip_weight(name, empty_expert_ids) for name in checkpoint_names):
+            process_group = None
+            extra_kwargs["tensor_filter"] = [
+                name
+                for name in checkpoint_names
+                if not should_skip_weight(name, local_expert_ids)
+            ]
+            extra_kwargs["tensor_name_match_mode"] = "exact"
+
     # copy=True yields tensors that own their memory, staying valid after the
     # context exits or InstantTensor reuses its buffer.
     with instanttensor.safe_open(
@@ -1387,6 +1405,7 @@ def instanttensor_weights_iterator(
         device=device,
         process_group=process_group,
         copy=True,
+        **extra_kwargs,
     ) as f:
         # Track bytes so the bar reports load throughput (GB/s).
         pbar = tqdm(