[l11y11z11]

Overview

ShardedStateLoader.save_model() can silently overwrite earlier checkpoint parts when a custom filename pattern renders the same name for multiple parts. Plan the parts and validate their rendered filenames before writing any files.

Related: #52110 added validation against PP/TP rank filename collisions. This fixes collisions between parts within one rank; it does not change parallel rank validation.

Claims

• Reject colliding part filenames before modifying an existing checkpoint or writing a partial export.
• Cover missing and escaped {part} fields, as well as lossy formatting that causes a later part to collide with an earlier one.
• Preserve default naming, shard boundaries, tensor contents, and valid single-file patterns that do not contain {part}.

Validation

Base: 852989f580820849113d67d13fd8053e279d0444.

Environment: Apple M4, macOS 27.0.1 arm64, Python 3.12.13, PyTorch 2.14.1, safetensors 0.8.0, pytest 9.1.1. CPU tensors and real safetensors files; only distributed rank/group accessors are mocked with the existing test helper. No model downloads or GPU are required for these cases.

Minimal reproduction on the base revision, with PP/TP rank 0 and world size 1:

model = torch.nn.Linear(1, 1)  # two independent float32 tensors, 4 bytes each
ShardedStateLoader.save_model(
    model,
    output_directory,
    pattern="model-rank-{rank}.safetensors",
    max_size=4,
)

Before: returns successfully, but only model-rank-0.safetensors remains, containing bias; weight has been overwritten. The default pattern produces two files and preserves both tensors.

After: the colliding pattern raises ValueError before any file is written. Regression cases also verify that a pre-existing checkpoint remains byte-for-byte untouched. Valid split and single-file exports are read back through ShardedStateLoader.load_weights() and every parameter is checked exactly.

.venv/bin/python -m pytest \
  tests/model_executor/model_loader/test_sharded_state_loader.py \
  -k 'filter_subtensors or render_default_pattern or get_default_pattern or save_model' -q

Result: 32 passed, 6 deselected. The three collision regressions failed on the unchanged production code (DID NOT RAISE ValueError); the four initial valid-save controls passed.

.venv/bin/pre-commit run ruff-check --files \
  vllm/model_executor/model_loader/sharded_state_loader.py \
  tests/model_executor/model_loader/test_sharded_state_loader.py
.venv/bin/pre-commit run ruff-format --files \
  vllm/model_executor/model_loader/sharded_state_loader.py \
  tests/model_executor/model_loader/test_sharded_state_loader.py
git diff --check

All three checks passed. Python compilation of both changed files passed. The full pre-commit run --files invocation for the two changed files also exited 0, including mypy, typos, and CI test-tethering checks.

Additional manual single-GPU integration validation passed on an RTX 3090 24 GB (driver 595.71.05, Linux x86_64, Python 3.12.3, PyTorch 2.14.0+cu130 / CUDA 13.0), using the same source revision and byte-identical patch. Qwen/Qwen3-0.6B revision c1899de289a04d12100db370d81485cdf75e47ca, float16, TP=1, eager execution, seed=0, max_model_len=2048, max_num_seqs=1, GPU memory utilization=0.7. For both the default pattern and custom-rank-{rank}-part-{part}.safetensors, a 16 MiB limit produced 86 actual parts. Each variant was reloaded in a fresh process: four fixed prompts, 32 tokens each, temperature=0 and ignore_eos=True, produced exactly the same 128 token IDs as before export. The targeted unit command above also passed on this host: 32 passed, 6 deselected (these unit cases use CPU tensors).

The manual validation script and commands are included below. The six existing GPU/model/distributed cases were not run; the additional Qwen check does not replace multi-process PP/TP or LoRA coverage. No performance benchmark or speedup is claimed. The added regression tests are in the existing model-loader suite already selected by distributed model Buildkite jobs.

Details

The exporter previously rendered and wrote each part immediately, without checking whether a later part would reuse the same filename. Checking only for a {part} placeholder is insufficient: {part!s:.1} maps both part 1 and part 10 to 1.

The change first builds the same part dictionaries as before, renders all filenames, and checks uniqueness. It then writes the planned parts. Planning retains additional Python dictionaries/references; it does not clone tensor storage. There is no new dependency, configuration option, or serving-path work.

This is not a transactional checkpoint writer: unrelated filesystem/write failures can still leave a partial export, as before. It does not change the policy for an existing output directory, recover an already overwritten export, or address arbitrary filesystem aliases. Integer-only formatting of {part} during loading is an independent existing limitation because the loader substitutes "*"; this PR does not expand pattern syntax.

Duplicate-work check on 2026-10-19: searched open and closed PRs by sharded-state titles, the loader name, {part}, filename/collision/overwrite terms, and inspected #52110's discussion. #51352 changes the offline example's output-directory validation, #61432 addresses strided tensors, and #62231 addresses attention-scale round trips. None of those is the same part-filename collision fix. No new public Issue was opened for this independently reproduced defect.

