Why is the context axis required for expert sharding when `shard_exp_on_fsdp` is enabled?

[pathfinder-pf]

Feature or Model Request

No response

Additional Context

No response

[NuojCheng]

Could you elaborate which lines of code and what errors you met? I don't think context axis is required for the shard_exp_on_fsdp flag.

[pathfinder-pf]

https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/configs/base.yml#L551 this line

[NuojCheng]

In MaxText, the context axis acts like FSDP to shard weight tensors. Because of this, context always shards the same dimensions as fsdp for weights and optimizer states. With the flag shard_exp_on_fsdp, this pattern is kept.

Alternatively, if you don't want the context axis to shard the expert dimension, you can define custom logical sharding rules. See the documentation for details.