Support TE Collective GEMM for Attention Layer Projections

[dinodeep]

Description

This PR adds support for performing collective GEMMS on projections within the attention layer: Q, K, V, and output projections when the user was previously only allowed to perform collective GEMMs on MLP GEMMs. Collective GEMMs reduce E2E execution time by better overlapping computation with exposed communication, and GEMMs in the attention layer can be large enough to benefit from such overlap (especially with QKV are fused).

This PR changes the way the user enables collectives from a boolean to using te_comm_gemm_overlap=(disabled|mlp|full). "disabled" means that there are no collective GEMMs enabled; "mlp" means that only the MLP GEMMs are turned into CGEMMs, and "full" means that the MLP GEMMs and the QKV/output GEMMs from the attention layer are overlapped. The default remains as disabled/off.

Also, some small pre-commit changes were included for formatting.

Tests

I tested this change by varying the flags across the three different configurations and validating via profiling that each of the three configurations produces collective GEMMs in the appropriate places.

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.

[codecov]

Codecov Report

❌ Patch coverage is 31.25000% with 11 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/quantizations.py	23.07%	10 Missing ⚠️
src/maxtext/trainers/pre_train/train.py	50.00%	0 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!