google / maxtext

A simple, performant and scalable Jax LLM!

Apache License 2.0

1.44k stars 263 forks source link

Add FSDP + Megablox #705

Closed RissyRan closed 2 months ago

RissyRan commented 2 months ago

Description

3x perf improvement with FSDP sharding

Test

Locally run: link
Internal benchmarks