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

.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 @@ else:
 
 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 @@ class 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 @@ class FastVideoArgs:
             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 @@ class FastVideoArgs:
                 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 @@ class FastVideoArgs:
         """
         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
+        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 UNIFIED_MEMORY_OFFLOAD_FLAGS:
-            if getattr(self, flag):
+            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 @@ from typing import Any, cast
 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 @@ class ComposedPipelineBase(ABC):
 
         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,
+            },
+        ),
+    ]

fastvideo/tests/platforms/test_unified_memory_offload.py

@@ -1,64 +1,180 @@
 # SPDX-License-Identifier: Apache-2.0
-"""CPU tests for applying unified-memory offload policy inside a worker."""
+"""CPU tests for worker-local offload policy on unified-memory devices."""
 from __future__ import annotations
 
+import dataclasses
 from unittest.mock import Mock
 
 import pytest
 
-from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.fastvideo_args import ExecutionMode, UNIFIED_MEMORY_OFFLOAD_FLAGS, FastVideoArgs
 
 
-def _args(**overrides) -> FastVideoArgs:
-    return FastVideoArgs(model_path="unused/for-this-test", **overrides)
+def _args_with_offloads(*enabled_flags: str, **overrides) -> FastVideoArgs:
+    kwargs = {flag: flag in enabled_flags for flag in UNIFIED_MEMORY_OFFLOAD_FLAGS}
+    kwargs.update(model_path="unused/for-this-test")
+    kwargs.update(overrides)
+    return FastVideoArgs(**kwargs)
 
 
-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:
+@pytest.fixture
+def as_unified_cuda(monkeypatch):
     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
+    monkeypatch.setattr("fastvideo.platforms.current_platform.is_mps", lambda: False)
+    return probe
 
 
-def test_discrete_device_keeps_text_encoder_offload(monkeypatch) -> None:
+def test_constructing_inference_args_defers_device_policy(monkeypatch) -> None:
+    probe = Mock(side_effect=AssertionError("device probe ran in the driver"))
+    monkeypatch.setattr("fastvideo.platforms.current_platform.has_unified_memory", probe)
+
+    args = FastVideoArgs(model_path="unused/for-this-test", use_fsdp_inference=True)
+
+    assert args.use_fsdp_inference is True
+    assert args.dit_layerwise_offload is True
+    assert args.dit_cpu_offload is True
+    probe.assert_not_called()
+
+
+def test_training_construction_retains_offload_conflict_normalization(monkeypatch) -> None:
+    monkeypatch.setattr("fastvideo.platforms.current_platform.is_mps", lambda: False)
+
+    args = FastVideoArgs(
+        model_path="unused/for-this-test",
+        mode=ExecutionMode.FINETUNING,
+        inference_mode=False,
+        sp_size=1,
+        hsdp_shard_dim=1,
+        use_fsdp_inference=True,
+    )
+
+    assert args.dit_layerwise_offload is True
+    assert args.dit_cpu_offload is False
+    assert args.use_fsdp_inference is False
+
+
+def test_policy_list_covers_every_declared_offload_flag() -> None:
+    declared = {field.name for field in dataclasses.fields(FastVideoArgs) if field.name.endswith("_offload")}
+
+    assert declared == set(UNIFIED_MEMORY_OFFLOAD_FLAGS)
+
+
+def test_unified_device_disables_every_offload_flag(as_unified_cuda) -> None:
+    args = FastVideoArgs(model_path="unused/for-this-test")
+
+    assert args.finalize_device_offload_policy(device_id=6) is True
+
+    as_unified_cuda.assert_called_once_with(6)
+    assert not [flag for flag in UNIFIED_MEMORY_OFFLOAD_FLAGS if getattr(args, flag)]
+
+
+@pytest.mark.parametrize("flag", UNIFIED_MEMORY_OFFLOAD_FLAGS)
+def test_each_offload_flag_is_independently_disabled(as_unified_cuda, flag: str) -> None:
+    args = _args_with_offloads(flag)
+    assert getattr(args, flag) is True
+
+    args.disable_offload_on_unified_memory(device_id=2)
+
+    assert getattr(args, flag) is False
+
+
+def test_discrete_device_classification_preserves_offload_requests(monkeypatch) -> None:
     probe = Mock(return_value=False)
     monkeypatch.setattr("fastvideo.platforms.current_platform.has_unified_memory", probe)
-    args = _args(text_encoder_cpu_offload=True)
+    args = FastVideoArgs(model_path="unused/for-this-test")
 
     assert args.disable_offload_on_unified_memory(device_id=3) is False
 
     probe.assert_called_once_with(3)
+    assert all(getattr(args, flag) for flag in UNIFIED_MEMORY_OFFLOAD_FLAGS)
+
+
+def test_discrete_device_finalization_retains_layerwise_precedence(monkeypatch) -> None:
+    monkeypatch.setattr("fastvideo.platforms.current_platform.has_unified_memory", lambda device_id: False)
+    monkeypatch.setattr("fastvideo.platforms.current_platform.is_mps", lambda: False)
+    args = FastVideoArgs(model_path="unused/for-this-test", use_fsdp_inference=True)
+
+    args.finalize_device_offload_policy(device_id=3)
+
+    assert args.dit_layerwise_offload is True
+    assert args.dit_cpu_offload is False
+    assert args.use_fsdp_inference is False
     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
+    assert args.vae_cpu_offload is True
+
+
+def test_workers_classify_their_own_device(monkeypatch) -> None:
+    seen_device_ids = []
+
+    def has_unified_memory(device_id):
+        seen_device_ids.append(device_id)
+        return device_id == 1
+
+    monkeypatch.setattr("fastvideo.platforms.current_platform.has_unified_memory", has_unified_memory)
+    monkeypatch.setattr("fastvideo.platforms.current_platform.get_device_name", lambda device_id: "NVIDIA GB10")
+    monkeypatch.setattr("fastvideo.platforms.current_platform.is_mps", lambda: False)
+    device_zero_args = FastVideoArgs(model_path="unused/for-this-test")
+    device_one_args = FastVideoArgs(model_path="unused/for-this-test")
+
+    device_zero_args.finalize_device_offload_policy(device_id=0)
+    device_one_args.finalize_device_offload_policy(device_id=1)
+
+    assert seen_device_ids == [0, 1]
+    assert device_zero_args.dit_layerwise_offload is True
+    assert device_zero_args.text_encoder_cpu_offload is True
+    assert not [flag for flag in UNIFIED_MEMORY_OFFLOAD_FLAGS if getattr(device_one_args, flag)]
+
+
+def test_repeated_policy_on_same_device_reuses_classification(as_unified_cuda) -> None:
+    args = FastVideoArgs(model_path="unused/for-this-test")
+
+    args.finalize_device_offload_policy(device_id=1)
+    args.finalize_device_offload_policy(device_id=1)
+
+    as_unified_cuda.assert_called_once_with(1)
+
+
+def test_mps_clears_offload_and_keeps_its_fsdp_rule(monkeypatch) -> None:
+    monkeypatch.setattr("fastvideo.platforms.current_platform.is_mps", lambda: True)
+    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: "mps")
+    args = FastVideoArgs(model_path="unused/for-this-test", use_fsdp_inference=True)
+
+    args.finalize_device_offload_policy()
+
+    assert args.use_fsdp_inference is False
+    assert not [flag for flag in UNIFIED_MEMORY_OFFLOAD_FLAGS if getattr(args, flag)]
+
+
+def test_cuda_unified_memory_preserves_realistic_fsdp_request(as_unified_cuda) -> None:
+    args = FastVideoArgs(model_path="unused/for-this-test", use_fsdp_inference=True)
+    assert args.dit_layerwise_offload is True
+    assert args.use_fsdp_inference is True
+
+    args.finalize_device_offload_policy()
+
+    assert args.use_fsdp_inference is True
+    assert not [flag for flag in UNIFIED_MEMORY_OFFLOAD_FLAGS if getattr(args, flag)]
+
+
+def test_pin_cpu_memory_is_not_a_host_offload_mode(as_unified_cuda) -> None:
+    args = FastVideoArgs(model_path="unused/for-this-test", pin_cpu_memory=True)
+
+    args.finalize_device_offload_policy()
+
+    assert "pin_cpu_memory" not in UNIFIED_MEMORY_OFFLOAD_FLAGS
+    assert args.pin_cpu_memory is True
+
+
+def test_already_disabled_flags_stay_disabled(as_unified_cuda) -> None:
+    args = _args_with_offloads()
+
+    args.finalize_device_offload_policy()
+
+    assert not [flag for flag in UNIFIED_MEMORY_OFFLOAD_FLAGS if getattr(args, flag)]
 
 
 @pytest.mark.parametrize("name_error", [NotImplementedError, ValueError, RuntimeError])
@@ -69,7 +185,7 @@ def test_platform_without_device_name_uses_generic_name(monkeypatch, name_error:
         raise name_error("device name unavailable")
 
     monkeypatch.setattr("fastvideo.platforms.current_platform.get_device_name", unsupported_name)
-    args = _args(text_encoder_cpu_offload=True)
+    args = _args_with_offloads("text_encoder_cpu_offload")
 
     assert args.disable_offload_on_unified_memory() is True
     assert args.text_encoder_cpu_offload is False

fastvideo/tests/worker/test_gpu_worker.py

@@ -1,3 +1,4 @@
+import os
 from types import SimpleNamespace
 from unittest.mock import Mock
 
@@ -42,7 +43,7 @@ def test_init_device_applies_offload_policy_after_binding_worker_device(monkeypa
     """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)))
+    args.finalize_device_offload_policy = 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")
@@ -64,6 +65,8 @@ def test_init_device_applies_offload_policy_after_binding_worker_device(monkeypa
         ("distributed", None),
         ("pipeline", None),
     ]
+    assert os.environ["LOCAL_RANK"] == "3"
+    assert worker.device == torch.device("cuda:3")
     assert worker.init_gpu_memory == 123
 
 

fastvideo/worker/gpu_worker.py

@@ -84,7 +84,7 @@ class Worker:
         # 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)
+        self.fastvideo_args.finalize_device_offload_policy(device_id)
 
         # Initialize the distributed environment.
         maybe_init_distributed_environment_and_model_parallel(self.fastvideo_args.tp_size, self.fastvideo_args.sp_size,