diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 48b59f560a..def6f5c4b5 100644 --- a/requirements/cuda.txt +++ b/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 diff --git a/requirements/test/cpu.txt b/requirements/test/cpu.txt index 86bd822110..d7d4692733 100644 --- a/requirements/test/cpu.txt +++ b/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 diff --git a/requirements/test/cuda.in b/requirements/test/cuda.in index 53a9b70da5..8c9b8570a6 100644 --- a/requirements/test/cuda.in +++ b/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 diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index e06e43a7f9..d464efaf67 100644 --- a/requirements/test/cuda.txt +++ b/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 diff --git a/requirements/test/nightly-torch.txt b/requirements/test/nightly-torch.txt index 66704d50d3..ec8ca4b9c9 100644 --- a/requirements/test/nightly-torch.txt +++ b/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 diff --git a/setup.py b/setup.py index 8958ab6dc7..abefefa35b 100644 --- a/setup.py +++ b/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", diff --git a/tests/model_executor/model_loader/instanttensor_loader/test_weight_utils.py b/tests/model_executor/model_loader/instanttensor_loader/test_weight_utils.py index 3d15f6f1c5..0a2ece6ca8 100644 --- a/tests/model_executor/model_loader/instanttensor_loader/test_weight_utils.py +++ b/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", diff --git a/vllm/model_executor/model_loader/default_loader.py b/vllm/model_executor/model_loader/default_loader.py index 6bedabfdc0..07283e5027 100644 --- a/vllm/model_executor/model_loader/default_loader.py +++ b/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"): diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index eecdf951e5..3262e94088 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/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(