AI assistance was used for investigation, implementation, tests, and this description (OpenAI Codex). The human submitter confirmed reviewing every changed line and personally running the relevant tests before publication.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill.
• AI assistance was used during the creation of this PR.
• Design Fit: Validates the existing export operation without changing its public API or serving code.
• Testing and Validation: Real CPU serialization/loading, targeted regression coverage, and manual single-GPU export/reload generation checks; remaining limitations stated above.
• Code Quality and Style: Ruff and formatting passed; no unrelated production changes.
• Pull Request Contents: Reproduction, root cause, compatibility, limitations, and duplicate-work checks included.

Reproduce the manual single-GPU export/reload check

Save the following as gpu_roundtrip.py outside the vLLM checkout. It is a manual integration check, not part of the proposed source change or a performance benchmark. Run from the patched repository root after installing vLLM and downloading Qwen/Qwen3-0.6B at revision c1899de289a04d12100db370d81485cdf75e47ca. Point --model at that local snapshot and --output at an empty results directory.

"""Manual single-GPU integration check; run each phase in a fresh process.

Not a performance benchmark. The export and reload phases must use the same
checkout, environment, model metadata, and deterministic generation settings.
"""

import argparse
import hashlib
import json
from pathlib import Path
import shutil
import subprocess

import torch

from vllm import LLM, SamplingParams


PROMPTS = [
    "Hello, my name is",
    "The capital of France is",
    "Write a Python function to add two numbers:",
    "Explain what a computer does in one sentence:",
]
PATTERNS = {
    "default": None,
    "custom": "custom-rank-{rank}-part-{part}.safetensors",
}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("phase", choices=["export", "reload"])
    parser.add_argument("--model", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--variant", choices=PATTERNS, default="default")
    args = parser.parse_args()
    args.output.mkdir(parents=True, exist_ok=True)
    checkpoint = args.output / args.variant
    extra = {}
    if args.phase == "export":
        model_path = args.model
    else:
        model_path = checkpoint
        extra = {
            "load_format": "sharded_state",
            "model_loader_extra_config": {"pattern": PATTERNS[args.variant]},
        }
    llm = LLM(
        model=str(model_path),
        dtype="float16",
        tensor_parallel_size=1,
        max_model_len=2048,
        max_num_seqs=1,
        gpu_memory_utilization=0.7,
        enforce_eager=True,
        seed=0,
        **extra,
    )
    outputs = llm.generate(
        PROMPTS, SamplingParams(temperature=0, max_tokens=32, ignore_eos=True)
    )
    tokens = [list(result.outputs[0].token_ids) for result in outputs]
    record = {
        "phase": args.phase,
        "variant": args.variant,
        "sha": subprocess.check_output(
            ["git", "rev-parse", "HEAD"], text=True
        ).strip(),
        "loader_sha256": hashlib.sha256(
            Path("vllm/model_executor/model_loader/sharded_state_loader.py").read_bytes()
        ).hexdigest(),
        "torch": torch.__version__,
        "cuda": torch.version.cuda,
        "gpu": torch.cuda.get_device_name(0),
        "token_ids": tokens,
        "texts": [result.outputs[0].text for result in outputs],
    }
    reference = args.output / f"{args.variant}-export.json"
    if args.phase == "export":
        checkpoint.mkdir(exist_ok=False)
        llm.llm_engine.engine_core.save_sharded_state(
            path=str(checkpoint),
            pattern=PATTERNS[args.variant],
            max_size=16 * 1024**2,
        )
        shards = sorted(checkpoint.glob("*.safetensors"))
        assert len(shards) > 1, "Workload must actually export multiple parts"
        record["shard_files"] = [file.name for file in shards]
        for file in args.model.iterdir():
            if file.is_file() and file.suffix in {".json", ".txt", ".model", ".jinja"}:
                shutil.copy2(file, checkpoint / file.name)
    else:
        expected = json.loads(reference.read_text())
        assert record["sha"] == expected["sha"]
        assert record["loader_sha256"] == expected["loader_sha256"]
        assert tokens == expected["token_ids"], "Reloaded greedy token IDs differ"
        record["exact_token_ids_match"] = True
    result_path = args.output / f"{args.variant}-{args.phase}.json"
    result_path.write_text(json.dumps(record, indent=2) + "\n")
    print(f"PASS: {result_path}", flush=True)


if __name__ == "__main__":
    main()

for variant in default custom; do
  for phase in export reload; do
    PYTHONPATH=. HF_HUB_OFFLINE=1 .venv/bin/python ../gpu_roundtrip.py "$phase" \
      --variant "$variant" \
      --model /path/to/local/Qwen3-0.6B \
      --output /path/to/empty/results
  done
done

Each phase runs in a separate Python process. The export phase generates the reference token IDs, writes actual shards, and copies model metadata. The reload phase loads those shards and asserts exact equality of generated token IDs.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