[diffusion] logging: log available gpu mem while loading and generating (#15936)

python/sglang/multimodal_gen/runtime/loader/component_loader.py

@@ -7,7 +7,6 @@ import glob
 import importlib.util
 import json
 import os
-import time
 import traceback
 from abc import ABC
 from collections.abc import Generator, Iterable
@@ -17,6 +16,7 @@ from typing import Any, cast
 import torch
 import torch.distributed as dist
 import torch.nn as nn
+from diffusers import AutoModel
 from safetensors.torch import load_file as safetensors_load_file
 from torch.distributed import init_device_mesh
 from transformers import AutoImageProcessor, AutoProcessor, AutoTokenizer
@@ -90,6 +90,26 @@ def _list_safetensors_files(model_path: str) -> list[str]:
     return sorted(glob.glob(os.path.join(str(model_path), "*.safetensors")))
 
 
+def get_memory_usage_of_component(module) -> float | None:
+    """
+    returned value is in GB, rounded to 2 decimal digits
+    """
+    if not isinstance(module, nn.Module):
+        return None
+    BYTES_PER_GB = 1024**3
+    if hasattr(module, "get_memory_footprint"):
+        usage = module.get_memory_footprint() / BYTES_PER_GB
+    else:
+        # manually
+        param_size = sum(p.numel() * p.element_size() for p in module.parameters())
+        buffer_size = sum(b.numel() * b.element_size() for b in module.buffers())
+
+        total_size_bytes = param_size + buffer_size
+        usage = total_size_bytes / (1024**3)
+
+    return round(usage, 2)
+
+
 class ComponentLoader(ABC):
     """Base class for loading a specific type of model component."""
 
@@ -118,7 +138,7 @@ class ComponentLoader(ABC):
         server_args: ServerArgs,
         module_name: str,
         transformers_or_diffusers: str,
-    ):
+    ) -> tuple[AutoModel, float]:
         """
         Template method that standardizes logging around the core load implementation.
         The priority of loading method is:
@@ -127,7 +147,13 @@ class ComponentLoader(ABC):
         If all of the above methods failed, an error will be thrown
 
         """
-        logger.info("Loading %s from %s", module_name, component_model_path)
+        gpu_mem_before_loading = current_platform.get_available_gpu_memory()
+        logger.info(
+            "Loading %s from %s. avail mem: %.2f GB",
+            module_name,
+            component_model_path,
+            gpu_mem_before_loading,
+        )
         try:
             component = self.load_customized(
                 component_model_path, server_args, module_name
@@ -159,20 +185,27 @@ class ComponentLoader(ABC):
 
         if component is None:
             logger.warning("Loaded %s returned None", module_name)
+            consumed = 0.0
         else:
+            current_gpu_mem = current_platform.get_available_gpu_memory()
+            consumed = get_memory_usage_of_component(component)
+            if consumed is None or consumed == 0.0:
+                consumed = gpu_mem_before_loading - current_gpu_mem
             logger.info(
-                f"Loaded %s: %s from: {source}",
+                f"Loaded %s: %s from {source}. avail mem: %.2f GB, %.2f GB consumed",
                 module_name,
                 component.__class__.__name__,
+                current_gpu_mem,
+                consumed,
             )
-        return component
+        return component, consumed
 
     def load_native(
         self,
         component_model_path: str,
         server_args: ServerArgs,
         transformers_or_diffusers: str,
-    ):
+    ) -> AutoModel:
         """
         Load the component using the native library (transformers/diffusers).
         """
@@ -273,9 +306,6 @@ class TextEncoderLoader(ComponentLoader):
         allow_patterns_overrides: list[str] | None = None
         """If defined, weights will load exclusively using these patterns."""
 
-    counter_before_loading_weights: float = 0.0
-    counter_after_loading_weights: float = 0.0
-
     def should_offload(self, server_args, model_config: ModelConfig | None = None):
         should_offload = server_args.text_encoder_cpu_offload
         if not should_offload:
@@ -355,8 +385,6 @@ class TextEncoderLoader(ComponentLoader):
         else:
             weights_iterator = pt_weights_iterator(hf_weights_files, to_cpu=to_cpu)
 
-        if self.counter_before_loading_weights == 0.0:
-            self.counter_before_loading_weights = time.perf_counter()
         # apply the prefix.
         return ((source.prefix + name, tensor) for (name, tensor) in weights_iterator)
 
@@ -443,12 +471,6 @@ class TextEncoderLoader(ComponentLoader):
             loaded_weights = model.load_weights(
                 self._get_all_weights(model, model_path, to_cpu=should_offload)
             )
-            self.counter_after_loading_weights = time.perf_counter()
-            logger.info(
-                "Loading weights took %.2f seconds",
-                self.counter_after_loading_weights
-                - self.counter_before_loading_weights,
-            )
 
             # Explicitly move model to target device after loading weights
             model = model.to(local_torch_device)