Training speed regression

Complete training seems to have regressed. See discussion here

• Write small benchmark for a training iteration
• Git bisect to isolate the regression commit
• Git bisect on GPU
• Check if a pure array index works better?
• Steal other things from torch? https://pytorch.org/docs/stable/data.html#single-and-multi-process-data-loading
  • prefetch
  • memory pinning
  • multiple workers
• seems like jax just doesn't do memory views. https://jax.readthedocs.io/en/latest/jax.numpy.html. Mem copies prob get optimized out in JIT. Way to JIT the training loop? Could be interesting.
  • if an inner fori should be used, clean up dataset._samples
  • clean up dataloader
  • fix prediction dataloader. right now cumlengths makes the +pred_offset wrap to the next array.
• check test_synth_data? seems to have pulled an old version

assigned to @hanikevi

mentioned in commit ce23ef3c

created branch 49-training-speed-regression to address this issue

Benchmark laptop example_data_train ddb78c3b, switched to example_data, 5 bags loaded, w10_b50_e20

056678b2, switch window/batches to match. 5 bags loaded, w10_b50_e20

3c171020 has nans after first iter, but first iteration is comparable

marked the checklist item Write small benchmark for a training iteration as completed

marked the checklist item Git bisect to isolate the regression commit as completed

GPU Training from main HEAD

056678b2, adjust run.train_data to use right folder, window/batches/epochs

Testing on 605a7e9a as suggested by Lisa, first get the 20 sec/iter:

But now I realize this is from the iterator being consumed in the first iteration, and I guess tqdm doesnt catch this error and makes the iter estimate from dividing by number of planned epochs. When I fix this bug we are back at 425 sec/iter :/

marked the checklist item Git bisect on GPU as completed

Hi @lis07128 I did some tests on the GPU today, didnt see any difference in the epoch time from commit 056678b2 to HEAD. I think theres two options: we had tested with different training parameters before or the regression happened before 056678b2. Do you have any instinct on that?

Unfortunately I never tested the optimized computation time right before/after it was merged in 605a7e9a. Maybe we achieved those mentioned 20sec per iteration in #36 (closed) by setting prep_batches to "True"?

On main with prep_batches=True in trainer:

removed casting of idxs to array in get_batch on dyn_slice_dataset. they use a yield operator here: https://github.com/BirkhoffG/jax-dataloader/blob/1842b49428e37de721e821bc31d78dc21c6d7390/jax_dataloader/loaders/jax.py#L18

Total training time (in .train()) at w10, b50, 2 epochs.

no switch in dyn_slice concatenates the trajectories then just adjusts index

# bags	prep_batches	train_jit	jit no switch in dyn slice
5	1.9	3.6	3.5
25	5.6	5.2	5.3
50	10.5	7.8	6.8
75	16.7	10.4	9.0
100	20.1	83.6	10.6
125	25.5	229.4	12.6

Now we're cooking

changed the description

marked the checklist item seems like jax just doesn't do memory views. https://jax.readthedocs.io/en/latest/jax.numpy.html. Mem copies prob get optimized out in JIT. Way to JIT the training loop? Could be interesting. as completed

changed the description

marked the checklist item Check if a pure array index works better? as completed

changed the description

marked the checklist item if an inner fori should be used, clean up dataset._samples as completed

changed the description

marked the checklist item fix prediction dataloader. right now cumlengths makes the +pred_offset wrap to the next array. as completed

changed the description

marked the checklist item clean up dataloader as completed

marked the checklist item check test_synth_data? seems to have pulled an old version as completed

mentioned in merge request !36 (merged)

closed with commit bcdba693

mentioned in commit bcdba693