From d265a454bf8d0eeda3225ba1bc9239b8a1bfaf84 Mon Sep 17 00:00:00 2001 From: Raghav K Date: Thu, 24 Sep 2026 17:05:59 -0700 Subject: [PATCH] [bugfix]: Fall back when large pinned output allocation fails (#1759) --- fastvideo/entrypoints/video_generator.py | 5 +-- .../minimax_h3/stages/minimax_h3_decoding.py | 7 ++- .../tests/entrypoints/test_video_generator.py | 34 +++++++++++++++ .../stages/test_minimax_h3_vae_streaming.py | 43 +++++++++++++++++++ .../test_minimax_h3_video_vae_streaming.py | 16 +++++++ fastvideo/utils.py | 32 ++++++++++++++ 6 files changed, 130 insertions(+), 7 deletions(-) diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index cf55c1e1f..65b6d8d4e 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -55,7 +55,7 @@ from fastvideo.api.sampling_param import SamplingParam from fastvideo.fastvideo_args import FastVideoArgs, WorkloadType from fastvideo.logger import init_logger from fastvideo.pipelines import ForwardBatch -from fastvideo.utils import align_to, shallow_asdict +from fastvideo.utils import allocate_cpu_tensor_with_pin_fallback, align_to, shallow_asdict from fastvideo.worker.executor import Executor fcntl: types.ModuleType | None @@ -819,9 +819,8 @@ class VideoGenerator: if skip_pixel_prealloc: samples = torch.empty(0, device='cpu') else: - samples = torch.empty( + samples = allocate_cpu_tensor_with_pin_fallback( (latent_batch_size, 3, sampling_param.num_frames, sampling_param.height, sampling_param.width), - device='cpu', pin_memory=fastvideo_args.pin_cpu_memory) thread.join() diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py index a12cca50d..65c1ed371 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py @@ -26,7 +26,7 @@ from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.pipelines.stages.base import PipelineStage from fastvideo.pipelines.stages.validators import StageValidators as V from fastvideo.pipelines.stages.validators import VerificationResult -from fastvideo.utils import is_pin_memory_available +from fastvideo.utils import allocate_cpu_tensor_with_pin_fallback logger = init_logger(__name__) @@ -135,11 +135,10 @@ class MiniMaxH3VideoDecodingStage(PipelineStage): output = None if is_output_rank: - output = torch.empty( + output = allocate_cpu_tensor_with_pin_fallback( self.vae.decoded_pixel_shape(latents.shape), - device="cpu", dtype=torch.float32, - pin_memory=fastvideo_args.pin_cpu_memory and is_pin_memory_available(), + pin_memory=fastvideo_args.pin_cpu_memory, ) # Attribute the streamed decoder computation while retaining # per-chunk device-to-host transfer and pinned-buffer reuse. diff --git a/fastvideo/tests/entrypoints/test_video_generator.py b/fastvideo/tests/entrypoints/test_video_generator.py index 0e7b940fa..4ccb0c922 100644 --- a/fastvideo/tests/entrypoints/test_video_generator.py +++ b/fastvideo/tests/entrypoints/test_video_generator.py @@ -8,6 +8,7 @@ import torch from einops import rearrange import fastvideo.entrypoints.video_generator as video_generator_module +import fastvideo.utils as fastvideo_utils from fastvideo.api import ( ConfigValidationError, GenerationRequest, @@ -273,6 +274,39 @@ def test_generate_single_video_return_frames_still_materializes_output(tmp_path) assert result["video_path"] is None +def test_generate_single_video_retries_pageable_output_when_pinned_allocation_fails(monkeypatch, tmp_path): + output = torch.ones((1, 3, 2, 16, 16), dtype=torch.float32) * 0.5 + output_batch = _single_video_output_batch(output) + fastvideo_args = _single_video_args() + fastvideo_args.pin_cpu_memory = True + generator = _single_video_generator(output_batch, fastvideo_args) + real_empty = torch.empty + allocation_attempts = [] + warnings = [] + + def fake_empty(size, **kwargs): + allocation_attempts.append(kwargs.get("pin_memory", False)) + if kwargs.get("pin_memory"): + raise RuntimeError("CUDA error: out of memory") + return real_empty(size, **kwargs) + + monkeypatch.setattr(fastvideo_utils, "is_pin_memory_available", lambda: True) + monkeypatch.setattr(fastvideo_utils.torch, "empty", fake_empty) + monkeypatch.setattr(fastvideo_utils.logger, "warning", lambda message, *args: warnings.append(message % args)) + + result = generator._generate_single_video( + prompt="pinned allocation fallback", + sampling_param=_small_sampling_param(save_video=False, return_frames=True), + fastvideo_args=fastvideo_args, + output_path=str(tmp_path / "unused.mp4"), + ) + + torch.testing.assert_close(result["samples"], output) + assert allocation_attempts[:2] == [True, False] + assert len(warnings) == 1 + assert "retrying with pageable memory" in warnings[0] + + def test_generate_single_video_frames_match_legacy_cpu_loop(tmp_path): """The on-device quantize path (#1362) must reproduce the legacy per-frame CPU loop (make_grid -> permute -> *255 -> uint8) bit-exactly diff --git a/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py b/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py index 6202e11d0..2e605139f 100644 --- a/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py +++ b/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py @@ -114,6 +114,49 @@ def test_decode_stage_uses_cpu_output_buffer(monkeypatch) -> None: assert torch.all(result.output == 0.25) +def test_decode_stage_requests_pin_fallback_output_buffer(monkeypatch) -> None: + latent_shape = (1, 4, 2, 4, 4) + latents = torch.randn(latent_shape) + rows = patchify_video_latents(latents, (1, 1, 1)) + batch = ForwardBatch(data_type="video", latents=rows, raw_latent_shape=latent_shape) + batch.extra[MINIMAX_H3_LAYOUT_KEY] = _layout(rows.shape[0], latent_shape) + observed = {} + + class VAE: + + def to(self, device): + return self + + def denormalize_latents(self, decoded_latents): + return decoded_latents + + def decoded_pixel_shape(self, shape): + return (1, 3, 5, 16, 16) + + def decode_to_pixels(self, decoded_latents, output): + output.zero_() + + def fake_allocate(size, *, dtype=None, pin_memory=False): + observed["size"] = size + observed["dtype"] = dtype + observed["pin_memory"] = pin_memory + return torch.empty(size, dtype=dtype) + + monkeypatch.setattr(minimax_h3_decoding, "get_local_torch_device", lambda: torch.device("cpu")) + monkeypatch.setattr(minimax_h3_decoding, "allocate_cpu_tensor_with_pin_fallback", fake_allocate) + + MiniMaxH3VideoDecodingStage(VAE(), SimpleNamespace(patch_size=(1, 1, 1))).forward( + batch, + SimpleNamespace(output_type="pil", pin_cpu_memory=True, vae_cpu_offload=False, vae_parallel_decode=False), + ) + + assert observed == { + "size": (1, 3, 5, 16, 16), + "dtype": torch.float32, + "pin_memory": True, + } + + def test_decode_stages_skip_vae_on_non_output_rank(monkeypatch) -> None: class VAE: diff --git a/fastvideo/tests/vaes/test_minimax_h3_video_vae_streaming.py b/fastvideo/tests/vaes/test_minimax_h3_video_vae_streaming.py index 5a2c8783e..68e70da3e 100644 --- a/fastvideo/tests/vaes/test_minimax_h3_video_vae_streaming.py +++ b/fastvideo/tests/vaes/test_minimax_h3_video_vae_streaming.py @@ -158,6 +158,22 @@ def test_decode_to_pixels_pinned_buffer_matches_dense_on_cuda() -> None: assert_close(pinned, expected, atol=0.0, rtol=0.0) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="pageable-buffer streaming requires CUDA") +@torch.inference_mode() +def test_decode_to_pixels_pageable_buffer_matches_dense_on_cuda() -> None: + """The allocation fallback's blocking chunk copies must equal dense decode.""" + torch.manual_seed(20260814) + vae = _tiny_vae().to("cuda") + latents = torch.randn(1, 4, 12, 4, 4, device="cuda") + expected = vae.denormalize_pixels(vae.decode(latents).sample.float()).clamp_(0, 1).cpu() + + pageable = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32) + vae.decode_to_pixels(latents, pageable) + + assert not pageable.is_pinned() + assert_close(pageable, expected, atol=0.0, rtol=0.0) + + def test_decode_to_pixels_rejects_incomplete_output(monkeypatch) -> None: vae = _tiny_vae() latents = torch.randn(1, 4, 2, 4, 4) diff --git a/fastvideo/utils.py b/fastvideo/utils.py index 79757a23a..af4d2acda 100644 --- a/fastvideo/utils.py +++ b/fastvideo/utils.py @@ -1306,3 +1306,35 @@ def _cached_pin_memory_available(pid: int) -> bool: def is_pin_memory_available() -> bool: return _cached_pin_memory_available(os.getpid()) + + +def allocate_cpu_tensor_with_pin_fallback( + size: tuple[int, ...] | torch.Size, + *, + dtype: torch.dtype | None = None, + pin_memory: bool = False, +) -> torch.Tensor: + """Allocate a CPU tensor, retrying pageable memory if pinning fails. + + ``is_pin_memory_available`` intentionally uses a small capability probe. + Large decoded-video buffers can still exhaust the CUDA host allocator after + that probe succeeds, especially on unified-memory systems. Treat pinning as + a performance preference: preserve the requested shape and dtype by + retrying the allocation without pinning. + """ + kwargs: dict[str, object] = {"device": "cpu"} + if dtype is not None: + kwargs["dtype"] = dtype + + if not pin_memory or not is_pin_memory_available(): + return torch.empty(size, **kwargs) + + try: + return torch.empty(size, pin_memory=True, **kwargs) + except RuntimeError as exc: + logger.warning( + "Pinned CPU allocation failed for shape %s; retrying with pageable memory: %s", + tuple(size), + exc, + ) + return torch.empty(size, **kwargs)