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

[diffusion] feat: generalize layerwise offloader to flux1 (#15633)

Browse Source
This commit is contained in:
ryang
2025-12-23 22:22:40 +08:00
committed by GitHub
parent b3f83cc1c5
commit ffc23ef877
3 changed files with 50 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
python/sglang/multimodal_gen/runtime/loader/component_loader.py
View File

@@ -712,16 +712,17 @@ class TransformerLoader(ComponentLoader):
model = model.eval()
if server_args.dit_layerwise_offload and hasattr(model, "blocks"):
if server_args.dit_layerwise_offload and hasattr(model, "dit_module_names"):
# TODO(will): support multiple module names
module_name = getattr(model, "dit_module_names", ["transformer_blocks"])[0]
try:
num_layers = len(getattr(model, "blocks"))
num_layers = len(getattr(model, module_name))
except Exception:
num_layers = None
if isinstance(num_layers, int) and num_layers > 0:
mgr = LayerwiseOffloadManager(
model,
module_list_attr="blocks",
module_list_attr=module_name,
num_layers=num_layers,
enabled=True,
pin_cpu_memory=server_args.pin_cpu_memory,