[dinodeep]

Description

This PR integrates the necessary changes for enabling TransformerEngine's collective GEMM API for MLP layers as well as adding a configurable flag (use_te_comm_gemm_overlap, defaults to false) to enable it.

Collective GEMMs are useful because they allow for fine-grained overlap of communication and computation of distributed GEMM computations, reducing the distributed GEMM latency and improving overall E2E step time when working with tensor sequence parallelism.

Enabling the flag requires the usage of tensor sequence parallelism, running on single-rank-per-GPU, and enabling a Transformer Engine quantization recipe.

It works by telling TransformerEngine which type of collective will be performed for specific parts of the MLP. Up-projections require all-gathers, and down projections require reduce-scatters under tensor sequence parallelism. Overall benefit is reduced exposed communication time for large distributed GEMMs.

The change is primarily in quantizations.py in which we also determine the type of collective the dense layer is performing if it is in the MLP.

Tests

Validated that enabling the collective GEMMs introduces them via profiling and looking at HLOs. Loss curves stay close when disabling/enabling collective GEMMs.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[codecov]

Codecov Report

❌ Patch coverage is 9.09091% with 20 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/quantizations.py	0.00%	13 Missing ⚠️
src/maxtext/trainers/pre_train/train.py	20.00%	2 Missing and 2 partials ⚠️
src/maxtext/utils/max_utils.py	25.00%	3 Missing ⚠️

📢 Thoughts on this report? Let us know!

[gobbleturk]

Can we include a test with the collective matmul enabled? In train_tests.py I think

[gobbleturk]

this test is getting skipped for two different reasons lol, unfortunately this feature is not protected by our CI