107 lines
2.8 KiB
Python
107 lines
2.8 KiB
Python
|
# SPDX-License-Identifier: Apache-2.0
|
||
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||
|
|
|
||
|
|
"""Shared model fixtures for diffusion offloader tests."""
|
||
|
|
|
||
|
|
from contextlib import contextmanager
|
||
|
|
|
||
|
|
import torch
|
||
|
|
from torch import nn
|
||
|
|
|
||
|
|
|
||
|
|
class DummyStream:
|
||
|
|
def wait_stream(self, _stream) -> None:
|
||
|
|
return None
|
||
|
|
|
||
|
|
def wait_event(self, _event) -> None:
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
class DummyEvent:
|
||
|
|
def record(self, _stream) -> None:
|
||
|
|
return None
|
||
|
|
|
||
|
|
def synchronize(self) -> None:
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
@contextmanager
|
||
|
|
def dummy_stream(_stream):
|
||
|
|
yield None
|
||
|
|
|
||
|
|
|
||
|
|
def patch_offload_runtime(monkeypatch, platform, *, synchronize: bool = False) -> None:
|
||
|
|
monkeypatch.setattr(platform, "Stream", DummyStream)
|
||
|
|
monkeypatch.setattr(platform, "Event", DummyEvent)
|
||
|
|
monkeypatch.setattr(platform, "current_stream", lambda: DummyStream())
|
||
|
|
monkeypatch.setattr(platform, "stream", dummy_stream)
|
||
|
|
if synchronize:
|
||
|
|
monkeypatch.setattr(platform, "synchronize", lambda: None)
|
||
|
|
monkeypatch.setattr(platform, "empty_cache", lambda: None)
|
||
|
|
|
||
|
|
|
||
|
|
class _DummyBlock(nn.Module):
|
||
|
|
def __init__(self):
|
||
|
|
super().__init__()
|
||
|
|
self.weight = nn.Parameter(torch.randn(10, 10))
|
||
|
|
|
||
|
|
|
||
|
|
class _SingleBlockModel(nn.Module):
|
||
|
|
_layerwise_offload_blocks_attrs = ["blocks"]
|
||
|
|
|
||
|
|
def __init__(self, num_blocks: int = 3):
|
||
|
|
super().__init__()
|
||
|
|
self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)])
|
||
|
|
|
||
|
|
|
||
|
|
class _StagedEncoder(nn.Module):
|
||
|
|
def __init__(self):
|
||
|
|
super().__init__()
|
||
|
|
self.vision = nn.Module()
|
||
|
|
self.vision.blocks = nn.ModuleList([_DummyBlock(), _DummyBlock()])
|
||
|
|
self.text_model = nn.Module()
|
||
|
|
self.text_model.layers = nn.ModuleList([_DummyBlock(), _DummyBlock()])
|
||
|
|
self.load_calls = 0
|
||
|
|
self.offload_calls = 0
|
||
|
|
self.to_calls = 0
|
||
|
|
|
||
|
|
def load_to_device(self):
|
||
|
|
self.load_calls += 1
|
||
|
|
|
||
|
|
def offload_to_cpu(self):
|
||
|
|
self.offload_calls += 1
|
||
|
|
for hook in getattr(self, "_omni_layerwise_hooks", []):
|
||
|
|
hook.offload_layer()
|
||
|
|
|
||
|
|
def to(self, *args, **kwargs):
|
||
|
|
self.to_calls += 1
|
||
|
|
return super().to(*args, **kwargs)
|
||
|
|
|
||
|
|
|
||
|
|
class _StagedVAE(nn.Module):
|
||
|
|
def __init__(self):
|
||
|
|
super().__init__()
|
||
|
|
self.proj = nn.Linear(2, 2)
|
||
|
|
self.offload_calls = 0
|
||
|
|
self.to_calls = 0
|
||
|
|
|
||
|
|
def load_to_device(self):
|
||
|
|
return None
|
||
|
|
|
||
|
|
def offload_to_cpu(self):
|
||
|
|
self.offload_calls += 1
|
||
|
|
return self.to("cpu")
|
||
|
|
|
||
|
|
def to(self, *args, **kwargs):
|
||
|
|
self.to_calls += 1
|
||
|
|
return super().to(*args, **kwargs)
|
||
|
|
|
||
|
|
|
||
|
|
class _PlainEncoder(nn.Module):
|
||
|
|
"""Standard encoder with no offload-specific lifecycle methods."""
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
super().__init__()
|
||
|
|
self.encoder = nn.Module()
|
||
|
|
self.encoder.block = nn.ModuleList([_DummyBlock(), _DummyBlock()])
|
||
|
|
self.final_norm = nn.Linear(2, 2)
|