[perf]: MiniMax H3 on GB10 - skip text encoder CPU offload on unified memory (5m49s to 30ms) (#1710)

Author: KyleNeverGivesUp

.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/ \

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
 

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

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:

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:

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

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)