Commit Graph

3 Commits

Author SHA1 Message Date
Richard Wong aca80720c8 Feat: t5_jax_simple_parallel implements a working example of fsdp 2024-09-20 23:42:51 +09:00
Richard Wong 429e1742ab Feat: flax pjit example 2024-09-16 12:19:07 +09:00
Richard Wong ad5cf7735f Feat: fsdp demo
Refactor: pulling dataloader code into dataload.py
2024-09-15 22:41:00 +09:00