code: def benchmark_memory(
      model,
      input_ids: torch.Tensor,
      masks: torch.Tensor,
      num_batches: int,
      is_decoder: bool,
      new_tokens: int,
      pad_token_id: int,
      memory_tracker: MemoryTracker,
  ):
      torch.cuda.empty_cache()
      gc.collect()
  
      print("Measuring peak memory...")
      with memory_tracker.track():
          gen_config = warmup(
              model,
              input_ids,
              masks,
              is_decoder,
              new_tokens,
              pad_token_id,
          )
  
          if is_decoder:
              _ = model.generate(
                  input_ids,
                  attention_mask=masks,
                  generation_config=gen_config,
              )
          else:
              _ = model(input_ids, masks)
  
          torch.cuda.synchronize()
  
      memory_stats = torch.cuda.memory_stats()
  
      peak_allocated_torch_mb = (
          memory_stats["allocated_bytes.all.peak"] * 1e-6
      )
      peak_reserved_torch_mb = (
          memory_stats["reserved_bytes.all.peak"] * 1e-6
      )
  
      peak_nvml_mb = memory_tracker.peak_memory
  
      # I am not sure whether we should substract here `inactive_split_bytes.all.peak` (not sure what it corresponds to, though it can get quite large, in the several GB).
      peak_external_mb = peak_nvml_mb - peak_reserved_torch_mb
      # assert peak_external_mb > 0
  
      # This formula is to confirm. We measure the actual allocated PyTorch memory, plus the additional non-PyTorch memory (as the CUDA context, CUDA extension device memory). We need to substract the PyTorch peak reserved memory since this one appears in the peak nvidia-smi/nvmlDeviceGetMemoryInfo.
  
      # NOTE: I verified this is only a ROUGH estimate. It may be better to use PYTORCH_NO_CUDA_MEMORY_CACHING=1 and just nvmlDeviceGetMemoryInfo.
      # We can actually doubt whether it make sense to try to estimate when we would OOM, given that different devices, CUDA version do have
      # a different CUDA context size.
      peak_memory_mb = peak_allocated_torch_mb + peak_external_mb
  
      print(
          f"DEBUG: peak allocated torch: {peak_allocated_torch_mb:.2f} MB"
      )
      print(f"DEBUG: peak nvidia-smi/nvml: {peak_nvml_mb:.2f} MB")
      print(
          f"DEBUG: peak reserved torch: {peak_reserved_torch_mb:.2f} MB"
      )
      print(f"DEBUG: peak external: {peak_external_mb:.2f} MB")
      print(f"DEBUG: global peak: {peak_memory_mb:.2f} MB")
  
      return peak_memory_mb
  
file_path: optimum/tests/benchmark/benchmark_gptq.py
imports: import argparse
  import gc
  import os
  import time
  import numpy
  import torch
  from auto_gptq.utils import Perplexity
  from memory_tracker import MemoryTracker
  from tqdm import tqdm
  from transformers import AutoModel, AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer, BitsAndBytesConfig, GenerationConfig, GPTQConfig
  from optimum.exporters import TasksManager
label_desc: Use image node_id 5 for calling a global function with example usage: benchmark_memory(model, input_ids, masks, num_batches, is_decoder, new_tokens, pad_token_id, memory_tracker) and returns: peak_memory_mb