The ordering of fsdp, ac, tp, pp and complie etc.

[aoyulong]

Based on the code, the ordering of parallelization and optimization appears to be: PP → TP → AC → Compile → FSDP/DDP.
Is it possible to modify this ordering? If not, could you explain the rationale for this specific sequence?

[tianyu-l]

PP -> SPMD
Making PP split first allows us to deal with clean SPMD regions.
[maybe not fundamental and can be reversed? @wconstab]

TP -> FSDP
because the semantics during training iterations would be the reverse.
[natural]

TP -> AC (citing @soulitzer)
Switching the ordering of apply_ac/apply_tp in torchtitan seems to result in a "got mixed torch.Tensor and DTensor, need to convert all torch.Tensor to DTensor before calling distributed operators" somewhere after backward.
[sounds not fundamental]

AC -> torch.compile (citing @soulitzer)
Today, compile wrapping AC is more recommended than AC wrapping compile. Compile already does recompute via the min-cut partitioner, so wrapping AC over a compiled region may lead to recomputing multiple times. If compile wraps AC, compile would incorporate the AC region information during the partitioner. It may be possible to improve this behavior though, e.g. detect any ambient AC contexts, and do something different in the partitioner and during runtime.
[recommended, sounds fundamental]

torch.compile -> FSDP
because torch.compile would graph break on FSDP2 wrapped modules.
[fundamental for pytorch FSDP2]

[tianyu-l]

we probably should document this somewhere @wwwjn