[JAX] Fix failure on pattern matching of FP8 GEMM when enabling FSDP.

[mingxu1067]

• Add a custom call, cast.
• Replace cast_and_transpose with cast to the kernel of layernorm_fp8_dot and the kernel_1 of layernrom_geglu_fp8_mlp to allow XLA handle transpose for avoiding unnecessary copy to break FP8 GEMM pattern matching.
• Replace cast_and_transpose with native XLA cast to the x and kernel of fp8_dot and the kernel_2 of layernrom_geglu_fp8_mlp to allow XLA handle transpose for avoiding unnecessary copy to break FP8 GEMM pattern matching.
• Fix a bug of enabling layernrom_geglu_fp8_mlp in flax.LayernormMLP.

[kaixih]

Basically, it seems the cast_transpose is replaced with either the custom cast or the native quantize. I feel the cast and quantize are functionally same. Do we have a rule or guide on when to use which?

[nouiz]

@denera to review.
@mingxu1067 Can you extent the all current fp8 sharding tests to verify that the gemm is in fp8?
This will test the failure you fix and make sure it doesn't regress and make sure no equivalent issue happen elsewhere.

[mingxu1067]

Basically, it seems the cast_transpose is replaced with either the custom cast or the native quantize. I feel the cast and quantize are functionally same. Do we have a rule or guide on when to use which?

We apply the native quantize and transpose when the tensors is split along its column. For example, a tensor in shape (M, N), and it is sharded along N. In this case, native quantize and transpose offer better flexibility for XLA to schedule all-gather and transpose, then avoid unnecessary copy.

However, there are few cases in backward, like x of fp8_dot, excluded from the above rule, but we found it also introduces unexpected copy. For this case, we also apply the native quantize and transpose to solve.

We also target on replacing all custom cast_fp8 and cast_transpose with the native implementation in the feature. Will schedule a time to evaluate the performance gap and work effort needed.

[mingxu1067]

@denera to review. @mingxu1067 Can you extent the all current fp8 sharding tests to verify that the gemm is in fp8? This will test the failure you fix and make sure it doesn't regress and make sure no equivalent issue happen elsewhere.

Currently, UTs does not include FSDP related tests, therefore the tests cannot capture this kind of failures. Extending the UTs to conver wide range of cases require some time. Had an internal discussion, will add this into TODO.

[denera]

LGTM!

[mingxu1067]

/te-ci jax

[nouiz]

This PR description tell that this fix a bug.
We need a test that make sure we don't regress.

[mingxu1067]

/te-ci jax

[mingxu1067]

/te-ci jax