mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(kompress): load merged.pt for the v2 checkpoint instead of the unmerged PEFT safetensors (#2716)
# PR draft: fix(kompress): load merged.pt for the v2 checkpoint instead of the unmerged PEFT safetensors Branch: `rnoz/fix-kompress-merged-checkpoint` (off `upstream/main`). Two commits. Issue: https://github.com/headroomlabs-ai/headroom/issues/2714 (filed, open). --- ## Description `_load_kompress_pytorch` in `headroom/transforms/kompress_compressor.py` downloaded `model.safetensors` from the default model repo `chopratejas/kompress-v2-base` and loaded it with `strict=False`, discarding the missing/unexpected key report. That file is the unmerged PEFT checkpoint (encoder keys prefixed `encoder.base_model.model...`), which never matches `HeadroomCompressorModel`'s plain `encoder.*` keys. The LoRA-adapted encoder weights were silently dropped while `token_head`/`span_conv` happened to match and loaded fine, so the model ran with a stock, non-adapted `answerdotai/ModernBERT-base` encoder feeding correctly trained decision heads, with no error and a healthy status reported everywhere. `scripts/export_kompress_v2_onnx.py` already documents this exact mismatch and loads the correct `merged.pt` sub-state-dicts for its own export path. This PR mirrors that same loading logic into the runtime PyTorch loader, with a fallback to the plain `model.safetensors` format for repos that never shipped a `merged.pt` (verified against the v1 `chopratejas/kompress-base` repo via the public HF API, which has no `merged.pt`). Both paths now check the missing/unexpected key report and raise instead of silently proceeding on a mismatch. A second commit fixes a gap an adversarial review caught in the first: the cache-only (`allow_download=False`, startup preload) path could not tell "this repo genuinely has no merged.pt" apart from "merged.pt exists but is not downloaded yet", so it would have fallen back to a stale `model.safetensors` left over from before this fix on exactly the upgrade path this PR is meant to close. It now uses `huggingface_hub`'s own `.no_exist` cache marker (via a new `hf_entry_known_absent()` helper in `headroom/onnx_runtime.py`) to make that distinction without a network call, and only falls back when absence is confirmed; otherwise it raises `KompressModelNotCached` so the caller defers instead of guessing. Closes #2714 ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) ## Changes Made - `headroom/transforms/kompress_compressor.py`: added `_load_merged_state_dict`, `_load_plain_state_dict`, and `_load_pytorch_weights`, replacing the inline `model.safetensors` download + `load_state_dict(strict=False)` call in `_load_kompress_pytorch`. `merged.pt` is tried first; the plain format is only used when its absence is confirmed. - `headroom/onnx_runtime.py`: added `hf_entry_known_absent()`, a thin wrapper around `huggingface_hub.try_to_load_from_cache()` that reads the on-disk `.no_exist` marker HF writes after a real 404, so cache-only code can distinguish "confirmed absent" from "never checked" without hitting the network. - `tests/test_transforms/test_kompress_compressor.py`: added `TestPytorchWeightLoading` (8 tests) covering the merged-checkpoint happy path, missing-section and key-mismatch failures, the plain-format fallback for repos without `merged.pt`, and the cache-only ambiguity fix (confirmed-absent vs unconfirmed). ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check`) - [x] Type checking passes (`mypy`) - [x] Manual testing performed (real download against the live `chopratejas/kompress-v2-base` repo, not just mocks) ### Test Output ```text $ .venv/bin/python3 -m pytest tests/ -k kompress -q 178 passed, 7 skipped in 22.39s $ .venv/bin/python3 -m ruff check headroom/onnx_runtime.py headroom/transforms/kompress_compressor.py tests/test_transforms/test_kompress_compressor.py All checks passed! $ .venv/bin/python3 -m ruff format --check headroom/onnx_runtime.py headroom/transforms/kompress_compressor.py tests/test_transforms/test_kompress_compressor.py 3 files already formatted $ .venv/bin/python3 -m mypy headroom/onnx_runtime.py headroom/transforms/kompress_compressor.py Success: no issues found in 2 source files ``` ## Real Behavior Proof Ran the actual fixed loader against the live `chopratejas/kompress-v2-base` HF repo (not a mock), before and after each fix: ```text # BEFORE (parsed the real cached model.safetensors header by hand, no safetensors lib needed): total tensors: 316 encoder.base_model.model.embeddings.norm.weight encoder.base_model.model.embeddings.tok_embeddings.weight ...(all 310 encoder tensors share this prefix)... span_conv.0.bias / span_conv.0.weight / span_conv.2.bias / span_conv.2.weight token_head.bias / token_head.weight exact prefix match count with plain "encoder.<rest>": 0 # This confirms the pre-fix code's model.load_state_dict(state_dict, strict=False) # silently dropped every encoder weight (0 keys match HeadroomCompressorModel.encoder), # while token_head/span_conv happened to match and loaded. # AFTER (commit 1, real merged.pt download + load): $ .venv/bin/python3 -c " import headroom.transforms.kompress_compressor as kmod model = kmod._get_model_class()() kmod._load_pytorch_weights(model, 'chopratejas/kompress-v2-base', allow_download=True) print('SUCCESS: 0 missing/unexpected keys across all three sections') " SUCCESS: 0 missing/unexpected keys across all three sections # AFTER (full pipeline, real end-to-end compression through the public API): $ .venv/bin/python3 -c " import headroom.transforms.kompress_compressor as kmod compressor = kmod.KompressCompressor() result = compressor.compress(sample_traceback_plus_boilerplate_text) print(result.original_tokens, result.compressed_tokens, result.tokens_saved) print('ValueError: bad input' in result.compressed) " 497 454 43 True # must-keep line (the actual error) survived compression # AFTER (commit 2, cache-only ambiguity): unit tests # test_cache_only_defers_instead_of_using_stale_plain_checkpoint: PASSED # test_cache_only_uses_plain_checkpoint_when_merged_pt_confirmed_absent: PASSED ``` ## Review Readiness - [x] I have performed a self-review - [x] An independent adversarial review pass was run on both commits before this PR was opened; its one finding (the cache-only ambiguity) is fixed in commit 2, verified with new regression tests - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation (N/A: internal loader behavior, no public API or config surface changed) - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective (regression tests for both the original silent-drop bug and the cache-only ambiguity found in review) - [x] New and existing unit tests pass locally with my changes - [x] I did not edit `CHANGELOG.md` ## Additional Notes - Both `merged.pt` and the plain `model.safetensors` fallback now raise loudly on any state-dict mismatch instead of proceeding with partially-loaded weights, closing the general silent-failure class this bug belonged to, not just this one instance of it. - No other call sites of `_load_kompress_pytorch` or its removed inline code exist; its public signature is unchanged.
This commit is contained in:
Raúl
GitHub
parent
6d5516dcb8
commit
46da91b2f1
3 changed files with 333 additions and 13 deletions
|
|
@ -117,6 +117,25 @@ def hf_hub_download_local_first(
|
|||
return str(hf_hub_download(repo_id, filename, revision=revision))
|
||||
|
||||
|
||||
def hf_entry_known_absent(repo_id: str, filename: str, *, revision: str | None = None) -> bool:
|
||||
"""True only if a prior network lookup already confirmed ``filename`` does
|
||||
not exist in ``repo_id`` at the resolved revision.
|
||||
|
||||
Backed by ``huggingface_hub``'s own cache of negative lookups (the
|
||||
``.no_exist`` marker written after a real 404), so this never makes a
|
||||
network call itself. Returns ``False`` both when the file is cached and
|
||||
when nothing is known yet about it, on purpose: callers in a cache-only
|
||||
(``allow_network=False``) code path can use this to tell "confirmed
|
||||
missing upstream, safe to use a fallback file" apart from "just never
|
||||
checked yet, do not guess."
|
||||
"""
|
||||
from huggingface_hub import _CACHED_NO_EXIST, try_to_load_from_cache
|
||||
|
||||
revision = _resolve_revision(repo_id, revision)
|
||||
result = try_to_load_from_cache(repo_id, filename, revision=revision)
|
||||
return result is _CACHED_NO_EXIST
|
||||
|
||||
|
||||
def create_cpu_session_options(
|
||||
ort: Any,
|
||||
*,
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ from ..config import TransformResult
|
|||
from ..onnx_runtime import (
|
||||
ONNX_CPU_ARENA_ENV,
|
||||
create_cpu_session_options,
|
||||
hf_entry_known_absent,
|
||||
hf_hub_download_local_first,
|
||||
trim_process_heap,
|
||||
)
|
||||
|
|
@ -721,6 +722,104 @@ def _load_modernbert_tokenizer(auto_tokenizer: Any, *, allow_download: bool) ->
|
|||
raise
|
||||
|
||||
|
||||
# Sub-state-dict keys inside a merged v2-style checkpoint (see
|
||||
# scripts/export_kompress_v2_onnx.py, which this mirrors).
|
||||
_MERGED_CHECKPOINT_KEYS = ("encoder_state_dict", "token_head_state_dict", "span_conv_state_dict")
|
||||
|
||||
|
||||
def _load_merged_state_dict(model: Any, ckpt_path: str, model_id: str) -> None:
|
||||
"""Load a merged v2-style checkpoint (LoRA already folded into the encoder).
|
||||
|
||||
The checkpoint is a dict of per-submodule state-dicts
|
||||
(``encoder_state_dict`` / ``token_head_state_dict`` / ``span_conv_state_dict``)
|
||||
rather than a single flat state-dict, so each piece is loaded into its
|
||||
matching submodule directly instead of via a single ``load_state_dict``
|
||||
call on the whole model.
|
||||
"""
|
||||
import torch
|
||||
|
||||
ckpt = torch.load(ckpt_path, map_location="cpu")
|
||||
missing_sections = [k for k in _MERGED_CHECKPOINT_KEYS if k not in ckpt]
|
||||
if missing_sections:
|
||||
raise RuntimeError(
|
||||
f"merged.pt for {model_id} is missing {missing_sections}; found keys: "
|
||||
f"{sorted(ckpt)}. This checkpoint format is not what the loader expects."
|
||||
)
|
||||
|
||||
for section, submodule in (
|
||||
("encoder_state_dict", model.encoder),
|
||||
("token_head_state_dict", model.token_head),
|
||||
("span_conv_state_dict", model.span_conv),
|
||||
):
|
||||
missing, unexpected = submodule.load_state_dict(ckpt[section], strict=False)
|
||||
if missing or unexpected:
|
||||
raise RuntimeError(
|
||||
f"{model_id} {section}: state_dict mismatch against {type(submodule).__name__} "
|
||||
f"(missing={list(missing)[:5]}, unexpected={list(unexpected)[:5]}). "
|
||||
"The checkpoint no longer matches HeadroomCompressorModel's architecture."
|
||||
)
|
||||
|
||||
|
||||
def _load_plain_state_dict(model: Any, weights_path: str, model_id: str) -> None:
|
||||
"""Load a plain, already-merged full state-dict (the pre-v2 / non-PEFT format)."""
|
||||
from safetensors.torch import load_file
|
||||
|
||||
state_dict = load_file(weights_path)
|
||||
missing, unexpected = model.load_state_dict(state_dict, strict=False)
|
||||
if missing or unexpected:
|
||||
raise RuntimeError(
|
||||
f"{model_id} model.safetensors: state_dict mismatch against "
|
||||
f"HeadroomCompressorModel (missing={list(missing)[:5]}, "
|
||||
f"unexpected={list(unexpected)[:5]}). Refusing to run with unloaded weights."
|
||||
)
|
||||
|
||||
|
||||
def _load_pytorch_weights(model: Any, model_id: str, *, allow_download: bool) -> None:
|
||||
"""Load PyTorch weights into ``model``, preferring the merged v2 checkpoint.
|
||||
|
||||
``merged.pt`` (when the repo ships one) holds LoRA-merged sub-state-dicts
|
||||
keyed by submodule name. In a PEFT-trained repo, ``model.safetensors`` is
|
||||
the *unmerged* adapter checkpoint (encoder keys prefixed
|
||||
``encoder.base_model.model...``) and does not map onto this module tree at
|
||||
all, so it is only used as a fallback for repos that never shipped a
|
||||
merged checkpoint (e.g. the original non-LoRA kompress-base).
|
||||
|
||||
In cache-only mode (``allow_download=False``) a ``merged.pt`` cache miss is
|
||||
ambiguous: it could mean the repo has no merged checkpoint (safe to use the
|
||||
plain fallback), or it could mean the repo has one but it just is not
|
||||
downloaded yet (in which case a *stale* cached ``model.safetensors`` from a
|
||||
prior run must not be used, since for a PEFT repo it is the wrong format).
|
||||
``hf_entry_known_absent`` disambiguates without a network call, using
|
||||
HuggingFace Hub's own cache of confirmed-404 lookups.
|
||||
"""
|
||||
try:
|
||||
ckpt_path = hf_hub_download_local_first(model_id, "merged.pt", allow_network=allow_download)
|
||||
except _NOT_CACHED_ERRORS as exc:
|
||||
if not allow_download:
|
||||
if not hf_entry_known_absent(model_id, "merged.pt"):
|
||||
raise KompressModelNotCached(model_id) from exc
|
||||
try:
|
||||
weights_path = hf_hub_download_local_first(
|
||||
model_id, "model.safetensors", allow_network=False
|
||||
)
|
||||
except _NOT_CACHED_ERRORS:
|
||||
raise KompressModelNotCached(model_id) from exc
|
||||
_load_plain_state_dict(model, weights_path, model_id)
|
||||
return
|
||||
if isinstance(exc, EntryNotFoundError):
|
||||
# merged.pt genuinely does not exist in this repo (confirmed by a
|
||||
# real network lookup, not just a cache miss) - fall back to the
|
||||
# plain format instead of treating it as a download failure.
|
||||
weights_path = hf_hub_download_local_first(
|
||||
model_id, "model.safetensors", allow_network=allow_download
|
||||
)
|
||||
_load_plain_state_dict(model, weights_path, model_id)
|
||||
return
|
||||
raise
|
||||
else:
|
||||
_load_merged_state_dict(model, ckpt_path, model_id)
|
||||
|
||||
|
||||
def _load_kompress_pytorch(
|
||||
model_id: str, device: str = "auto", *, allow_download: bool = True
|
||||
) -> tuple[Any, Any, str]:
|
||||
|
|
@ -738,22 +837,10 @@ def _load_kompress_pytorch(
|
|||
|
||||
logger.info("Downloading Kompress PyTorch model from %s ...", model_id)
|
||||
|
||||
try:
|
||||
weights_path = hf_hub_download_local_first(
|
||||
model_id, "model.safetensors", allow_network=allow_download
|
||||
)
|
||||
except _NOT_CACHED_ERRORS as exc:
|
||||
if not allow_download:
|
||||
raise KompressModelNotCached(model_id) from exc
|
||||
raise
|
||||
|
||||
HeadroomCompressorModel = _get_model_class()
|
||||
model = HeadroomCompressorModel()
|
||||
|
||||
from safetensors.torch import load_file
|
||||
|
||||
state_dict = load_file(weights_path)
|
||||
model.load_state_dict(state_dict, strict=False)
|
||||
_load_pytorch_weights(model, model_id, allow_download=allow_download)
|
||||
|
||||
if device == "auto":
|
||||
if torch.cuda.is_available():
|
||||
|
|
|
|||
|
|
@ -539,3 +539,217 @@ class TestOnnxBackendPrefixGating:
|
|||
|
||||
with pytest.raises(AttributeError):
|
||||
compressor._timed_canary(model, pt_tokenizer, "pytorch")
|
||||
|
||||
|
||||
class TestPytorchWeightLoading:
|
||||
"""_load_pytorch_weights must load the merged v2 checkpoint format correctly,
|
||||
fall back to the plain format only when the repo genuinely has no merged.pt,
|
||||
and refuse to run on a state-dict mismatch instead of silently ignoring it.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _make_model(torch):
|
||||
import torch.nn as nn
|
||||
|
||||
model = nn.Module()
|
||||
model.encoder = nn.Linear(4, 4)
|
||||
model.token_head = nn.Linear(4, 2)
|
||||
model.span_conv = nn.Sequential(nn.Conv1d(4, 4, 1), nn.GELU())
|
||||
return model
|
||||
|
||||
def test_merged_checkpoint_loads_into_matching_submodules(self, tmp_path, monkeypatch) -> None:
|
||||
import pytest
|
||||
|
||||
torch = pytest.importorskip("torch")
|
||||
import headroom.transforms.kompress_compressor as kmod
|
||||
|
||||
model = self._make_model(torch)
|
||||
ckpt_path = tmp_path / "merged.pt"
|
||||
torch.save(
|
||||
{
|
||||
"encoder_state_dict": model.encoder.state_dict(),
|
||||
"token_head_state_dict": model.token_head.state_dict(),
|
||||
"span_conv_state_dict": model.span_conv.state_dict(),
|
||||
},
|
||||
ckpt_path,
|
||||
)
|
||||
|
||||
fresh_model = self._make_model(torch)
|
||||
monkeypatch.setattr(kmod, "hf_hub_download_local_first", lambda *a, **k: str(ckpt_path))
|
||||
|
||||
kmod._load_pytorch_weights(fresh_model, "some/repo", allow_download=True)
|
||||
|
||||
for name, param in model.encoder.state_dict().items():
|
||||
assert torch.equal(param, fresh_model.encoder.state_dict()[name])
|
||||
|
||||
def test_merged_checkpoint_missing_section_raises(self, tmp_path, monkeypatch) -> None:
|
||||
import pytest
|
||||
|
||||
torch = pytest.importorskip("torch")
|
||||
import headroom.transforms.kompress_compressor as kmod
|
||||
|
||||
ckpt_path = tmp_path / "merged.pt"
|
||||
torch.save({"encoder_state_dict": {}}, ckpt_path)
|
||||
|
||||
fresh_model = self._make_model(torch)
|
||||
monkeypatch.setattr(kmod, "hf_hub_download_local_first", lambda *a, **k: str(ckpt_path))
|
||||
|
||||
with pytest.raises(RuntimeError, match="missing"):
|
||||
kmod._load_pytorch_weights(fresh_model, "some/repo", allow_download=True)
|
||||
|
||||
def test_merged_checkpoint_key_mismatch_raises_instead_of_silently_dropping(
|
||||
self, tmp_path, monkeypatch
|
||||
) -> None:
|
||||
"""Regression test for the bug this loader used to have: loading a
|
||||
state-dict that does not match the module tree (e.g. an unmerged PEFT
|
||||
checkpoint) must fail loudly, not silently skip the mismatched keys.
|
||||
"""
|
||||
import pytest
|
||||
|
||||
torch = pytest.importorskip("torch")
|
||||
import headroom.transforms.kompress_compressor as kmod
|
||||
|
||||
ckpt_path = tmp_path / "merged.pt"
|
||||
torch.save(
|
||||
{
|
||||
# Wrong prefix, mimics the unmerged PEFT structure documented
|
||||
# in scripts/export_kompress_v2_onnx.py.
|
||||
"encoder_state_dict": {"base_model.model.weight": torch.zeros(4, 4)},
|
||||
"token_head_state_dict": {},
|
||||
"span_conv_state_dict": {},
|
||||
},
|
||||
ckpt_path,
|
||||
)
|
||||
|
||||
fresh_model = self._make_model(torch)
|
||||
monkeypatch.setattr(kmod, "hf_hub_download_local_first", lambda *a, **k: str(ckpt_path))
|
||||
|
||||
with pytest.raises(RuntimeError, match="state_dict mismatch"):
|
||||
kmod._load_pytorch_weights(fresh_model, "some/repo", allow_download=True)
|
||||
|
||||
def test_missing_merged_pt_falls_back_to_plain_safetensors(self, tmp_path, monkeypatch) -> None:
|
||||
import pytest
|
||||
|
||||
torch = pytest.importorskip("torch")
|
||||
safetensors_torch = pytest.importorskip("safetensors.torch")
|
||||
import headroom.transforms.kompress_compressor as kmod
|
||||
|
||||
model = self._make_model(torch)
|
||||
weights_path = tmp_path / "model.safetensors"
|
||||
safetensors_torch.save_file(dict(model.state_dict()), str(weights_path))
|
||||
|
||||
fresh_model = self._make_model(torch)
|
||||
|
||||
def fake_download(model_id, filename, *, allow_network=True, **kwargs): # noqa: ANN001
|
||||
if filename == "merged.pt":
|
||||
raise kmod.EntryNotFoundError("no merged.pt in this repo")
|
||||
assert filename == "model.safetensors"
|
||||
return str(weights_path)
|
||||
|
||||
monkeypatch.setattr(kmod, "hf_hub_download_local_first", fake_download)
|
||||
|
||||
kmod._load_pytorch_weights(fresh_model, "some/v1/repo", allow_download=True)
|
||||
|
||||
for name, param in model.state_dict().items():
|
||||
assert torch.equal(param, fresh_model.state_dict()[name])
|
||||
|
||||
def test_plain_safetensors_key_mismatch_raises(self, tmp_path, monkeypatch) -> None:
|
||||
import pytest
|
||||
|
||||
torch = pytest.importorskip("torch")
|
||||
safetensors_torch = pytest.importorskip("safetensors.torch")
|
||||
import headroom.transforms.kompress_compressor as kmod
|
||||
|
||||
weights_path = tmp_path / "model.safetensors"
|
||||
safetensors_torch.save_file({"totally.unrelated.key": torch.zeros(2)}, str(weights_path))
|
||||
|
||||
fresh_model = self._make_model(torch)
|
||||
|
||||
def fake_download(model_id, filename, *, allow_network=True, **kwargs): # noqa: ANN001
|
||||
if filename == "merged.pt":
|
||||
raise kmod.EntryNotFoundError("no merged.pt in this repo")
|
||||
return str(weights_path)
|
||||
|
||||
monkeypatch.setattr(kmod, "hf_hub_download_local_first", fake_download)
|
||||
|
||||
with pytest.raises(RuntimeError, match="state_dict mismatch"):
|
||||
kmod._load_pytorch_weights(fresh_model, "some/v1/repo", allow_download=True)
|
||||
|
||||
def test_cache_only_miss_raises_kompress_model_not_cached(self, monkeypatch) -> None:
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("torch")
|
||||
import headroom.transforms.kompress_compressor as kmod
|
||||
|
||||
def fake_download(model_id, filename, *, allow_network=True, **kwargs): # noqa: ANN001
|
||||
raise kmod.LocalEntryNotFoundError("not cached")
|
||||
|
||||
monkeypatch.setattr(kmod, "hf_hub_download_local_first", fake_download)
|
||||
monkeypatch.setattr(kmod, "hf_entry_known_absent", lambda *a, **k: False)
|
||||
|
||||
with pytest.raises(kmod.KompressModelNotCached):
|
||||
kmod._load_pytorch_weights(SimpleNamespace(), "some/repo", allow_download=False)
|
||||
|
||||
def test_cache_only_defers_instead_of_using_stale_plain_checkpoint(self, monkeypatch) -> None:
|
||||
"""Regression test: if merged.pt is not cached yet and we have no
|
||||
confirmation it is genuinely absent upstream, a stale model.safetensors
|
||||
left over from a previous (pre-fix) run must NOT be used as a silent
|
||||
fallback - that would reintroduce the original bug for exactly the
|
||||
upgrade scenario that motivated this fix. It must defer instead.
|
||||
"""
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("torch")
|
||||
import headroom.transforms.kompress_compressor as kmod
|
||||
|
||||
plain_download_calls: list[str] = []
|
||||
|
||||
def fake_download(model_id, filename, *, allow_network=True, **kwargs): # noqa: ANN001
|
||||
if filename == "merged.pt":
|
||||
raise kmod.LocalEntryNotFoundError("merged.pt not cached yet")
|
||||
plain_download_calls.append(filename)
|
||||
return "/fake/cached/model.safetensors"
|
||||
|
||||
monkeypatch.setattr(kmod, "hf_hub_download_local_first", fake_download)
|
||||
# Nothing has ever confirmed merged.pt is absent from this repo -
|
||||
# simulates a v2-style repo mid-upgrade, not a v1-style repo.
|
||||
monkeypatch.setattr(kmod, "hf_entry_known_absent", lambda *a, **k: False)
|
||||
|
||||
with pytest.raises(kmod.KompressModelNotCached):
|
||||
kmod._load_pytorch_weights(
|
||||
SimpleNamespace(), "chopratejas/kompress-v2-base", allow_download=False
|
||||
)
|
||||
|
||||
assert plain_download_calls == [], (
|
||||
"must not fall back to model.safetensors without confirming merged.pt is absent"
|
||||
)
|
||||
|
||||
def test_cache_only_uses_plain_checkpoint_when_merged_pt_confirmed_absent(
|
||||
self, tmp_path, monkeypatch
|
||||
) -> None:
|
||||
import pytest
|
||||
|
||||
torch = pytest.importorskip("torch")
|
||||
safetensors_torch = pytest.importorskip("safetensors.torch")
|
||||
import headroom.transforms.kompress_compressor as kmod
|
||||
|
||||
model = self._make_model(torch)
|
||||
weights_path = tmp_path / "model.safetensors"
|
||||
safetensors_torch.save_file(dict(model.state_dict()), str(weights_path))
|
||||
|
||||
fresh_model = self._make_model(torch)
|
||||
|
||||
def fake_download(model_id, filename, *, allow_network=True, **kwargs): # noqa: ANN001
|
||||
if filename == "merged.pt":
|
||||
raise kmod.LocalEntryNotFoundError("merged.pt not cached")
|
||||
return str(weights_path)
|
||||
|
||||
monkeypatch.setattr(kmod, "hf_hub_download_local_first", fake_download)
|
||||
# A prior real network lookup already confirmed this repo has no
|
||||
# merged.pt at all (the v1-style case), so the plain fallback is safe.
|
||||
monkeypatch.setattr(kmod, "hf_entry_known_absent", lambda *a, **k: True)
|
||||
|
||||
kmod._load_pytorch_weights(fresh_model, "chopratejas/kompress-base", allow_download=False)
|
||||
|
||||
for name, param in model.state_dict().items():
|
||||
assert torch.equal(param, fresh_model.state_dict()[name])
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue