fix(lora): materialize balanced-offload modules before touching weights

Under a pressed balanced offload, dispatched modules hold meta tensors
whose data lives in the accelerate offload map. The factor path raised
trying to move a meta svd tensor and aborted activation mid-pass; the
legacy requantize path silently skipped those layers. Both left the
model with a partially applied network.

Rebuild the offload state with apply_balanced_offload(force) at
activate and deactivate entry: modules come back real on cpu with
hooks intact and the execution device unchanged, so both paths see
usable tensors and the next forward re-onloads under the watermark.
This commit is contained in:
CalamitousFelicitousness
2026-07-17 12:41:26 +01:00
parent 5b694d3300
commit 28bd5e8d74
+4
View File
@@ -24,6 +24,8 @@ def network_activate(include=None, exclude=None):
if shared.opts.diffusers_offload_mode == "sequential":
sd_models.disable_offload(sd_model)
sd_models.move_model(sd_model, device=devices.cpu)
elif shared.opts.diffusers_offload_mode == "balanced":
sd_model = sd_models.apply_balanced_offload(sd_model, force=True) # dispatched modules hold meta tensors backed by the offload map; rebuild them real on cpu with hooks intact before touching weights
device = None
modules = {}
components = include if len(include) > 0 else default_components
@@ -111,6 +113,8 @@ def network_deactivate(include=None, exclude=None):
if shared.opts.diffusers_offload_mode == "sequential":
sd_models.disable_offload(sd_model)
sd_models.move_model(sd_model, device=devices.cpu)
elif shared.opts.diffusers_offload_mode == "balanced":
sd_model = sd_models.apply_balanced_offload(sd_model, force=True) # dispatched modules hold meta tensors backed by the offload map; rebuild them real on cpu with hooks intact before touching weights
modules = {}
components = include if len(include) > 0 else ['text_encoder', 'text_encoder_2', 'text_encoder_3', 'unet', 'transformer', 'llm_adapter']