Non-SPMD DiLoCo

src/maxtext/configs/base.yml

@@ -896,6 +896,15 @@ dcn_bandwidth_latency: "50ms"
 # The network interface to apply throttling rules to.
 dcn_bandwidth_interface: "eth0"
 
+# Streaming DiLoCo params
+enable_streaming_diloco: false
+enable_non_spmd_diloco: false
+num_diloco_fragments: 16
+use_sequential_layers: false
+num_communication_overlapping_steps: 0
+communication_overlapping_alpha: 0.0
+
+
 # You may disable clipping by setting gradient_clipping_threshold to zero.
 gradient_clipping_threshold: 1.0
 

src/maxtext/configs/types.py

@@ -1521,14 +1521,45 @@ class DilocoParams(BaseModel):
   diloco_outer_lr: float = Field(0.3, description="learning rate for outer optimizer.")
   diloco_outer_momentum: float = Field(0.9, description="momentum for outer optimizer.")
   dcn_bandwidth_limit: str = Field(
-      "", description="Programmatic DCN egress bandwidth limit (e.g., '28gbit'). Empty means no limit."
+      "", description="Programmatic DCN egress bandwidth limit per VM (e.g., '28gbit'). Empty means no limit."
   )
   dcn_bandwidth_burst: str = Field("10mb", description="Burst size for Token Bucket Filter (TBF) traffic shaping.")
   dcn_bandwidth_latency: str = Field(
       "50ms", description="Latency threshold for Token Bucket Filter (TBF) traffic shaping."
   )
   dcn_bandwidth_interface: str = Field("eth0", description="Network interface to apply bandwidth limits on.")
 
+  # Streaming DiLoCo parameters
+  enable_streaming_diloco: bool = Field(False, description="Enable streaming DiLoCo parallelism.")
+  enable_non_spmd_diloco: bool = Field(False, description="Enable non-SPMD, multi-threaded streaming DiLoCo.")
+  num_diloco_fragments: int = Field(16, description="Number of fragments to partition the model layers into.")
+  use_sequential_layers: bool = Field(False, description="Whether to sync layers sequentially (or interleaved).")
+  num_communication_overlapping_steps: int = Field(
+      0, description="Steps of communication overlap with computation. \\tau from the paper."
+  )
+  communication_overlapping_alpha: float = Field(
+      0.0,
+      description=(
+          "Interpolation factor between local and global parameters. alpha=1"
+          " means no communication between islands, alpha=0 means discards any"
+          " updates done in the inner optimizer in the first"
+          " `num_communication_overlapping_steps` steps. alpha=0.5 does a"
+          " uniform average between the local fragment parameters and the"
+          " globally shared one."
+      ),
+  )
+
+  @model_validator(mode="after")
+  def validate_overlap_steps(self) -> "DilocoParams":
+    if self.enable_streaming_diloco:
+      if self.num_communication_overlapping_steps >= self.diloco_sync_period:
+        raise ValueError(
+            "num_communication_overlapping_steps must be strictly less than"
+            f" diloco_sync_period. Got {self.num_communication_overlapping_steps=}"
+            f" and {self.diloco_sync_period=}"
+        )
+    return self
+
 
 class Optimizer(BaseModel):
   """Configuration for the optimizer and learning rate schedule."""
@@ -3025,6 +3056,8 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
           "Colocated python data input is only supported with Pathways (single"
           " controller) enabled (`enable_single_controller=True`)."
       )
