huggingface / huggingface/diffusers

RMSNorm's weight not registered as submodules when initializing

Aperta
#11,938 3 commenti 0 reazioni 0 assegnatari Vedi su GitHub
bug stale
Lingua principale
Python
Stelle
34.5k
Fork
7.3k
Merge medio
3g 3h
PR unite (30g)
91

Descrizione

### 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_

Guida per i contributori

Apri la guida per i contributori

Direzione di ricerca

Inizia in models/normalization.py, all'interno di RMSNorm, e riproduci il problema tramite FluxPipeline.from_pretrained usando la configurazione mostrata. Esamina come l'inizializzazione sul meta-device e la registrazione dei parametri interagiscono con il modulo stampato, quindi verifica che il peso sia visibile a FSDP e che la riproduzione di Flux continui a essere eseguita correttamente.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Valutazione

Stack tecnologico
python, pytorch
Ambito
distributed-systems, machine-learning
Tipo di issue
Bug
Difficoltà
4/5
Tempo stimato
3-5 giorni
Stato di attività
Ferma
Chiarezza
Abbastanza chiara
Idoneità per principianti
35/100

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.