Memory-mapped caching for image translation training

[ziw-liu]

#195 implemented in-RAM caching for image translation trainings. However it does not scale to datasets larger than available system memory. This PR implements a node-local disk cache via tensordict's memory-mapped tensor.

[ziw-liu]

Progress summary: previously training on 2.3 TB of datasets takes 20 s/iter, now training on 3 TB takes 10 s/iter.

Lessons learned:

• Default local scratch configuration on our compute nodes was not optimal for sustained I/O (ZFS and very small sector size). This was fixed on select nodes and will be rolled out more broadly later.
• Optimizations/mitigations done for system memory caching might hurt performance in the mmap setup. Moving some transforms back to the CPU improved end-to-end timing. This might be because MONAI transforms are not batched (executed in a loop), and CPU/GPU sync could be taking much longer than the actual compute.

To be investigated:

• Moving augmentations back to the CPU recreated the CPU compute bottleneck (removing augmentations further reduces end-to-end training time to 3s/iter, or a 3x reduction). This is potentially fixable by using batched augmentation or distributing the compute better across devices.
• Precomputing the normalization to simplify training-time logic (d2cd340).

[ziw-liu]

This might be because MONAI transforms are not batched (executed in a loop), and CPU/GPU sync could be taking much longer than the actual compute.

Benchmark of 3D random affine in 6a88ec4 (10 runs, milliseconds):

Device	MONAI (sequential)	Kornia (batched)	Relative
Zen 2 CPU (1 thread)	9160	3800	2.4
Zen 2 CPU (16 threads)	7320	556	13.2
A40 GPU	2620	210	12.5