fix: remove redundant torch.cuda.empty_cache()

recipe/prime/prime_fsdp_workers.py

@@ -247,8 +247,6 @@ def init_model(self):
                                                         lr_scheduler=self.reward_lr_scheduler,
                                                         tokenizer=self.tokenizer)
 
-        torch.cuda.empty_cache()
-
     @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)
     def compute_rm_score(self, data: DataProto):
         data = data.to('cuda')
@@ -321,7 +319,6 @@ def update_rm(self, data: DataProto):
             offload_fsdp_model_to_cpu(self.ref_module)
         if self._is_offload_optimizer:
             offload_fsdp_optimizer(optimizer=self.reward_optimizer)
-        torch.cuda.empty_cache()
         output = output.to('cpu')
         return output
 

verl/workers/fsdp_workers.py

@@ -411,8 +411,6 @@ def init_model(self):
                 lr_scheduler=self.actor_lr_scheduler,
                 processing_class=self.processor if self.processor is not None else self.tokenizer)
 
-        torch.cuda.empty_cache()
-
     @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)
     def update_actor(self, data: DataProto):
         # Support all hardwares
@@ -498,7 +496,6 @@ def generate_sequences(self, prompts: DataProto):
         output = output.to('cpu')
 
         # clear kv cache
-        torch.cuda.empty_cache()
         log_gpu_memory_usage('After recompute log prob', logger=logger)
         return output
 
@@ -561,7 +558,6 @@ def compute_ref_log_prob(self, data: DataProto):
         if self.world_size > 1:
             self.ref_policy.actor_module._handle.reshard(True)
 
-        torch.cuda.empty_cache()
         return output
 
     @register(dispatch_mode=Dispatch.ONE_TO_ALL)
@@ -784,8 +780,6 @@ def init_model(self):
             lr_scheduler=self.critic_lr_scheduler,
             processing_class=self.processor if self.processor is not None else self.tokenizer)
 
-        torch.cuda.empty_cache()
-
     @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)
     def compute_values(self, data: DataProto):
 
@@ -981,7 +975,6 @@ def init_model(self):
         # This is used to import external_lib into the huggingface systems
         import_external_libs(self.config.model.get('external_lib', None))
         self.reward_module = self._build_model(config=self.config)
-        torch.cuda.empty_cache()
 
     def _forward_micro_batch(self, micro_batch):
         from flash_attn.bert_padding import pad_input, unpad_input, index_first_axis, rearrange
@@ -1154,5 +1147,4 @@ def compute_rm_score(self, data: DataProto):
         self.reward_module._handle.reshard(True)
 
         output = output.to('cpu')
-        torch.cuda.empty_cache()
         return output

verl/workers/sharding_manager/fsdp_vllm.py

@@ -70,6 +70,15 @@ def __init__(self,
             self.gen_random_states = None
 
     def __enter__(self):
+        # NOTE: Basically, we only need `torch.cuda.empty_cache()` before vllm wake_up and
+        # after vllm sleep, since vllm has its own caching memory allocator CuMemAllocator.
+        # Out of vllm scope, we should avoid empty cache to let pytorch using caching memory
+        # to speed up memory allocations.
+        #
+        # pytorch: https://pytorch.org/docs/stable/notes/cuda.html#memory-management
+        # vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/device_allocator/cumem.py#L103
+        torch.cuda.empty_cache()
+
         log_gpu_memory_usage('Before state_dict() in sharding manager memory', logger=logger)
         params = self.module.state_dict()
         log_gpu_memory_usage('After state_dict() in sharding manager memory', logger=logger)
@@ -89,7 +98,6 @@ def __enter__(self):
         log_gpu_memory_usage('After sync model weights in sharding manager', logger=logger)
 
         del params
-        torch.cuda.empty_cache()
         log_gpu_memory_usage('After del state_dict and empty_cache in sharding manager', logger=logger)
 
         # TODO: offload FSDP model weights