Implement batch bounds to even out resource usage during training

[frostedoyster]

This PR implements even load management between different batches, based on the total number of atoms in a batch. This is implemented by letting the user specify a lower and upper bound of atoms in a batch. In practice, training is done as normal but, if a batch has too few or too many atoms, it is skipped.

This is useful for two main reasons:

• avoids OOM issues due to unlucky large batches, allowing for much larger batch sizes which potentially increase training efficiency
• evens out the load on different processes in distributed settings, where all processes have to wait for the slowest batch across all processes at every training step

For PET, this complements the adaptive cutoff strategies to fully even out memory and computation usage across batches

---

📚 Documentation preview 📚: https://metatrain--990.org.readthedocs.build/en/990/

[pfebrer]

Hey, I have some questions:

1. Why not skipping the batch in the data loader itself, instead of creating a None batch and then checking for it in the trainer? Apart from the code being nicer, the current implementation seems to have the possibility of dropping a lot of data in distributed environments.
2. Does data get shuffled at every epoch or will this result in some structures never appearing in training because their batch is too large/too small?
3. Would it make sense to implement load balancing by interfering in the batch creation, instead of naively creating batches and then dropping?

[frostedoyster]

1. If you skip just a single batch in one dataloader, the dataloaders fall out of sync, because, in my understanding, the length of each dataloader is determined in advance and it's the same for all parallel dataloaders. This is the standard way to skip batches in distributed environments in GNNs, as far as I know
2. The data is shuffled at each epoch
3. Maybe but the implementation would be much more difficult. When I say much, I mean much. This is the standard, as far as I know

[pfebrer]

1. Ok, but this doesn't mean that you can't synchronize the skipping inside the dataloaders/datasets no?

If this is the way to go, could we at least keep track of how many batches were skipped and print it somewhere in the log file (e.g. as a percentage of the full dataset)?

[pfebrer]

If we can synchronize the skipping at the beggining of the collate function, apart from the code (this is just an opinion) being nicer, we would save a lot of wasted computation.