chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] fix: fix suppressing error log on non-main ranks (#17712)

Browse Source
This commit is contained in:
Mick
2026-01-28 09:29:19 +08:00
committed by GitHub
parent 331a22427c
commit 1507dc6cdf
4 changed files with 45 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files

54
python/sglang/multimodal_gen/runtime/managers/gpu_worker.py
View File

@@ -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