huggingface / huggingface/diffusers

HiDream auxiliary loss for MoE experts not tied to computation graph

オープン
#11,301 コメント 2 件 リアクション 0 件 担当者 0 名 GitHub で見る
bug stale
主要言語
Python
スター
34.5k
フォーク
7.3k
平均マージ
3日 3時間
マージ済み PR(30日)
91

説明

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

コントリビューションガイド

コントリビューションガイドを開く

調査の方向性

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.

索引モデルが issue の本文から書いたものです。

評価

技術スタック
python, pytorch
領域
machine-learning
issue の種類
バグ
難易度
4/5
見積もり時間
3〜5日
活発さ
停滞
明瞭さ
説明が足りない
初心者へのやさしさ
35/100

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。