Rebase ALPA to JAX 0.4.10

alpa-projects / alpa

Training and serving large-scale neural networks with auto parallelization.

https://alpa.ai

Apache License 2.0

3.08k stars 360 forks source link

Open yhtang opened 1 year ago

yhtang commented 1 year ago

This is still a work-in-progress.

[ ] Migrate from tensorflow to OpenXLA
- [ ] migrate all delta between tensorflow-alpa and tensorflow upstream to xla-alpa
- [ ] build ALPA's jaxlib wheel
[ ] Rebase ALPA's Python code
- [ ] Fix frontend-backend compatibility

yhtang commented 1 year ago

Adding @nouiz @ZYHowell @gjoliver @merrymercy @mjsML for viz