[Bugfix] Force reduce in bwd for MoE router weights

[acisseJZhong]

solve #2387 where user reported a bug that MoE router weights diff on each rank due to a missing all reduce under certain senarios:

• self.score_before_experts=False for both ETP=1 and ETP=TP
• self.score_before_experts=True and ETP=1

Bug analysis

Case1:
When ETP=1 < TP, MoE is using TP2EP(etp_mesh is None, there is no ETP), gate is Replicated(via NoParallel), experts weight is Sharded(via EP mesh), MoE input is Replicated.

Root Cause: ReordererSequenceParallell is applied to split top_scores across TP ranks. Since each TP rank process its local slice of tokens through the experts, d_routed_input is Partial and there is no forced all-reduce.

Case2:
When ETP = TP, MoE is using DP2EP, gate is Replicated(via NoParallel), experts weight is Sharded(via ETP which force all reduce in bwd #1878), MoE input is Replicated on TP mesh.

Root Cause: The gradient flowing back to gate.weight is Partial when self.score_before_experts=False, which we incorrectly marked as Replicated and therefore results in wrong numerics.

Explanation:
When self.score_before_experts=True, d_routed_input is already all reduced because of the fix #1878 we added to ETP.

routed_input = routed_input * top_scores_experts_sorted
routed_output = self.experts(routed_input, ...) 
out_experts = routed_output_unsorted.sum(dim=1)    

When self.score_before_experts=False, gradient of top_scores from bmm is Partial. The previous all reduce fix only force all reduce on d_routed_input, but not d_top_scores which is wrong.

routed_output = self.experts(routed_input, ...)
out_experts = bmm(top_scores, routed_output_unsorted) 

The fix we propose that could fix both Case1 and 2

• We revert the fix in fix MoE TP backward #1878 to keep Partial dx flowing back, and delay the all reduce later at the boundary of MoE.
• We choose to not use ColWiseParallel/RowwiseParallel for shared expert as they force all reduce dx; but we want to keep Partial dx. So we introduced MoEColWiseParallel/MoERowwiseParallel which removed the input and output hooks. See the reasoning here for more details.
• We annotate gradient of router to be partial in bwd through adding a field to NoParallel(), so that DTensor will handle the all reduce

Caveat: The graph is plotted from Claude generated computational graph , please take it with a grain of salt.

cc @volcacius @fatih-uzlmz

[wconstab]

is doing 2 DTensor operations (to, from) really the best way here? why not just insert one autograd.function where the backwards is an all_reduce?

[acisseJZhong]

yeah I guess we could also insert a customized function where fwd=no_op and bwd does all_reduce? I was following a similar fix #1878
cc @tianyu-l if you know why we might prefer this over customized function

[wconstab]

ok. well short term it seems harmless to land this, other than it would have more overhead. (worth confirming, but I assume it does not do any extra collectives or operations?) anyway, we could fix them all to use a dedicated function from spmd_types once that's available.

[acisseJZhong]

it should not have any extra collectives other than the all reduce in bwd. yes I think we could directly use the function from spmd_types once it's ready!

[fatih-uzlmz]

@acisseJZhong
Thanks for the ping and for putting together the great deep dive, very fascinating to see how the bug split into two distinct cases based on the ETP/TP config

I'm glad @volcacius and I were on the right track with the NoParallel wrapper being the correct conceptual boundary for the reduction, but your fix handling the activations directly is a really clean solution.

Really happy to have helped kick off this investigation and grateful for the opportunity to learn from the team here :)

Thanks everyone once again