add support for simplefsdp+ep

[ruisizhang123]

As titled, this pr adds support for simplefsdp+ep.

In a SimpleFSDP + EP tlparse, you can see the all-to-all co-exists with replicated shared experts (_grouped_mm), which means we could potentially reorder them for overlapping.

Profiler Trace & Correctness (eager-mode)

The following results are benchmarks on 8 H100

1. FSDP + EP(degree=2)

• Loss: As seen the losses almost match between FSDP2+EP and SimpleFSDP+EP.

• Trace: The first stream is FSDP (on dp_shard_cp dim), the second stream is all_to_all in token dispatch, and the third stream is FSDP (on dp_shard_mod_ep). It's on different streams probably because its submesh names are different, but it should be or we hope they will be on the same....

2. FSDP + TP(degree=2)

• Loss: As seen the losses almost match between FSDP2+TP and SimpleFSDP+TP.

• Trace: The first stream is FSDP, the second stream is TP.

3. FSDP + TP(degree=2) + EP(degree=2)

• Loss: As seen the losses almost match between FSDP2+TP+EP and SimpleFSDP+TP+EP.

• Trace: The first stream is FSDP communication (on dp_shard_cp dim), the second stream is TP communication, the third stream is all-to-all for token dispatch, and the fourth stream is FSDP communication (on dp_shard_mod_ep dim).

4. HSDP + TP(degree=2) + EP(degree=2)

• Loss: As seen the losses almost match between FSDP2(HSDP)+TP+EP and SimpleFSDP(HSDP)+TP+EP.

• Trace: The first stream is FSDP communication (on dp_shard_cp dim), the second stream is TP communication, the third stream is all-to-all for token dispatch, the fourth stream is FSDP communication (on dp_shard_mod_ep dim), and the fifth stream is DDP communication.

What is not working:

1. AC + EP has three graph breaks:

• one in buffer mutation in self.input_splits = num_tokens_per_expert.view(device_mesh.shape[0], -1).sum(dim=1) in expert_parallel.py after confirming with @xmfan LINK;
• one in the _A2A class as the temporary fix for AC leak... LINK
• there is also one graph break in all_to_all_single_autograd's input, which is converted using tolist(). But for some magical reason, I found if inputs to all_to_all_single is parsed as output_split_sizes.tolist() instead of output_split_sizes, there will be no graph break.... After looking into the tlparse, output_split_sizes is still treated as a tensor, but when copy data out of output_split_sizes, the triton code will do an additional .item(). LINK

2. All-gather in the backward --> This also happens to FSDP + TP(without EP). There should be sth wrong with checkpoint in deepseek_v3 in general after some discussion in with @tianyu-l
3. The losses are quite close, but not perfectly match. This is potentially because the checkpoint implementing reshard_after_forward is not working. Thus, the behavior of two are different.

cc. @anijain2305