From 2218fd091d3e4e1279fd2f7fb5668489b5ff60fc Mon Sep 17 00:00:00 2001 From: Kyle Date: Fri, 14 Aug 2026 21:03:41 -0700 Subject: [PATCH 1/4] [perf]: skip text encoder CPU offload on unified memory On integrated CUDA and MPS devices, host and accelerator allocations share one physical pool. Text-encoder CPU offload therefore frees no memory and can briefly retain both copies; on DGX Spark this made the 63 GB MiniMax H3 encoder spend 5m49s sharding before the process was killed. Apply the decision to the worker-local FastVideoArgs only after that worker binds its selected device. This avoids creating a CUDA context in the parent process and avoids classifying every worker from device 0. Clearing the text_encoder_cpu_offload flag keeps load placement, FSDP, generic text encoding, and MiniMax H3 conditioning consistent. Keep a role-aware loader fallback for direct text-encoder callers. ImageEncoderLoader shares this implementation, so its role remains unchanged here and can opt into the same policy separately. CUDA classification follows cudaDeviceProp integrated, while discrete and unknown platforms preserve existing behavior. --- fastvideo/fastvideo_args.py | 37 ++++++ fastvideo/models/loader/component_loader.py | 17 ++- fastvideo/platforms/cuda.py | 57 +++++++++ fastvideo/platforms/interface.py | 14 +++ fastvideo/platforms/mps.py | 5 + ...est_text_encoder_unified_memory_offload.py | 108 ++++++++++++++++++ .../tests/platforms/test_unified_memory.py | 80 +++++++++++++ .../platforms/test_unified_memory_offload.py | 75 ++++++++++++ fastvideo/tests/worker/test_gpu_worker.py | 29 +++++ fastvideo/worker/gpu_worker.py | 7 ++ 10 files changed, 427 insertions(+), 2 deletions(-) create mode 100644 fastvideo/tests/loader/test_text_encoder_unified_memory_offload.py create mode 100644 fastvideo/tests/platforms/test_unified_memory.py create mode 100644 fastvideo/tests/platforms/test_unified_memory_offload.py 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..a2cbc4a205 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -347,11 +347,23 @@ 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", ): + runtime_device = get_local_torch_device() + device_id = runtime_device.index if runtime_device.index is not None else 0 + disable_cpu_offload = fastvideo_args.disable_offload_on_unified_memory(device_id, + offload_flag=offload_flag) + if cpu_offload is None: - cpu_offload = fastvideo_args.text_encoder_cpu_offload + cpu_offload = getattr(fastvideo_args, offload_flag) + elif cpu_offload and disable_cpu_offload: + # Direct loader callers can override the args-level flag. Apply the + # same device-local policy here before choosing the load target so + # the model is never constructed on the host in the first place. + logger.info("Disabling explicit %s on unified-memory device %d", offload_flag, device_id) + cpu_offload = False + target_device = runtime_device use_cpu_offload = (cpu_offload and len(getattr(model_config, "_fsdp_shard_conditions", [])) > 0) - runtime_device = get_local_torch_device() from fastvideo.platforms import current_platform @@ -555,6 +567,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..ea886b517d --- /dev/null +++ b/fastvideo/tests/loader/test_text_encoder_unified_memory_offload.py @@ -0,0 +1,108 @@ +# 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", [None, True]) +def test_unified_memory_uses_worker_device_before_model_construction(monkeypatch, tmp_path, cpu_offload) -> 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) + requested_target = torch.device("cpu") if cpu_offload is True else torch.device("cuda:5") + + 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..acc821354b 100644 --- a/fastvideo/tests/worker/test_gpu_worker.py +++ b/fastvideo/tests/worker/test_gpu_worker.py @@ -37,6 +37,35 @@ def test_cuda_device_uuid_receipt_identifies_profiled_worker(monkeypatch) -> Non ) +def test_init_device_applies_offload_policy_after_binding_worker_device(monkeypatch) -> None: + """The runtime probe must see this worker's device, never driver device 0.""" + events = [] + args = FastVideoArgs(model_path="test", num_gpus=1) + 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.setattr("fastvideo.worker.gpu_worker.get_local_torch_device", lambda: torch.device("cuda:3")) + 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..cf518ffd65 100644 --- a/fastvideo/worker/gpu_worker.py +++ b/fastvideo/worker/gpu_worker.py @@ -76,6 +76,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) From ebd8a468c58dee6069d36e602db5083ee2784169 Mon Sep 17 00:00:00 2001 From: William Lin <8941107+SolitaryThinker@users.noreply.github.com> Date: Tue, 25 Aug 2026 22:10:22 -0700 Subject: [PATCH 2/4] [bugfix]: apply unified-memory policy before loader placement --- fastvideo/models/loader/component_loader.py | 15 ++++++++------- .../test_text_encoder_unified_memory_offload.py | 13 ++++++++++--- ...est_lingbot_video_dense_text_encoder_parity.py | 1 + .../zimage/test_zimage_encoder_parity.py | 10 +++++----- 4 files changed, 24 insertions(+), 15 deletions(-) diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index a2cbc4a205..6a9e9b29c1 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -351,18 +351,19 @@ def load_model( ): 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 cpu_offload is None: - cpu_offload = getattr(fastvideo_args, offload_flag) - elif cpu_offload and disable_cpu_offload: - # Direct loader callers can override the args-level flag. Apply the - # same device-local policy here before choosing the load target so - # the model is never constructed on the host in the first place. - logger.info("Disabling explicit %s on unified-memory device %d", offload_flag, device_id) + 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 diff --git a/fastvideo/tests/loader/test_text_encoder_unified_memory_offload.py b/fastvideo/tests/loader/test_text_encoder_unified_memory_offload.py index ea886b517d..73b9707d1d 100644 --- a/fastvideo/tests/loader/test_text_encoder_unified_memory_offload.py +++ b/fastvideo/tests/loader/test_text_encoder_unified_memory_offload.py @@ -28,8 +28,16 @@ def _model_config(): return SimpleNamespace(architectures=["PassthroughEncoder"], _fsdp_shard_conditions=[], quant_config=None) -@pytest.mark.parametrize("cpu_offload", [None, True]) -def test_unified_memory_uses_worker_device_before_model_construction(monkeypatch, tmp_path, cpu_offload) -> 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")) @@ -40,7 +48,6 @@ def test_unified_memory_uses_worker_device_before_model_construction(monkeypatch lambda architectures: (_PassthroughEncoder, None), ) args = FastVideoArgs(model_path=str(tmp_path), text_encoder_cpu_offload=True) - requested_target = torch.device("cpu") if cpu_offload is True else torch.device("cuda:5") model = TextEncoderLoader().load_model( str(tmp_path), 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), From 7d7f7eff5570c6d18f5a7768080d32931340e52b Mon Sep 17 00:00:00 2001 From: William Lin <8941107+SolitaryThinker@users.noreply.github.com> Date: Tue, 25 Aug 2026 22:21:20 -0700 Subject: [PATCH 3/4] [ci]: collect loader tests in the unit lane --- .buildkite/scripts/unit_test.sh | 1 + 1 file changed, 1 insertion(+) 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/ \ From 450aad5ae0cff4de953f2633c0aa5f4bb5a1821b Mon Sep 17 00:00:00 2001 From: William Lin <8941107+SolitaryThinker@users.noreply.github.com> Date: Tue, 25 Aug 2026 22:26:06 -0700 Subject: [PATCH 4/4] [bugfix]: bind Ray workers to their local device --- fastvideo/tests/worker/test_gpu_worker.py | 7 ++++--- fastvideo/worker/gpu_worker.py | 7 +++++-- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/fastvideo/tests/worker/test_gpu_worker.py b/fastvideo/tests/worker/test_gpu_worker.py index acc821354b..f4436cec99 100644 --- a/fastvideo/tests/worker/test_gpu_worker.py +++ b/fastvideo/tests/worker/test_gpu_worker.py @@ -37,14 +37,15 @@ def test_cuda_device_uuid_receipt_identifies_profiled_worker(monkeypatch) -> Non ) -def test_init_device_applies_offload_policy_after_binding_worker_device(monkeypatch) -> None: +@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) + 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.setattr("fastvideo.worker.gpu_worker.get_local_torch_device", lambda: torch.device("cuda:3")) + 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))) diff --git a/fastvideo/worker/gpu_worker.py b/fastvideo/worker/gpu_worker.py index cf518ffd65..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)