[diffusion] chore: minor cleanups (#19123)

This commit is contained in:
Mick
2026-02-22 19:07:25 +08:00
committed by GitHub
parent 1a1c768d44
commit 87823722b3
6 changed files with 128 additions and 113 deletions
@@ -135,22 +135,21 @@ class GPUWorker:
# otherwise empty offloaded weights could fail lora converting
if self.server_args.dit_layerwise_offload:
# enable layerwise offload if possible
for dit in filter(
None,
[
self.pipeline.get_module("transformer"),
self.pipeline.get_module("transformer_2"),
self.pipeline.get_module("video_dit"),
self.pipeline.get_module("video_dit_2"),
self.pipeline.get_module("audio_dit"),
],
):
if isinstance(dit, OffloadableDiTMixin):
dit.configure_layerwise_offload(self.server_args)
else:
logger.info(
f"Module {type(dit).__name__} does not support layerwise offload. Skipping."
)
for module_name in [
"transformer",
"transformer_2",
"video_dit",
"video_dit_2",
"audio_dit",
]:
dit = self.pipeline.get_module(module_name)
if dit:
if isinstance(dit, OffloadableDiTMixin):
dit.configure_layerwise_offload(self.server_args)
else:
logger.info(
f"Module {type(dit).__name__} does not support layerwise offload. Skipping."
)
logger.info(
f"Worker {self.rank}: Initialized device, model, and distributed environment."
@@ -199,7 +198,7 @@ class GPUWorker:
logger.info(
f"Peak GPU memory: {peak_reserved_gb:.2f} GB, "
f"Peak allocated: {peak_allocated_gb:.2f} GB, "
f"Memory pool overhead: {pool_overhead_gb:.2f} GB ({pool_overhead_gb/peak_reserved_gb*100:.1f}%), "
f"Memory pool overhead: {pool_overhead_gb:.2f} GB ({pool_overhead_gb / peak_reserved_gb * 100:.1f}%), "
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}"