Reduce vRAM usage during generation by allowing to transfer logits to CPU

memory_test.py

@@ -0,0 +1,123 @@
+"""
+Test script to evaluate memory efficiency of CPU logits offloading.
+"""
+
+import gc
+import time
+
+import torch
+
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+def get_gpu_memory_mb():
+    """Get current GPU memory usage in MB."""
+    return torch.cuda.memory_allocated() / 1024 / 1024
+
+
+def clear_memory():
+    """Clear GPU memory."""
+    torch.cuda.empty_cache()
+    gc.collect()
+
+
+def test_memory_scenario():
+    """Test memory efficiency of CPU offloading with gpt2-large."""
+
+    model_name = "gpt2-large"
+    max_new_tokens = 1000
+    prompt_text = "This essay will start with a detailed and verbose introduction to deep learning. To start with,"
+
+    print(f"Testing model: {model_name}")
+    print(f"Tokens to generate: {max_new_tokens}")
+    print(f"Device: {torch.cuda.get_device_name()}")
+
+    # Load model
+    print("Loading model...")
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name, dtype=torch.float16).cuda()
+
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    inputs = tokenizer(prompt_text, return_tensors="pt")
+    inputs = {k: v.cuda() for k, v in inputs.items()}
+
+
+    # Test WITHOUT CPU offloading
+    clear_memory()
+    print("\nTesting without CPU offloading...")
+
+    initial_memory = get_gpu_memory_mb()
+    start_time = time.time()
+
+    with torch.no_grad():
+        outputs_gpu = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            output_scores=True,
+            output_logits=True,
+            return_dict_in_generate=True,
+            offload_logits_to_cpu=False,
+            do_sample=False,
+            pad_token_id=tokenizer.eos_token_id,
+        )
+
+    gpu_time = time.time() - start_time
+    gpu_peak = get_gpu_memory_mb()
+    gpu_increase = gpu_peak - initial_memory
+
+    print(f"Initial memory: {initial_memory:.1f} MB")
+    print(f"Peak memory: {gpu_peak:.1f} MB")
+    print(f"Memory increase: {gpu_increase:.1f} MB")
+    print(f"Generation time: {gpu_time:.2f}s")
+    print(f"Tokens generated: {len(outputs_gpu.logits) if outputs_gpu.logits else 0}")
+
+    # Store for comparison
+    gpu_sequences = outputs_gpu.sequences.clone()
+    del outputs_gpu
+    clear_memory()
+
+    # Test WITH CPU offloading
+    print("\nTesting with CPU offloading...")
+
+    initial_memory = get_gpu_memory_mb()
+    start_time = time.time()
+
+    with torch.no_grad():
+        outputs_cpu = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            output_scores=True,
+            output_logits=True,
+            return_dict_in_generate=True,
+            offload_logits_to_cpu=True,
+            do_sample=False,
+            pad_token_id=tokenizer.eos_token_id,
+        )
+
+    cpu_time = time.time() - start_time
+    cpu_peak = get_gpu_memory_mb()
+    cpu_increase = cpu_peak - initial_memory
+
+    print(f"Initial memory: {initial_memory:.1f} MB")
+    print(f"Peak memory: {cpu_peak:.1f} MB")
+    print(f"Memory increase: {cpu_increase:.1f} MB")
+    print(f"Generation time: {cpu_time:.2f}s")
+    print(f"Tokens generated: {len(outputs_cpu.logits) if outputs_cpu.logits else 0}")
+
+    # Analysis
+    memory_saved = gpu_increase - cpu_increase
+    reduction_pct = (memory_saved / gpu_increase * 100) if gpu_increase > 0 else 0
+    time_overhead = ((cpu_time - gpu_time) / gpu_time * 100) if gpu_time > 0 else 0
+    sequences_match = torch.equal(gpu_sequences, outputs_cpu.sequences)
+
+    print("\nResults:")
+    print("Number of tokens generated:", outputs_cpu.sequences.shape[1])
+    print(f"Memory saved: {memory_saved:.1f} MB")
+    print(f"Memory reduction: {reduction_pct:.1f}%")
+    print(f"Time overhead: {time_overhead:+.1f}%")
+    print(f"Sequences match: {sequences_match}")
+
+if __name__ == "__main__":
+    test_memory_scenario()

src/transformers/generation/configuration_utils.py

@@ -256,6 +256,12 @@ class GenerationConfig(PushToHubMixin):
         output_logits (`bool`, *optional*):
             Whether or not to return the unprocessed prediction logit scores. See `logits` under returned tensors for
             more details.
