Optimize Computation Time in Training

Hey @hanikevi ,

this issue is related to 35-out-of-memory for optimizing the computation time in training.

Current state is:

• jitted build_view_batch_fori function with fori_loop()
• jitted get_traj_idx function
• testing is added in dyn_slice_dataset.py

Looking forward to work together on this issue. Validation activities are not related to this.

---

• Pad arrays
• So the view_batch and view_batch_fori are working
• Write minimal benchmark including slicing and memory transfer w/ trivial computation on GPU
• Can we just vmap to get batches?
• move cumsum into default _get_idx
• Possible to do slicing on GPU? If it's JIT-ted is this possible?
• If we have pytree on CPU, generate batch, then feed that into network, when does that go to GPU?
• Why not faster on smaller batches? Prep_batches does a small batch in 13sec.
  • Check control flow - I suspect calling get_batch(idxs) in training loop might be less efficient than using an iterator / generator
  • If so; split out dataset and dataloader

marked this issue as related to #35 (closed)

created branch 36-optimize-computation-time-in-training to address this issue

mentioned in issue #33 (closed)

assigned to @hanikevi

mentioned in merge request !21 (merged)

created branch 36-optimize-computation-time-in-training-2 to address this issue

Looking at loading speeds, b/c we suspect that is dominating training time right now (huge variation in training speed depending on data loader).

• Reference values
  • SSD -> CPU Memory: ~1-6 Gb/sec (samsung 870 evo)
  • CPU -> GPU: 64 Gb/sec (nvidia 4090)
  • Dataset:
    • Unsliced: ~2 Gb
    • Sliced: sizeof(unsliced) * window_size, ~160 Gb

From this, the naïve slicing on CPU, transfer each slice should have memory transfer costs of 3 sec/epoch. This suggests the slicing time may be critical? Or somehow a bottleneck in slicing process?

changed the description

Tracing memory locations:

• JIT is associated with a device. Can make explicit w/ device= kwarg in JIT fn.
• Display array's device via y.device_buffer.device()

changed the description

marked the checklist item Pad arrays as completed

marked the checklist item So the view_batch and view_batch_fori are working as completed

marked the checklist item Write minimal benchmark including slicing and memory transfer w/ trivial computation on GPU as completed

changed the description

condition	speed	cuda speed
nojit	0.060
jit	0.016
fori jit	0.011	0.002
vmap jit	0.011	0.0005

first call on fori takes 2.5 sec to compile. vmap is no overhead. both pass all tests

marked the checklist item move cumsum into default _get_idx as completed

marked the checklist item Can we just vmap to get batches? as completed

changed the description