Description

This PR introduces vocabulary tiling, a new feature for MaxText designed to significantly reduce peak memory consumption during training.

This optimization is particularly beneficial for two key scenarios:

1. Training large-vocabulary models on limited devices: It makes it possible to train models like Gemma 2B and Gemma-2 2B, which have large vocabulary sizes, on hardware with limited memory.
2. Training with long sequence lengths: It makes training on long context lengths more feasible by avoiding the need to store the entire, large logits tensor in memory for the final loss computation.

The core idea of vocabulary tiling is to avoid explicitly materializing the full final logits tensor. Instead, the logits activation is chunked (or "tiled") along the batch-sequence dimension.

As illustrated in the diagram below, the forward and backward passes are repeated num_vocab_tiling times. In each iteration, a small slice of the logits is computed, used to calculate the loss, and the Vector-Jacobian product is immediately backpropagated. This iterative process avoids holding the complete, memory-intensive logits tensor.

Figure 1: The vocabulary tiling process. The forward and backward passes are repeated for each tile, preventing the full logits tensor from being stored in memory.

For a more in-depth technical explanation, please see the design document.

Doc: go/maxtext-vocab-tiling
FIXES: b/429255841

Tests

Correctness Tests

Test losses and embedding table gradient differences in MaxText/tests/vocab_tiling_test.py with 1% relative error tolerance in following cases:

• Default non-tied embedding (logits_via_embedding=False).
• Tied embedding and default sharding (FSDP).
• Data parallelism.
• Tensor parallelism.
• Context parallelism.

Performance Tests

See doc.

Checklist

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

• I have performed a self-review of my code.
• 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.