Improve sharding propagation for triangle updates outgoing

[wujingyue]

Repro: #5890

$ mpirun -np 1 -x NVFUSER_DUMP=pre_segmenter_logging pytest tests/python/multidevice/test_alphafold3.py -k outgoing --only-mpi -vs

The code of interest:

Fuser/tests/python/direct/test_alphafold3.py

Lines 123 to 133 in 77abd29

	match direction:
	case Direction.OUTGOING:
	# z_out = einsum("bikc,bjkc->bijc", a, b)
	a = fd.ops.permute(a, [0, 3, 1, 2]) # [b, c, i, k]
	b = fd.ops.permute(b, [0, 3, 2, 1]) # [b, c, k, j]
	case Direction.INCOMING:
	# z_out = einsum("bkic,bkjc->bijc", a, b)
	a = fd.ops.permute(a, [0, 3, 2, 1]) # [b, c, i, k]
	b = fd.ops.permute(b, [0, 3, 1, 2]) # [b, c, k, j]
	z = fd.ops.matmul(a, b) # [b, c, i, j]
	z = fd.ops.permute(z, [0, 2, 3, 1]) # [b, i, j, c]

The current heuristic for the forward propagation is to prefer the second input (usually the weight). Therefore, in the einsum output, j is sharded by DIDy not DIDx. This breaks the backprop from z_in (i by DIDy and j by DIDx) to the einsum output, because z_in wants j to be sharded by DIDx instead.

By the way, this is not a problem for "incoming" mode. Following the current heuristic, the einsum output does have j sharded on DIDx.

cc @DejunL

[wujingyue]

I discussed with @Priya2698 offline. We think the right way to fix this is to make sharding propagation more incremental and iterative.

The current algorithm tries to decide the sharding for the einsum output too soon. If it waited for the sharding to be propagated from z_in through gating and back through layernorm, the einsum output would naturally get the right sharding.

Fuser/tests/python/direct/test_alphafold3.py

Lines 135 to 136 in 77abd29

	z = layer_norm(fd, z, w_norm_out, b_norm_out)
	z = gating(fd, z, w_p_out, z_in, w_g_out)

More concretely, we'd like sharding propagation

1. defer conflicts and solve easy problems first.
2. be flexible about direction. Currently, the algorithm enforces one forward pass followed by one backward pass, an artificial propagation order.

[DejunL]

Thanks for tagging me! There is probably some domain language here I don't really understand, e.g., what are "dy", "dx" in the context of the sharding by the "^" branching symbol? What about "slice"? How are all these related to sharding propagation?

[wujingyue]

Did I show you https://github.com/NVIDIA/Fuser/blob/main/doc/reading/multigpu.md?

dy stands for DIDy and dx for DIDx. The ^ branching symbol is the Split op on IterDomains.

Sorry, it's indeed very domain specific.

[DejunL]

^ah, that's new but very interesting. I'll take a look

[wujingyue]

cc @chaserileyroberts

Wonder your thoughts on this. I'm sure you know GSPMD and Shardy more than I do.

[wujingyue]

To show the full picture, the figure below is the designed sharding for outgoing mode:

The figure below is the suboptimal sharding we currently generate:

[wujingyue]

More concretely, we'd like sharding propagation

1. defer conflicts and solve easy problems first.
2. be flexible about direction. Currently, the algorithm enforces one forward pass followed by one backward pass, an artificial propagation order.

This won't work out of the box for sequence parallel: https://github.com/NVIDIA/Fuser/blob/main/doc/reading/multigpu.md#sharding-propagation-1. Gelu's output of shape [t, 4h] will easily get sharded on t instead of 4h. Size information could help find the right heuristics and we don't have full size information during preseg.

[Priya2698]

More concretely, we'd like sharding propagation

1. defer conflicts and solve easy problems first.
2. be flexible about direction. Currently, the algorithm enforces one forward pass followed by one backward pass, an artificial propagation order.

This won't work out of the box for sequence parallel: https://github.com/NVIDIA/Fuser/blob/main/doc/reading/multigpu.md#sharding-propagation-1. Gelu's output of shape [t, 4h] will easily get sharded on t instead of 4h. Size information could help find the right heuristics and we don't have full size information during preseg.

I was reading Shardy concepts and priority based sharding: https://openxla.org/shardy/propagation#detailed_design appears useful.
What we are observing is that for different models, different ops/dimensions are more important. For AlphaFold, sequence parallel is important, where-as in SP+TP, we would want to propagate weights first. So incorporating the notion of priority can help us provide hints on how to resolve conflicts in different models.