| """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) |
|
|
