huggingface / huggingface/diffusers

RMSNorm's weight not registered as submodules when initializing

Open
#11,938 3 comments 0 reactions 0 assignees View on GitHub
bug stale
Dominant language
Python
Stars
34.5k
Forks
7.3k
Avg merge
3d 3h
Merged PRs (30d)
91

Description

### Describe the bug

RMSNorm() look like this, which is in models/normalization.py
```python
class RMSNorm(nn.Module):
def __init__(self, dim, eps: float, elementwise_affine: bool = True, bias: bool = False):
super().__init__()

self.eps = eps
self.elementwise_affine = elementwise_affine

if isinstance(dim, numbers.Integral):
dim = (dim,)

self.dim = torch.Size(dim)

self.weight = None
self.bias = None

if elementwise_affine:
self.weight = nn.Parameter(torch.ones(dim))
if bias:
self.bias = nn.Parameter(torch.zeros(dim))
print(self.weight)
print(self)
```
This module is used in Flux.1-dev, and maybe other models I don't know.
When I print the model, it looks like `(norm_q): RMSNorm()`, which means the self.weight is not registered as submodules.
I've checked that self.weight does exist. Actually, when initialized from Attention, elementwise_affine and bias are never modified, so they are using the default values.
The output looks like this:
```
tensor(..., device='meta', size=(128,), requires_grad=True)
RMSNorm()
```
LLM told me it's because the self.weight is assigned in a conditional branch, AND it's initialized on meta device with from_pretrained(). But it's not really the cause. I removed the branch but the bug persisted.
This bug might not be fatal in most cases, but it prevents me from using FSDP with Flux.

### Reproduction

```python
import torch
from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained(
'black-forest-labs/FLUX.1-dev', torch_dtype=torch.bfloat16
).to('cuda')

prompt = "a cat holding a paper with word prompt on it"
image = pipe(
prompt,
height=1024,
width=1024,
).images[0]
image.save("flux.png")
'''

### Logs

```shell

```

### System Info

diffusers["torch"]==0.32.2

### Who can help?

_No response_

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.