How long should training take?

[wigging]

I know it depends on the hardware specs like number of GPUs and CPUs and amount of memory, but how long should it take to train a model? I tried training with MP-20 using two Nvidia A100s where each one has 80 GB of memory. I canceled the training after 46 hours and 408 epochs. I used the command shown below for the training. At this rate, it seems like it would take 2-3 days to train this model. Is that a reasonable amount of time or am I doing something wrong with the configuration settings? How long did it take to train the provided models in the repository and what was the hardware specs used to generate the pre-trained models?

mattergen-train \
    data_module=mp_20 \
    ~trainer.logger \
    trainer.devices=2 \
    trainer.accumulate_grad_batches=2 \
    data_module.num_workers.train=0 \
    data_module.num_workers.val=0 \
    data_module.num_workers.test=0

Edit:

In the Supplementary Information for the paper I found the following:

D.1.1 Hyperparameters for training base model
The base unconditional model was trained for 1.74 million steps with a batch size of
64 per GPU over eight A100 GPUs using the Adam optimizer [18]. The learning rate
was initialized at 0.0001 and was decayed using the ReduceLROnPlateau scheduler
with decay factor 0.6, patience 100 and minimum learning rate 10↓6. For the training
loss we use ϱcoord = 0.1 and ϱcell = ϱtypes = 1, as well as the recommended value
ϱCE = 0.01 for the D3PM cross-entropy loss.

So how long did it take to train with eight A100s?

[danielzuegner]

Hi @wigging,

Training took around 5 hours (one epoch took a bit less than 20 seconds typically).

[danielzuegner]

It trained for 900 epochs, as is specified in the config. Wandb is short-hand for Weights & Biases, the system we used to log the training run & the related metadata.

[wigging]

Is this the command you used to run the training? Other than devices, are there other config settings you changed or is everything using the default values?

mattergen-train \
    data_module=mp_20 \
    ~trainer.logger \
    trainer.devices=8

[danielzuegner]

Yes, so I didn't remove the logger in order to use the Wandb logger but otherwise that's what I was using. Have you checked whether dropping trainer.accumulate_grad_batches or setting it to 1 leads to OOM? If not, it will for sure speed up the training if your GPU RAM is sufficient.

[wigging]

@danielzuegner Can you update the training section in the README with more information about the hardware specs and training time for mp-20 and alex mp-20? The information you shared in this discussion would be useful to have in the README.

[rongaoli]

Took me ~7 days to train 2200 epochs on Alex_MP_20.
32 x NVIDIA V100, 5 min per epoch, batch_size=512.

FYI :)