Merge upstream/main into report-signal-terminated-workers

Both sides added a bullet to the GB10 tuning list and git could not tell
they were about different things. #1715 documents that FastVideo now
disables the offload modes once a worker binds its device; this branch
documents that earlyoom prefers Python and that a worker's SIGTERM
traceback shows where it was interrupted, not why it was chosen. Neither
replaces the other, so both are kept, offload first.

Author: KyleNeverGivesUp

.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.
 - Some Spark images run `earlyoom` with a preference for terminating Python
   processes under memory pressure. A worker's SIGTERM log and traceback show
   where it was interrupted, not why it was selected; confirm the cause in the

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/models/loader/fsdp_load.py

@@ -545,7 +545,14 @@ def load_model_from_full_model_state_dict(
     sharded_sd = {}
     custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict(full_sd_iterator,
                                                                            param_names_mapping)  # type: ignore
-    for target_param_name, full_tensor in custom_param_sd.items():
+    # Drain rather than iterate. Production safetensors values may retain
+    # memory-mapped shard storage, while mapped or merged parameters can own
+    # ordinary allocations. Keeping the dict retains all of that source
+    # storage until loading finishes; popping releases each reference as soon
+    # as its conversion completes and lowers the host/unified-memory working
+    # set.
+    for target_param_name in list(custom_param_sd):
+        full_tensor = custom_param_sd.pop(target_param_name)
         meta_sharded_param = meta_sd.get(target_param_name)
         if meta_sharded_param is None:
             # Some checkpoints include extra entries that are not part of the

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_fsdp_load_releases_checkpoint.py

@@ -0,0 +1,163 @@
+# SPDX-License-Identifier: Apache-2.0
+"""The loader must not hold the whole checkpoint alive while it copies it.
+
+``hf_to_custom_state_dict`` drains the weight iterator into one dict before a
+single parameter is placed. Production safetensors values may retain
+memory-mapped shard storage, and mapped or merged parameters can own ordinary
+allocations. Keeping every source reference until loading finishes raises the
+host or unified-memory working set enough to kill a large-model load.
+
+Checking after the call proves nothing, because the dict is a local and dies
+with the frame. So these tests keep a reference to that dict from the outside
+and assert it comes back empty: if the loader iterated instead of draining,
+every source tensor would still be reachable through the reference we hold.
+"""
+from __future__ import annotations
+
+import gc
+import weakref
+
+import pytest
+import torch
+from torch import nn
+
+from fastvideo.models.loader import fsdp_load
+from fastvideo.models.loader.fsdp_load import load_model_from_full_model_state_dict
+
+PARAM_NAMES = ("a.weight", "b.weight", "c.weight")
+
+
+class _TinyModel(nn.Module):
+    """Plain module, no device mesh, so the loader takes its unsharded path."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.a = nn.Linear(8, 8, bias=False)
+        self.b = nn.Linear(8, 8, bias=False)
+        self.c = nn.Linear(8, 8, bias=False)
+
+
+def _identity_mapping(name: str) -> tuple[str, None, None]:
+    return name, None, None
+
+
+def _source_tensors(scale: bool = False) -> dict[str, torch.Tensor]:
+    # FP32 on purpose: the loader casts to param_dtype, and a cast is what makes
+    # the copy a real copy. Handing it tensors that already match would let
+    # `.to()` return the same object, and the model would then legitimately keep
+    # the source alive for reasons that have nothing to do with this fix.
+    return {
+        name: torch.ones(8, 8, dtype=torch.float32) * (index + 1 if scale else 1)
+        for index, name in enumerate(PARAM_NAMES)
+    }
+
+
+def _capture_state_dict(monkeypatch) -> dict:
+    """Hold the loader's internal dict from outside so we can inspect it after."""
+    captured: dict = {}
+    real = fsdp_load.hf_to_custom_state_dict
+
+    def spy(*args, **kwargs):
+        custom_param_sd, reverse = real(*args, **kwargs)
+        captured["sd"] = custom_param_sd
+        return custom_param_sd, reverse
+
+    monkeypatch.setattr(fsdp_load, "hf_to_custom_state_dict", spy)
+    return captured
+
+
+def test_the_state_dict_is_drained_not_iterated(monkeypatch) -> None:
+    captured = _capture_state_dict(monkeypatch)
+
+    load_model_from_full_model_state_dict(
+        _TinyModel(),
+        iter(list(_source_tensors().items())),
+        torch.device("cpu"),
+        torch.bfloat16,
+        strict=False,
+        param_names_mapping=_identity_mapping,
+        training_mode=False,
+    )
+
+    assert captured["sd"] == {}, ("the loader finished with the checkpoint still in hand; on a real model that is the "
+                                  "whole file held resident for the length of the copy")
+
+
+@pytest.mark.parametrize(
+    ("skipped_name", "strict"),
+    (("metadata._extra_state", True), ("unexpected.weight", False)),
+)
+def test_skipped_source_entry_is_popped_before_continue(monkeypatch, skipped_name: str, strict: bool) -> None:
+    """Both skip branches must drop their source before advancing the loop."""
+    captured = _capture_state_dict(monkeypatch)
+    warning_names = []
+
+    def assert_popped_before_warning(_message, warned_name) -> None:
+        assert warned_name not in captured["sd"]
+        warning_names.append(warned_name)
+
+    monkeypatch.setattr(fsdp_load.logger, "warning", assert_popped_before_warning)
+    sources = {**_source_tensors(), skipped_name: torch.ones(8, 8)}
+
+    load_model_from_full_model_state_dict(
+        _TinyModel(),
+        iter(sources.items()),
+        torch.device("cpu"),
+        torch.bfloat16,
+        strict=strict,
+        param_names_mapping=_identity_mapping,
+        training_mode=False,
+    )
+
+    assert warning_names == [skipped_name]
+    assert captured["sd"] == {}
+
+
+def test_source_tensors_become_collectable(monkeypatch) -> None:
+    """The reason the drain matters: the tensors have to actually go."""
+    captured = _capture_state_dict(monkeypatch)
+    sources = _source_tensors()
+    refs = {name: weakref.ref(tensor) for name, tensor in sources.items()}
+
+    # Mirror safetensors_weights_iterator, which drops each tensor as it yields.
+    def iterator():
+        while sources:
+            yield sources.popitem()
+
+    load_model_from_full_model_state_dict(
+        _TinyModel(),
+        iterator(),
+        torch.device("cpu"),
+        torch.bfloat16,
+        strict=False,
+        param_names_mapping=_identity_mapping,
+        training_mode=False,
+    )
+
+    # captured["sd"] is still in scope here on purpose: it is the reference that
+    # would keep them alive if the loader had not popped.
+    gc.collect()
+    alive = sorted(name for name, ref in refs.items() if ref() is not None)
+    assert not alive, f"still reachable through the loader's state dict: {alive}"
+
+
+def test_weights_still_land_in_the_model() -> None:
+    """Releasing early must not cost correctness."""
+    sources = _source_tensors(scale=True)
+    expected = {name: tensor[0, 0].item() for name, tensor in sources.items()}
+
+    model = _TinyModel()
+    load_model_from_full_model_state_dict(
+        model,
+        iter(list(sources.items())),
+        torch.device("cpu"),
+        torch.bfloat16,
+        strict=False,
+        param_names_mapping=_identity_mapping,
+        training_mode=False,
+    )
+
+    loaded = dict(model.named_parameters())
+    for name, value in expected.items():
+        assert loaded[name].dtype == torch.bfloat16
+        assert loaded[name][0, 0].item() == value

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)