fattorib / ZeRO-transformer

Two implementations of ZeRO-1 optimizer sharding in JAX

fattorib/ZeRO-transformer Issues

Replace naive loops with jax.lax.scan
Closed a year ago
Add support for converting to PyTorch models
Closed a year ago1
Benchmark code for ```pjit``` meshes
Closed 2 years ago2
Train loss increasing during training
Closed 2 years ago5
Refactor attention to be ```pjit``` friendly
Closed 2 years ago1