Malte lig 4952 add test for all components using gatherlayer

Tests for losses

Moved the test_dist__gather.py to test_dist__gather__losses.py.
Extended it to not only test NTXentLoss, but also TiCoLoss, VICRegLoss and DCLLoss
Not included VICRegLLoss, as it requires the global and local features and grids as input, which does not fit into the current framework. I think the effort to also test it is not worth it for the moment. If we want to do it, we should do it in a follow-up issue.
This covers all losses using gather(...)

Aligning the import style is another topic. Here it is a bit more difficult, as there is the name clash between lightly.utils.dist and torch.dist. Currently some files import the one and some the other as dist.

Test for Benchmark module

Added test_dist__gather__benchmark_module.py for benchmarking the BenchmarkModule. Also fixed bug in the BenchmarkModule.on_validation_epoch_start: As the self.dataloader_kNN is not split by the DDP, all devices have the full dataloader and no gathering is necessary.

Other changes

Some typing updates

Missing

No missing test for `lightly.utils.dist.gather`

All usages of lightly.utils.dist.gather are covered now (except VICRegLLoss):

Missing test for other DDP functions

We also have more gather-like functions, e.g. concat_all_gather in lightly/models/utils.py.

And the kNN also use some gather functions:

However, these are out of scope for this issue.

lightly-ai / lightly