[WIP] Compile for dp2ep

[xmfan]

Status

• 🚧 NEED pytorch/pytorch changes 🚧
  • [c10d] support dynamic shapes for all_to_all_single_autograd pytorch#157521
  • [inductor] use aten grouped_mm by default, move template autotuning to max-autotune pytorch#158045
• Graph breaks/Recompiles for MoE. tlparse: https://fburl.com/amorfehs, still looking through the graph breaks
  • Mostly due to experts being wrapped by MoE. Since we will not be tracing FSDP wrapper, these graph breaks are considered fundamental until we migrate to SimpleFSDP.
  • staticmethod on user-defined classes can not be generically supported, I moved those out.
  • TODO: token dispatch a2a overlap with expert

Repro

tested on debug model NGPU=2 CONFIG_FILE="./torchtitan/experiments/llama4/train_configs/debug_model.toml" ./run_train.sh --parallelism.data_parallel_shard_degree=2 --parallelism.expert_parallel_degree=2 --training.compile
logs: https://gist.github.com/xmfan/41b822d9f09eb07fee62d684a061cec1

memory: 2.20GiB -> 1.42GiB
speedup: no big change, need to check with actual model