+        offload_logits_to_cpu (`bool`, *optional*, defaults to `False`):
+            Whether or not to transfer logits, scores, and attentions to CPU memory after each generation step to reduce GPU memory
+            usage. Only effective when `output_scores=True`, `output_logits=True`, or `output_attentions=True`. This can significantly reduce
+            GPU memory consumption during long generations at the cost of some CPU-GPU transfer overhead.
+            NB: if `offload_logits_to_cpu=True`, the returned `scores`, `logits`, and `attentions` tensors in the generation output
+            will be located on CPU memory, not on the same device as the model.
         return_dict_in_generate (`bool`, *optional*, defaults to `False`):
             Whether or not to return a [`~utils.ModelOutput`], as opposed to returning exclusively the generated
             sequence. This flag must be set to `True` to return the generation cache (when `use_cache` is `True`)
@@ -328,7 +334,13 @@ class GenerationConfig(PushToHubMixin):
             need to use this flag.
     """
 
-    extra_output_flags = ("output_attentions", "output_hidden_states", "output_scores", "output_logits")
+    extra_output_flags = (
+        "output_attentions",
+        "output_hidden_states",
+        "output_scores",
+        "output_logits",
+        "offload_logits_to_cpu",
+    )
 
     def __init__(self, **kwargs):
         # Parameters that control the length of the output
@@ -390,6 +402,7 @@ def __init__(self, **kwargs):
         self.output_hidden_states = kwargs.pop("output_hidden_states", False)
         self.output_scores = kwargs.pop("output_scores", False)
         self.output_logits = kwargs.pop("output_logits", None)
+        self.offload_logits_to_cpu = kwargs.pop("offload_logits_to_cpu", False)
         self.return_dict_in_generate = kwargs.pop("return_dict_in_generate", False)
 
         # Special tokens that can be used at generation time

src/transformers/generation/utils.py

@@ -143,6 +143,22 @@
 }
 
 
+def _maybe_offload_to_cpu(tensor: torch.Tensor, offload_to_cpu: bool) -> torch.Tensor:
+    """
+    Conditionally transfer a tensor to CPU memory if offloading is enabled.
+
+    Args:
+        tensor (torch.Tensor): The tensor to potentially transfer to CPU.
+        offload_to_cpu (bool): Whether to transfer the tensor to CPU.
+
+    Returns:
+        torch.Tensor: The tensor on CPU if offloading is enabled, otherwise unchanged.
+    """
+    if offload_to_cpu and tensor.device != torch.device("cpu"):
+        return tensor.cpu()
+    return tensor
+
+
 @dataclass
 class GenerateDecoderOnlyOutput(ModelOutput):
     """
