[NVIDIA] Move collective pipeliner after post-layout opts

[kaixih]

To perform the fp8 dot, we insert quant and dequant ops before the dot and rely on the gemm-rewriter to pattern match and rewrite to cublaslt calls for GPUs.

When this fp8 pattern appears in while loops that is qualified for the collective pipeliner (CP), the CP might disrupts the pattern by accidentally peeling the quant or dequant ops off the loop.

So, this PR moves the CP after the gemm-rewriter (when the newly-introduced xla_gpu_collective_pipeliner_post_layout_optimization is set), since this can save us from redesigning the CP to support this special case.

Note, for the long-term solution, we should move the collectives to the "safer" places before applying the CP. For example, in the above use case, we can move the collectives in-between the quant and dequant and then the CP should be able to preserve the dequant->dot pattern. This is actually already supported as an integral part in the gemm-rewriter.

cc. @Tixxx @hx89 @nluehr @philipphack

[kaixih]

@frgossen Can you help review?

[frgossen]

Thanks for explaining the motivation. From the description, it is not quite clear why this is needed. Can you expand the motivation, maybe with an example? What fp8 pattern appears in a while loop and how is it transformed vs should be transformed?

Can you add test cases with the pattern variations you care about? I think that could also help understand the motivation behind this change.

[kaixih]

Sure. Basically, the original graph looks like

while->q->dq->AG->dot

where q takes bf16 and outputs fp8 and dq takes fp8 and outputs bf16.
Then our gemm rewriter moves the AG and pattern matches dq->dot to cublaslt_fp8gemm, like:

# move AG:
while->q->AG->dq->dot
# pattern match:
while->q->AG->custom_call(=cublaslt-fp8)

Now, with the CP pass, we see that the original graph is transformed into:

q->dq->AG->while->dot 

where you can see the q->dq->AG is moved out of the loop. Then, the gemm-rewriter can no longer rewrites this dot to the custom call of cublaslt-fp8.

So, that is why this PR moves the CP after the gemm rewriter. So, it doesn't need to worry about breaking any patterns.

Please let me know if the above makes sense to you. At the same time, I am trying to extract the above from our real model training as a unit test to showcase the issue.

[kaixih]

@frgossen Any update?

We have recently tested this in some real models and we can see the fp8 gemms are back.

[bixia1]

Do you know why moving CP to post layout is not generally good?

[bixia1]

It is better to rename and make it clear this is post layout CP: RunPostLayoutCollectivePipelinerPasses

[kaixih]

Done. PTAL.

[kaixih]

Do you know why moving CP to post layout is not generally good?

I think theoretically it should be fine to move CP to post layout, since the CP is supposed to be agnostic of non-collective ops. The main concern (from @Tixxx ) might be some passes in layout (e.g. layout normalization) might alter the collectives through inserting copies or transposes, for instance. We are not 100% confident that this would pose a problem or not.

So, our plan is to keep the flag and set it to True by default. If no complaints are received, we will remove the flag. How do you think of the plan?

[frgossen]

Thanks!

[frgossen]

Can you address the build failures?

[kaixih]

@frgossen Rebased. PTAL.

(Only AMD ROCm is failing but the log implies it is irrelevant to this PR).

[frgossen]

Looks like this runs into a merge conflict. Can you rebase?