Announcement: big_vision is transitioning from jax.pmap to jax.jit. #66
akolesnikoff
announced in
Announcements
Replies: 1 comment
-
|
Transition is done |
Beta Was this translation helpful? Give feedback.
0 replies
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
-
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.
Beta Was this translation helpful? Give feedback.
All reactions