[FSDP2] enable per-param mesh FSDP2 for MoE

[weifengpy]

command: NGPU=8 MODULE=deepseek_v3 CONFIG=deepseek_v3_16b ./run_train.sh --training.steps 20 --parallelism.expert-parallel-degree 4

fsdp2 support per-param mesh: pytorch/pytorch#173509

this PR applies fully_shard on transformer_block, sharding experts on edp_mesh, and other params on dp_mesh. FSDPModule schedule 2 all-gather sequentially: 1st on transformer blocks, 2nd on experts

def _shard_placement_fn(param: nn.Parameter) -> ShardPlacementResult:
    if param in expert_params:
        # Expert parameters: use Shard(1) on edp_mesh
        return ShardPlacementResult(
            placement=Shard(1), mesh_info=edp_mesh_info
        )
    else:
        # Non-expert parameters: use Shard(0) on dp_mesh
        return ShardPlacementResult(
            placement=Shard(0), mesh_info=dp_mesh_info
        )

this make it possible for apply torch.compile on each transformer_block. I didn't enable compile per block yet becuase there is still a gap in torch.compile + ac + MoE: #2341

AG order in forward are exactly the same before and after this change

AG order in backward are different but is better

Explicit Backward AllGather Order                                                                                                                                                                                         
  layers.7       @ 118.83ms   (attention/ffn params)                                                                 
  layers.6       @ 121.52ms   (attention/ffn params)                                                                 
  layers.6.moe   @ 122.04ms   (MoE expert params)                                                                    
  layers.7.moe   @ 125.81ms   (MoE expert params)  ← delayed!                                                        
                                                                                                                     
  Per-param Backward AllGather Order                                                                                 
  layers.7       @ 114.30ms   (first FSDP unit)                                                                      
  layers.7       @ 115.14ms   (second FSDP unit, includes MoE)                                                       
  layers.6       @ 117.42ms   (first FSDP unit)                                                                      
  layers.6       @ 117.89ms   (second FSDP unit, includes MoE)   

Numerics remains bitwise equal with/without this change

 Loss Comparison                                                                                                                                                                                                                           
  ┌──────┬───────────────┬───────────────┬───────┐                                                                                                                                                                                          
  │ Step │ Old (0d93c63) │ New (e1c47c8) │ Match │                                                                                                                                                                                          
  ├──────┼───────────────┼───────────────┼───────┤                                                                                                                                                                                          
  │ 1    │ 8.01151657    │ 8.01151657    │ ✓     │                                                                                                                                                                                          
  ├──────┼───────────────┼───────────────┼───────┤                                                                                                                                                                                          
  │ 5    │ 3.85572004    │ 3.85572004    │ ✓     │                                                                                                                                                                                          
  ├──────┼───────────────┼───────────────┼───────┤                                                                                                                                                                                          
  │ 10   │ 3.15517211    │ 3.15517211    │ ✓     │                                                                                                                                                                                          
  ├──────┼───────────────┼───────────────┼───────┤                                                                                                                                                                                          
  │ 15   │ 3.07873583    │ 3.07873583    │ ✓     │                                                                                                                                                                                          
  ├──────┼───────────────┼───────────────┼───────┤                                                                                                                                                                                          
  │ 20   │ 2.92206621    │ 2.92206621    │ ✓     │                                                                                                                                                                                          
  ├──────┼───────────────┼───────────────┼───────┤                                                                                                                                                                                          
  │ 25   │ 2.89102936    │ 2.89102936    │ ✓     │                                                                                                                                                                                          
  ├──────┼───────────────┼───────────────┼───────┤                                                                                                                                                                                          
  │ 30   │ 2.81201696    │ 2.81201696    │ ✓     │                                                                                                                                                                                          
  ├──────┼───────────────┼───────────────┼───────┤                                                                                                                                                                                          
  │ 35   │ 2.84123349    │ 2.84123349    │ ✓     │                                                                                                                                                                                          
  ├──────┼───────────────┼───────────────┼───────┤                                                                                                                                                                                          
  │ 40   │ 2.76206398    │ 2.76206398    │ ✓     │                                                                                                                                                                                          
  ├──────┼───────────────┼───────────────┼───────┤                                                                                                                                                                                          
  │ 45   │ 2.82969308    │ 2.82969308    │ ✓     │                                                                                                                                                                                          
  ├──────┼───────────────┼───────────────┼───────┤                                                                                                                                                                                          
  │ 50   │ 2.77560568    │ 2.77560568    │ ✓     │                                                                                                                                                                                          
  ├──────┼───────────────┼───────────────┼───────┤                                                                                                                                                                                          
  │ 55   │ 2.75578761    │ 2.75578761    │ ✓     │                                                                                                                                                                                          
  ├──────┼───────────────┼───────────────┼───────┤                                                                                                                                                                                          
  │ 60   │ 2.75143075    │ 2.75143075    │ ✓     │                                                                                                                                                                                          
  ├──────┼───────────────┼───────────────┼───────┤                                                                                                                                                                                          
  │ 65   │ 2.74203372    │ 2.74203372    │ ✓     │                                                                                                                                                                                          
  ├──────┼───────────────┼───────────────┼───────┤                                                                                                                                                                                          
  │ 70   │ 2.71638918    │ 2.71638918    │ ✓     │                                                                                                                                                                                          
  ├──────┼───────────────┼───────────────┼───────┤                                                                                                                                                                                          
  │ 75   │ 2.74999237    │ 2.74999237    │ ✓     │                                                                                                                                                                                          
  ├──────┼───────────────┼───────────────┼───────┤                                                                                                                                                                                          
  │ 80   │ 2.75584078    │ 2.75584078    │ ✓     │                                                                                                                                                                                          
  ├──────┼───────────────┼───────────────┼───────┤                                                                                                                                                                                          
  │ 85   │ 2.74837303    │ 2.74837303    │ ✓     │                                                                                                                                                                                          
  ├──────┼───────────────┼───────────────┼───────┤                                                                                                                                                                                          
  │ 90   │ 2.72101045    │ 2.72101045    │ ✓     │                                                                                                                                                                                          
  ├──────┼───────────────┼───────────────┼───────┤                                                                                                                                                                                          
  │ 95   │ 2.73645735    │ 2.73645735    │ ✓     │                                                                                                                                                                                          
  ├──────┼───────────────┼───────────────┼───────┤                                                                                                                                                                                          
  │ 100  │ 2.70604038    │ 2.70604038    │ ✓     │                                                                                                                                                                                          
  └──────┴───────────────┴───────────────┴───────┘                                                                                                                                                                                          

[weifengpy]

those CI errors because we need to land pytorch/pytorch#173509 first

[tianyu-l]

SGTM, is there a plan for solving #2341?

Also, not sure if it's blocked by the issue, but we should modify apply_compile to remove the previous workaround fine-grained compilation code.

[xmfan]

@tianyu-l We're still gonna need the fine-grained workarounds at least for mxfp8, due to #2250 (comment). Until that issue is fixed (dynamo bwd tracing of autograd functions), we will trace wrong bwd graph.

[tianyu-l]

@xmfan oh, I didn't know. I marked it as high priority for now.

[weifengpy]

Also, not sure if it's blocked by the issue, but we should modify apply_compile to remove the previous workaround fine-grained compilation code.

I didn't change apply_compile because of #2341. #2250 (comment) is new to me

[tianyu-l]

sgtm