huggingface / huggingface/diffusers

HiDream auxiliary loss for MoE experts not tied to computation graph

Offen
#11,301 2 Kommentare 0 Reaktionen 0 zugewiesene Personen Auf GitHub ansehen
bug stale
Vorherrschende Sprache
Python
Sterne
34.5k
Forks
7.3k
Ø Merge
3 T. 3 Std.
Gemergte PRs (30 T.)
91

Beschreibung

### Describe the bug

The MOEFeedForward for HiDream has the auxiliary loss commented out in the upstream prototype code.

Additionally, the MoEGate has a memory leak in it.

### Reproduction

- Train HiDream
- Observe outOfMemory on backward pass
- Resolve MoEGate OOM by implementing gradient checkpointing
- Observe extraordinarily high loss values

### Logs

```shell
2025-04-12 10:03:59,851 [INFO] cls: , settings: {'betas': (0.9, 0.999), 'weight_decay': 0.01, 'eps': 1e-06}
2025-04-12 10:03:59,855 [INFO] Optimizer arguments={'lr': 4e-05, 'betas': (0.9, 0.999), 'weight_decay': 0.01, 'eps': 1e-06}
2025-04-12 10:03:59,855 [INFO] Loading constant learning rate scheduler with 100 warmup steps
2025-04-12 10:03:59,855 [INFO] Using generic 'constant' learning rate scheduler.
2025-04-12 10:03:59,857 [INFO] Preparing models..
2025-04-12 10:03:59,858 [INFO] Loading our accelerator...
2025-04-12 10:03:59,875 [INFO] Resuming from checkpoint checkpoint-8000
2025-04-12 10:04:00,033 [INFO] Previous checkpoint had 0 exhausted buckets.
2025-04-12 10:04:00,034 [INFO] Previous checkpoint was on epoch 471.
2025-04-12 10:04:00,034 [INFO] Previous checkpoint had 10 seen images.
2025-04-12 10:04:00,034 [INFO] Resuming from global_step 8000.
2025-04-12 10:04:00,034 [INFO]
(Rank: 0) -> Number of seen images: 10
(Rank: 0) -> Number of unseen images: 7
(Rank: 0) -> Current Bucket: None
(Rank: 0) -> 1 Buckets: ['1.0']
(Rank: 0) -> 0 Exhausted Buckets: []
2025-04-12 10:04:00,093 [INFO]
***** Running training *****
- Num batches = 17
- Num Epochs = 589
- Current Epoch = 471
- Total train batch size (w. parallel, distributed & accumulation) = 1
- Instantaneous batch size per device = 1
- Gradient Accumulation steps = 1
- Total optimization steps = 10000
- Steps completed: 8000
- Total optimization steps remaining = 2000
Epoch 478/589, Steps: 81%|████████████▏ | 8114/10000 [03:30<57:25, 1.83s/it, grad_absmax=0.00149, lr=4e-5, step_loss=1.13]
```

### System Info

Diffusers git main

### Who can help?

_No response_

Beitragsleitfaden

Beitragsleitfaden öffnen

Rechercherichtung

Start with the HiDream MOEFeedForward and MoEGate components referenced in the report, then reproduce HiDream training and observe the backward-pass OOM and unusually high loss. Trace whether the auxiliary loss is connected to the computation graph and investigate gradient checkpointing for the gate. Done means training avoids the memory leak and the auxiliary loss contributes correctly without the reported extreme values.

Vom Indexierungsmodell aus dem Issue-Text verfasst.

Bewertung

Tech-Stack
python, pytorch
Bereich
machine-learning
Issue-Typ
Bug
Schwierigkeit
4/5
Geschätzter Aufwand
3-5 Tage
Aktivitätsstatus
Veraltet
Klarheit
Muss geklärt werden
Anfängerfreundlichkeit
35/100

Neue Issues direkt in Ihr Postfach

Eine kurze Übersicht über anfängerfreundliche GitHub-Issues.