Add configurable pin_memory and prefetch_factor for DataLoader

[stephengreen]

Summary

• Adds pin_memory and prefetch_factor options to training settings under local section
• Allows users to tune memory vs speed tradeoffs during training
• Reduces default prefetch_factor from 2 to 1 for lower memory usage

Resolves #308.

Details

Investigation of memory usage during training revealed that pin_memory=True combined with the prefetch queue can significantly increase memory consumption. This PR:

• pin_memory (default True): Uses pinned memory for faster GPU transfers but increases memory usage. Can be set to False to save memory at cost of slower data transfer.
• prefetch_factor (default 1, reduced from 2): Controls how many batches each worker prefetches. Lower values reduce memory usage.

Memory usage scales roughly as: num_workers × prefetch_factor × batch_size × data_size, and is roughly doubled when using pinned memory.

Performance testing

In a realistic training setup:

• pin_memory=False increases training time by ~20%, so keeping the default True is recommended
• prefetch_factor=1 (reduced from default 2) does not slow down training at all, while reducing memory usage

Test plan

• Verified training runs with default settings (pin_memory=True, prefetch_factor=1)
• Verified pin_memory=False works (slower but lower memory)
• Verified prefetch_factor=1 vs prefetch_factor=2 (no speed difference, lower memory)

🤖 Generated with Claude Code

[stephengreen]

Added documentation for the new settings in docs/source/training.md, including the memory usage formula and performance findings.

[stephengreen]

@srashtig

[annalena-k]

I have a few questions, but these are minor things.
I haven't tested the implementation myself, so I can't judge how much this change reduces the memory footprint.

[annalena-k]

I tested the following with 64 CPUs and 63 workers:

• pin_memory=True & prefetch_factor=1 required 128GB RAM and the GPU load looked stable (roughly constant over batches)
• pin_memory=False & prefetch_factor=1 required 95 GB RAM, but the GPU load was unstable (see screenshot, maybe longer training would have solved this)

• pin_memory=True & prefetch_factor=2 required ~270 GB RAM and looked stable (after longer training time)

I found the way the new options are included in the settings file confusing.
This is conceptually different from how we previously included comments in the example settings files, e.g. the early stopping is commented (disabled) and the user can uncomment it to enable it. Similar for local_cache_path. Defaults like leave_waveforms_on_disk: True are not commented.

I would recommend to stay consistent and not comment defaults.

[annalena-k]

When I was testing this, it was confusing that this line is commented out.
As far as I understood, pin_memory=True by default, right?
Therefore, uncommenting this line actually does nothing. This is conceptually different from how we previously included comments in the example settings files, e.g. the early stopping is commented (disabled) and the user can uncomment it to enable it. Similar for local_cache_path. Defaults like leave_waveforms_on_disk: True are not commented.

I would recommend to stay consistent and not comment defaults to make it easy to understand for the user.

[annalena-k]

The same applies to prefetch_factor and all example files.