GPU Memory Requirement for Stage2

[LaVieEnRoseSMZ]

Hi, Thank you for your great work! I’m wondering about the GPU memory requirements for the second stage. I’m encountering an OOM issue when training at 1080p resolution on an 80GB A100. Do you have any recommendations to optimize memory usage?

Also, I’m looking forward to any updates on the training code.

[jshilong]

Thank you for your interest in our work.
Training with a single GPU is not sufficient, as the parameters need to be distributed across multiple GPUs during training to save memory usage.

I will organize and release the training code soon; currently, it is on another branch , but you are welcome to use it as a reference for now.
#12

[LaVieEnRoseSMZ]

Hi, thanks for your prompt response! I understand the idea of distributing different models across multiple GPUs. In this case, would multiple 80GB A100 still be feasible, or is H100 necessary for training? I assume running 1080p resolution on the transformer — even after offloading other parts to different GPUs — still demands significant memory.

I’ve already incorporated the latest pull request into my implementation. I really appreciate the continued work and look forward to the upcoming release!

[jshilong]

There should be no difference in training between the A100 and H100, as both have 80GB of memory.