From 14adfe2edce76f4c19ab6cfd8c271e7b101d26ad Mon Sep 17 00:00:00 2001 From: Wenxuan Tan Date: Wed, 9 Jul 2025 16:45:20 -0500 Subject: [PATCH] Remove all unnecessary torch.cuda.empty_cache (#606) --- fastvideo/v1/distributed/parallel_state.py | 10 ---------- fastvideo/v1/entrypoints/video_generator.py | 3 --- fastvideo/v1/models/vaes/common.py | 3 +-- fastvideo/v1/pipelines/stages/encoding.py | 1 - fastvideo/v1/pipelines/stages/image_encoding.py | 3 --- fastvideo/v1/training/training_pipeline.py | 3 --- fastvideo/v1/worker/gpu_worker.py | 8 -------- 7 files changed, 1 insertion(+), 30 deletions(-) diff --git a/fastvideo/v1/distributed/parallel_state.py b/fastvideo/v1/distributed/parallel_state.py index bdbb807e5..0bcdbea42 100644 --- a/fastvideo/v1/distributed/parallel_state.py +++ b/fastvideo/v1/distributed/parallel_state.py @@ -23,7 +23,6 @@ If you only need to use the distributed environment without model parallelism, you can skip the model parallel initialization and destruction steps. """ import contextlib -import gc import os import pickle import weakref @@ -1016,15 +1015,6 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False): if shutdown_ray: import ray # Lazy import Ray ray.shutdown() - gc.collect() - from fastvideo.v1.platforms import current_platform - if not current_platform.is_cpu(): - torch.cuda.empty_cache() - try: - torch._C._host_emptyCache() - except AttributeError: - logger.warning( - "torch._C._host_emptyCache() only available in Pytorch >=2.5") def in_the_same_node_as(pg: ProcessGroup | StatelessProcessGroup, diff --git a/fastvideo/v1/entrypoints/video_generator.py b/fastvideo/v1/entrypoints/video_generator.py index 9b0726775..1a2f47554 100644 --- a/fastvideo/v1/entrypoints/video_generator.py +++ b/fastvideo/v1/entrypoints/video_generator.py @@ -6,7 +6,6 @@ This module provides a consolidated interface for generating videos using diffusion models. """ -import gc import math import os import time @@ -277,5 +276,3 @@ class VideoGenerator: """ self.executor.shutdown() del self.executor - gc.collect() - torch.cuda.empty_cache() diff --git a/fastvideo/v1/models/vaes/common.py b/fastvideo/v1/models/vaes/common.py index 11f124c1d..6975a5b38 100644 --- a/fastvideo/v1/models/vaes/common.py +++ b/fastvideo/v1/models/vaes/common.py @@ -239,7 +239,6 @@ class ParallelTiledVAE(ABC): results = torch.cat(local_results, dim=0).contiguous() del local_results - torch.cuda.empty_cache() # first gather size to pad the results local_size = torch.tensor([results.size(0)], device=results.device, @@ -253,7 +252,7 @@ class ParallelTiledVAE(ABC): padded_results = torch.zeros(max_size, device=results.device) padded_results[:results.size(0)] = results del results - torch.cuda.empty_cache() + # Gather all results gathered_dim_metadata = [None] * world_size gathered_results = torch.zeros_like(padded_results).repeat( diff --git a/fastvideo/v1/pipelines/stages/encoding.py b/fastvideo/v1/pipelines/stages/encoding.py index a4c0a04ad..a99bf9c2f 100644 --- a/fastvideo/v1/pipelines/stages/encoding.py +++ b/fastvideo/v1/pipelines/stages/encoding.py @@ -136,7 +136,6 @@ class EncodingStage(PipelineStage): self.maybe_free_model_hooks() self.vae.to("cpu") - torch.cuda.empty_cache() return batch diff --git a/fastvideo/v1/pipelines/stages/image_encoding.py b/fastvideo/v1/pipelines/stages/image_encoding.py index 11e9bbf77..62370c3e0 100644 --- a/fastvideo/v1/pipelines/stages/image_encoding.py +++ b/fastvideo/v1/pipelines/stages/image_encoding.py @@ -5,8 +5,6 @@ Image encoding stages for I2V diffusion pipelines. This module contains implementations of image encoding stages for diffusion pipelines. """ -import torch - from fastvideo.v1.distributed import get_local_torch_device from fastvideo.v1.fastvideo_args import FastVideoArgs from fastvideo.v1.forward_context import set_forward_context @@ -68,7 +66,6 @@ class ImageEncodingStage(PipelineStage): if fastvideo_args.use_cpu_offload: self.image_encoder.to('cpu') - torch.cuda.empty_cache() return batch diff --git a/fastvideo/v1/training/training_pipeline.py b/fastvideo/v1/training/training_pipeline.py index 30dd27de5..750e0368f 100644 --- a/fastvideo/v1/training/training_pipeline.py +++ b/fastvideo/v1/training/training_pipeline.py @@ -1,5 +1,4 @@ # SPDX-License-Identifier: Apache-2.0 -import gc import math import os import time @@ -706,5 +705,3 @@ class TrainingPipeline(ComposedPipelineBase, ABC): # Re-enable gradients for training training_args.inference_mode = False transformer.train() - gc.collect() - torch.cuda.empty_cache() diff --git a/fastvideo/v1/worker/gpu_worker.py b/fastvideo/v1/worker/gpu_worker.py index fb2e82783..797cc6d48 100644 --- a/fastvideo/v1/worker/gpu_worker.py +++ b/fastvideo/v1/worker/gpu_worker.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 import contextlib import faulthandler -import gc import multiprocessing as mp import os import signal @@ -69,8 +68,6 @@ class Worker: torch.cuda.set_device(self.device) # _check_if_gpu_supports_dtype(self.model_config.dtype) - gc.collect() - torch.cuda.empty_cache() self.init_gpu_memory = torch.cuda.mem_get_info()[0] os.environ["MASTER_ADDR"] = "localhost" @@ -102,9 +99,6 @@ class Worker: if hasattr(self, 'pipeline') and self.pipeline is not None: # Clean up pipeline resources if needed pass - # Release CUDA resources - if torch.cuda.is_available(): - torch.cuda.empty_cache() # Destroy the distributed environment cleanup_dist_env_and_memory(shutdown_ray=False) @@ -133,8 +127,6 @@ class Worker: # Handle regular RPC calls if method_name == 'execute_forward': - gc.collect() - torch.cuda.empty_cache() forward_batch = recv_rpc['kwargs']['forward_batch'] fastvideo_args = recv_rpc['kwargs']['fastvideo_args'] output_batch = self.execute_forward(forward_batch,