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.
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.