Comfy-Org / Comfy-Org/ComfyUI

DoRA `weight_decompose` broadcasts incorrectly for non-square weights

Open Beginner friendly
#12,938 2 comments 0 reactions 0 assignees View on GitHub
Potential Bug
Dominant language
Python
Stars
133k
Forks
15.7k
Avg merge
1d 6h
Merged PRs (30d)
155

Description

### Custom Node Testing

- [x] I have tried disabling custom nodes and the issue persists (see [how to disable custom nodes](https://docs.comfy.org/troubleshooting/custom-node-issues#step-1%3A-test-with-all-custom-nodes-disabled) if you need help)

### Expected Behavior

`weight_decompose()` in `comfy/weight_adapter/base.py` produces wrong results (or crashes) when applying DoRA to non-square weight matrices (e.g. MLP layers where `d_ff != d_model`).

### Actual Behavior

## Root Cause

On [line 301](https://github.com/Comfy-Org/ComfyUI/blob/master/comfy/weight_adapter/base.py#L301):

```python
weight_calc *= (dora_scale / weight_norm).type(weight.dtype)
```

`dora_scale` is 1D `[N]` and `weight_norm` is 2D `[N, 1]`. PyTorch's broadcasting aligns dimensions from the right, so:

- `dora_scale` `[N]` becomes `[1, N]`
- `weight_norm` stays `[N, 1]`
- Result: `[1, N] / [N, 1]` = `[N, N]` (outer-product-shaped)

This `[N, N]` tensor then fails to multiply with `weight_calc` of shape `[N, M]` when `N != M`.

### Steps to Reproduce

Train a DoRA (LoRA with `use_dora=True`) that targets MLP layers (any non-square `nn.Linear`), then load it via the `LoraLoader` node.

Example error:

```
ERROR lora diffusion_model.blocks.0.mlp.layer1.weight The size of tensor a (2048) must match the size of tensor b (8192) at non-singleton dimension 1
```

This affects any model architecture with non-square MLP weights (Anima, Cosmos Predict2, etc). Square weights (most attention Q/K/V/O projections) are silently unaffected since `[N, N] * [N, N]` doesn't raise a shape error.

### Debug Logs

```powershell
...
ERROR lora diffusion_model.blocks.1.mlp.layer1.weight The size of tensor a (2048) must match
the size of tensor b (8192) at non-singleton dimension 1
ERROR lora diffusion_model.blocks.0.mlp.layer2.weight The size of tensor a (8192) must match
the size of tensor b (2048) at non-singleton dimension 1
ERROR lora diffusion_model.blocks.0.mlp.layer1.weight The size of tensor a (2048) must match
the size of tensor b (8192) at non-singleton dimension 1
...
```

### Other

## Fix

Reshape `dora_scale` to match `weight_norm`'s dimensionality before dividing:

```python
if wd_on_output_axis:
dora_scale = dora_scale.reshape(weight.shape[0], *[1] * (weight.dim() - 1))
else:
dora_scale = dora_scale.reshape(*[1] * (weight.dim() - 1), weight.shape[-1])
weight_calc *= (dora_scale / weight_norm).type(weight.dtype)
```

This ensures `dora_scale / weight_norm` produces `[N, 1]` (element-wise per output neuron) instead of `[N, N]`.

Contributor guide

Open the contributing guide

Research direction

Start in comfy/weight_adapter/base.py at weight_decompose(), especially line 301 and the wd_on_output_axis handling. Reproduce through the LoraLoader node with a DoRA targeting a non-square MLP weight, then verify the change works for both output- and input-axis weights without breaking square projections. Done means non-square MLP layers load without the reported tensor-size error.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
2/5
Estimated time
1-3 hours
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.