File size: 3,477 Bytes
2ea2835 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | """Regression coverage for the missing disabled-DINO offload entry."""
from types import SimpleNamespace
from unittest.mock import patch
import pytest
import torch
from model.wan_kv_mmd import WanKVMMD
def _model(names, device="cpu"):
model = object.__new__(WanKVMMD)
torch.nn.Module.__init__(model)
model.device = torch.device(device)
model.generator_replay_frozen_encoder_cpu_offload = True
model._generator_replay_frozen_encoder_offload_state = None
model.dual_branch_registry = {
name: {"encoder_attribute": f"{name}_encoder"} for name in names}
model.vae = torch.nn.Linear(3, 3).to(device).eval().requires_grad_(False)
for name in names:
setattr(model, f"{name}_encoder", torch.nn.Linear(3, 3).to(device).eval().requires_grad_(False))
return model
@pytest.mark.parametrize("names", [
("vjepa2", "videomae_v2"), ("dinov3", "vjepa2", "videomae_v2")])
def test_active_offload_restore_repeats_without_changing_weights(names):
model = _model(names)
modules = ("vae",) + tuple(f"{name}_encoder" for name in names)
before = {name: {k: v.clone() for k, v in getattr(model, name).state_dict().items()} for name in modules}
with patch.object(torch.cuda, "synchronize"), patch.object(torch.cuda, "empty_cache"):
for _ in range(2):
moved = model.offload_frozen_modules_for_generator_replay()
assert moved > 0
assert model._generator_replay_frozen_encoder_offload_state["module_names"] == modules
with pytest.raises(RuntimeError, match="already offloaded"):
model.offload_frozen_modules_for_generator_replay()
assert model.restore_frozen_modules_before_feature_replay() == moved
assert model._generator_replay_frozen_encoder_offload_state is None
for name in modules:
module = getattr(model, name)
assert not module.training and not any(p.requires_grad for p in module.parameters())
for key, value in module.state_dict().items():
assert torch.equal(value, before[name][key])
if "dinov3" not in names:
assert not hasattr(model, "dinov3_encoder")
def test_missing_active_encoder_still_fails_closed_and_rolls_back():
model = _model(("vjepa2", "videomae_v2"))
del model.vjepa2_encoder
with patch.object(torch.cuda, "synchronize"), patch.object(torch.cuda, "empty_cache"):
with pytest.raises(RuntimeError, match="vjepa2_encoder is unavailable"):
model.offload_frozen_modules_for_generator_replay()
assert model._generator_replay_frozen_encoder_offload_state is None
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA GPU")
def test_cuda_two_step_offload_restore_is_bitwise_exact():
model = _model(("vjepa2", "videomae_v2"), "cuda:0")
modules = (model.vae, model.vjepa2_encoder, model.videomae_v2_encoder)
inputs = torch.randn(4, 3, device="cuda:0")
expected = [module(inputs).detach() for module in modules]
for _ in range(2):
model.offload_frozen_modules_for_generator_replay()
assert all(next(module.parameters()).device.type == "cpu" for module in modules)
model.restore_frozen_modules_before_feature_replay()
assert all(next(module.parameters()).device.type == "cuda" for module in modules)
for module, baseline in zip(modules, expected):
assert torch.equal(module(inputs), baseline)
|