From a0f28d73a83bbf0491428efb0f11c2b8f9dc9067 Mon Sep 17 00:00:00 2001 From: Ning Xie Date: Sat, 10 Oct 2026 02:08:31 +0800 Subject: [PATCH] [sharded state loader] support pp in sharded state loader (#52110) Signed-off-by: Andy Xie Co-authored-by: Simon Mo --- .../save_sharded_state_offline.py | 7 +- .../model_loader/test_sharded_state_loader.py | 438 ++++++++++++++++++ vllm/distributed/parallel_state.py | 5 + .../model_loader/sharded_state_loader.py | 152 ++++-- 4 files changed, 572 insertions(+), 30 deletions(-) diff --git a/examples/features/sharded_state/save_sharded_state_offline.py b/examples/features/sharded_state/save_sharded_state_offline.py index 7ed0de1f10..44b68cecf0 100644 --- a/examples/features/sharded_state/save_sharded_state_offline.py +++ b/examples/features/sharded_state/save_sharded_state_offline.py @@ -38,8 +38,11 @@ def parse_args(): parser.add_argument( "--file-pattern", type=str, - default=ShardedStateLoader.DEFAULT_PATTERN, - help="string pattern of saved filenames", + help=f"string pattern of saved filenames." + f"{ShardedStateLoader.TP_ONLY_PATTERN!r} for tensor parallel only " + f"(pipeline parallelism is 1). " + f"{ShardedStateLoader.PP_AND_TP_PATTERN!r} for pipeline and " + f"tensor parallel.", ) parser.add_argument( "--max-file-size", diff --git a/tests/model_executor/model_loader/test_sharded_state_loader.py b/tests/model_executor/model_loader/test_sharded_state_loader.py index 06354fe3c5..6ec5f1ce43 100644 --- a/tests/model_executor/model_loader/test_sharded_state_loader.py +++ b/tests/model_executor/model_loader/test_sharded_state_loader.py @@ -6,6 +6,7 @@ import multiprocessing as mp import os import shutil from tempfile import TemporaryDirectory +from types import SimpleNamespace import pytest import torch @@ -49,6 +50,318 @@ def test_filter_subtensors(): assert tensor is state_dict[key] +def _mock_pp_group( + monkeypatch: pytest.MonkeyPatch, + world_size: int, + *, + pp_rank: int = 0, + tp_rank: int = 0, + tp_world_size: int = 1, +) -> None: + """Stub the parallel-state accessors used by the sharded state loader.""" + monkeypatch.setattr( + "vllm.distributed.get_pp_group", + lambda: SimpleNamespace(world_size=world_size, rank_in_group=pp_rank), + ) + monkeypatch.setattr( + "vllm.distributed.get_tp_group", + lambda: SimpleNamespace(world_size=tp_world_size, rank_in_group=tp_rank), + ) + monkeypatch.setattr( + "vllm.distributed.get_tensor_model_parallel_rank", + lambda: tp_rank, + ) + monkeypatch.setattr( + "vllm.distributed.get_pipeline_model_parallel_rank", + lambda: pp_rank, + ) + + +@pytest.mark.parametrize( + "world_size,expected", + [ + (1, ShardedStateLoader.TP_ONLY_PATTERN), + (2, ShardedStateLoader.PP_AND_TP_PATTERN), + ], +) +def test_get_default_pattern(monkeypatch, world_size, expected): + """The default pattern is chosen from the pipeline parallel group size.""" + _mock_pp_group(monkeypatch, world_size) + assert ShardedStateLoader._get_default_pattern() == expected + + +def test_render_default_pattern_default_tp(monkeypatch): + """Without a custom pattern and PP=1, legacy TP-only names are rendered.""" + _mock_pp_group(monkeypatch, 1) + assert ( + ShardedStateLoader._render_default_pattern(part=0) + == "model-rank-0-part-0.safetensors" + ) + + +def test_render_default_pattern_default_pp(monkeypatch): + """Without a custom pattern and PP>1, pp/tp-aware names are rendered.""" + _mock_pp_group(monkeypatch, 2) + assert ( + ShardedStateLoader._render_default_pattern(part=0) + == "model-pp-0-tp-0-part-0.safetensors" + ) + + +def test_render_default_pattern_custom_tp_only(monkeypatch): + """A custom TP-only pattern is honored when pipeline parallelism is off.""" + _mock_pp_group(monkeypatch, 1) + assert ( + ShardedStateLoader._render_default_pattern( + ShardedStateLoader.TP_ONLY_PATTERN, part=0 + ) + == "model-rank-0-part-0.safetensors" + ) + + +def test_render_default_pattern_tp_only_with_pp_fails(monkeypatch): + """TP-only filenames would collide across pipeline stages; fail fast.""" + _mock_pp_group(monkeypatch, 2) + with pytest.raises(ValueError, match="does not embed the pipeline parallel"): + ShardedStateLoader._render_default_pattern( + ShardedStateLoader.TP_ONLY_PATTERN, part=0 + ) + + +def test_render_default_pattern_pp_pattern_without_pp_fails(monkeypatch): + """A pp/tp-aware pattern cannot be rendered when pipeline parallelism is off.""" + _mock_pp_group(monkeypatch, 1) + with pytest.raises(ValueError, match="requires pipeline parallelism"): + ShardedStateLoader._render_default_pattern( + ShardedStateLoader.PP_AND_TP_PATTERN, part=0 + ) + + +def test_render_default_pattern_custom_pp(monkeypatch): + """A custom pp/tp-aware pattern is honored when pipeline parallelism is on.""" + _mock_pp_group(monkeypatch, 2) + assert ( + ShardedStateLoader._render_default_pattern( + ShardedStateLoader.PP_AND_TP_PATTERN, part=0 + ) + == "model-pp-0-tp-0-part-0.safetensors" + ) + + +@pytest.mark.parametrize( + "pattern,match", + [ + # tp field present but pp field missing: pipeline stages with the + # same tp rank would write the same file. + ( + "model-tp-{tp_rank}-part-{part}.safetensors", + "does not embed the pipeline parallel", + ), + # pp field present but tp field missing: TP workers would overwrite + # each other's shards. + ( + "model-pp-{pp_rank}-part-{part}.safetensors", + "does not embed the tensor parallel", + ), + ], +) +def test_render_default_pattern_incomplete_custom_patterns_fail( + monkeypatch, pattern, match +): + """A custom pattern must distinguish every active rank (pp and tp).""" + _mock_pp_group(monkeypatch, 2, pp_rank=1, tp_rank=0, tp_world_size=2) + with pytest.raises(ValueError, match=match): + ShardedStateLoader._render_default_pattern(pattern, part=0) + + +def test_render_default_pattern_no_rank_field_with_tp_fails(monkeypatch): + """Without any rank field, TP workers would collide even when PP is off.""" + _mock_pp_group(monkeypatch, 1, tp_world_size=2) + with pytest.raises(ValueError, match="does not embed the tensor parallel"): + ShardedStateLoader._render_default_pattern( + "model-part-{part}.safetensors", part=0 + ) + + +def test_render_default_pattern_custom_pp_tp1_allows_pp_only(monkeypatch): + """With TP=1 a pp-only custom pattern is safe and must be accepted.""" + _mock_pp_group(monkeypatch, 2, pp_rank=1, tp_world_size=1) + assert ( + ShardedStateLoader._render_default_pattern( + "model-pp-{pp_rank}-part-{part}.safetensors", part=0 + ) + == "model-pp-1-part-0.safetensors" + ) + + +@pytest.mark.parametrize( + "pattern,match", + [ + # {{pp_rank}} renders as literal text: both stages would write the + # same file, so it must not count as a pp field. + ( + "model-{{pp_rank}}-tp-{tp_rank}-part-{part}.safetensors", + "does not embed the pipeline parallel", + ), + # Same for an escaped tp field: TP workers would overwrite each + # other's shards. + ( + "model-pp-{pp_rank}-tp-{{tp_rank}}-part-{part}.safetensors", + "does not embed the tensor parallel", + ), + ], +) +def test_render_default_pattern_escaped_field_not_rank_field( + monkeypatch, pattern, match +): + """Escaped braces are literal text and cannot pose as rank fields.""" + _mock_pp_group(monkeypatch, 2, pp_rank=1, tp_rank=0, tp_world_size=2) + with pytest.raises(ValueError, match=match): + ShardedStateLoader._render_default_pattern(pattern, part=0) + + +@pytest.mark.parametrize( + "pattern,world_size,tp_world_size,pp_rank,tp_rank,expected", + [ + # Mixed tp alias works when pipeline parallelism is on. + ( + "model-pp-{pp_rank}-rank-{rank}-part-{part}.safetensors", + 2, + 2, + 1, + 0, + "model-pp-1-rank-0-part-0.safetensors", + ), + # tp_rank alias works for TP-only checkpoints too. + ( + "model-tp-{tp_rank}-part-{part}.safetensors", + 1, + 2, + 0, + 3, + "model-tp-3-part-0.safetensors", + ), + ], +) +def test_render_default_pattern_tp_aliases_accepted( + monkeypatch, pattern, world_size, tp_world_size, pp_rank, tp_rank, expected +): + """Both tp aliases are supplied at render time, so either works.""" + _mock_pp_group( + monkeypatch, + world_size, + pp_rank=pp_rank, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ) + assert ShardedStateLoader._render_default_pattern(pattern, part=0) == expected + + +def test_render_default_pattern_unsupported_field_fails(monkeypatch): + """Unknown placeholders are rejected instead of KeyError-ing later.""" + _mock_pp_group(monkeypatch, 2, pp_rank=1, tp_world_size=2) + with pytest.raises(ValueError, match="unsupported placeholder"): + ShardedStateLoader._render_default_pattern( + "model-{pp_rank}-{tp_rank}-{epoch}-part-{part}.safetensors", part=0 + ) + + +@pytest.mark.parametrize( + "pp_rank,tp_rank,expected", + [ + (0, 0, "model-pp-0-tp-0-part-0.safetensors"), + (1, 2, "model-pp-1-tp-2-part-0.safetensors"), + (3, 5, "model-pp-3-tp-5-part-0.safetensors"), + ], +) +def test_render_default_pattern_non_zero_ranks(monkeypatch, pp_rank, tp_rank, expected): + """The rendered filename must embed the actual pp/tp ranks, not zeros.""" + _mock_pp_group(monkeypatch, 2, pp_rank=pp_rank, tp_rank=tp_rank) + assert ( + ShardedStateLoader._render_default_pattern( + ShardedStateLoader.PP_AND_TP_PATTERN, part=0 + ) + == expected + ) + + +def test_render_default_pattern_part_wildcard(monkeypatch): + """The glob wildcard '*' passes through as the part token for file lookup.""" + _mock_pp_group(monkeypatch, 1, tp_rank=4) + assert ( + ShardedStateLoader._render_default_pattern(part="*") + == "model-rank-4-part-*.safetensors" + ) + + +def test_save_model_pp_aware_filenames(monkeypatch, tmp_path): + """save_model writes pp/tp-aware filenames for pipeline-parallel models.""" + _mock_pp_group(monkeypatch, 2, pp_rank=1, tp_rank=2) + + saved: list[str] = [] + + def _capture(tensors, path, **kwargs): + saved.append(os.path.basename(path)) + + monkeypatch.setattr("safetensors.torch.save_file", _capture) + + model = SimpleNamespace( + state_dict=lambda: {"w": torch.zeros(4, dtype=torch.float32)} + ) + ShardedStateLoader.save_model(model, str(tmp_path)) + assert saved == ["model-pp-1-tp-2-part-0.safetensors"] + + +def test_save_model_part_split(monkeypatch, tmp_path): + """save_model rolls over to a new part when max_size is exceeded.""" + _mock_pp_group(monkeypatch, 1, tp_rank=0) + + saved: list[str] = [] + + def _capture(tensors, path, **kwargs): + saved.append(os.path.basename(path)) + + monkeypatch.setattr("safetensors.torch.save_file", _capture) + + # Two float32[4] tensors (16 B each); max_size=20 fits one but not two. + model = SimpleNamespace( + state_dict=lambda: { + "a": torch.zeros(4, dtype=torch.float32), + "b": torch.zeros(4, dtype=torch.float32), + } + ) + ShardedStateLoader.save_model(model, str(tmp_path), max_size=20) + assert saved == [ + "model-rank-0-part-0.safetensors", + "model-rank-0-part-1.safetensors", + ] + + +def test_save_model_rejects_bad_pattern_before_overwrite(monkeypatch, tmp_path): + """An unsafe pattern is rejected before any shard file is written.""" + _mock_pp_group(monkeypatch, 2, pp_rank=1, tp_rank=0, tp_world_size=2) + + # With the escaped pp field this pattern renders both stages to the + # same literal name; the validation must fire before any write. + colliding = tmp_path / "model-{pp_rank}-tp-0-part-0.safetensors" + colliding.write_bytes(b"existing shard") + before = sorted(p.name for p in tmp_path.iterdir()) + + model = SimpleNamespace( + state_dict=lambda: {"w": torch.zeros(4, dtype=torch.float32)} + ) + with pytest.raises(ValueError, match="does not embed the pipeline parallel"): + ShardedStateLoader.save_model( + model, + str(tmp_path), + pattern="model-{{pp_rank}}-tp-{tp_rank}-part-{part}.safetensors", + ) + + # The pre-existing shard is untouched and nothing new was written. + assert colliding.read_bytes() == b"existing shard" + assert sorted(p.name for p in tmp_path.iterdir()) == before + + @pytest.fixture(scope="module") def llama_3p2_1b_files(): input_dir = hf_api().snapshot_download( @@ -163,3 +476,128 @@ def test_sharded_state_loader( queue.join_thread() assert out_before == out_after + + +def _run_sharded_state_round_trip( + input_dir: str, + *, + tp_size: int, + pp_size: int, + num_gpus_available: int, +) -> None: + """Save a TP/PP sharded checkpoint, reload it, and compare outputs. + + Args: + input_dir: directory containing the original (un-sharded) model. + tp_size: tensor parallel size used for both save and load. + pp_size: pipeline parallel size used for both save and load. + num_gpus_available: number of GPUs on the current machine. + + """ + if num_gpus_available < tp_size * pp_size: + pytest.skip(f"Not enough GPUs for TP={tp_size}, PP={pp_size}") + + weights_patterns = ("*.safetensors",) + gpu_memory_utilization = 0.8 + ctx = mp.get_context("spawn") + + # Keep batching deterministic: this test compares exact greedy outputs + # across separate engine processes, whose schedulers may otherwise form + # different batches depending on prompt-processing timing. + platform_args = {"max_num_seqs": 1} + parallel_args = dict( + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + ) + + with TemporaryDirectory() as output_dir: + p = ctx.Process( + target=_run_writer, + args=(input_dir, output_dir, weights_patterns), + kwargs=dict( + **parallel_args, + gpu_memory_utilization=gpu_memory_utilization, + enforce_eager=True, + **platform_args, + ), + ) + p.start() + p.join() + assert p.exitcode == 0, f"writer process failed with exit code {p.exitcode}" + + if pp_size > 1: + # PP checkpoints must use the pp/tp-aware naming scheme so that + # different pipeline stages write distinct files. + sharded_files = [ + f for f in os.listdir(output_dir) if f.endswith(".safetensors") + ] + assert sharded_files, "no sharded checkpoint files were written" + assert any( + fnmatch.fnmatch(f, "model-pp-*-tp-*-part-*.safetensors") + for f in sharded_files + ), f"expected PP-aware filenames, got {sharded_files}" + + queue = ctx.Queue() + p = ctx.Process( + target=_run_generate, + args=(input_dir, queue), + kwargs=dict( + enable_lora=False, + gpu_memory_utilization=gpu_memory_utilization, + **parallel_args, + **platform_args, + ), + ) + p.start() + # Call queue.get() before p.join() to prevent deadlock: + # If p.join() is called before queue.get() and the queue is full, + # the child process may block while writing to the queue and never + # terminate, causing the parent to wait indefinitely on p.join(). + # See: https://github.com/vllm-project/vllm/pull/22371#discussion_r2257773814 + out_before = queue.get() + p.join() + assert p.exitcode == 0, f"generate process failed with exit code {p.exitcode}" + queue.close() + queue.join_thread() + + queue = ctx.Queue() + p = ctx.Process( + target=_run_generate, + args=(output_dir, queue), + kwargs=dict( + enable_lora=False, + gpu_memory_utilization=gpu_memory_utilization, + load_format="sharded_state", + **parallel_args, + **platform_args, + ), + ) + p.start() + out_after = queue.get() + p.join() + assert p.exitcode == 0, f"generate process failed with exit code {p.exitcode}" + queue.close() + queue.join_thread() + + assert out_before == out_after + + +def test_sharded_state_loader_pp(num_gpus_available, llama_3p2_1b_files): + """Round-trip sharded state save/load with pipeline parallelism (PP=2).""" + _run_sharded_state_round_trip( + llama_3p2_1b_files, + tp_size=1, + pp_size=2, + num_gpus_available=num_gpus_available, + ) + + +def test_sharded_state_loader_pp_tp(num_gpus_available, llama_3p2_1b_files): + """Round-trip sharded state save/load with pipeline + tensor parallelism + (PP=2, TP=2).""" + _run_sharded_state_round_trip( + llama_3p2_1b_files, + tp_size=2, + pp_size=2, + num_gpus_available=num_gpus_available, + ) diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index feb8f25d1e..fea0c4e3c7 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -2444,6 +2444,11 @@ def get_tensor_model_parallel_rank() -> int: return get_tp_group().rank_in_group +def get_pipeline_model_parallel_rank() -> int: + """Return my rank for the pipeline model parallel group.""" + return get_pp_group().rank_in_group + + def get_node_count() -> int: """Return the total number of nodes in the distributed environment.""" assert _NODE_COUNT is not None, "distributed environment is not initialized" diff --git a/vllm/model_executor/model_loader/sharded_state_loader.py b/vllm/model_executor/model_loader/sharded_state_loader.py index d17bfdc6d4..97a31e8e87 100644 --- a/vllm/model_executor/model_loader/sharded_state_loader.py +++ b/vllm/model_executor/model_loader/sharded_state_loader.py @@ -4,6 +4,7 @@ import collections import glob import os +import string import time from collections.abc import Generator from copy import copy @@ -28,13 +29,14 @@ logger = init_logger(__name__) class ShardedStateLoader(BaseModelLoader): """Model loader that directly loads each worker's model state dict, which - enables a fast load path for large tensor-parallel models where each worker - only needs to read its own shard rather than the entire checkpoint. See - `examples/features/sharded_state/save_sharded_state_offline.py` for creating - a sharded checkpoint. + enables a fast load path for large tensor-parallel/pipeline-parallel models + where each worker only needs to read its own shard rather than the entire + checkpoint. See `examples/features/sharded_state/save_sharded_state_offline.py` + for creating a sharded checkpoint. """ - DEFAULT_PATTERN = "model-rank-{rank}-part-{part}.safetensors" + TP_ONLY_PATTERN = "model-rank-{rank}-part-{part}.safetensors" + PP_AND_TP_PATTERN = "model-pp-{pp_rank}-tp-{tp_rank}-part-{part}.safetensors" def __init__(self, load_config: LoadConfig): super().__init__(load_config) @@ -44,14 +46,94 @@ class ShardedStateLoader(BaseModelLoader): if load_config.model_loader_extra_config is None else copy(load_config.model_loader_extra_config) ) - self.pattern = extra_config.pop("pattern", self.DEFAULT_PATTERN) + # A custom shard filename pattern can be supplied via + # model_loader_extra_config; when omitted (None), the default pattern + # is selected at render time based on the parallel configuration. + self.pattern = extra_config.pop("pattern", None) if extra_config: raise ValueError( f"Unexpected extra config keys for load format " - f"{load_config.load_format}: " - f"{load_config.model_loader_extra_config.keys()}" + f"{load_config.load_format}: {extra_config.keys()}" ) + @staticmethod + def _get_default_pattern() -> str: + from vllm.distributed import get_pp_group + + return ( + ShardedStateLoader.PP_AND_TP_PATTERN + if get_pp_group().world_size > 1 + else ShardedStateLoader.TP_ONLY_PATTERN + ) + + @staticmethod + def _render_default_pattern(pattern: str | None = None, *, part: int | str) -> str: + from vllm.distributed import ( + get_pipeline_model_parallel_rank, + get_pp_group, + get_tensor_model_parallel_rank, + get_tp_group, + ) + + # _get_default_pattern is called here because _render_default_pattern + # can be invoked from save_model without instantiating a + # ShardedStateLoader instance. + if pattern is None: + pattern = ShardedStateLoader._get_default_pattern() + + # Parse the real format fields instead of matching placeholder + # text: escaped braces like {{pp_rank}} render as literal text + # and must not count as a rank field. + fields = { + field + for _, field, _, _ in string.Formatter().parse(pattern) + if field is not None + } + unsupported = fields - {"pp_rank", "tp_rank", "rank", "part"} + if unsupported: + raise ValueError( + f"Pattern '{pattern}' uses unsupported placeholder(s) " + f"{sorted(unsupported)}; supported fields are pp_rank, " + "tp_rank, rank and part" + ) + tp_rank = get_tensor_model_parallel_rank() + pp_rank = get_pipeline_model_parallel_rank() + pp_world_size = get_pp_group().world_size + tp_world_size = get_tp_group().world_size + # The pattern must carry a rank field for every active + # parallelism dimension, otherwise shards from different ranks + # collide on the same filename and silently overwrite each other. + has_pp_field = "pp_rank" in fields + has_tp_field = "tp_rank" in fields or "rank" in fields + if pp_world_size > 1: + if not has_pp_field: + raise ValueError( + f"Pattern '{pattern}' does not embed the pipeline parallel " + "rank; with pipeline parallelism every stage would write the " + f"same file. Use '{ShardedStateLoader.PP_AND_TP_PATTERN}' " + "instead." + ) + if tp_world_size > 1 and not has_tp_field: + raise ValueError( + f"Pattern '{pattern}' does not embed the tensor parallel " + "rank; TP workers would overwrite each other's shards." + ) + # Supply both tp aliases so mixed-alias patterns render. + return pattern.format( + pp_rank=pp_rank, tp_rank=tp_rank, rank=tp_rank, part=part + ) + if has_pp_field: + raise ValueError( + f"Pattern '{pattern}' requires pipeline parallelism, but the " + "pipeline parallel size is 1" + ) + if tp_world_size > 1 and not has_tp_field: + raise ValueError( + f"Pattern '{pattern}' does not embed the tensor parallel rank; " + "TP workers would overwrite each other's shards." + ) + return pattern.format(tp_rank=tp_rank, rank=tp_rank, part=part) + @staticmethod def _filter_subtensors( tensors: dict[str, torch.Tensor], @@ -106,30 +188,36 @@ class ShardedStateLoader(BaseModelLoader): self._prepare_weights(model_config.model, model_config.revision) def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: - from vllm.distributed import get_tensor_model_parallel_rank - model_weights = model_config.model if model_weights_override := model_config.model_weights: model_weights = model_weights_override local_model_path = model_weights - rank = get_tensor_model_parallel_rank() - pattern = os.path.join( - local_model_path, - self.pattern.format(rank=rank, part="*"), + rendered_pattern = ShardedStateLoader._render_default_pattern( + self.pattern, part="*" ) - filepaths = [] if is_s3(local_model_path): - file_pattern = f"*{self.pattern.format(rank=rank, part='*')}" + file_pattern = f"*{rendered_pattern}" filepaths = s3_glob(path=local_model_path, allow_pattern=[file_pattern]) else: - filepaths = glob.glob(pattern) + filepaths = glob.glob( + os.path.join( + local_model_path, + rendered_pattern, + ) + ) if not filepaths: # TODO: support un-sharded checkpoints too raise ValueError( - f"Could not find checkpoint files '{pattern}', only " - f"pre-sharded checkpoints are currently supported!" + f"Could not find checkpoint files " + f"'{ + os.path.join( + local_model_path, + rendered_pattern, + ) + }', " + f"only pre-sharded checkpoints are currently supported!" ) state_dict = self._filter_subtensors(model.state_dict()) counter_before_loading_weights = time.perf_counter() @@ -137,8 +225,17 @@ class ShardedStateLoader(BaseModelLoader): # If loading with LoRA enabled, additional padding may # be added to certain parameters. We only load into a # narrowed view of the parameter data. - param_data = state_dict[key].data - param_shape = state_dict[key].shape + param = state_dict[key] + if param.dtype != tensor.dtype: + logger.warning( + "Attempted to load weight %s with dtype %s into " + "parameter with dtype %s", + key, + tensor.dtype, + param.dtype, + ) + param_data = param.data + param_shape = param.shape for dim, size in enumerate(tensor.shape): if size < param_shape[dim]: param_data = param_data.narrow(dim, 0, size) @@ -182,11 +279,6 @@ class ShardedStateLoader(BaseModelLoader): ) -> None: from safetensors.torch import save_file - from vllm.distributed import get_tensor_model_parallel_rank - - if pattern is None: - pattern = ShardedStateLoader.DEFAULT_PATTERN - rank = get_tensor_model_parallel_rank() part_idx = 0 total_size = 0 state_dict = ShardedStateLoader._filter_subtensors(model.state_dict()) @@ -194,7 +286,9 @@ class ShardedStateLoader(BaseModelLoader): for key, tensor in state_dict.items(): param_size = tensor.nelement() * tensor.element_size() if max_size is not None and total_size + param_size > max_size: - filename = pattern.format(rank=rank, part=part_idx) + filename = ShardedStateLoader._render_default_pattern( + pattern, part=part_idx + ) save_file( state_dict_part, os.path.join(path, filename), @@ -205,7 +299,9 @@ class ShardedStateLoader(BaseModelLoader): state_dict_part[key] = tensor total_size += param_size if len(state_dict_part) > 0: - filename = pattern.format(rank=rank, part=part_idx) + filename = ShardedStateLoader._render_default_pattern( + pattern, part=part_idx + ) save_file( state_dict_part, os.path.join(path, filename),