[WIP][SimpleFSDP] add manual bucketing pass

[ruisizhang123]

This PR adds support for aten-level manual bucketing in SimpleFSDP+aot_eager backend. Dependent on PyTorch PR

TODO List:

• We should have better way of handling region info other than a list of str FQNs in current manual_bucketed_modules. It would be very easy to miss some of model modules. (cc. @xmfan @SherlockNoMad )
• Currently, the reordering happens under the hood and overlap with last/next compute. We should allow users to specify which module they want to reorder.

1. Performance (FSDP2 under eager mode, SimpleFSDP uses aot_eager backend)

Llama 3-8B

• Single Node, 8 H100, Performance. (The slower TPS on Single Node is sort of as expected, since FSDP2 handles copy-in/out in two different streams, whereas SimpleFSDP handles copy-in/out in the same stream)

Method	Parallelism	Memory	TPS	Trace
SimpleFSDP	FSDP=8	40.96GiB(43.12%)	7,227	LINK
FSDP2-eager	FSDP=8	47.82GiB(50.35%)	7,380	LINK
FSDP2-aot_eager	FSDP=8
SimpleFSDP	FSDP=4 TP=2
FSDP2	FSDP=4 TP=2

Example SimpleFSDP 1D overlapping trace:

• Bitwise Loss:

FSDP-only:

FSDP+TP:

DeepSeekV3-16B