GPU Requirements for Model Training

[yanliang3612]

Hi, thanks for your great work!
I have a question regarding the hardware requirements for training the model:

1. What types of GPUs are recommended or required for training this model?
   (e.g., NVIDIA A100, V100, RTX 3090, etc.)

2. How many GPUs are typically needed to successfully train the model as reported in your paper or repo?

Any guidance or details about the GPU configuration (e.g., memory per GPU, distributed setup) would be greatly appreciated!

Thank you!

[lantunes]

Hi,

Thank you for your interest in this project. The large and small models were trained using a single A100 GPU with 80 GB of memory. On an A100 GPU, the small model requires 3 seconds per training iteration. Therefore, 100,000 small model training iterations requires 83.3 hours, or approximately 3.5 days. The large model requires 16 seconds per training iteration on an A100 GPU. Therefore, 48,000 large model training iterations requires 213.3 hours, or approximately 8.9 days. While training times could be decreased with distributed training, the current codebase does not support distributed training. Adding this support is currently a work in progress, here.

For inference, a T4 or A10G GPU can be used.