Comfy-Org / Comfy-Org/ComfyUI

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

Open
#14,046 1 comment 1 reaction 0 assignees View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.