ScatterMoE Gradient Norm Needs to be Properly Computed When Used With FSDP

[fabianlim]

When ScatterMoE is used together with FSDP, HF Accelerate will call FSDP's clip_grad_norm, see here, which does not know how to properly compute the grad norms for ScatterMoE shards.

We need to be able to hook in some logic to handle the grad norm properly.

[fabianlim]

@kmehant do you have any ideas on how to fix this issue? cc: @anhuong @willmj

[kmehant]

@fabianlim Would need to think given its a mix of flat params and DTensors. Thank you for looping me in.

[fabianlim]

we will need to patch the clip_grad_norm function in the FSDP case, and have some logic to handle