Announcement: big_vision is transitioning from jax.pmap to jax.jit.
akolesnikoff opened this issue · 1 comments
akolesnikoff commented
In the coming 1-2 weeks big_vision is expected to transition from pmap-based parallelism to jit-parallelism. This will enable more flexible parallelisation strategies, including, but not limited to, ZERO* and fully-sharded data-parallel (aka fsdp) training.
This transition may temporarily break project-specific code (or we will just remove such code). If you want to read/run the old code, please see the table at the end of the README for the project-specific commits to sync to.
akolesnikoff commented
Transition is done