Data loader bottlenecking training

[JakobLindscheid]

Hi,
Thank you for publishing the pretraining and finetuning scripts! They are really helpful.
For a university project, we are trying to reproduce the results from the paper. However, running the pretrain script, we observe very slow training speeds (~1 minute per epoch) on our hardware.
Running the pytorch profiler for 16 training batches, we see the following:

FIT Profiler Report (relevant lines)

Action	Mean duration (s)	Num calls	Total time (s)	Percentage %
Total	-	1397	99.734	100 %
run_training_epoch	91.657	1	91.657	91.901
[_TrainingEpochLoop].train_dataloader_next	5.2931	16	84.689	84.915
[_EvaluationLoop].val_next	0.246	19	4.674	4.6865
[LightningModule]LagLlamaLightningModule.optimizer_step	0.10731	16	1.717	1.7216
run_training_batch	0.10731	16	1.717	1.7216
[Strategy]SingleDeviceStrategy.training_step	0.091875	16	1.47	1.4739
[Strategy]SingleDeviceStrategy.validation_step	0.044368	19	0.843	0.84525
[Strategy]SingleDeviceStrategy.backward	0.0135	16	0.216	0.21658
[Callback]ModelCheckpoint{'monitor': None, 'mode': 'min', 'every_n_train_steps': 0, 'every_n_epochs': 1, 'train_time_interval': None}.on_train_epoch_end	0.141	1	0.141	0.14138
[Callback]ModelCheckpoint{'monitor': 'val_loss', 'mode': 'min', 'every_n_train_steps': 0, 'every_n_epochs': 1, 'train_time_interval': None}.on_train_epoch_end	0.093	1	0.093	0.093248
[LightningModule]LagLlamaLightningModule.transfer_batch_to_device	0.0022286	35	0.078	0.078208
[Strategy]SingleDeviceStrategy.batch_to_device	0.0022286	35	0.078	0.078208
[LightningModule]LagLlamaLightningModule.on_validation_model_train	0.008	2	0.016	0.016043
[Callback]ModelSummary.on_fit_start	0.015	1	0.015	0.01504
[Callback]TQDMProgressBar.on_validation_batch_end	0.00078947	19	0.015	0.01504
[LightningModule]LagLlamaLightningModule.optimizer_zero_grad	0.0009375	16	0.015	0.01504

Apparently the data loader needs 5 seconds for each batch, which is 84% of the full time of the training step.
After some further investigation, we found that the train data loader does the following:

1. Apply the transformation to a full time series.
2. Sample a window from the transformed data (inside the InstanceSplitter).
3. Extract the window from the transformed data (InstanceSplitter).
4. Create the batches of data according to the batch size.

This means a full timeseries gets transformed and then most of the transformed data is not used. This is then done for each item in a batch. We observed ~10 ms for transforming a full timeseries and with a batch size of 512, we get the >5 seconds reported by the profiler.

The order of execution is partly given by the gluonts package. So I am not aware of an obvious solution without addressing it there.

Now my question. Did you face the same issue during your experiments? How can we solve the problem we observe?

[ashok-arjun]

Hi @JakobLindscheid !

Thanks for the detailed issue!

I was not aware of this issue as I never checked the data loading speed in my experiments.

Can I check this on my end and get back to you soon?

[JakobLindscheid]

Sure, thank you for having a look!
For now, I added a data = list(data) before the instance splitter is applied. This forces the transformation to be done before training starts. Obviously it's not the nicest solution ever since it takes a few minutes before training starts, but the total training time is improved a lot.

[ashok-arjun]

That's useful to know, thanks for sharing.