[perf] Disable every offload path on unified memory, unblocking MiniMax H3 generation on one GB10

.buildkite/scripts/unit_test.sh

@@ -8,6 +8,8 @@ exec pytest \
   ./fastvideo/tests/workflow/ \
   ./fastvideo/tests/entrypoints/ \
   ./fastvideo/tests/loader/ \
+  ./fastvideo/tests/pipelines/ \
+  ./fastvideo/tests/platforms/ \
   ./fastvideo/tests/train/ \
   ./fastvideo/tests/stages/ \
   ./fastvideo/tests/ops/ \

docs/getting_started/installation/spark_performance.md

@@ -156,8 +156,11 @@ is power-cycled. To avoid it:
 
 - **Builds** (flash-attn, kernel): `nice -n 19`, `MAX_JOBS=2`, `nohup`. Never a
   bare foreground high-parallelism build.
-- Leave `*_cpu_offload` at the example defaults — "CPU" offload is the *same*
-  unified RAM on the GB10, so the win is tiling + sane resolution, not offloading.
+- FastVideo automatically disables DiT layerwise/CPU offload and encoder/VAE CPU
+  offload after each worker binds its GB10 device. Do not force those modes back
+  on: "CPU" offload uses the same unified RAM. Multi-GPU FSDP sharding remains
+  available because it partitions weights without parking them in a separate
+  host pool.
 
 ## Gotchas specific to the GB10
 

docs/inference/offloading.md