+    if self.enable_non_spmd_diloco and self.colocated_python_data_input:
+      raise ValueError("Non-SPMD DiLoCo does not support colocated python data input.")
     if self.grain_use_elastic_iterator and self.grain_file_type != "arrayrecord":
       raise ValueError(
           "`grain_use_elastic_iterator=True` only supports `grain_file_type=arrayrecord`. "

src/maxtext/input_pipeline/grain_data_processing.py

@@ -422,6 +422,8 @@ def make_grain_train_iterator(
     config: ml_collections.ConfigDict,
     global_mesh,
     process_indices,
+    learner_idx: int = 0,
+    num_learners: int = 1,
 ):
   """Load, preprocess dataset and return iterators"""
   assert (
@@ -478,9 +480,11 @@ def make_grain_train_iterator(
   if 0 < config.expansion_factor_real_data < 1:
     num_dataloader_to_restore = int(1 / config.expansion_factor_real_data)
     train_dataloader_list = []
-    dataloading_host_count = len(process_indices) * num_dataloader_to_restore
+    dataloading_host_count = len(process_indices) * num_dataloader_to_restore * num_learners
     for i in range(num_dataloader_to_restore):
-      dataloading_host_index = len(process_indices) * i + process_indices.index(jax.process_index())
+      dataloading_host_index = (
+          len(process_indices) * i + process_indices.index(jax.process_index())
+      ) * num_learners + learner_idx
       train_ds = get_ds_fn(dataloading_host_index=dataloading_host_index, dataloading_host_count=dataloading_host_count)
       train_dataloader = preprocessing_fn(dataset=train_ds)
       train_dataloader_list.append(train_dataloader)
@@ -490,8 +494,8 @@ def make_grain_train_iterator(
     ]
 
   # Default non-colocated, non-expansion path
-  shard_index = process_indices.index(jax.process_index())
-  shard_count = len(process_indices)
+  shard_index = process_indices.index(jax.process_index()) * num_learners + learner_idx
+  shard_count = len(process_indices) * num_learners
   train_ds = get_ds_fn(
       dataloading_host_index=shard_index,
       dataloading_host_count=shard_count,
@@ -523,6 +527,8 @@ def make_grain_eval_iterator(
     config: ml_collections.ConfigDict,
     global_mesh,
     process_indices,
+    learner_idx: int = 0,
+    num_learners: int = 1,
 ):
   """Load, preprocess dataset and return iterators"""
   assert (
@@ -556,8 +562,8 @@ def make_grain_eval_iterator(
 
   if not config.colocated_python_data_input:
     eval_ds = get_ds_fn(
-        dataloading_host_index=process_indices.index(jax.process_index()),
-        dataloading_host_count=len(process_indices),
+        dataloading_host_index=process_indices.index(jax.process_index()) * num_learners + learner_idx,
+        dataloading_host_count=len(process_indices) * num_learners,
     )
     eval_dataloader = preprocessing_fn(dataset=eval_ds)
     return multihost_dataloading.MultiHostDataLoadIterator(

src/maxtext/input_pipeline/input_pipeline_interface.py

@@ -14,6 +14,7 @@
 
 """Input pipeline"""
 import functools
+import inspect
 
 import jax
 from jax.sharding import PartitionSpec as P
@@ -46,20 +47,33 @@ def get_process_loading_real_data(
   return list(process_loading_real_data)
 
 
-def create_process_specific_iterator(config: pyconfig.HyperParameters, mesh, process_indices, input_iterator):
+def create_process_specific_iterator(
+    config: pyconfig.HyperParameters,
+    mesh,
+    process_indices,
+    input_iterator,
+    learner_idx: int = 0,
+    num_learners: int = 1,
+):
   """
   If the current process's index is among the `process_indices`, a real
   data iterator is created. Otherwise, a placeholder iterator is returned.
   """
   if jax.process_index() in process_indices:
-    iterator_fn = functools.partial(input_iterator, config, mesh, process_indices)
+    sig = inspect.signature(input_iterator)
+    kwargs = {}
+    if "learner_idx" in sig.parameters:
+      kwargs["learner_idx"] = learner_idx
+    if "num_learners" in sig.parameters:
+      kwargs["num_learners"] = num_learners
+    iterator_fn = functools.partial(input_iterator, config, mesh, process_indices, **kwargs)
     output_iterator = iterator_fn()
   else:
     output_iterator = PlaceHolderDataIterator(config, mesh)
   return output_iterator
 
 
-def create_data_iterator(config: pyconfig.HyperParameters, mesh):
+def create_data_iterator(config: pyconfig.HyperParameters, mesh, learner_idx: int = 0, num_learners: int = 1):
   """Create train and eval data iterators given configs and mesh."""
 
   # Return synthetic dataset if selected
@@ -99,7 +113,9 @@ def create_data_iterator(config: pyconfig.HyperParameters, mesh):
       config.max_target_length,
       mesh,
   )
-  output_train_iterator = create_process_specific_iterator(config, mesh, process_indices_train, train_iterator)
+  output_train_iterator = create_process_specific_iterator(
+      config, mesh, process_indices_train, train_iterator, learner_idx=learner_idx, num_learners=num_learners
+  )
   if config.expansion_factor_real_data > 1:  # assert number of hosts loading real data
     assert len(process_indices_train) == jax.process_count() // config.expansion_factor_real_data
 
@@ -116,5 +132,7 @@ def create_data_iterator(config: pyconfig.HyperParameters, mesh):
 
     if config.expansion_factor_real_data > 1:
       assert len(process_indices_eval) == jax.process_count() // config.expansion_factor_real_data
-    output_eval_iterator = create_process_specific_iterator(config, mesh, process_indices_eval, eval_iterator)
+    output_eval_iterator = create_process_specific_iterator(
+        config, mesh, process_indices_eval, eval_iterator, learner_idx=learner_idx, num_learners=num_learners
+    )
   return output_train_iterator, output_eval_iterator

src/maxtext/input_pipeline/tfds_data_processing.py

@@ -185,6 +185,8 @@ def make_tfds_train_iterator(
     config: ml_collections.ConfigDict,
     global_mesh,
     process_indices_train,
+    learner_idx: int = 0,
+    num_learners: int = 1,
 ):
   """load dataset, preprocess and return iterators"""
   assert (
@@ -200,8 +202,8 @@ def make_tfds_train_iterator(
   }
   if not config.colocated_python_data_input:
     train_ds = get_datasets(
-        dataloading_host_index=process_indices_train.index(jax.process_index()),
-        dataloading_host_count=len(process_indices_train),
+        dataloading_host_index=process_indices_train.index(jax.process_index()) * num_learners + learner_idx,
+        dataloading_host_count=len(process_indices_train) * num_learners,
         **get_datasets_kwargs,
     )
     train_dataloader = preprocessing_pipeline(
@@ -254,6 +256,8 @@ def make_tfds_eval_iterator(
     config: ml_collections.ConfigDict,
     global_mesh,
     process_indices_eval,
+    learner_idx: int = 0,
+    num_learners: int = 1,
 ):
   """load eval dataset, preprocess and return iterators"""
   assert (
@@ -266,8 +270,8 @@ def make_tfds_eval_iterator(
         data_split=config.eval_split,
         shuffle_files=False,
         shuffle_seed=config.data_shuffle_seed,
-        dataloading_host_index=process_indices_eval.index(jax.process_index()),
-        dataloading_host_count=len(process_indices_eval),
+        dataloading_host_index=process_indices_eval.index(jax.process_index()) * num_learners + learner_idx,
+        dataloading_host_count=len(process_indices_eval) * num_learners,
     )
     eval_dataloader = preprocessing_pipeline(
         dataset=eval_ds,

src/maxtext/trainers/diloco/decomposed_transport.py

@@ -0,0 +1,148 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Thread-safe transport layer for non-SPMD multi-threaded DiLoCo."""
+
+import queue
+from typing import Any
+from concurrent.futures import ThreadPoolExecutor
+import traceback
+import jax
+import jax.numpy as jnp
+from jax.experimental import colocated_python
+from maxtext.utils import max_logging
+
+
+
+class ThreadedTransportManager:
+  """Manages in-memory communication between learner threads and the syncer thread."""
+
+  def __init__(self, num_learners: int):
+    self.num_learners = num_learners
+
+    # Thread-safe FIFO queues for each learner.
+    # Stores tuples of (step, fragment_id, data).
+    self._learner_to_syncer_queues = [queue.Queue() for _ in range(num_learners)]
+    self._syncer_to_learner_queues = [queue.Queue() for _ in range(num_learners)]
+
+    # Local buffers for out-of-order storage.
+    # Since only the receiving thread (syncer for learner_to_syncer, and respective
+    # learner for syncer_to_learner) accesses these buffers, we do not need locks.
+    self._learner_to_syncer_buffers = [{} for _ in range(num_learners)]
+    self._syncer_to_learner_buffers = [{} for _ in range(num_learners)]
+
+  def send_to_syncer(self, learner_idx: int, step: int, fragment_id: int, data: Any):
+    """Learner sends data to the syncer."""
+    self._learner_to_syncer_queues[learner_idx].put((step, fragment_id, data))
+
+  def recv_from_learner(self, learner_idx: int, step: int, fragment_id: int) -> Any:
+    """Syncer receives data from a specific learner. Blocks if not available."""
+    key = (step, fragment_id)
+    buffer = self._learner_to_syncer_buffers[learner_idx]
+    if key in buffer:
+      return buffer.pop(key)
+
+    while True:
+      rec_step, rec_frag, data = self._learner_to_syncer_queues[learner_idx].get()
+      if rec_step == step and rec_frag == fragment_id:
+        return data
+      buffer[(rec_step, rec_frag)] = data
+
+  def send_to_learner(self, learner_idx: int, step: int, fragment_id: int, data: Any):
+    """Syncer sends data to a specific learner."""
+    self._syncer_to_learner_queues[learner_idx].put((step, fragment_id, data))
+
+  def recv_from_syncer(self, learner_idx: int, step: int, fragment_id: int) -> Any:
+    """Learner receives data from the syncer. Blocks if not available."""
+    key = (step, fragment_id)
+    buffer = self._syncer_to_learner_buffers[learner_idx]
+    if key in buffer:
+      return buffer.pop(key)
+
+    while True:
+      rec_step, rec_frag, data = self._syncer_to_learner_queues[learner_idx].get()
+      if rec_step == step and rec_frag == fragment_id:
+        return data
+      buffer[(rec_step, rec_frag)] = data
+
+
+class LearnerTransport:
+  """Wrapper for learner threads to communicate with the syncer."""
+
+  def __init__(
+      self,
+      manager: ThreadedTransportManager,
+      learner_idx: int,
+      local_cpu_mesh: jax.sharding.Mesh,
+  ):
+    self.manager = manager
+    self.learner_idx = learner_idx
+    self.local_cpu_mesh = local_cpu_mesh
+    self._executor = ThreadPoolExecutor(max_workers=1)
+
+  def send_to_syncer_async(self, step: int, fragment_id: int, data: Any):
+    """Asynchronously offloads TPU data to local CPU mesh and sends to syncer."""
+    # 1. Asynchronously offload to CPU colocated mesh (non-blocking on main thread)
+    cpu_sharding = jax.tree_util.tree_map(
+        lambda s: jax.sharding.NamedSharding(self.local_cpu_mesh, s.spec),
+        jax.tree_util.tree_map(lambda x: x.sharding, data),
+    )
+    frag_cpu = jax.device_put(data, cpu_sharding)
+
+    # 2. Block and send in the background executor thread
+    def _send():
+      try:
+        max_logging.log(f"Learner {self.learner_idx}: async send starting for step {step} frag {fragment_id}")
+        jax.block_until_ready(frag_cpu)
+        max_logging.log(f"Learner {self.learner_idx}: async send block_until_ready done")
+        self.manager.send_to_syncer(self.learner_idx, step, fragment_id, frag_cpu)
+        max_logging.log(f"Learner {self.learner_idx}: async send sent to syncer")
+      except Exception as e:
+        max_logging.error(f"Learner {self.learner_idx}: async send failed: {e}")
+        max_logging.error(traceback.format_exc())
+        raise e
+
+    self._executor.submit(_send)
+
+  def send_to_syncer(self, step: int, fragment_id: int, data: Any):
+    """Synchronously offloads TPU data to local CPU mesh and sends to syncer."""
+    cpu_sharding = jax.tree_util.tree_map(
+        lambda s: jax.sharding.NamedSharding(self.local_cpu_mesh, s.spec),
+        jax.tree_util.tree_map(lambda x: x.sharding, data),
+    )
+    frag_cpu = jax.device_put(data, cpu_sharding)
+    jax.block_until_ready(frag_cpu)
+    self.manager.send_to_syncer(self.learner_idx, step, fragment_id, frag_cpu)
+
+  def recv_from_syncer(self, step: int, fragment_id: int) -> Any:
+    return self.manager.recv_from_syncer(self.learner_idx, step, fragment_id)
+
+  def close(self):
+    """Shutdown background thread executor."""
+    self._executor.shutdown(wait=True)
+
+
+class SyncerTransport:
+  """Wrapper for the syncer thread to communicate with learners."""
+
+  def __init__(self, manager: ThreadedTransportManager):
+    self.manager = manager
+
+  def send_to_learner(self, learner_idx: int, step: int, fragment_id: int, data: Any):
+    self.manager.send_to_learner(learner_idx, step, fragment_id, data)
+
+  def recv_from_learner(self, learner_idx: int, step: int, fragment_id: int) -> Any:
+    return self.manager.recv_from_learner(learner_idx, step, fragment_id)
+
+