Repository navigation
Way to skip self.model.to(device_to) inside ModelPatcher.unpatch_model #14142
Description
Feature Idea
ModelPatcher.unpatch_model(unpatch_weights=True) ends its weight-unpatch block with self.model.to(device_to) (master comfy/model_patcher.py around the current_weight_patches_uuid = None / backup.clear() lines). That walk uses nn.Module.to which iterates every parameter, ignoring comfy_cast_weights = True on custom layers.
For custom ops that intentionally keep weights on offload_device and only cast on the fly through comfy_cast_weights, the walk hits parameters that are not safe to .to(...) (mmap-backed, lazily quantized, etc.) and faults. ComfyUI-GGUF hits this with Windows EXCEPTION_ACCESS_VIOLATION on mmap-backed quantized tensors (city96/ComfyUI-GGUF#444).
The only existing escape hatch is unpatch_weights=False, but that block also runs unpatch_hooks, unpin_all_weights, lowvram cleanup, non-quantized backup restoration, model_loaded_weight_memory = 0, model_offload_buffer_memory = 0, and comfy_patched_weights deletion. Custom patchers that want to keep all of that but skip the device walk end up mirroring the whole block by hand. city96/ComfyUI-GGUF#445 is the latest iteration of that workaround there.
Two shapes that would work:
- Add
unpatch_device: bool = Truetounpatch_model, gating just theself.model.to(device_to)line. Defaults preserve behavior. - Replace
self.model.to(device_to)with a per-module walk that skips modules withcomfy_cast_weights = True(their weights are managed by the op layer, not byModelPatcher).
Either would let custom patchers drop their unpatch_model overrides entirely.
Existing Solutions
- ComfyUI-GGUF currently maintains a multi-iteration downstream workaround in
nodes.py:GGUFModelPatcher.unpatch_model(see history at6dbb4ba->717a0e1-> Option for nodes to output any signals they have #355 -> Added a queue for the queue action #357 -> Allow reroute to have in/out in different orientation #445). The maintainer flagged this as needing an upstream fix in717a0e1's commit message in September 2024. - Other quantization custom nodes (e.g. bitsandbytes-based loaders) avoid the issue because their weights survive
.to, but they would still benefit from the cleaner contract.
Other
The nn.Module.to walk is also redundant on the partially_load call path (comfy/model_patcher.py partially_load calls unpatch_model(self.offload_device, ...) and then immediately re-loads), so option (2) would also avoid wasted work on that hot path. The detach path is the only caller that genuinely needs the weights on offload_device after unpatch.
Activity
coderabbitai commentedon May 27, 2026
🔗 Related PRs
#13959 - implement/fix dynamic clip saving [merged]
#14031 - Fix reshaping lora application [merged]
#14108 - multigpu: use unet_manual_cast for SelectModelDevice compute dtype [merged]
🧪 Issue enrichment is currently in open beta.
You can configure auto-planning by selecting labels in the issue_enrichment configuration.
To disable automatic issue enrichment, add the following to your
.coderabbit.yaml:💬 Have feedback or questions? Drop into our discord!
nodes: also zero memory counters and run hook/pin cleanup in unpatch …
woodenriver05 commentedon Aug 20, 2026
Follow-up correction to option (2), based on a full-load/full-unload measurement.
For GGUF,
comfy_cast_weights = Truedoes not imply that a module has no CUDA-resident weight storage. After a full load, the quantizedGGMLLayerweights can be physically resident on CUDA. An unconditional module walk that skips everycomfy_cast_weightsmodule would therefore fail the postcondition of a real detach/full unload.Single-host A/B result on ComfyUI 0.33.1 + PyTorch 2.13.0+cu130:
19039.68 MB6861.54 MBavailable; it loaded partially with16766.75 MBoffloaded25918.22 MBavailable and a complete23276.97 MBloadThe upstream walk did not reproduce the historical Windows access violation on this current stack. This is one current-stack sample, not evidence that the older crash cannot occur.
So option (2) as currently written is not safe for the
detachpath. Option (1) also needs call-site semantics: disabling the move can be appropriate for thepartially_loadunload/reload hot path, but not for a full detach unless a custom-op hook still moves/releases the storage safely.A useful upstream contract would distinguish those two intents, or let a custom patcher provide a safe offload operation while preserving the requirement that full detach ends with the model on
offload_device. This would avoid both the mmap-sensitive generic walk and the physical-VRAM/accounting divergence observed with the downstream mirrored block.Booyaka101 commentedon Aug 20, 2026
Follow-up from the downstream side. On ComfyUI v0.33.1, torch 2.11.0+cu130, Windows 11 I can't reproduce the original access violation either, five unpatch scenario shapes all clean. Probing the mechanism with the real gguf loader: nn.Module.to over still-mmap'd tensors doesn't fault here (same-device is identity, cross-device is a plain read), but any in-place write into one is an instant AV, since torch.from_numpy on the read-only mapping yields a tensor torch treats as writable.
city96/ComfyUI-GGUF#445 now handles this without needing anything upstream: super().unpatch_model(device_to=None, ...) so the walk is skipped, then an explicit move of only the tensors that report a different device. That keeps VRAM accounting and full unload correct, measured 10512MB back to 0 on detach where the previous shape stranded it. @woodenriver05's correction stands: unconditionally skipping the walk for comfy_cast_weights modules would not be unload-correct upstream either. If anything is worth changing here it would be moving per-tensor by residency inside the walk, but from our side this can be closed if you'd rather keep unpatch_model as is.