Add activation checkpoint offloading
Open
- Dominant language
- Python
- Stars
- 200
- Forks
- 45
- PR merge metrics
- No merged PRs in 30d
Description
JAX now supports saved activation offloading. These policies should be expressible in praxis and thus downstream libraries like paxml.
Maxtext already has such integration.
https://github.com/google/maxtext/blob/ebd39aa64d670fa13a313b6f776e01ad9e450321/MaxText/layers/models.py#L231
Contributor guide
Assessment
This issue has not been assessed yet.