diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 90b8c0ba9..3f5a15827 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -111,6 +111,41 @@ class GPUWorker: f"Worker {self.rank}: Initialized device, model, and distributed environment." ) + def do_mem_analysis(self, output_batch: OutputBatch): + peak_memory_bytes = torch.cuda.max_memory_allocated() + output_batch.peak_memory_mb = peak_memory_bytes / (1024**2) + peak_memory_gb = peak_memory_bytes / (1024**3) + remaining_gpu_mem_gb = ( + current_platform.get_device_total_memory() / (1024**3) - peak_memory_gb + ) + can_stay_resident = self.get_can_stay_resident_components(remaining_gpu_mem_gb) + suggested_args = set() + component_to_arg = { + "vae": "--vae-cpu-offload", + "text_encoder": "--text-encoder-cpu-offload", + "text_encoder_2": "--text-encoder-cpu-offload", + "image_encoder": "--image-encoder-cpu-offload", + } + + for component in can_stay_resident: + if component == "transformer": + if self.server_args.dit_layerwise_offload: + suggested_args.add("--dit-layerwise-offload") + elif self.server_args.dit_cpu_offload: + suggested_args.add("--dit-cpu-offload") + elif component in component_to_arg: + suggested_args.add(component_to_arg[component]) + + suggested_args_str = ( + ", ".join(sorted(suggested_args)) if suggested_args else "None" + ) + logger.info( + f"Peak GPU memory: {peak_memory_gb:.2f} GB, " + f"Remaining GPU memory at peak: {remaining_gpu_mem_gb:.2f} GB. " + f"Components that could stay resident (based on the last request workload): {can_stay_resident}. " + f"Related offload server args to disable: {suggested_args_str}" + ) + def execute_forward(self, batch: List[Req]) -> OutputBatch: """ Execute a forward pass. @@ -138,23 +173,8 @@ class GPUWorker: else: output_batch = result - if self.rank == 0: - peak_memory_bytes = torch.cuda.max_memory_allocated() - output_batch.peak_memory_mb = peak_memory_bytes / (1024**2) - peak_memory_gb = peak_memory_bytes / (1024**3) - remaining_gpu_mem_gb = ( - current_platform.get_device_total_memory() / (1024**3) - - peak_memory_gb - ) - can_stay_resident = self.get_can_stay_resident_components( - remaining_gpu_mem_gb - ) - if not req.suppress_logs: - logger.info( - f"Peak GPU memory: {peak_memory_gb:.2f} GB, " - f"Remaining GPU memory at peak: {remaining_gpu_mem_gb:.2f} GB. " - f"Components that can stay resident: {can_stay_resident}" - ) + if self.rank == 0 and not req.suppress_logs: + self.do_mem_analysis(output_batch) duration_ms = (time.monotonic() - start_time) * 1000 output_batch.timings.total_duration_ms = duration_ms diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py index 01c9af0f4..d4b1c71d6 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py @@ -138,7 +138,7 @@ class InputValidationStage(PipelineStage): scale = max(ow / iw, oh / ih) img = img.resize((round(iw * scale), round(ih * scale)), Image.LANCZOS) - logger.info("resized img height: %s, img width: %s", img.height, img.width) + logger.debug("resized img height: %s, img width: %s", img.height, img.width) # center-crop x1 = (img.width - ow) // 2 diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index 9363d4dd9..e56115157 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args.py @@ -12,7 +12,6 @@ import os import random import sys import tempfile -from contextlib import contextmanager from dataclasses import field from enum import Enum from typing import Any, Optional @@ -1090,8 +1089,6 @@ class PortArgs: ) -# TODO: not sure what _current_server_args is for, using a _global_server_args instead -_current_server_args = None _global_server_args = None @@ -1103,29 +1100,9 @@ def prepare_server_args(argv: list[str]) -> ServerArgs: ServerArgs.add_cli_args(parser) raw_args = parser.parse_args(argv) server_args = ServerArgs.from_cli_args(raw_args) - global _current_server_args - _current_server_args = server_args return server_args -@contextmanager -def set_current_server_args(server_args: ServerArgs): - """ - Temporarily set the current sgl_diffusion config. - Used during model initialization. - We save the current sgl_diffusion config in a global variable, - so that all modules can access it, e.g. custom ops - can access the sgl_diffusion config to determine how to dispatch. - """ - global _current_server_args - old_server_args = _current_server_args - try: - _current_server_args = server_args - yield - finally: - _current_server_args = old_server_args - - def set_global_server_args(server_args: ServerArgs): """ Set the global sgl_diffusion config for each process @@ -1134,17 +1111,7 @@ def set_global_server_args(server_args: ServerArgs): _global_server_args = server_args -def get_current_server_args() -> ServerArgs | None: - if _current_server_args is None: - # in ci, usually when we test custom ops/modules directly, - # we don't set the sgl_diffusion config. In that case, we set a default - # config. - # TODO(will): may need to handle this for CI. - raise ValueError("Current sgl_diffusion args is not set.") - return _current_server_args - - -def get_global_server_args() -> ServerArgs | None: +def get_global_server_args() -> ServerArgs: if _global_server_args is None: # in ci, usually when we test custom ops/modules directly, # we don't set the sgl_diffusion config. In that case, we set a default diff --git a/python/sglang/multimodal_gen/runtime/utils/logging_utils.py b/python/sglang/multimodal_gen/runtime/utils/logging_utils.py index 4be477cc5..18af7e71f 100644 --- a/python/sglang/multimodal_gen/runtime/utils/logging_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/logging_utils.py @@ -142,6 +142,7 @@ def get_is_local_main_process(): def _log_process_aware( + server_log_level: int, level: int, logger_self: Logger, msg: object, @@ -153,12 +154,12 @@ def _log_process_aware( """Helper function to log a message if the process rank matches the criteria.""" is_main_process = get_is_main_process() is_local_main_process = get_is_local_main_process() - should_log = ( not main_process_only and not local_main_process_only or (main_process_only and is_main_process) or (local_main_process_only and is_local_main_process) + or server_log_level <= logging.DEBUG ) if should_log: @@ -234,6 +235,8 @@ def init_logger(name: str) -> _SGLDiffusionLogger: logger = logging.getLogger(name) + server_log_level = logger.getEffectiveLevel() + # Patch instance methods setattr(logger, "info_once", MethodType(_print_info_once, logger)) setattr(logger, "warning_once", MethodType(_print_warning_once, logger)) @@ -252,6 +255,7 @@ def init_logger(name: str) -> _SGLDiffusionLogger: **kwargs: Any, ) -> None: _log_process_aware( + server_log_level, level, self, msg, @@ -281,7 +285,7 @@ def init_logger(name: str) -> _SGLDiffusionLogger: setattr( logger, "error", - MethodType(_create_patched_method(logging.ERROR, False, True), logger), + MethodType(_create_patched_method(logging.ERROR, False, False), logger), ) return cast(_SGLDiffusionLogger, logger)