DoRA `weight_decompose` broadcasts incorrectly for non-square weights
- 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
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