[Bugfix][Model Loader] Run post-load processing for raw Tensorizer weights

[nicole-lihui]

Purpose

Raw/Hugging Face Tensorizer checkpoints can load their weights successfully but
return a model with state required by forward still uninitialized. For example,
a CPU UnquantizedLinearMethod fails with
AttributeError: 'Linear' object has no attribute 'cpu_linear'.

Route the non-vLLM-tensorized branch through BaseModelLoader.load_model and
remove its separate loading helper. This reuses the existing weight iterator
and shared post-load lifecycle before returning the model. The native
vLLM-tensorized path is unchanged; its state-restoration and eval behavior are
outside this PR's scope.

Refs #55198.

The shared path also honors model_config and load_config.device, including
online-quantization finalization when applicable. Calling the existing
load_weights adds one format-detection read on the raw path.

Duplicate-work checks on 2026-09-04 found no open PR implementing this fix
(55198 in:body, Tensorizer post-load, and TensorizerLoader). #42823 excludes
Tensorizer; #37677 addresses storage URIs. The issue has a contributor's
expression of interest; this PR is limited to the reproduced raw/HF failure.

Test Plan

The regression writes a real Tensorizer file, loads it through the public loader,
and compares the first forward with the source module. It replaces full-model
construction with a tiny module and uses vLLM's real CPU linear method. CPU
selection is local to the test so accelerator CI hosts can also execute it.

VLLM_TARGET_DEVICE=cpu .venv/bin/python -m pytest \
  tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py \
  -k 'deserialized_hf_weights_are_ready_for_forward or tensorize_vllm_model_shuts_down_engine' -q

.venv/bin/pre-commit run --files \
  vllm/model_executor/model_loader/tensorizer_loader.py \
  tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py

.venv/bin/pre-commit run mypy-3.12 --hook-stage manual --files \
  vllm/model_executor/model_loader/tensorizer_loader.py \
  tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py

Existing full-model comparison for the standard test environment:

.venv/bin/python -m pytest \
  tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py::test_deserialized_hf_model_has_same_outputs -q

Test Result

• Targeted regression and five existing shutdown cases: 6 passed.
• The new regression against baseline loader 9cd956c7e6: 1 failed at
  forward with missing cpu_linear. The baseline loader was loaded in memory,
  with its model factory forwarded to the same tiny-model test factory.
• Pre-commit and the Python 3.12 mypy hook: passed.
• Full-model comparison: 1 passed on macOS arm64 / CPU / FP16, torch 2.13.0,
  Tensorizer 2.10.1, cached facebook/opt-125m. HF and raw Tensorizer-loaded vLLM
  generated identical outputs for four prompts, up to 50 greedy tokens each.
  The existing test used the public resource settings shown below.
• The local CPU extension was rebuilt against torch 2.13.0 to replace a stale
  binary. No C++ source or dependency changes are included.
• Native round-trip and GPU/quantized-model validation were not run.

Reproduce the local CPU full-model comparison

The default memory fraction exceeded this machine's available memory. The local
run used 256 MiB KV cache, eight sequences, a 1,024-token batch limit, and eager
execution. The test body, loading, kernels, and output assertions were unchanged.

The executed local helper was
.venv/bin/python dist/contribute-pr/issue-55198/run_cpu_model_regression.py.
That investigation file is excluded from the commit. For another checkout, save
the equivalent helper below as run_cpu_tensorizer_test.py at the repository root
and run .venv/bin/python run_cpu_tensorizer_test.py after caching OPT-125m:

import os
import sys
from pathlib import Path

os.environ["VLLM_TARGET_DEVICE"] = "cpu"
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"
root = Path(__file__).resolve().parent
sys.path.insert(0, str(root))
os.chdir(root)

import pytest
import tests.conftest as test_config


class CpuVllmRunner(test_config.VllmRunner):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault("kv_cache_memory_bytes", 256 * 1024**2)
        kwargs.setdefault("max_num_seqs", 8)
        kwargs.setdefault("max_num_batched_tokens", 1024)
        kwargs.setdefault("enforce_eager", True)
        super().__init__(*args, **kwargs)


if __name__ == "__main__":
    test_config.VllmRunner = CpuVllmRunner
    raise SystemExit(pytest.main([
        "tests/model_executor/model_loader/tensorizer_loader/"
        "test_tensorizer.py::test_deserialized_hf_model_has_same_outputs",
        "-q",
    ]))

AI assistance: Codex assisted with investigation, implementation, regression-test
design, and local validation.

---

Essential elements

• Purpose and related issue, with the scope stated explicitly.
• Test plan and commands.
• Actual regression and model-output comparison results.
• Documentation impact considered; no public API or configuration is added.

[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.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Team

Run ID: a2772e6b-c47f-4bae-abc0-aab792d0116b

📥 Commits

Reviewing files that changed from the base of the PR and between fd4a151 and 08c7d79.

📒 Files selected for processing (2)

• tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py
• vllm/model_executor/model_loader/tensorizer_loader.py

Included review availability: Your plan provides up to 10 included reviews per hour; 9 remain after this review.

---

📝 Summary

Summary by CodeRabbit

• Bug Fixes

  • Improved loading of serialized tensorized models so their weights are prepared correctly before the first inference.
  • Enhanced compatibility for loading models across CPU and accelerator environments.
  • Preserved support for non-native serialized models through the standard model-loading flow.

• Tests

  • Added coverage confirming that loaded models produce the same outputs as their original versions.

Walkthrough

TensorizerLoader now delegates non-vLLM tensorized model loading to BaseModelLoader. A new test verifies that deserialized weights are CPU-processed before the first forward pass.

Changes

Tensorizer loading

Layer / File(s)	Summary
Delegate fallback loading to the base loader vllm/model_executor/model_loader/tensorizer_loader.py	Removes the CPU-only serialized loading helper and delegates non-vLLM tensorized models to BaseModelLoader.
Validate deserialized weights before forward tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py	Adds a test that serializes a small model, loads it with TensorizerLoader, and compares the loaded output with the source output.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Merge Risk: ⚪ Minimal · up to 08c7d

Raw Tensorizer checkpoints now receive the shared post-load preparation required for CPU linear inference, with regression coverage confirming the first forward matches the source model. No current merge-blocking risk remains.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 10.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 10 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly identifies the bugfix and the main change: applying post-load processing to raw Tensorizer weights.
Description check	✅ Passed	The description directly explains the raw Tensorizer loading issue, the implementation approach, scope, regression tests, and validation results.

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Create stacked PR
• Commit on current branch

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[codetalker7]

This looks totally LLM generated. By the way, even for the raw Huggingface path, why did you remove the call to _load_model_serialized_cpu? Also, it'll be great if you have results for real huggingface weights that you tested (it looks like your test contains a custom nn module). This will be helpful for me for the native path!