Move Stage to second axis

google / maxtext

A simple, performant and scalable Jax LLM!

Apache License 2.0

1.39k stars 247 forks source link

Closed gobbleturk closed 1 month ago

gobbleturk commented 1 month ago

Move stage axis to second in list - orbax currently makes some assumptions that the data parallel axis is first for a few new features.

Tested via dcn XPK logs