[MoE][compile][full ac] weave torch.compile around the FSDP(GroupedExperts) graph break

[xmfan]

Stacked PRs:

• ->[MoE][compile][full ac] weave torch.compile around the FSDP(GroupedExperts) graph break #1895

---

This PR changes how we compile MoE layers to work around Compile + AC limitations. When you AC(Compile(block)) or Compile(AC(block)) and there is a graph break in block, we fall back the entire block to eager. For llama3, we've worked around this problem by addressing all graph breaks. With MoE models particularly dp2ep, we need to wrapFSDP(block.moe.experts), meaning that we will have graph breaks when tracing block.moe.experts.__call__, meaning that whenever AC was enabled, the entire block for MoE would fallback to eager: https://gist.github.com/xmfan/50f4de1e89d789cd63a21aca9e600132 (Note in the tlparse, graph 0/1 is empty and it corresponds to the block containing the MoE).

The workaround in this PR is to avoid tracing block.moe.experts.__call__. This is done by individually wrapping torch.compile on submodules of TransformerBlock. Note that we are leaving some perf on the table as this might exclude some ops in TransformerBlock.forward and MoE.forward. This is an API limitation, as we have no way to acquire those ops while decoupling the wrapper from model code. This workaround will no longer be necessity when either:

• We can do Compile + AC with graph breaks
• We remove the FSDP graph break

This change introduces a small regression to the non-AC configuration. You can see a small perf dip from before this PR and after this PR. Given that AC is a necessity to run non-toy configurations of these models, I chose to stick to this implementation to make comparisons easier.

Validated on DSv3 debug model:

• dp2ep, no AC, no compile: https://gist.github.com/xmfan/927f354158ad36f4c5c1ffedde4e4ebe
• dp2ep, no AC, compile: https://gist.github.com/xmfan/11561b5406b3f92ecd08da94bc5ee4e3
  • before this PR (compile w/ nested graph break): https://gist.github.com/xmfan/0b32e95980d263cf3f62869fa4d85921
• dp2ep, full AC, compile: https://gist.github.com/xmfan/6ed5b48aa51ce0ac2b6bfceb86a0c482
  • before this PR (whole moe block in eager): https://gist.github.com/xmfan/50f4de1e89d789cd63a21aca9e600132
• dp2ep, full AC, no compile: https://gist.github.com/xmfan/2308355c2aa4814fe3d12243445555fa
• dp2ep, pp, full AC, compile: https://gist.github.com/xmfan/5a1ac23f00abdf93dbcc1539f552e840
• dp2ep, pp, full AC, no compile: https://gist.github.com/xmfan/302cda7191e53ffad5c4dc1e4b8f02de