set_attr_param re-wraps Parameter subclasses, breaking bnb Params4bit (NF4/FP4) at partially_unload

[j2gg0s]

Summary

Loading any bitsandbytes-quantized model (NF4/FP4) with a custom-OPS loader such as CheckpointLoaderNF4 (from comfyanonymous/ComfyUI_bitsandbytes_NF4) crashes at KSampler with:

RuntimeError: Creating a Parameter from an instance of type Params4bit requires that
detach() returns an instance of the same type, but return type Tensor was found
instead. To use the type as a Parameter, please correct the detach() semantics
defined by its __torch_dispatch__() implementation.

The model loads successfully; the crash fires the first time the sampler triggers model_management.load_models_gpu → partially_load → partially_unload.

Reproduction

1. Install comfyanonymous/ComfyUI_bitsandbytes_NF4 custom node.
2. Download lllyasviel/flux1-dev-bnb-nf4's flux1-dev-bnb-nf4-v2.safetensors into models/checkpoints/.
3. Build a minimal txt2img workflow: CheckpointLoaderNF4 → CLIPTextEncode ×2 → KSamplerAdvanced → VAEDecode → SaveImage.
4. Queue the prompt. PyTorch ≥ 2.6 + ComfyUI ≥ v0.3.x (any version with dynamic-vbar / partially_unload) reproduces.

Traceback (abridged, full chain)

File "comfy/sample.py", sample
File "comfy/samplers.py", outer_sample
File "comfy/sampler_helpers.py", prepare_sampling
File "comfy/model_management.py", load_models_gpu
File "comfy/model_management.py", LoadedModel.model_load
File "comfy/model_management.py", model_use_more_vram
File "comfy/model_patcher.py", partially_load
File "comfy/model_patcher.py:1194", detach()
File "comfy/model_patcher.py:1843", unpatch_model(unpatch_weights=True)
File "comfy/model_patcher.py:1804", partially_unload(None, 1e32)
File "comfy/utils.py:1005", set_attr_param(self.model, key, bk.weight)
        return set_attr(obj, attr, torch.nn.Parameter(value, requires_grad=False))
File "torch/nn/parameter.py:62", Parameter.__new__
        raise RuntimeError("Creating a Parameter from an instance of type ...")

Root Cause

set_attr_param (comfy/utils.py:1000) unconditionally wraps value in torch.nn.Parameter(). When value is already a Parameter subclass with a non-trivial detach semantic — notably bitsandbytes.nn.Params4bit / ForgeParams4bit — PyTorch's Parameter.__new__ strict check rejects the construction.

The bnb side cannot reasonably fix this from their end:

• Params4bit.__new__ constructs via torch.Tensor._make_subclass(cls, data, ...), never round-trips through detach(), so detach() is not part of bnb's API contract.
• Params4bit.__torch_function__ only intercepts torch.chunk / torch.split (added in Fix Params4bit tensor subclass handling bitsandbytes-foundation/bitsandbytes#1719 for FSDP2). Extending it to detach/clone/copy_ is push-back-prone given upstream's preference to minimize __torch_function__ for torch.compile reasons.
• Forge's ForgeParams4bit also does not override detach.

The natural fix is in ComfyUI: skip the rewrap when value is already a Parameter.

Proposed Fix

def set_attr_param(obj, attr, value):
    if (not torch.is_inference_mode_enabled()) and value.is_inference():
        value = value.clone()
    # Already a Parameter (incl. bnb Params4bit / ForgeParams4bit subclasses).
    # Re-wrapping triggers Parameter.__new__'s detach-roundtrip check which
    # fails for subclasses whose detach() returns a plain Tensor.
    if isinstance(value, torch.nn.Parameter):
        return set_attr(obj, attr, value)
    return set_attr(obj, attr, torch.nn.Parameter(value, requires_grad=False))

Behavior is unchanged for plain tensors (the dominant path), and the inference-mode clone above still applies before the early-return because a Parameter wrapping an inference-mode tensor would already have hit is_inference() semantics during its own construction. The branch is a no-op rewrap elision.

I have applied this patch locally and confirmed the previously crashing partially_unload(...) step completes without the RuntimeError. Happy to follow with a PR (with the same patch and a brief test) if maintainers agree on the direction.

Downstream reports observed in the wild

These all show the same Params4bit × Parameter.__new__ traceback but were filed on consumer repos rather than upstream — they are surfacing the same root cause:

• T2v and i2v workflows fail after latest ComfyUI_windows_portable update Lightricks/ComfyUI-LTXVideo#437 — T2v/i2v workflows fail after ComfyUI update (Gemma 4-bit prompt enhancer path)
• 🅛🅣🅧 Gemma 3 Prompt Enhancer does not work with 4-bit gemma Lightricks/ComfyUI-LTXVideo#425 — Gemma-3-12b-it-qat-bnb-4bit enhancer broken
• RuntimeError: Creating a Parameter from an instance of type Params4bit requires that detach() returns an instance of the same type, but return type Tensor was found instead. To use the type as a Parameter, please correct the detach() semantics defined by its __torch_dispatch__() implementation. baichuan-inc/Baichuan2#418 — lmdeploy + Baichuan2 int4 (cross-ecosystem confirmation that this is a Parameter.__new__ × bnb subclass issue, not a ComfyUI-only bug)

Related upstream activity

• PR Fix OOM regression in _apply() for quantized models during inference #13372 (merged) — _apply() OOM fix for FP8/quantized models. Same family of small fixes around quantized weight handling.
• PR Fix QuantizedTensor weight restore failure causing progressive drift across generations #12978 (open) — set_attr_param semantics on fp8_scaled QuantizedTensor with LoRA. Different bug, same function; flagged for context.

[coderabbitai]

🔗 Related PRs

#13699 - fix: Hunyuan3D 2.1 batch size crashes in attention and forward pass [merged]
#13735 - Reduce LTX2.3 peak VRAM when guide_mask is in use (CORE-166) [merged]
#13801 - fix: Fix bug when mask not on same device (CORE-181) [merged]
#13959 - implement/fix dynamic clip saving [merged]
#14031 - Fix reshaping lora application [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:

issue_enrichment:
  auto_enrich:
    enabled: false

💬 Have feedback or questions? Drop into our discord!