[DiLoCo] Add Docs for Streaming DiLoCo

[Dr-Left]

Description

Adds comprehensive reference documentation and a production tutorial for DiLoCo (Distributed Low-Communication) and Streaming DiLoCo training in MaxText.

Changes Made

1. DiLoCo Theory & Mathematics Reference (docs/reference/core_concepts/diloco.md):

   • Details the mathematical formulation of bi-level optimization (Inner local AdamW, pseudo-gradient computation, outer Nesterov momentum).
   • Derives operational Arithmetic Intensity ($AI = \frac{\text{FLOPs}}{\text{Bytes}}$) and physical network rooflines demonstrating how DiLoCo scales effective AI by a factor of $H$.
   • Formulates Streaming DiLoCo parameter fragmentation ($P$ fragments), staggered sync scheduling ($\Delta h = \lfloor H/P \rfloor$), asynchronous apply delay ($V$), and soft parameter blending ($\alpha$).
   • Explains the native JAX SPMD + NNX design (DiLoCoTrainState, drjax.program multi-placement, FragmentedTreeManipulator).
   • Clarifies the distinction between current SPMD algorithmic simulation and future MPMD multi-threading hardware overlap for $V$ and $\alpha$.
   • Integrated into docs/reference/core_concepts.md.

2. DiLoCo Pre-training Tutorial (docs/tutorials/diloco_pretraining.md):

   • Provides a clear comparison between Vanilla DiLoCo (periodic full-model barrier) and Streaming DiLoCo (continuous pipelined fragment communication).
   • Includes production recipes for:
     • Recipe 1: Vanilla DiLoCo Multi-Slice Pre-training (Qwen3-8B on 2x v5p-128).
     • Recipe 2: Streaming DiLoCo Dense Pre-training (Qwen3-8B on 2x v5p-128 with $H=37, P=37$).
     • Recipe 3: Streaming DiLoCo MoE Pre-training (Qwen3-30B-A3B with Grain data loading).
   • Documents multi-replica Orbax checkpointing, automatic resumption, and single-slice parameter bootstrapping.
   • Provides concrete hyperparameter tuning tables and heuristics ($H = P = N_{\text{layers}} + 1$, Outer LR 0.3–0.9 inverse scaling rule).
   • Integrated as a standalone tutorial in docs/tutorials.md.

3. Homepage & Index Updates:

   • Added DiLoCo and Streaming DiLoCo support announcement to ## 🔥 Latest news 🔥 in README.md (rendered on docs/index.md).
   • Updated pre-training overview references in README.md.

Future Improvements

• In the current JAX SPMD design, num_communication_overlapping_steps and communication_overlapping_alpha simulate the algorithmic behavior of delayed weight merging without cross-step asynchronous hardware overlap. Future MPMD multi-threading designs will leverage these parameters for true non-blocking hardware compute/communication overlap.

Tests

• Sphinx Documentation Build: Verified clean HTML compilation with zero errors:

  source maxtext_venv/bin/activate
  cd docs && sphinx-build -b html . _build/html

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.

[gemini-code-assist]

Code Review

This pull request adds comprehensive documentation and tutorials for Vanilla and Streaming DiLoCo training in MaxText, detailing the theoretical background, arithmetic intensity, and production recipes. The review feedback highlights critical formatting issues in the LaTeX math blocks that prevent proper rendering, as well as a mismatch in the Vanilla DiLoCo recipe where the referenced shell script does not support the ENABLE_STREAMING_DILOCO environment variable.

[huytransformer]

bi level optimization!!