Integrate Sonic-Moe in FSDP

[ChangyiYang]

As titled
issue #1144

[junjzhang]

Hi, happended to see this pr, I guess the numerical accuracy is not correct. Sonic-moe will require a interleaved weight like below

 I = w1.shape[1] // 2
# Interleave gate/up: [E, 2*I, H] concat -> [E, 2*I, H] interleaved
w1 = torch.stack([w1[:, :I, :], w1[:, I:, :]], dim=2).reshape(w1.shape[0], 2 * I, w1.shape[2])

We tested this in our own training framework, I suppose this is the correct way. BTW, we do not inspect significant improvement in MFU. Happy to discuss your results.

[ChangyiYang]

Thanks for the detailed feedback! This makes sense — I’ve indeed observed some nonsensical outputs and am currently debugging the numerical correctness. I’ll try the interleaved gate/up weight layout you suggested and will update the PR description with performance results once ready.

[faresobeid]

Any updates on this?