VRAM usage with Visison-TinyEngram?

[2211427551]

Hello, thanks for this excellent work!
I'm trying to adopt Vision-TinyEngram to ERNIE-Image, but my approach encountering severe VRAM usage issue. So I would like to have a VRAM usage baseline with Vision-TinyEngram.

[Alichua]

Hi @2211427551

Thanks for raising this.

Currently we do not have a measured peak VRAM number for TinyEngram-Vision. (We did not treat it as a separately worth-reporting metric before, sry about that.)

However, the thing we can confirm is that: SD3.5 training already requires quite huge VRAM, and ERNIE-Image is expected to be even more memory demanding.

TinyEngram module itself is very small. However, the frozen heavy backbone still participates in the forward/backward computation. Freezing the UNet/DiT does remove weight gradients and optimizer states, but it does not make the backbone free.

For your question about VRAM, I can offer an approximate static residency lower bound, using public Hugging Face checkpoint file sizes, as a basic reference:

Pipeline	Static lower-bound
SD1.5	~1.99 GiB
SD3.5 Medium	~15.15 GiB
SD3.5 Large	~25.73 GiB
ERNIE-Image, use_pe=False	~22.29 GiB
ERNIE-Image, use_pe=True	~29.43 GiB

Lower bound is sad because it only takes resident model weights into consideration. They exclude activations, CUDA workspace, allocator overhead, ...so the actual training peak VRAM must be higher.

So why your ERNIE-Image so huge?

From the model recipe, we can see that the text side is already heavy: the Mistral3 text encoder is about 7 GiB, and the Prompt Enhancer, if enabled and kept on GPU, adds another ~7 GiB. And of course: the 8B DiT backbone is about 15 GiB. (Bong!) And this is only the lower bound.

Our training was performed on an H20 machine with multiple devices and relatively redundant VRAM, so we did not record a peak VRAM baseline at that time. However, if you are trying to perform training on consumer-level devices, even 32GB-class cards such as RTX 5090 may still be quite tricky, especially with use_pe=True(i.e. enabling Prompt Enhancer).

Possible directions would be:

• disabling or offloading the Prompt Enhancer
• caching frozen/base states where possible
• using memory-efficient attention, enabling gradient checkpointing
• reducing resolution/batch size
• using model offloading.

[2211427551]

I think disabling the Prompt Enhancer is reasonable, since in practise, I usually use a API like Gemini 3 Flash, so Prompt Enhancer can be ignored. And gradient checkpointing is a must for me.

[Alichua]

I think disabling the Prompt Enhancer is reasonable, since in practise, I usually use a API like Gemini 3 Flash, so Prompt Enhancer can be ignored. And gradient checkpointing is a must for me.

Indeed, and if possible, please share your device specs and real experiment results, not only VRAM usage but also visual performance. I am personally very interested in ERNIE-Image training and would love to see how TinyEngram-Vision behaves on it.

[2211427551]

Seems like Vision-TinyEngram, with my implementation, the convergence speed of ERNIE-Image are quite slow, batch size = 2 with RTX Pro 6000, loss decreased from 1.1 to 0.8 uses 1200 steps. My dataset, with repeat, has 210.

[Alichua]

Seems like Vision-TinyEngram, with my implementation, the convergence speed of ERNIE-Image are quite slow, batch size = 2 with RTX Pro 6000, loss decreased from 1.1 to 0.8 uses 1200 steps. My dataset, with repeat, has 210.

Thanks for sharing the numbers. This setup is quite different from the original TinyEngram Vision recipe.

In our SD3.5 experiment, the target binding was very narrow: 5 images and one synthetic trigger name, "Aldric Vortex-9 CyberNebula". With max n-gram size 7 and batch size 1, we got useful convergence around 3500 steps, which is roughly 700 exposures per image. If your repeated dataset length is 210 and batch size is 2, then 1200 steps is only about 11 epochs. So the slower loss drop may simply mean each training item has been seen much less often.

Another important difference is concept entropy. TinyEngram Vision only trains the trigger-specific memory vectors and scale parameters, while keeping the image backbone and text encoders frozen. This works best for a compact subject/concept that can be bound to one trigger phrase(e.g. One person with specific gear and iconic appearance.). Hence, if your ERNIE-Image setup is learning a broader style or a more diverse concept, such as costume design or a specific artist style, the current narrow engram allocation may not have enough capacity.

Before tuning heavily, I would suggest a few sanity checks:

1. Try to overfit a tiny subset first, e.g. 5-10 images with the same trigger phrase in every caption. If it cannot overfit that, it is more likely an implementation/injection issue than a data-size issue.
2. Print the tokenized trigger and the number of registered n-grams for the ERNIE-Image tokenizer. Also count actual n-gram matches per batch. If the trigger does not match exactly, the engram parameters may receive little or no gradient(This may sounds funny but it is a potential risk since you mentioned forming prompts via Gemini.).
3. If the target is a style rather than a specific subject, consider increasing trigger capacity, using multiple style-specific trigger slots, or using general engram as we implemented in Qwen experiments.

Looking forward to seeing how it goes! If you can share more logs or implementation details later, we’d be happy to help take another look.