diff --git a/.buildkite/scripts/unit_test.sh b/.buildkite/scripts/unit_test.sh index 8f6ac456f8..440082c4e3 100644 --- a/.buildkite/scripts/unit_test.sh +++ b/.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/ \ diff --git a/docs/getting_started/installation/spark_performance.md b/docs/getting_started/installation/spark_performance.md index 2687f073b3..a2a3a14526 100644 --- a/docs/getting_started/installation/spark_performance.md +++ b/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 diff --git a/docs/inference/offloading.md b/docs/inference/offloading.md index 373b28d17b..61773f710b 100644 --- a/docs/inference/offloading.md +++ b/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 diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index aaca40c71f..7ff44f5ad1 100644 --- a/fastvideo/fastvideo_args.py +++ b/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) diff --git a/fastvideo/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py index bc2a0f9a14..ec1acb963c 100644 --- a/fastvideo/pipelines/composed_pipeline_base.py +++ b/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 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 73b9707d1d..c955358a3c 100644 --- a/fastvideo/tests/loader/test_text_encoder_unified_memory_offload.py +++ b/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) diff --git a/fastvideo/tests/pipelines/test_direct_pipeline_unified_memory.py b/fastvideo/tests/pipelines/test_direct_pipeline_unified_memory.py new file mode 100644 index 0000000000..0201188fc7 --- /dev/null +++ b/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, + }, + ), + ] diff --git a/fastvideo/tests/platforms/test_unified_memory_offload.py b/fastvideo/tests/platforms/test_unified_memory_offload.py index fdf00dcab5..aec74b0960 100644 --- a/fastvideo/tests/platforms/test_unified_memory_offload.py +++ b/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: +@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") + monkeypatch.setattr("fastvideo.platforms.current_platform.is_mps", lambda: False) + return probe + + +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 = _args(text_encoder_cpu_offload=True) + args = FastVideoArgs(model_path="unused/for-this-test", use_fsdp_inference=True) - assert args.text_encoder_cpu_offload is 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_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) +def test_training_construction_retains_offload_conflict_normalization(monkeypatch) -> None: + monkeypatch.setattr("fastvideo.platforms.current_platform.is_mps", lambda: False) - assert args.disable_offload_on_unified_memory(device_id=6, offload_flag="text_encoder_cpu_offload") is True + 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, + ) - probe.assert_called_once_with(6) - assert args.text_encoder_cpu_offload is False + 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_discrete_device_keeps_text_encoder_offload(monkeypatch) -> None: +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 + assert args.image_encoder_cpu_offload is True + assert args.vae_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) +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") - args = _args(text_encoder_cpu_offload=True, image_encoder_cpu_offload=True) + 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") - applies_to_image_encoder = args.disable_offload_on_unified_memory( - device_id=0, - offload_flag="image_encoder_cpu_offload", - ) + device_zero_args.finalize_device_offload_policy(device_id=0) + device_one_args.finalize_device_offload_policy(device_id=1) - assert applies_to_image_encoder is False - assert args.text_encoder_cpu_offload is False - assert args.image_encoder_cpu_offload is True + 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 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) + 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 diff --git a/fastvideo/tests/worker/test_gpu_worker.py b/fastvideo/tests/worker/test_gpu_worker.py index f4436cec99..56cfd9e1e1 100644 --- a/fastvideo/tests/worker/test_gpu_worker.py +++ b/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 diff --git a/fastvideo/worker/gpu_worker.py b/fastvideo/worker/gpu_worker.py index 72bd030c55..3e6d18250d 100644 --- a/fastvideo/worker/gpu_worker.py +++ b/fastvideo/worker/gpu_worker.py @@ -84,7 +84,7 @@ def init_device(self) -> None: # 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,