huggingface / huggingface/diffusers
Ideogram4: allow the text conditioning projection to be hoisted out of the denoising loop
- Vorherrschende Sprache
- Python
- Sterne
- 34.5k
- Forks
- 7.3k
- Ø Merge
- 3 T. 3 Std.
- Gemergte PRs (30 T.)
- 91
Beschreibung
I was profiling Ideogram 4 inference on a 24 GB card and found that most of the text conditioning
work repeats on every denoising step.
`Ideogram4Transformer2DModel.forward` normalizes and projects `encoder_hidden_states` each step:
```python
encoder_hidden_states = self.llm_cond_norm(encoder_hidden_states)
encoder_hidden_states = self.llm_cond_proj(encoder_hidden_states) * llm_token_mask
```
But `encoder_hidden_states` is fixed for the whole generation, and the projection runs at the full
packed width before masking most of the result away. The tensor is 53,248 rows wide and 512 carry
text, so at 2048² that is 16,384 of 16,896 rows projected and discarded. Classifier-free guidance
adds a second branch that is zero throughout. Under a 4 bit `llm_cond_proj` that is a wide NF4 GEMM
per step over rows that cannot affect the output.
Hoisting it out of the loop, with bitsandbytes NF4 against `7685bffe8904`:
| | baseline | hoisted | |
|---|---|---|---|
| peak allocation, 2048² (A100) | 26.835333 GiB | 17.378107 GiB | **−9.457 GiB** |
| peak allocation, 1024² (L40S) | 18.189322 GiB | 16.053430 GiB | −2.136 GiB |
Byte identical output, same SHA-256 on latents, pixels, and OCR, across 18 generations over three
prompt lengths and three seeds. At 2048² the unmodified path OOMs on an A10's 24 GiB. The hoisted
path completes.
## Why It Needs an API Change
I tried to avoid one. The caller builds the `(1, 16896, 53248)` tensor before `forward` is entered,
so the transformer cannot unspend it. And skipping the masked rows internally changes the GEMM's M,
which changes which NF4 kernel runs: projecting only the 512 real rows changes pixels. I measured
that. Preserving the row count is what keeps it exact.
## Proposed Shape
One public method and one keyword only flag, both opt-in, default path untouched:
```python
projected = transformer.project_text_conditioning(
text_features[:, indicator[0] == LLM_TOKEN_INDICATOR],
# Pads back to the row count the default path would have used, so a quantized llm_cond_proj
# selects the same kernel and the result stays bitwise identical.
minimum_projection_rows=text_features.shape[0] * text_features.shape[1],
)
# then per step: transformer(..., encoder_hidden_states=projected, encoder_hidden_states_projected=True)
```
An unconditional branch with no text positions can pass `encoder_hidden_states=None` instead of a
wide all zero tensor.
Would you take a change of this shape? Implemented against `7685bffe8904` with tests and docs,
`make style` / `make quality` / `make fix-copies` clean, and the model's tests passing on CPU:
Two things I would rather hear your view on. The opt-in path is not `torch.compile(fullgraph=True)`
compatible, because rejecting a ragged batch means reading `indicator`. The default path still
compiles. And `minimum_projection_rows` is an odd argument for a public signature, but without it
the change is no longer exact, which removes the reason to make it.
The code was written by an AI agent, per the agentic contribution guidelines. I have run
`self-review` and can post the notes.
@yiyixuxu @DN6
Beitragsleitfaden
Rechercherichtung
Start at Ideogram4Transformer2DModel.forward and trace how encoder_hidden_states is built and consumed across the denoising loop. Review the proposed project_text_conditioning and encoder_hidden_states_projected entry points, then run the model's CPU tests and the mentioned quality checks. Done means the opt-in path preserves default behavior and the reported memory and output equivalence are covered by tests and docs.
Vom Indexierungsmodell aus dem Issue-Text verfasst.
Bewertung
- Tech-Stack
- python, pytorch
- Bereich
- api, machine-learning, performance
- Issue-Typ
- Feature
- Schwierigkeit
- 5/5
- Geschätzter Aufwand
- Über eine Woche
- Aktivitätsstatus
- Ruhig
- Klarheit
- Größtenteils klar
- Anfängerfreundlichkeit
- 45/100