@@ -14,6 +14,13 @@ vae_cpu_offload: bool = True
 pin_cpu_memory: bool = True
 ```
 
+On unified-memory accelerators such as NVIDIA GB10 and Apple silicon, FastVideo
+detects the selected device inside each worker and disables all five host-offload
+modes before loading modules. Host and accelerator allocations share one physical
+pool there, so offload adds transfers and duplicate residency instead of freeing
+memory. CUDA FSDP sharding remains enabled when requested; MPS continues to
+disable FSDP. `pin_cpu_memory` is not an offload mode and is left unchanged.
+
 ## Behavior Explanation
 
 !!! note

fastvideo/fastvideo_args.py

@@ -25,9 +25,16 @@
 
 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", )
+# Offload flags that trade device memory for host memory. All of them are a loss
+# on a device where the two are the same physical pool. Keeping the policy
+# centralized lets every loader and stage share one worker-local decision.
+UNIFIED_MEMORY_OFFLOAD_FLAGS = (
+    "dit_layerwise_offload",
+    "dit_cpu_offload",
+    "text_encoder_cpu_offload",
+    "image_encoder_cpu_offload",
+    "vae_cpu_offload",
+)
 
 
 class ExecutionMode(str, Enum):
@@ -854,20 +861,6 @@ def from_kwargs(cls, **kwargs: Any) -> "FastVideoArgs":
 
     def check_fastvideo_args(self) -> None:
         """Validate inference arguments for consistency"""
-        from fastvideo.platforms import current_platform
-
-        if current_platform.is_mps():
-            self.use_fsdp_inference = False
-            self.dit_layerwise_offload = False
-
-        if self.dit_layerwise_offload:
-            if self.use_fsdp_inference:
-                logger.warning("dit_layerwise_offload is enabled, automatically disabling use_fsdp_inference.")
-                self.use_fsdp_inference = False
-            if self.dit_cpu_offload:
-                logger.warning("dit_layerwise_offload is enabled, automatically disabling dit_cpu_offload.")
-                self.dit_cpu_offload = False
-
         # Validate mode and inference_mode consistency
         assert isinstance(self.mode, ExecutionMode), f"Mode must be an ExecutionMode enum, got {type(self.mode)}"
         assert self.mode in ExecutionMode.choices(), f"Invalid execution mode: {self.mode}"
@@ -884,6 +877,14 @@ def check_fastvideo_args(self) -> None:
             logger.warning("Mode is '%s' but inference_mode is False. Setting inference_mode to True.", self.mode)
             self.inference_mode = True
 
+        # Inference policy must wait until a worker owns and binds its device:
+        # a unified-memory device disables layerwise offload before conflicts
+        # are resolved, preserving an explicit FSDP request. Training does not
+        # pass through the inference worker boundary, so retain its historical
+        # constructor-time normalization.
+        if not self.inference_mode:
+            self._resolve_device_offload_conflicts()
+
         if not self.inference_mode:
             assert self.hsdp_replicate_dim != -1, "hsdp_replicate_dim must be set for training"
             assert self.hsdp_shard_dim != -1, "hsdp_shard_dim must be set for training"
@@ -918,6 +919,28 @@ def check_fastvideo_args(self) -> None:
                 self.pipeline_config.vae_config.load_encoder = True
             self.preprocess_config.check_preprocess_config()
 
+    def _resolve_device_offload_conflicts(self) -> None:
+        """Resolve offload modes after device-local policy has been applied."""
+        from fastvideo.platforms import current_platform
+
+        if current_platform.is_mps():
+            self.use_fsdp_inference = False
+            self.dit_layerwise_offload = False
+
+        if self.dit_layerwise_offload:
+            if self.use_fsdp_inference:
+                logger.warning("dit_layerwise_offload is enabled, automatically disabling use_fsdp_inference.")
+                self.use_fsdp_inference = False
+            if self.dit_cpu_offload:
+                logger.warning("dit_layerwise_offload is enabled, automatically disabling dit_cpu_offload.")
+                self.dit_cpu_offload = False
+
+    def finalize_device_offload_policy(self, device_id: int = 0) -> bool:
+        """Apply device-local memory policy, then resolve incompatible modes."""
+        has_unified_memory = self.disable_offload_on_unified_memory(device_id)
+        self._resolve_device_offload_conflicts()
+        return has_unified_memory
+
     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.
 
@@ -931,20 +954,29 @@ def disable_offload_on_unified_memory(self, device_id: int = 0, *, offload_flag:
         """
         from fastvideo.platforms import current_platform
 
-        if not current_platform.has_unified_memory(device_id):
+        cached_device_id = getattr(self, "_unified_memory_device_id", None)
+        cached_result = getattr(self, "_unified_memory_result", None)
+        if cached_device_id != device_id or cached_result is None:
+            cached_result = current_platform.has_unified_memory(device_id)
+            self._unified_memory_device_id = device_id
+            self._unified_memory_result = cached_result
+
+        if not cached_result:
             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):
+        enabled_flags = [flag for flag in UNIFIED_MEMORY_OFFLOAD_FLAGS if getattr(self, flag)]
+        if enabled_flags:
+            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 enabled_flags:
                 logger.info(
                     "Disabling %s: %s has unified memory, so moving weights to the host duplicates "
                     "them rather than freeing device memory.", flag, device_name)

fastvideo/pipelines/composed_pipeline_base.py

@@ -13,7 +13,11 @@
 import torch
 
 from fastvideo.configs.pipelines import PipelineConfig
-from fastvideo.distributed import (maybe_init_distributed_environment_and_model_parallel, get_world_group)
+from fastvideo.distributed import (
+    get_local_torch_device,
+    get_world_group,
+    maybe_init_distributed_environment_and_model_parallel,
+)
 from fastvideo.distributed.communication_op import (warmup_sequence_parallel_communication)
 from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
 from fastvideo.hooks.activation_trace import attach_activation_trace, detach_activation_trace
@@ -90,6 +94,14 @@ def __init__(self,
 
         maybe_init_distributed_environment_and_model_parallel(fastvideo_args.tp_size, fastvideo_args.sp_size)
 
+        # VideoGenerator applies this in each Worker before building the
+        # pipeline. Keep direct from_pretrained/build_pipeline callers aligned,
+        # but only after distributed setup has selected this process's device.
+        if fastvideo_args.inference_mode:
+            local_device = get_local_torch_device()
+            device_id = local_device.index if local_device.index is not None else 0
+            fastvideo_args.finalize_device_offload_policy(device_id)
+
         # Torch profiler. Enabled and configured through env vars:
         # FASTVIDEO_TORCH_PROFILER_DIR=/path/to/save/trace
         trace_dir = envs.FASTVIDEO_TORCH_PROFILER_DIR

fastvideo/tests/loader/test_text_encoder_unified_memory_offload.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-"""Regression tests for text-encoder placement on unified memory."""
+"""Regression tests for encoder placement on unified memory."""
 from __future__ import annotations
 
 from types import SimpleNamespace
@@ -84,7 +84,7 @@ def test_discrete_memory_preserves_explicit_cpu_target(monkeypatch, tmp_path) ->
     assert _PassthroughEncoder.loaded_device == torch.device("cpu")
 
 
-def test_text_policy_does_not_change_inherited_image_encoder_path(monkeypatch, tmp_path) -> None:
+def test_image_encoder_explicit_offload_resets_cpu_target(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"))
@@ -109,7 +109,7 @@ def test_text_policy_does_not_change_inherited_image_encoder_path(monkeypatch, t
         offload_flag="image_encoder_cpu_offload",
     )
 
-    assert _PassthroughEncoder.loaded_device == torch.device("cpu")
+    assert _PassthroughEncoder.loaded_device == torch.device("cuda:4")
     assert args.text_encoder_cpu_offload is False
-    assert args.image_encoder_cpu_offload is True
+    assert args.image_encoder_cpu_offload is False
     probe.assert_called_once_with(4)

fastvideo/tests/pipelines/test_direct_pipeline_unified_memory.py

@@ -0,0 +1,70 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Direct pipeline construction applies offload policy after device setup."""
+from __future__ import annotations
+
+from contextlib import nullcontext
+from types import SimpleNamespace
+
+import torch
+
+import fastvideo.pipelines.composed_pipeline_base as composed_pipeline_base
+from fastvideo.fastvideo_args import UNIFIED_MEMORY_OFFLOAD_FLAGS, FastVideoArgs
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+
+
+class _Profiler:
+
+    def region(self, name):
+        del name
+        return nullcontext()
+
+
+class _Pipeline(ComposedPipelineBase):
+    events = []
+
+    def load_modules(self, fastvideo_args, loaded_modules=None):
+        del loaded_modules
+        policy_state = {flag: getattr(fastvideo_args, flag) for flag in UNIFIED_MEMORY_OFFLOAD_FLAGS}
+        policy_state["use_fsdp_inference"] = fastvideo_args.use_fsdp_inference
+        self.events.append(("load_modules", policy_state))
+        return {}
+
+    def create_pipeline_stages(self, fastvideo_args):
+        del fastvideo_args
+
+
+def test_direct_pipeline_applies_policy_after_device_initialization(monkeypatch) -> None:
+    events = []
+    monkeypatch.setattr(_Pipeline, "events", events)
+    args = FastVideoArgs(model_path="unused", use_fsdp_inference=True)
+
+    def classify_device(device_id):
+        events.append(("offload_policy", device_id))
+        return True
+
+    monkeypatch.setattr(
+        composed_pipeline_base,
+        "maybe_init_distributed_environment_and_model_parallel",
+        lambda *args: events.append(("distributed", None)),
+    )
+    monkeypatch.setattr(composed_pipeline_base, "get_local_torch_device", lambda: torch.device("cuda:4"))
+    monkeypatch.setattr(composed_pipeline_base, "get_world_group", lambda: SimpleNamespace(local_rank=4))
+    monkeypatch.setattr(composed_pipeline_base, "get_or_create_profiler", lambda trace_dir: _Profiler())
+    monkeypatch.setattr("fastvideo.platforms.current_platform.has_unified_memory", classify_device)
+    monkeypatch.setattr("fastvideo.platforms.current_platform.get_device_name", lambda device_id: "NVIDIA GB10")
+    monkeypatch.setattr("fastvideo.platforms.current_platform.is_mps", lambda: False)
+
+    pipeline = _Pipeline("unused", args, required_config_modules=[])
+
+    assert pipeline.modules == {}
+    assert events == [
+        ("distributed", None),
+        ("offload_policy", 4),
+        (
+            "load_modules",
+            {
+                **{flag: False for flag in UNIFIED_MEMORY_OFFLOAD_FLAGS},
+                "use_fsdp_inference": True,
+            },
+        ),
+    ]