set_attr_param re-wraps Parameter subclasses, breaking bnb Params4bit (NF4/FP4) at partially_unload
- Dominant language
- Python
- Stars
- 133k
- Forks
- 15.7k
- Avg merge
- 1d 7h
- Merged PRs (30d)
- 158
Description
## 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 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
```python
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:
- Lightricks/ComfyUI-LTXVideo#437 — T2v/i2v workflows fail after ComfyUI update (Gemma 4-bit prompt enhancer path)
- Lightricks/ComfyUI-LTXVideo#425 — Gemma-3-12b-it-qat-bnb-4bit enhancer broken
- 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 #13372 (merged) — `_apply()` OOM fix for FP8/quantized models. Same family of small fixes around quantized weight handling.
- PR #12978 (open) — `set_attr_param` semantics on fp8_scaled `QuantizedTensor` with LoRA. Different bug, same function; flagged for context.
Contributor guide
Assessment
This issue has not been assessed yet.