A simple notebook demonstrating capabilities of a low rank gradient approximation algorithm.
You will also find an implementation for a DDP communication hook which approximates gradients and all gathers the results.
In the future, we expect to add a JAX implementation.
Please read the notebook for more information: notebook.