iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[sharded state loader] support pp in sharded state loader (#52110)

Browse Source 
Signed-off-by: Andy Xie <andy.xning@gmail.com>
Co-authored-by: Simon Mo <simon.mo@hey.com>
This commit is contained in:
Ning Xie
2026-10-09 11:08:31 -07:00
committed by GitHub
co-authored by Simon Mo
parent c811943137
commit a0f28d73a8
4 changed files with 572 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/features/sharded_state/save_sharded_state_offline.py
+5 -2
View File
@@ -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",
tests/model_executor/model_loader/test_sharded_state_loader.py
+438
View File
@@ -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,
)
vllm/distributed/parallel_state.py
+5
View File
@@ -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"
vllm/model_executor/model_loader/sharded_state_loader.py
+124 -28
View File
@@ -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),