From 00adb5203c3234a267f410640e164281d07f3982 Mon Sep 17 00:00:00 2001 From: Chris Zuo Date: Wed, 19 Aug 2026 05:56:35 +0000 Subject: [PATCH] Add Docs for DiLoCo Fix LaTeX formatting in DiLoCo docs and support ENABLE_STREAMING_DILOCO in runner script Update Recipe 1 in DiLoCo tutorial to direct train command and keep run_spmd_streaming_diloco.sh hardcoded for streaming --- README.md | 1 + docs/reference/core_concepts.md | 8 + docs/reference/core_concepts/diloco.md | 215 ++++++++++++++++++++++++ docs/tutorials.md | 8 + docs/tutorials/diloco_pretraining.md | 221 +++++++++++++++++++++++++ 5 files changed, 453 insertions(+) create mode 100644 docs/reference/core_concepts/diloco.md create mode 100644 docs/tutorials/diloco_pretraining.md diff --git a/README.md b/README.md index 49e422b34c..be8f3abffe 100644 --- a/README.md +++ b/README.md @@ -48,6 +48,7 @@ See our guide on running MaxText in decoupled mode, without any GCP dependencies ## 🔥 Latest news 🔥 +- [August 19, 2026] Distributed low-communication training with **Streaming DiLoCo** is now supported in MaxText! Scale LLMs and MoEs across multi-slice TPU clusters and low-bandwidth DCN/WAN networks. See the [DiLoCo training tutorial](https://maxtext.readthedocs.io/en/latest/tutorials/diloco_pretraining.html) and [DiLoCo theory reference](https://maxtext.readthedocs.io/en/latest/reference/core_concepts/diloco.html) for details. - [August 5, 2026] DeepSeek V4 Flash (284B) is now supported. See the [user guide](https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/end_to_end/tpu/deepseek/Run_DeepSeek.md) for checkpoint conversion and evaluation details. - [August 3, 2026] Tokamax GMM v2 is now available for MoE kernel. See the [MoE configuration guide](https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/reference/core_concepts/moe_configuration.md) for details on how to enable it. - [July 15, 2026] RL is now supported for Qwen3 30B and GPT-OSS 20B. See the [Qwen3 30B RL tutorial](https://maxtext.readthedocs.io/en/latest/tutorials/posttraining/rl_qwen3_30b.html) and [GPT-OSS 20B RL tutorial](https://maxtext.readthedocs.io/en/latest/tutorials/posttraining/rl_gptoss_20b.html) for details. diff --git a/docs/reference/core_concepts.md b/docs/reference/core_concepts.md index 78a61f19a5..bff295cbac 100644 --- a/docs/reference/core_concepts.md +++ b/docs/reference/core_concepts.md @@ -60,6 +60,13 @@ How MaxText leverages JAX, XLA, and Pallas for efficiency. Configuring Mixture of Experts (MoE) models. ::: + +:::{grid-item-card} 🌐 DiLoCo & Streaming DiLoCo Theory +:link: core_concepts/diloco +:link-type: doc + +Bi-level optimization theory, arithmetic intensity analysis, and pipelined communication overlapping. +::: :::: ```{toctree} @@ -74,4 +81,5 @@ core_concepts/quantization.md core_concepts/tiling.md core_concepts/jax_xla_and_pallas.md core_concepts/moe_configuration.md +core_concepts/diloco.md ``` diff --git a/docs/reference/core_concepts/diloco.md b/docs/reference/core_concepts/diloco.md new file mode 100644 index 0000000000..37ef8d51b8 --- /dev/null +++ b/docs/reference/core_concepts/diloco.md @@ -0,0 +1,215 @@ + + +(diloco-theory)= + +# DiLoCo and Streaming DiLoCo Theory + +This document provides a comprehensive theoretical reference for **DiLoCo (Distributed Low-Communication)** and **Streaming DiLoCo** in MaxText. + +```{note} +For step-by-step launch commands and production recipes, see the [DiLoCo Training Tutorial](../../tutorials/diloco_pretraining.md). +``` + +______________________________________________________________________ + +## 1. Overview & Motivation + +Standard distributed training paradigms (such as Distributed Data Parallelism / FSDP) rely on high-frequency, synchronous collective communications (`all-reduce` or `reduce-scatter`) executed at **every single optimization step**. When scaling across multiple datacenter pods, geographically distributed clusters, or preemptible multi-region compute pools, inter-cluster network bandwidth becomes the primary bottleneck that throttles accelerator compute efficiency. + +DiLoCo addresses this challenge through **bi-level optimization**: + +1. **Local Inner Loop**: Multiple computing islands (e.g., TPU slices) train independently on their local data shards for $H$ steps (the *inner loop*) using fast local interconnects (such as TPU Inter-Chip Interconnect, ICI). +2. **Global Outer Loop**: Every $H$ steps, islands communicate pseudo-gradients over the slower inter-island network (Data Center Network, DCN) to execute a centralized outer momentum update. + +### Comparison of Distributed Training Paradigms + +| Feature | Synchronous DDP / FSDP | Vanilla DiLoCo | Streaming DiLoCo | +| :-------------------------------- | :----------------------------------- | :---------------------------------------------------------- | :------------------------------------------------------- | +| **Communication Frequency** | Every step ($H = 1$) | Periodic burst every $H$ steps (e.g., $H = 36$–$500$) | Pipelined every $\Delta h = \lfloor H/P \rfloor$ steps | +| **Network Bandwidth Requirement** | High (DCN/WAN bottleneck) | Low ($100\times$–$500\times$ reduction in volume over time) | Low + constant bandwidth profile | +| **Compute Idle Bubbles** | Frequent stalls on DCN collectives | Periodic barrier stall at step $H$ | **Zero stall** (computation overlaps with communication) | +| **Fault Tolerance** | Single host failure hangs global job | Isolated for $H$ steps; elastic failure recovery | Isolated + smooth staggered weight blending | +| **Primary Use Cases** | Single high-speed pod/slice (ICI) | Multi-slice, multi-cluster, WAN | High-throughput multi-slice & cross-datacenter training | + +______________________________________________________________________ + +## 2. Arithmetic Intensity & Hardware Rooflines + +The communication-to-computation trade-off is governed by **Arithmetic Intensity ($AI$)**: + +$$\text{Arithmetic Intensity } (AI) = \frac{\text{Computation FLOPs}}{\text{Communication Bytes}}$$ + +To prevent hardware from stalling on network transfers, the model's operational arithmetic intensity must exceed the physical network's hardware arithmetic intensity: + +$$AI_{\text{model}} \ge AI_{\text{hardware}} = \frac{\text{Peak Compute FLOP/s}}{\text{Network Bandwidth (Bytes/s)}}$$ + +Consider scaling DeepSeek-V3 ($671\text{B}$ MoE) or Qwen3-30B on Google TPU Ironwood / v5p with a global batch size ($GBS$) of $64\text{M}$ tokens across DCN-connected slices: + +$$\text{Max DCN Slices} = \frac{GBS}{146{,}000 \times S_{\text{seq\_len}}} \approx 14 \text{ slices}$$ + +Under standard synchronous data parallelism, scaling beyond 14 DCN slices causes the model's arithmetic intensity to drop below the hardware threshold, making communication the dominant bottleneck. + +**DiLoCo bypasses this limit by a factor of $H$**: because inter-slice communication occurs only once every $H$ steps, the effective operational arithmetic intensity scales linearly with $H$: + +$$AI_{\text{DiLoCo}} = H \times AI_{\text{standard}}$$ + +This allows scaling across dozens of TPU slices or low-bandwidth WAN connections without degrading accelerator utilization. + +______________________________________________________________________ + +## 3. Vanilla DiLoCo Algorithm + +DiLoCo splits optimization into two distinct levels: + +```text ++-----------------------------------------------------------------------------------+ +| Initial Outer Model Weights (θ_outer) | ++-----------------------------------------------------------------------------------+ + | | + | Broadcast | Broadcast + v v ++-------------------------------+ +------------------------------------+ +| Island 1 (TPU Slice 1) | | Island 2 (TPU Slice 2) | +| Local Inner Loop (Fast ICI) | | Local Inner Loop (Fast ICI) | +| Train H steps with AdamW | | Train H steps with AdamW | +| θ_1 ← local updates | | θ_2 ← local updates | +| Δθ_1 = θ_outer - θ_1 | | Δθ_2 = θ_outer - θ_2 | ++-------------------------------+ +------------------------------------+ + | | + | Send Pseudo-Gradient (Δθ_1) | Send Pseudo-Gradient (Δθ_2) + +-------------------------------+-------------------------------+ + | + v ++-----------------------------------------------------------------------------------+ +| Global Outer Optimizer (Slow Inter-Island DCN / WAN) | +| | +| 1. Global All-Reduce Average: Δθ = (Δθ_1 + Δθ_2) / 2 | +| 2. Outer Nesterov Momentum: v ← β · v + Δθ | +| θ_outer ← θ_outer - η_outer · (Δθ + β · v) | ++-----------------------------------------------------------------------------------+ + | | + +-------------------------------+-------------------------------+ + | + Broadcast New θ_outer to all Islands +``` + +### Mathematical Formulation + +1. **Inner Optimization (Local per island $k \in \{1, \dots, K\}$)**: + For local steps $t = 1, \dots, H$: + + $$\theta_{k, t} = \theta_{k, t-1} - \eta_{\text{inner}} \cdot \text{AdamW}(\nabla \mathcal{L}_k(\theta_{k, t-1}))$$ + +2. **Pseudo-Gradient Computation**: + + $$\Delta \theta_k = \theta_{\text{outer}} - \theta_{k, H}$$ + +3. **Global All-Reduce**: + + $$\overline{\Delta \theta} = \frac{1}{K} \sum_{k=1}^K \Delta \theta_k$$ + +4. **Outer Optimizer Step (Nesterov Momentum)**: + + $$v \leftarrow \beta \cdot v + \overline{\Delta \theta}$$ + + $$\theta_{\text{outer}} \leftarrow \theta_{\text{outer}} - \eta_{\text{outer}} \cdot (\overline{\Delta \theta} + \beta \cdot v)$$ + +5. **Broadcast**: + + $$\theta_{k, 0} \leftarrow \theta_{\text{outer}} \quad \forall k$$ + +______________________________________________________________________ + +## 4. Streaming DiLoCo: Pipelined Communication Overlapping + +While Vanilla DiLoCo reduces total communication volume, it introduces a **periodic barrier**: all workers must pause local training every $H$ steps to exchange full-model weights over DCN. + +**Streaming DiLoCo** ([Cherepanov et al., 2025](https://arxiv.org/abs/2501.18512)) eliminates this idle bubble through **pipelined parameter fragmentation and communication overlapping**: + +### 1. Parameter Fragmentation ($P$ Fragments) + +The model parameters $\Theta$ are partitioned into $P$ disjoint subsets via [`FragmentedTreeManipulator`](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/trainers/diloco/utils/fragmenter.py): + +- **Fragment 0**: Non-scanned parameters (token embeddings, final RMS norm, output projection head). +- **Fragments $1, \dots, P-1$**: Scanned transformer decoder layers partitioned either sequentially or interleaved. + +### 2. Staggered Synchronization Schedule + +Instead of synchronizing all parameters every $H$ steps, one fragment is synchronized every $\Delta h = \lfloor H / P \rfloor$ steps: + +$$\text{Synchronize Fragment } f = \left(\frac{t \bmod H}{\Delta h}\right) \quad \text{when } t > 0 \text{ and } t \bmod \Delta h == 0$$ + +When $H = P$, $\Delta h = 1$, synchronizing exactly one parameter fragment on **every single step**. + +### 3. Asynchronous Apply Delay ($V$) + +To overlap the cross-island DCN collective with inner-step computation, the newly updated outer weights for fragment $f$ are merged back into the local replica after a delay of $V$ steps (`num_communication_overlapping_steps`): + +$$\text{Apply Fragment } f \text{ at Step } t \quad \text{where } (t - V) \bmod \Delta h == 0$$ + +### 4. Soft Weight Blending ($\alpha$) & Delayed Merging + +An optional interpolation parameter $\alpha$ (`communication_overlapping_alpha`) smoothly blends the local replica weights with the outer synchronized weights: + +$$\theta_{\text{inner}}^{(f)} \leftarrow \alpha \cdot \theta_{\text{inner}}^{(f)} + (1 - \alpha) \cdot \theta_{\text{outer}}^{(f)}$$ + +Setting $\alpha = 0.0$ applies an exact replacement. + +```{important} +**SPMD vs. Future MPMD Multi-Threading Design**: +* `num_communication_overlapping_steps` ($V$) and `communication_overlapping_alpha` ($\alpha$) are coupled in defining the asynchronous weight merging policy. +* **Current SPMD Design**: Because JAX SPMD compiles the entire step (compute, collective reduction, and weight update) into a single synchronous XLA graph per step, setting $V > 0$ and $\alpha > 0$ **does not enhance hardware training efficiency or hide network latency**. However, it allows researchers to faithfully **simulate the algorithmic convergence behavior** of delayed weight merging and soft parameter blending on real workloads. +* **Future MPMD Multi-Threading Design**: In future MPMD architectures featuring dedicated background communication threads running independently of the TPU/GPU compute engine, $V$ and $\alpha$ will provide true, non-blocking hardware compute/communication overlap. +``` + +``` +Vanilla DiLoCo: +[------- H steps Compute -------][ Full All-Reduce ][------- H steps Compute -------] + ^^^^^^^^^^^^^^^ + Idle Bubble + +Streaming DiLoCo: +[ Compute Step 1..k ][ Compute Step k+1..2k ][ Compute Step 2k+1..3k ] ... + └─ Sync Frag 0 ────┘ └─ Sync Frag 1 ─────┘ └─ Sync Frag 2 ─────┘ (Continuous Fragment Pipelining) +``` + +______________________________________________________________________ + +## 5. Pure JAX SPMD & NNX Architecture + +MaxText integrates DiLoCo natively with **JAX SPMD and NNX** without relying on external sidecars or multi-controller processes: + +### State Representation (`DiLoCoTrainState`) + +In [`src/maxtext/trainers/diloco/diloco.py`](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/trainers/diloco/diloco.py): + +- `inner_state`: Per-replica NNX `TrainState` containing sharded weights and AdamW states with a leading `diloco` mesh axis of dimension $K$ (`dcn_diloco_parallelism`). +- `params`: Global synchronized outer model parameters (PyTree of `Param` leaves). +- `outer_opt_state`: Optax Nesterov momentum state `(TraceState(trace=...), EmptyState())`. +- `step`: Global step tensor. + +### Multi-Placement Execution with `drjax` + +Local training steps are mapped across replicas using `@drjax.program(placements={"diloco": K})` and `drjax.map_fn`. Collectives across islands use `drjax.reduce_mean` and `drjax.broadcast`. + +______________________________________________________________________ + +## 6. References + +1. **DiLoCo**: Douillard, A., Su, Y., Roberts, A., et al. *DiLoCo: Distributed Low-Communication Training of Language Models*. [arXiv:2311.08105](https://arxiv.org/abs/2311.08105), 2023. +2. **Streaming DiLoCo**: Cherepanov, A., et al. *Streaming DiLoCo with overlapping communication: Towards a Distributed Free Lunch*. [arXiv:2501.18512](https://arxiv.org/abs/2501.18512), 2025. +3. **MaxText Sharding & Arithmetic Intensity Guide**: [Sharding on TPUs](../../guides/optimization/sharding.md). diff --git a/docs/tutorials.md b/docs/tutorials.md index 666c38e9d8..3ddffdc495 100644 --- a/docs/tutorials.md +++ b/docs/tutorials.md @@ -37,6 +37,13 @@ Step-by-step guide for building MaxText docker images. Step-by-step guides for pre-training with real datasets like C4 using HuggingFace, Grain, or TFDS. ``` +```{grid-item-card} 🌐 DiLoCo Training +:link: tutorials/diloco_pretraining +:link-type: doc + +Step-by-step production recipes for distributed low-communication multi-slice training with DiLoCo and Streaming DiLoCo. +``` + ```{grid-item-card} 🧩 Post-training :link: tutorials/post_training_index :link-type: doc @@ -59,6 +66,7 @@ maxdepth: 1 --- tutorials/build_maxtext.md tutorials/pretraining.md +tutorials/diloco_pretraining.md tutorials/post_training_index.md tutorials/inference.md ``` diff --git a/docs/tutorials/diloco_pretraining.md b/docs/tutorials/diloco_pretraining.md new file mode 100644 index 0000000000..1ee5812bdf --- /dev/null +++ b/docs/tutorials/diloco_pretraining.md @@ -0,0 +1,221 @@ + + +(diloco-pretraining)= + +# DiLoCo (Distributed Low-Communication) Training + +This tutorial guides you through configuring and running **DiLoCo** and **Streaming DiLoCo** training in MaxText across multi-slice TPU clusters, multi-datacenter pods, and low-bandwidth DCN/WAN networks. + +```{note} +For theoretical background, arithmetic intensity calculations, and in-depth architectural details, refer to the [DiLoCo Theory & Mathematics Reference](../reference/core_concepts/diloco.md). +``` + +______________________________________________________________________ + +## 1. Vanilla DiLoCo vs. Streaming DiLoCo + +MaxText supports two modes of distributed low-communication training: + +``` +Vanilla DiLoCo: +[------- H steps Compute (Local ICI) -------][ Full Model All-Reduce (DCN) ][------- H steps Compute -------] + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + Periodic Barrier Pause + +Streaming DiLoCo: +[ Step 1 Compute ][ Step 2 Compute ][ Step 3 Compute ] ... + └─ Sync Frag 0 ─┘ └─ Sync Frag 1 ─┘ └─ Sync Frag 2 ─┘ (Continuous Pipelined Synchronization) +``` + +### Key Differences: + +- **Vanilla DiLoCo (`enable_streaming_diloco=false`)**: + + - **How it works**: Each computing island trains independently for $H$ inner steps (e.g., $H=100$). At every $H$-th step, training pauses for a global collective all-reduce where the entire model's pseudo-gradient ($\Delta \theta = \theta_{\text{outer}} - \theta_{\text{inner}}$) is averaged across all islands over DCN and updated using outer Nesterov momentum. + - **When to use**: Simpler baseline, ideal when $H$ is large (e.g. $H \ge 500$) and the periodic all-reduce pause represents a negligible fraction of total training time. + +- **Streaming DiLoCo (`enable_streaming_diloco=true`)**: + + - **How it works**: The model parameters are partitioned into $P$ fragments (typically $P = N_{\text{layers}} + 1$). By setting $H = P$, exactly 1 fragment is synchronized on every single local inner step ($\Delta h = 1$). + - **When to use**: Optimal for high-throughput scaling across lower-bandwidth DCN/WAN networks, as it eliminates bursty communication spikes and removes the periodic step-$H$ idle barrier. + +______________________________________________________________________ + +## 2. Prerequisites + +1. **MaxText Environment**: Follow the [installation guide](../install_maxtext.md) to set up your environment (`maxtext[tpu]` or `maxtext[cuda12]`). +2. **Compute Resources**: A Google Kubernetes Engine (GKE) cluster with TPU slices managed via [XPK](https://github.com/AI-Hypercomputer/xpk). +3. **Storage**: A Google Cloud Storage (GCS) bucket for logging and Orbax checkpoints (`gs://`). + +______________________________________________________________________ + +## 3. Production Recipe 1: Vanilla DiLoCo Multi-Slice Pre-training + +In this recipe, we train a model (e.g., **Qwen3-8B**) across **2 TPU v5p-128 slices** using Vanilla DiLoCo with periodic synchronization every $H=100$ steps: + +```bash +python3 -m maxtext.trainers.pre_train.train \ + maxtext/configs/base.yml \ + run_name="vanilla-dlco-8b-01" \ + base_output_directory="gs://your-bucket/maxtext-logs" \ + dataset_path="gs://your-bucket/maxtext-datasets" \ + dataset_name='c4/en:3.0.1' \ + model_name="qwen3-8b" \ + per_device_batch_size=8 \ + max_target_length=2048 \ + enable_diloco=true \ + enable_streaming_diloco=false \ + dcn_diloco_parallelism=2 \ + diloco_sync_period=100 \ + diloco_outer_lr=0.7 \ + diloco_outer_momentum=0.9 \ + pure_nnx=true \ + steps=1000 \ + enable_checkpointing=true \ + checkpoint_period=100 +``` + +### Configuration Breakdown: + +- `enable_diloco=true`: Enables outer optimization and multi-slice Low-Communication training across `dcn_diloco_parallelism=2` slices. +- `enable_streaming_diloco=false`: Disables parameter fragmentation and performs full-model pseudo-gradient all-reduce. +- `diloco_sync_period=100`: Islands execute 100 local AdamW steps independently before pausing to sync. +- `diloco_outer_lr=0.7` and `diloco_outer_momentum=0.9`: Outer Nesterov momentum parameters. + +______________________________________________________________________ + +## 4. Production Recipe 2: Streaming DiLoCo Dense Pre-training (Qwen3-8B) + +In this recipe, we train **Qwen3-8B** with Streaming DiLoCo across **2 TPU v5p-128 slices** with $H=P=37$ (synchronizing 1 fragment every step) via the SPMD runner script: + +```bash +CLUSTER="mlperf-v5p" \ +ZONE="europe-west4-b" \ +PROJECT="cloud-tpu-multipod-dev" \ +DEVICE_TYPE="v5p-128" \ +NUM_SLICES="2" \ +RUNNAME="stream-dlco-8b-01" \ +XPK_WORKLOAD="stream-dlco-01" \ +BASE_OUTPUT_DIRECTORY="gs://your-bucket/maxtext-logs" \ +DATASET_PATH="gs://your-bucket/maxtext-datasets" \ +MODEL_NAME="qwen3-8b" \ +STEPS="1000" \ +CHECKPOINT_PERIOD="100" \ +DILOCO_SYNC_PERIOD="37" \ +DILOCO_NUM_FRAGMENTS="37" \ +DILOCO_NUM_COMM_OVERLAP_STEPS="0" \ +DILOCO_OUTER_LR="0.7" \ +DILOCO_OUTER_MOMENTUM="0.9" \ +bash src/maxtext/trainers/diloco/scripts/run_spmd_streaming_diloco.sh +``` + +______________________________________________________________________ + +## 5. Production Recipe 3: Streaming DiLoCo MoE Pre-training (Qwen3-30B-A3B) + +For large Mixture-of-Experts (MoE) architectures, this recipe demonstrates Streaming DiLoCo pre-training with the **OLMo Grain** data pipeline across 2x `v5p-128` TPU slices: + +```bash +XPK_CLUSTER="mlperf-v5p" \ +XPK_ZONE="europe-west4-b" \ +XPK_PROJECT="cloud-tpu-multipod-dev" \ +XPK_DEVICE_TYPE="v5p-128" \ +XPK_NUM_SLICES="2" \ +RUN_NAME="qw3-olmo-dlco-01" \ +WORKLOAD_NAME="qw3-olmo-01" \ +BASE_OUTPUT_DIRECTORY="gs://your-bucket/maxtext-logs" \ +OLMO_GCS_BASE="gs://your-bucket/datasets" \ +MODEL_NAME="qwen3-30b-a3b" \ +ENABLE_STREAMING_DILOCO="true" \ +DILOCO_SYNC_PERIOD="49" \ +DILOCO_NUM_FRAGMENTS="49" \ +DILOCO_OUTER_LR="0.7" \ +DILOCO_OUTER_MOMENTUM="0.9" \ +bash src/maxtext/trainers/diloco/scripts/run_olmo_qwen3_30b_streaming_diloco.sh +``` + +- **49 Fragments**: 48 MoE transformer decoder layers + 1 embedding/head fragment ($H=49, P=49$). +- **Grain Pipeline**: Deterministic streaming from pre-tokenized numpy arrays. + +______________________________________________________________________ + +## 6. Checkpointing & Resumption + +### Automatic Resumption + +To resume an interrupted DiLoCo pre-training run, submit the workload with the same `RUNNAME` and `BASE_OUTPUT_DIRECTORY`: + +```bash +RUNNAME="stream-dlco-8b-01" \ +XPK_WORKLOAD="dlco-resm-01" \ +STEPS="2000" \ +bash src/maxtext/trainers/diloco/scripts/run_spmd_streaming_diloco.sh +``` + +MaxText detects existing Orbax checkpoints, restores both the per-replica inner optimizer moments and the outer Nesterov momentum state, and continues training seamlessly. + +### Bootstrapping from Single-Slice Weights + +To initialize a multi-slice DiLoCo run from standard pre-trained single-slice weights, specify `LOAD_FULL_STATE_PATH`: + +```bash +LOAD_FULL_STATE_PATH="gs://your-bucket/checkpoints/base_model/0/items" \ +bash src/maxtext/trainers/diloco/scripts/run_spmd_streaming_diloco.sh +``` + +MaxText broadcasts the single-slice model weights across all islands and initializes a clean outer optimizer state automatically. + +______________________________________________________________________ + +## 7. Tuning Guidelines + +| Hyperparameter | Recommended Setting | Description | +| :------------------------------------------- | :--------------------------- | :----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `diloco_sync_period` ($H$) | **`= num_diloco_fragments`** | Sync period. Setting $H = P$ ensures exactly 1 fragment is synchronized every local step ($\Delta h = 1$). | +| `num_diloco_fragments` ($P$) | **`num_layers + 1`** | Partition count (1 for non-scanned embeddings/head + 1 per transformer decoder layer). | +| `diloco_outer_lr` ($\eta_{\text{outer}}$) | **`0.3` – `0.9`** | Outer learning rate. Start from `0.3` – `0.9` (e.g. `0.7`) and tune based on inner LR. | +| `diloco_outer_momentum` ($\beta$) | `0.9` | Nesterov momentum coefficient for the outer optimizer. | +| `num_communication_overlapping_steps` ($V$) | `0` (or `1`) | Delay in inner steps before applying outer weights. In the current SPMD design, this does not enhance hardware efficiency but simulates the algorithmic behavior of delayed weight merging; it will provide non-blocking hardware overlap in future MPMD multi-threading. Coupled with $\alpha$. | +| `communication_overlapping_alpha` ($\alpha$) | `0.0` | Soft parameter blending factor ($\theta_{\text{inner}} \leftarrow \alpha \theta_{\text{inner}} + (1 - \alpha) \theta_{\text{outer}}$). Simulates soft weight interpolation in current SPMD; will be effective alongside $V$ in future MPMD. `0.0` applies direct replacement. | + +### Practical Tuning Heuristics + +- **Synchronize Every Step ($H = P = N_{\text{layers}} + 1$)**: + Setting `diloco_sync_period` equal to `num_diloco_fragments` with $P = N_{\text{layers}} + 1$ (e.g., $H=37, P=37$ for 36-layer models like Qwen3-8B, or $H=49, P=49$ for 48-layer models like Qwen3-30B) ensures a steady, constant stream of background communications by syncing 1 fragment on every local step. + +- **Outer Learning Rate Tuning & Inverse Scaling Rule**: + Outer LR should be tuned alongside the inner optimizer learning rate. As a core heuristic: + + $$\text{Higher Inner Learning Rate} \implies \text{Lower Outer Learning Rate}$$ + + $$\text{Lower Inner Learning Rate} \implies \text{Higher Outer Learning Rate}$$ + + When using standard AdamW inner optimization, starting with `diloco_outer_lr: 0.7` is a strong baseline. + +- **Overlapping Steps ($V$) and Alpha ($\alpha$) in SPMD vs. MPMD**: + `num_communication_overlapping_steps` ($V$) and `communication_overlapping_alpha` ($\alpha$) are coupled in defining the asynchronous weight merging policy: + + - **Current SPMD Design**: Because JAX SPMD compiles each step into a synchronous XLA graph, setting $V > 0$ or $\alpha > 0$ **does not enhance hardware training efficiency or hide network latency**. However, it allows researchers to accurately **simulate the algorithmic convergence behavior** of delayed weight merging and soft parameter blending. Setting $V=0$ is standard for performance. + - **Future MPMD Multi-Threading Design**: In upcoming MPMD architectures with independent background communication threads, $V$ and $\alpha$ will provide true, non-blocking hardware compute/communication overlap. + +______________________________________________________________________ + +## 8. Next Steps + +- Deep dive into mathematical foundations: [DiLoCo Theory & Mathematics Reference](../reference/core_concepts/diloco.md). +- Explore input pipeline options: [Data Input Pipeline Guides](../guides/data_input_pipeline.md). +- Learn about sharding strategies on TPUs: [Sharding on TPUs](../guides/optimization/sharding.md).