@@ -2712,6 +2728,7 @@ def _sample(
         output_hidden_states = generation_config.output_hidden_states
         output_scores = generation_config.output_scores
         output_logits = generation_config.output_logits
+        offload_logits_to_cpu = generation_config.offload_logits_to_cpu
         return_dict_in_generate = generation_config.return_dict_in_generate
         has_eos_stopping_criteria = any(hasattr(criteria, "eos_token_id") for criteria in stopping_criteria)
         do_sample = generation_config.do_sample
@@ -2790,15 +2807,25 @@ def _sample(
             # Store scores, attentions and hidden_states when required
             if return_dict_in_generate:
                 if output_scores:
-                    scores += (next_token_scores,)
+                    scores += (_maybe_offload_to_cpu(next_token_scores, offload_logits_to_cpu),)
                 if output_logits:
-                    raw_logits += (next_token_logits,)
+                    raw_logits += (_maybe_offload_to_cpu(next_token_logits, offload_logits_to_cpu),)
                 if output_attentions:
-                    decoder_attentions += (
-                        (outputs.decoder_attentions,) if self.config.is_encoder_decoder else (outputs.attentions,)
-                    )
                     if self.config.is_encoder_decoder:
-                        cross_attentions += (outputs.cross_attentions,)
+                        decoder_attentions += (
+                            tuple(
+                                _maybe_offload_to_cpu(att, offload_logits_to_cpu) for att in outputs.decoder_attentions
+                            ),
+                        )
+                        cross_attentions += (
+                            tuple(
+                                _maybe_offload_to_cpu(att, offload_logits_to_cpu) for att in outputs.cross_attentions
+                            ),
+                        )
+                    else:
+                        decoder_attentions += (
+                            tuple(_maybe_offload_to_cpu(att, offload_logits_to_cpu) for att in outputs.attentions),
+                        )
 
                 if output_hidden_states:
                     decoder_hidden_states += (
@@ -3145,6 +3172,7 @@ def _beam_search(
         output_hidden_states = generation_config.output_hidden_states
         output_scores = generation_config.output_scores
         output_logits = generation_config.output_logits
+        offload_logits_to_cpu = generation_config.offload_logits_to_cpu
         return_dict_in_generate = generation_config.return_dict_in_generate
         do_sample = generation_config.do_sample
         early_stopping = generation_config.early_stopping
@@ -3273,18 +3301,30 @@ def _beam_search(
             # Store logits, attentions and hidden_states when required
             if return_dict_in_generate:
                 if output_logits:
-                    raw_logits += (logits.clone(),)
+                    raw_logits += (_maybe_offload_to_cpu(logits.clone(), offload_logits_to_cpu),)
                 if return_dict_in_generate and output_scores:
-                    all_scores += (log_probs.clone(),)
+                    all_scores += (_maybe_offload_to_cpu(log_probs.clone(), offload_logits_to_cpu),)
 
                 if output_attentions:
-                    decoder_attentions += (
-                        (model_outputs.decoder_attentions,)
-                        if self.config.is_encoder_decoder
-                        else (model_outputs.attentions,)
-                    )
                     if self.config.is_encoder_decoder:
-                        cross_attentions += (model_outputs.cross_attentions,)
+                        decoder_attentions += (
+                            tuple(
+                                _maybe_offload_to_cpu(att, offload_logits_to_cpu)
+                                for att in model_outputs.decoder_attentions
+                            ),
+                        )
+                        cross_attentions += (
+                            tuple(
+                                _maybe_offload_to_cpu(att, offload_logits_to_cpu)
+                                for att in model_outputs.cross_attentions
+                            ),
+                        )
+                    else:
+                        decoder_attentions += (
+                            tuple(
+                                _maybe_offload_to_cpu(att, offload_logits_to_cpu) for att in model_outputs.attentions
+                            ),
+                        )
 
                 if output_hidden_states:
                     decoder_hidden_states += (
@@ -3514,6 +3554,7 @@ def _assisted_decoding(
         output_hidden_states = generation_config.output_hidden_states
         output_scores = generation_config.output_scores
         output_logits = generation_config.output_logits
+        offload_logits_to_cpu = generation_config.offload_logits_to_cpu
         return_dict_in_generate = generation_config.return_dict_in_generate
 
         # init attention / hidden states / scores tuples
@@ -3653,9 +3694,15 @@ def _assisted_decoding(
             if return_dict_in_generate:
                 newly_added_length = n_matches + 1
                 if output_scores:
-                    scores += tuple(new_logits[:, i, :] for i in range(newly_added_length))
+                    scores += tuple(
+                        _maybe_offload_to_cpu(new_logits[:, i, :], offload_logits_to_cpu)
+                        for i in range(newly_added_length)
+                    )
                 if output_logits:
-                    raw_logits += tuple(next_token_logits[:, i, :] for i in range(newly_added_length))
+                    raw_logits += tuple(
+                        _maybe_offload_to_cpu(next_token_logits[:, i, :], offload_logits_to_cpu)
+                        for i in range(newly_added_length)
+                    )
 
                 newly_added_length = new_cur_len if is_first_iteration else newly_added_length
                 if output_attentions:
@@ -3670,6 +3717,18 @@ def _assisted_decoding(
                             newly_added_length,
                             is_decoder_attention=True,
                         )
+                        # Apply CPU offloading to the newly added attention tensors
+                        if offload_logits_to_cpu:
+                            # Offload the last `newly_added_length` tuples in cross_attentions
+                            cross_attentions = cross_attentions[:-newly_added_length] + tuple(
+                                tuple(_maybe_offload_to_cpu(att, True) for att in attention_tuple)
+                                for attention_tuple in cross_attentions[-newly_added_length:]
+                            )
+                            # Offload the last `newly_added_length` tuples in decoder_attentions
+                            decoder_attentions = decoder_attentions[:-newly_added_length] + tuple(
+                                tuple(_maybe_offload_to_cpu(att, True) for att in attention_tuple)
+                                for attention_tuple in decoder_attentions[-newly_added_length:]
+                            )
                     # some (V)LLMs have hard requirement on SDPA and thus never return attn
                     elif outputs.attentions[0] is not None:
                         decoder_attentions = _split_model_outputs(
@@ -3679,6 +3738,13 @@ def _assisted_decoding(
                             newly_added_length,
                             is_decoder_attention=True,
                         )
+                        # Apply CPU offloading to the newly added attention tensors
+                        if offload_logits_to_cpu:
+                            # Offload the last `newly_added_length` tuples in decoder_attentions
+                            decoder_attentions = decoder_attentions[:-newly_added_length] + tuple(
+                                tuple(_maybe_offload_to_cpu(att, True) for att in attention_tuple)
+                                for attention_tuple in decoder_attentions[-newly_added_length:]
+                            )
                 if output_hidden_states:
                     if self.config.is_encoder_decoder:
                         decoder_hidden_states = _split_model_outputs(