Supporting DDP / Lightning Fabric

[genemerewether]

Thanks for the sweet library! Any update on DDP / lightning fabric compatibility?

[Charl-AI]

Hey! Thanks for reaching out, and thanks for the kind words. I do intend on adding this eventually, but honestly I can't see myself getting round to it for a little while. I'll give you a quick explainer of the situation in the hope that it may help you find a workaround.

[This may not be 100% correct, it was a little while since I last looked into this.]

As it stands, it is possible to use the library with DDP, but it's just less efficient than it should be because (IIRC) each of the distributed processes gets an independent cache. In single GPU training, the cache is created in the main process and shared across the worker processes spawned by the data loader. However, I think that in DDP, we essentially get multiple main processes, each spawning their own data loader children processes too.

Now, it turns out that Python has two different ways of creating arrays backed by shared memory, mp.Array (which we use currently), and mp.shared_memory. The latter is more 'decentralised', meaning that if we switch to using it, then we may be able to connect each of the DDP processes to the same cache. There is probably an extra wrinkle to iron out when it comes to multi node, because then you would need to have separate caches per machine and would need to make sure that each machine gets the same subset of data each epoch.

A diagram to visualise the processes:

I think that a migration to the mp.shared_memory API is the best way to support DDP in the future, but I really don't want to break the current SharedCache abstraction (which I think works really nicely for single GPU training). The best solution may just be to make a new class SharedCacheDDP, which can handle all this.

Anyway, this was a bit of an information dump. I hope it helps you to solve your problem. If you find a nice solution, please feel free to open a PR and I'd be happy to take a look and merge it.