DistributedSampler with modification will give you the almost equal-sized splits.

I don’t know how you defined your model but you should also use DDP to maximally parallelize the models with multiple GPUs & use DistributedSampler with multiple processes.
make sure to customize the sampler so that there is no overlap between the different ranks (processes).
you should communicate between different processes to collect loss or accuracy metrics.

You may want to take a look at my github repository 16 for an example.