[Platform] Enable layerwise offload on all hardware (#1492)
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
This commit is contained in:
1 parent
43a9e48ff7
commit
3b0c21f94d
3 files changed
+10
-12
No files matched your search
@@ -7,7 +7,6 @@ from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
|
||||
|
||||
from tests.utils import GPUMemoryMonitor
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
# ruff: noqa: E402
|
||||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
@@ -68,7 +67,6 @@ def run_inference(
|
||||
return peak
|
||||
|
||||
|
||||
@pytest.mark.skipif(current_omni_platform.is_npu() or current_omni_platform.is_rocm(), reason="Hardware not supported")
|
||||
@pytest.mark.parametrize("model_name", MODELS_SAVED_MEMORY_MB.keys())
|
||||
def test_layerwise_offload_diffusion_model(model_name: str):
|
||||
"""Test that layerwise offloading reduces GPU memory usage.
|
||||
|
||||
Reference in new issue
Block a user