[diffusion] chore: minor cleanups (#19123)
This commit is contained in:
@@ -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}"
|
||||
|
||||
Reference in New Issue
Block a user