diff --git a/.buildkite/scripts/unit_test.sh b/.buildkite/scripts/unit_test.sh index 1d5c4475fa..8f6ac456f8 100644 --- a/.buildkite/scripts/unit_test.sh +++ b/.buildkite/scripts/unit_test.sh @@ -7,6 +7,7 @@ exec pytest \ ./fastvideo/tests/dataset/ \ ./fastvideo/tests/workflow/ \ ./fastvideo/tests/entrypoints/ \ + ./fastvideo/tests/loader/ \ ./fastvideo/tests/train/ \ ./fastvideo/tests/stages/ \ ./fastvideo/tests/ops/ \ diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 5bb9e0fce5..aaca40c71f 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -25,6 +25,10 @@ logger = init_logger(__name__) +# Offload flags that trade device memory for host memory. Keeping the policy +# centralized lets components share the same worker-local device decision. +UNIFIED_MEMORY_OFFLOAD_FLAGS = ("text_encoder_cpu_offload", ) + class ExecutionMode(str, Enum): """ @@ -914,6 +918,39 @@ def check_fastvideo_args(self) -> None: self.pipeline_config.vae_config.load_encoder = True self.preprocess_config.check_preprocess_config() + def disable_offload_on_unified_memory(self, device_id: int = 0, *, offload_flag: str | None = None) -> bool: + """Disable host offload after a worker has selected its device. + + CUDA's unified-memory probe reads runtime device properties and may + initialize a CUDA context. Callers must therefore use this only inside + a device-owning process, after selecting and binding ``device_id``. + Returning the classification lets direct component-loader callers + apply the same policy to explicit per-call overrides. When + ``offload_flag`` is given, the return value says whether this policy + covers that component role. + """ + from fastvideo.platforms import current_platform + + if not current_platform.has_unified_memory(device_id): + return False + + try: + device_name = current_platform.get_device_name(device_id) + except Exception: + # Device naming is diagnostic only. NVML can be unavailable on an + # integrated GPU (for example Jetson), and its physical-ordinal + # lookup cannot interpret CUDA_VISIBLE_DEVICES UUID/MIG selectors. + # Neither case should undo an authoritative driver classification. + device_name = current_platform.device_name + + for flag in UNIFIED_MEMORY_OFFLOAD_FLAGS: + if getattr(self, flag): + logger.info( + "Disabling %s: %s has unified memory, so moving weights to the host duplicates " + "them rather than freeing device memory.", flag, device_name) + setattr(self, flag, False) + return offload_flag is None or offload_flag in UNIFIED_MEMORY_OFFLOAD_FLAGS + _current_fastvideo_args = None diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 6d6aadc79b..6a9e9b29c1 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -347,11 +347,24 @@ def load_model( dtype: str = "fp16", use_text_encoder_override: bool = False, # prevent subclasses from misusing cpu_offload: bool | None = None, + offload_flag: str = "text_encoder_cpu_offload", ): - if cpu_offload is None: - cpu_offload = fastvideo_args.text_encoder_cpu_offload - use_cpu_offload = (cpu_offload and len(getattr(model_config, "_fsdp_shard_conditions", [])) > 0) runtime_device = get_local_torch_device() + device_id = runtime_device.index if runtime_device.index is not None else 0 + requested_cpu_offload = getattr(fastvideo_args, offload_flag) if cpu_offload is None else cpu_offload + disable_cpu_offload = fastvideo_args.disable_offload_on_unified_memory(device_id, + offload_flag=offload_flag) + + if requested_cpu_offload and disable_cpu_offload: + # Direct loader callers can choose a CPU target before the worker + # applies its device-local policy. Reset both the request and the + # target so the model is never constructed on the host first. + logger.info("Disabling %s on unified-memory device %d", offload_flag, device_id) + cpu_offload = False + target_device = runtime_device + else: + cpu_offload = requested_cpu_offload + use_cpu_offload = (cpu_offload and len(getattr(model_config, "_fsdp_shard_conditions", [])) > 0) from fastvideo.platforms import current_platform @@ -555,6 +568,7 @@ def load(self, model_path: str, fastvideo_args: FastVideoArgs): fastvideo_args, encoder_precision, cpu_offload=fastvideo_args.image_encoder_cpu_offload, + offload_flag="image_encoder_cpu_offload", ) diff --git a/fastvideo/platforms/cuda.py b/fastvideo/platforms/cuda.py index 08c35afe5b..b6f8a9cb3f 100644 --- a/fastvideo/platforms/cuda.py +++ b/fastvideo/platforms/cuda.py @@ -4,6 +4,7 @@ pynvml. However, it should not initialize cuda context. """ +import ctypes import os from collections.abc import Callable from functools import lru_cache, wraps @@ -24,6 +25,56 @@ pynvml = import_pynvml() # type: ignore[no-untyped-call] +_CUDA_SUCCESS = 0 +# Stable value from CUDA's public CUdevice_attribute enum. +_CU_DEVICE_ATTRIBUTE_INTEGRATED = 18 +_CUDA_DRIVER_LIBRARY = "nvcuda.dll" if os.name == "nt" else "libcuda.so.1" + + +def _cuda_driver_device_is_integrated(device_id: int) -> bool: + """Query ``CU_DEVICE_ATTRIBUTE_INTEGRATED`` without creating a context. + + ``cuInit`` loads the CUDA driver, while these device-management calls only + inspect the logical device ordinal. They neither create nor retain a CUDA + context. Driver initialization is not pre-fork safe, so this probe must + remain worker-local after process creation and device binding. Using the + driver ordinal preserves CUDA_VISIBLE_DEVICES ordering, including UUID and + MIG selectors. + """ + try: + driver = ctypes.CDLL(_CUDA_DRIVER_LIBRARY) + driver.cuInit.argtypes = [ctypes.c_uint] + driver.cuInit.restype = ctypes.c_int + driver.cuDeviceGet.argtypes = [ctypes.POINTER(ctypes.c_int), ctypes.c_int] + driver.cuDeviceGet.restype = ctypes.c_int + driver.cuDeviceGetAttribute.argtypes = [ + ctypes.POINTER(ctypes.c_int), + ctypes.c_int, + ctypes.c_int, + ] + driver.cuDeviceGetAttribute.restype = ctypes.c_int + + if driver.cuInit(0) != _CUDA_SUCCESS: + return False + + device = ctypes.c_int() + if driver.cuDeviceGet(ctypes.byref(device), device_id) != _CUDA_SUCCESS: + return False + + is_integrated = ctypes.c_int() + if driver.cuDeviceGetAttribute( + ctypes.byref(is_integrated), + _CU_DEVICE_ATTRIBUTE_INTEGRATED, + device, + ) != _CUDA_SUCCESS: + return False + return bool(is_integrated.value) + except Exception: + # Missing/incompatible driver libraries and unavailable devices must + # preserve the established discrete-memory offload policy. + return False + + # pytorch 2.5 uses cudnn sdpa by default, which will cause crash on some models # see https://github.com/huggingface/diffusers/issues/9704 for details torch.backends.cuda.enable_cudnn_sdp(False) @@ -79,6 +130,12 @@ def get_device_name(cls, device_id: int = 0) -> str: def get_device_total_memory(cls, device_id: int = 0) -> int: raise NotImplementedError + @classmethod + def has_unified_memory(cls, device_id: int = 0) -> bool: + # This is cudaDeviceProp::integrated's driver-level source of truth. It + # is true on parts such as GB10 and Jetson whose GPU reads host memory. + return _cuda_driver_device_is_integrated(device_id) + @classmethod def is_async_output_supported(cls, enforce_eager: bool | None) -> bool: if enforce_eager: diff --git a/fastvideo/platforms/interface.py b/fastvideo/platforms/interface.py index f97cd6af55..ea55f77803 100644 --- a/fastvideo/platforms/interface.py +++ b/fastvideo/platforms/interface.py @@ -118,6 +118,20 @@ def is_mps(self) -> bool: def is_npu(self) -> bool: return self._enum == PlatformEnum.NPU + @classmethod + def has_unified_memory(cls, device_id: int = 0) -> bool: + """Whether host and device allocations come out of one physical pool. + + Where this is true, moving a tensor between host and device frees + nothing: both ends are the same RAM. Anything that offloads to save + memory needs to know, because on such a device the copy is at best a + no-op and at worst holds two copies at once. + + Implementations may need runtime device properties. Call this only + after the current worker has selected and initialized ``device_id``. + """ + return False + @classmethod def get_attn_backend_cls(cls, selected_backend: AttentionBackendEnum | None, head_size: int, dtype: torch.dtype) -> str: diff --git a/fastvideo/platforms/mps.py b/fastvideo/platforms/mps.py index 461360b80e..6f8df1ff1a 100644 --- a/fastvideo/platforms/mps.py +++ b/fastvideo/platforms/mps.py @@ -16,6 +16,11 @@ class MpsPlatform(Platform): dispatch_key: str = "MPS" device_control_env_var: str = "MPS_VISIBLE_DEVICES" + @classmethod + def has_unified_memory(cls, device_id: int = 0) -> bool: + # Apple silicon shares one pool between CPU and GPU. + return True + @classmethod def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None: raise NotImplementedError diff --git a/fastvideo/tests/loader/test_text_encoder_unified_memory_offload.py b/fastvideo/tests/loader/test_text_encoder_unified_memory_offload.py new file mode 100644 index 0000000000..73b9707d1d --- /dev/null +++ b/fastvideo/tests/loader/test_text_encoder_unified_memory_offload.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regression tests for text-encoder placement on unified memory.""" +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +import torch.nn as nn + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.models.loader.component_loader import ImageEncoderLoader, TextEncoderLoader + + +class _PassthroughEncoder(nn.Module): + supports_hf_from_pretrained = True + loaded_device: torch.device | None = None + + @classmethod + def from_pretrained_local(cls, model_path, model_config, *, dtype, device): + del model_path, model_config, dtype + cls.loaded_device = torch.device(device) + return cls() + + +def _model_config(): + return SimpleNamespace(architectures=["PassthroughEncoder"], _fsdp_shard_conditions=[], quant_config=None) + + +@pytest.mark.parametrize( + ("cpu_offload", "requested_target"), + [ + (None, torch.device("cuda:5")), + (None, torch.device("cpu")), + (True, torch.device("cpu")), + ], +) +def test_unified_memory_uses_worker_device_before_model_construction(monkeypatch, tmp_path, cpu_offload, + requested_target) -> None: + probe = Mock(return_value=True) + monkeypatch.setattr("fastvideo.models.loader.component_loader.get_local_torch_device", + lambda: torch.device("cuda:5")) + monkeypatch.setattr("fastvideo.platforms.current_platform.has_unified_memory", probe) + monkeypatch.setattr("fastvideo.platforms.current_platform.get_device_name", lambda device_id: "NVIDIA GB10") + monkeypatch.setattr( + "fastvideo.models.loader.component_loader.ModelRegistry.resolve_model_cls", + lambda architectures: (_PassthroughEncoder, None), + ) + args = FastVideoArgs(model_path=str(tmp_path), text_encoder_cpu_offload=True) + + model = TextEncoderLoader().load_model( + str(tmp_path), + _model_config(), + requested_target, + args, + cpu_offload=cpu_offload, + ) + + assert isinstance(model, _PassthroughEncoder) + assert _PassthroughEncoder.loaded_device == torch.device("cuda:5") + assert args.text_encoder_cpu_offload is False + probe.assert_called_once_with(5) + + +def test_discrete_memory_preserves_explicit_cpu_target(monkeypatch, tmp_path) -> None: + monkeypatch.setattr("fastvideo.models.loader.component_loader.get_local_torch_device", + lambda: torch.device("cuda:2")) + monkeypatch.setattr("fastvideo.platforms.current_platform.has_unified_memory", lambda device_id: False) + monkeypatch.setattr( + "fastvideo.models.loader.component_loader.ModelRegistry.resolve_model_cls", + lambda architectures: (_PassthroughEncoder, None), + ) + args = FastVideoArgs(model_path=str(tmp_path), text_encoder_cpu_offload=False) + + TextEncoderLoader().load_model( + str(tmp_path), + _model_config(), + torch.device("cuda:2"), + args, + cpu_offload=True, + ) + + assert _PassthroughEncoder.loaded_device == torch.device("cpu") + + +def test_text_policy_does_not_change_inherited_image_encoder_path(monkeypatch, tmp_path) -> None: + probe = Mock(return_value=True) + monkeypatch.setattr("fastvideo.models.loader.component_loader.get_local_torch_device", + lambda: torch.device("cuda:4")) + monkeypatch.setattr("fastvideo.platforms.current_platform.has_unified_memory", probe) + monkeypatch.setattr("fastvideo.platforms.current_platform.get_device_name", lambda device_id: "NVIDIA GB10") + monkeypatch.setattr( + "fastvideo.models.loader.component_loader.ModelRegistry.resolve_model_cls", + lambda architectures: (_PassthroughEncoder, None), + ) + args = FastVideoArgs( + model_path=str(tmp_path), + text_encoder_cpu_offload=True, + image_encoder_cpu_offload=True, + ) + + ImageEncoderLoader().load_model( + str(tmp_path), + _model_config(), + torch.device("cpu"), + args, + cpu_offload=True, + offload_flag="image_encoder_cpu_offload", + ) + + assert _PassthroughEncoder.loaded_device == torch.device("cpu") + assert args.text_encoder_cpu_offload is False + assert args.image_encoder_cpu_offload is True + probe.assert_called_once_with(4) diff --git a/fastvideo/tests/platforms/test_unified_memory.py b/fastvideo/tests/platforms/test_unified_memory.py new file mode 100644 index 0000000000..f1bf7a1e3a --- /dev/null +++ b/fastvideo/tests/platforms/test_unified_memory.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU tests for device-local unified-memory classification.""" +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from fastvideo.platforms.cuda import (CudaPlatformBase, + _CU_DEVICE_ATTRIBUTE_INTEGRATED) +from fastvideo.platforms.interface import Platform +from fastvideo.platforms.mps import MpsPlatform + + +def test_base_platform_reports_separate_pools() -> None: + # Discrete accelerators are the default, so the base must answer False and + # leave the existing offload path alone. + assert Platform.has_unified_memory(device_id=3) is False + + +def test_mps_is_unified() -> None: + assert MpsPlatform.has_unified_memory(device_id=2) is True + + +def _fake_cuda_driver(integrated: bool = True): + driver = SimpleNamespace( + cuInit=Mock(return_value=0), + cuDeviceGet=Mock(), + cuDeviceGetAttribute=Mock(), + ) + + def get_device(device, device_id): + device._obj.value = device_id + 40 + return 0 + + def get_attribute(value, attribute, device): + value._obj.value = int(integrated) + return 0 + + driver.cuDeviceGet.side_effect = get_device + driver.cuDeviceGetAttribute.side_effect = get_attribute + return driver + + +@pytest.mark.parametrize("integrated", [True, False]) +def test_cuda_follows_driver_integrated_attribute(monkeypatch, integrated: bool) -> None: + # The driver attribute is cudaDeviceProp::integrated's source of truth and + # does not initialize a PyTorch CUDA context. + driver = _fake_cuda_driver(integrated) + get_device_properties = Mock(side_effect=AssertionError("torch fallback used")) + monkeypatch.setattr("fastvideo.platforms.cuda.ctypes.CDLL", lambda library: driver) + monkeypatch.setattr("torch.cuda.get_device_properties", get_device_properties) + + assert CudaPlatformBase.has_unified_memory(device_id=7) is integrated + + driver.cuInit.assert_called_once_with(0) + assert driver.cuDeviceGet.call_args.args[1] == 7 + assert driver.cuDeviceGetAttribute.call_args.args[1] == _CU_DEVICE_ATTRIBUTE_INTEGRATED + assert driver.cuDeviceGetAttribute.call_args.args[2].value == 47 + get_device_properties.assert_not_called() + + +@pytest.mark.parametrize("failing_call", ["cuInit", "cuDeviceGet", "cuDeviceGetAttribute"]) +def test_cuda_driver_call_failure_assumes_separate_pools(monkeypatch, failing_call: str) -> None: + driver = _fake_cuda_driver() + getattr(driver, failing_call).side_effect = None + getattr(driver, failing_call).return_value = 1 + monkeypatch.setattr("fastvideo.platforms.cuda.ctypes.CDLL", lambda library: driver) + + assert CudaPlatformBase.has_unified_memory(device_id=4) is False + + +def test_cuda_without_driver_library_assumes_separate_pools(monkeypatch) -> None: + def unavailable(library): + raise OSError("CUDA driver unavailable") + + monkeypatch.setattr("fastvideo.platforms.cuda.ctypes.CDLL", unavailable) + + assert CudaPlatformBase.has_unified_memory(device_id=5) is False diff --git a/fastvideo/tests/platforms/test_unified_memory_offload.py b/fastvideo/tests/platforms/test_unified_memory_offload.py new file mode 100644 index 0000000000..fdf00dcab5 --- /dev/null +++ b/fastvideo/tests/platforms/test_unified_memory_offload.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU tests for applying unified-memory offload policy inside a worker.""" +from __future__ import annotations + +from unittest.mock import Mock + +import pytest + +from fastvideo.fastvideo_args import FastVideoArgs + + +def _args(**overrides) -> FastVideoArgs: + return FastVideoArgs(model_path="unused/for-this-test", **overrides) + + +def test_constructing_args_does_not_probe_runtime_device_properties(monkeypatch) -> None: + probe = Mock(side_effect=AssertionError("device probe ran in the driver")) + monkeypatch.setattr("fastvideo.platforms.current_platform.has_unified_memory", probe) + + args = _args(text_encoder_cpu_offload=True) + + assert args.text_encoder_cpu_offload is True + probe.assert_not_called() + + +def test_unified_device_disables_text_encoder_offload_for_selected_device(monkeypatch) -> None: + probe = Mock(return_value=True) + monkeypatch.setattr("fastvideo.platforms.current_platform.has_unified_memory", probe) + monkeypatch.setattr("fastvideo.platforms.current_platform.get_device_name", lambda device_id: "NVIDIA GB10") + args = _args(text_encoder_cpu_offload=True) + + assert args.disable_offload_on_unified_memory(device_id=6, offload_flag="text_encoder_cpu_offload") is True + + probe.assert_called_once_with(6) + assert args.text_encoder_cpu_offload is False + + +def test_discrete_device_keeps_text_encoder_offload(monkeypatch) -> None: + probe = Mock(return_value=False) + monkeypatch.setattr("fastvideo.platforms.current_platform.has_unified_memory", probe) + args = _args(text_encoder_cpu_offload=True) + + assert args.disable_offload_on_unified_memory(device_id=3) is False + + probe.assert_called_once_with(3) + assert args.text_encoder_cpu_offload is True + + +def test_policy_does_not_claim_unlisted_component_role(monkeypatch) -> None: + monkeypatch.setattr("fastvideo.platforms.current_platform.has_unified_memory", lambda device_id: True) + monkeypatch.setattr("fastvideo.platforms.current_platform.get_device_name", lambda device_id: "NVIDIA GB10") + args = _args(text_encoder_cpu_offload=True, image_encoder_cpu_offload=True) + + applies_to_image_encoder = args.disable_offload_on_unified_memory( + device_id=0, + offload_flag="image_encoder_cpu_offload", + ) + + assert applies_to_image_encoder is False + assert args.text_encoder_cpu_offload is False + assert args.image_encoder_cpu_offload is True + + +@pytest.mark.parametrize("name_error", [NotImplementedError, ValueError, RuntimeError]) +def test_platform_without_device_name_uses_generic_name(monkeypatch, name_error: type[Exception]) -> None: + monkeypatch.setattr("fastvideo.platforms.current_platform.has_unified_memory", lambda device_id: True) + + def unsupported_name(device_id): + raise name_error("device name unavailable") + + monkeypatch.setattr("fastvideo.platforms.current_platform.get_device_name", unsupported_name) + args = _args(text_encoder_cpu_offload=True) + + assert args.disable_offload_on_unified_memory() is True + assert args.text_encoder_cpu_offload is False diff --git a/fastvideo/tests/worker/test_gpu_worker.py b/fastvideo/tests/worker/test_gpu_worker.py index 77850595d3..f4436cec99 100644 --- a/fastvideo/tests/worker/test_gpu_worker.py +++ b/fastvideo/tests/worker/test_gpu_worker.py @@ -37,6 +37,36 @@ def test_cuda_device_uuid_receipt_identifies_profiled_worker(monkeypatch) -> Non ) +@pytest.mark.parametrize("executor_backend", ["mp", "ray"]) +def test_init_device_applies_offload_policy_after_binding_worker_device(monkeypatch, executor_backend: str) -> None: + """The runtime probe must see this worker's device, never driver device 0.""" + events = [] + args = FastVideoArgs(model_path="test", num_gpus=1, distributed_executor_backend=executor_backend) + args.disable_offload_on_unified_memory = Mock(side_effect=lambda device_id: events.append(("policy", device_id))) + worker = Worker(args, local_rank=3, rank=3, distributed_init_method="env://") + + monkeypatch.setenv("LOCAL_RANK", "0") + monkeypatch.setattr("fastvideo.platforms.current_platform.is_cuda_alike", lambda: True) + monkeypatch.setattr("fastvideo.platforms.current_platform.is_cuda", lambda: False) + monkeypatch.setattr(torch.cuda, "set_device", lambda device: events.append(("set_device", device.index))) + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda device: (123, 456)) + monkeypatch.setattr( + "fastvideo.worker.gpu_worker.maybe_init_distributed_environment_and_model_parallel", + lambda *args: events.append(("distributed", None)), + ) + monkeypatch.setattr("fastvideo.worker.gpu_worker.build_pipeline", lambda args: events.append(("pipeline", None))) + + worker.init_device() + + assert events == [ + ("set_device", 3), + ("policy", 3), + ("distributed", None), + ("pipeline", None), + ] + assert worker.init_gpu_memory == 123 + + def _worker_returning(output_batch: ForwardBatch) -> Worker: worker = Worker.__new__(Worker) worker.fastvideo_args = SimpleNamespace() diff --git a/fastvideo/worker/gpu_worker.py b/fastvideo/worker/gpu_worker.py index dcaac6e438..72bd030c55 100644 --- a/fastvideo/worker/gpu_worker.py +++ b/fastvideo/worker/gpu_worker.py @@ -56,8 +56,11 @@ def init_device(self) -> None: # Set environment variables BEFORE calling get_local_torch_device() # so that each worker uses the correct device - if self.fastvideo_args.distributed_executor_backend == "mp": - os.environ["LOCAL_RANK"] = str(self.local_rank) + # Both multiprocessing and Ray pass the worker-local rank explicitly. + # Ray deliberately excludes LOCAL_RANK from the copied driver + # environment and exposes all GPUs assigned to the node, so leaving an + # inherited or missing value here would bind every Ray actor to cuda:0. + os.environ["LOCAL_RANK"] = str(self.local_rank) os.environ["RANK"] = str(self.rank) os.environ["WORLD_SIZE"] = str(self.fastvideo_args.num_gpus) @@ -76,6 +79,13 @@ def init_device(self) -> None: # For MPS, we can't get memory info the same way self.init_gpu_memory = 0 + # CUDA's unified-memory classification reads runtime device + # properties, so make this decision only after this worker has bound + # its own device. The worker-local args object is what every loader and + # pipeline stage below will consume. + device_id = self.device.index if self.device.index is not None else 0 + self.fastvideo_args.disable_offload_on_unified_memory(device_id) + # Initialize the distributed environment. maybe_init_distributed_environment_and_model_parallel(self.fastvideo_args.tp_size, self.fastvideo_args.sp_size, self.distributed_init_method) diff --git a/tests/local_tests/encoders/test_lingbot_video_dense_text_encoder_parity.py b/tests/local_tests/encoders/test_lingbot_video_dense_text_encoder_parity.py index edb631af4c..1089f6cc3f 100644 --- a/tests/local_tests/encoders/test_lingbot_video_dense_text_encoder_parity.py +++ b/tests/local_tests/encoders/test_lingbot_video_dense_text_encoder_parity.py @@ -64,6 +64,7 @@ def _load_native(device: torch.device, checkpoint: Path) -> LingBotVideoQwen3VLT override_text_encoder_quant=None, override_text_encoder_safetensors=None, pin_cpu_memory=False, + disable_offload_on_unified_memory=lambda device_id, offload_flag=None: False, ) model = TextEncoderLoader().load_model( str(checkpoint / "text_encoder"), diff --git a/tests/local_tests/zimage/test_zimage_encoder_parity.py b/tests/local_tests/zimage/test_zimage_encoder_parity.py index dfc3be3d9c..caf04902f6 100644 --- a/tests/local_tests/zimage/test_zimage_encoder_parity.py +++ b/tests/local_tests/zimage/test_zimage_encoder_parity.py @@ -138,6 +138,7 @@ def _loader_args(cpu_offload: bool) -> SimpleNamespace: override_text_encoder_quant=None, override_text_encoder_safetensors=None, pin_cpu_memory=False, + disable_offload_on_unified_memory=lambda device_id, offload_flag=None: False, ) @@ -264,6 +265,7 @@ def _print_diag(label: str, ref: torch.Tensor, fv: torch.Tensor) -> torch.Tensor ) def test_qwen3_production_loader_avoids_device_context_and_honors_placement( monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, cpu_offload: bool, target_device: torch.device, expected_device: torch.device, @@ -315,17 +317,15 @@ def fail_causal_lm_from_pretrained(*args, **kwargs): ), ) monkeypatch.setattr(torch.cuda, "is_available", lambda: True) - monkeypatch.setattr( - "fastvideo.distributed.get_local_torch_device", - lambda: torch.device("cuda"), - ) + monkeypatch.setattr("fastvideo.models.loader.component_loader.get_local_torch_device", + lambda: torch.device("cuda")) config = Qwen3TextConfig() # The pinned checkpoint names the causal wrapper even though the official # Z-Image loader intentionally asks Transformers for the body-only model. config.arch_config.architectures = ["Qwen3ForCausalLM"] loaded = TextEncoderLoader().load_model( - "unused-by-mock", + str(tmp_path), config, target_device, _loader_args(cpu_offload=cpu_offload),