Enable 2D shardign with `minibatch=True` for SPMD

[miladm]

🐛 Bug

Enable 2D shardign with minibatch=True for SPMD. The current implementation supports 1D sharding. Here is a link to this topic on the original implementation of multi-host SPMD data loading.

cc @qihqi @tengyifei

[miladm]

@tengyifei we took a deep dive into this topic when enabling llama3.1 405B. Any context to share on your experience?

[tengyifei]

My experience with minibatch in Llama 3.1 405B:

• I set minibatch=True under 2D sharding (FSDP, TP) of the model, while using FSDP to shard the input data. The model loss was NaN.
• I set minibatch=False and everything trains correctly

I'm not sure why this produces NaN. NaNs are generally very hard to debug. I think something should at least throw an error if this is not a supported configuration.

[miladm]

@lsy323 do we have a target timeline for this issue?