Merge pull request #4759 from AI-Hypercomputer:chris/spmd-streaming-dlco

PiperOrigin-RevId: 966865365

Author: Google-ML-Automation

src/maxtext/common/checkpointing.py

@@ -24,6 +24,7 @@
 
 from etils import epath
 from flax import nnx
+from flax import struct
 
 
 from flax.training import train_state
@@ -36,6 +37,7 @@
 from maxtext.input_pipeline.multihost_dataloading import MultiHostDataLoadIterator
 from maxtext.input_pipeline.multihost_dataloading import RemoteIteratorWrapper
 from maxtext.input_pipeline.synthetic_data_processing import PlaceHolderDataIterator
+from maxtext.trainers.diloco.utils import spmd_diloco_checkpointing as diloco_checkpoint_utils
 from maxtext.utils import elastic_utils
 from maxtext.utils import exceptions
 from maxtext.utils import gcs_utils
@@ -186,6 +188,16 @@ def _load_linen_checkpoint_into_nnx(
   present, else keep their fresh init value. A genuinely-missing weight raises.
   """
   max_logging.log(f"Restoring Linen-layout checkpoint into NNX state at {path}")
+  if config and getattr(config, "enable_diloco", False):
+    return diloco_checkpoint_utils.restore_diloco_checkpoint(
+        path,
+        abstract_nnx_state,
+        checkpoint_storage_concurrent_gb,
+        use_ocdbt=use_ocdbt,
+        use_zarr3=use_zarr3,
+        config=config,
+    )
+
   linen_abstract = train_state_nnx.to_checkpoint_dict(abstract_nnx_state)
   if config and getattr(getattr(config, "lora", None), "enable_lora", False):
     linen_abstract = _filter_lora_trainable_state(linen_abstract)
@@ -536,9 +548,21 @@ def map_to_pspec(data):
         )
         ocp.type_handlers.register_type_handler(jax.Array, array_handler, override=True)
 
-      restore_target = (
-          train_state_nnx.to_checkpoint_dict(abstract_unboxed_pre_state) if is_nnx else abstract_unboxed_pre_state
-      )
+      is_diloco = bool(maxtext_config and getattr(maxtext_config, "enable_diloco", False))
+
+      # Map the expected training state to the on-disk checkpoint dictionary layout:
+      # - DiLoCo: DiLoCoTrainState (wrapping NNX or Linen inner state + outer params + opt state).
+      # - Standard non-DiLoCo NNX: TrainStateNNX (converted to Linen collection layout for storage).
+      # - Standard non-DiLoCo Linen: TrainState dataclass (used directly).
+      if is_diloco:
+        restore_target = diloco_checkpoint_utils.to_diloco_checkpoint_dict(
+            abstract_unboxed_pre_state, config=maxtext_config
+        )
+      elif is_nnx:
+        restore_target = train_state_nnx.to_checkpoint_dict(abstract_unboxed_pre_state)
+      else:
+        restore_target = abstract_unboxed_pre_state
+
       if maxtext_config and getattr(getattr(maxtext_config, "lora", None), "enable_lora", False):
         restore_target = _filter_lora_trainable_state(restore_target)
       restore_args = jax.tree_util.tree_map(map_to_pspec, restore_target)
@@ -560,7 +584,11 @@ def map_to_pspec(data):
             ),
         ):
           restored = checkpoint_manager.restore(step, args=Composite(state=checkpoint_args)).state
-          if is_nnx:
+          if is_diloco:
+            restored = diloco_checkpoint_utils.from_diloco_checkpoint_dict(
+                restored, abstract_unboxed_pre_state, config=maxtext_config
+            )
+          elif is_nnx:
             restored = _restored_linen_to_nnx(restored, abstract_unboxed_pre_state, config=maxtext_config)
           return (
               restored,
@@ -585,15 +613,25 @@ def map_to_pspec(data):
               checkpoint_args,
               expansion_factor_real_data,
           )
-          if is_nnx:
+          if is_diloco:
+            restored_items = diloco_checkpoint_utils.from_diloco_checkpoint_dict(
+                restored["items"], abstract_unboxed_pre_state, config=maxtext_config
+            )
+            restored = {"items": restored_items}
+          elif is_nnx:
             restored_items = _restored_linen_to_nnx(restored["items"], abstract_unboxed_pre_state, config=maxtext_config)
             restored = {"items": restored_items}
           return (restored, iterator)
         # Case 3: Default/Fallback case.
         # This case acts as a wildcard ('_') and matches if none of the preceding cases were met.
         case _:
           restored = checkpoint_manager.restore(step, args=Composite(items=checkpoint_args))
-          if is_nnx:
+          if is_diloco:
+            restored_items = diloco_checkpoint_utils.from_diloco_checkpoint_dict(
+                restored["items"], abstract_unboxed_pre_state, config=maxtext_config
+            )
+            restored = {"items": restored_items}
+          elif is_nnx:
             restored_items = _restored_linen_to_nnx(restored["items"], abstract_unboxed_pre_state, config=maxtext_config)
             restored = {"items": restored_items}
           return (restored, None)
@@ -848,7 +886,9 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step
     _handle_post_checkpoint_preemption(checkpoint_manager, actual_step, force_ckpt_save)
     return
 
-  if latest_step(checkpoint_manager) == actual_step:
+  # Skip if step directory already exists (e.g. step 0 or prior checkpoints in all_steps())
+  # to prevent Orbax OCDBT UUID collisions during auto-resume / continuation runs for DiLoCo.
+  if latest_step(checkpoint_manager) == actual_step or actual_step in checkpoint_manager.all_steps():
     max_logging.log(f"Checkpoint for step {actual_step} already exists, skipping save.")
     return
 
@@ -903,18 +943,18 @@ def _filter_dict(val, path=()):
 
 def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator=None, force=False):
   """Wrapper for saving checkpoint."""
-  if not isinstance(state, (dict, nnx.State, train_state.TrainState)):
+  # Allow struct.PyTreeNode so Flax dataclass states (e.g. DiLoCoTrainState) aren't cleared to empty dicts ({})
+  if not isinstance(state, (dict, nnx.State, train_state.TrainState, struct.PyTreeNode)):
     if isinstance(state, train_state_nnx.TrainStateNNX):
       state = nnx.state(state)
     elif not isinstance(state, (dict, nnx.State)):
       state = {}
 
-  if config and getattr(config, "pure_nnx", False) and isinstance(state, nnx.State):
+  if config and getattr(config, "enable_diloco", False):
+    state = diloco_checkpoint_utils.to_diloco_checkpoint_dict(state, config)
+  elif config and getattr(config, "pure_nnx", False):
     # Save in the Linen on-disk layout so pure_nnx and Linen checkpoints are interchangeable.
-    if getattr(config, "enable_diloco", False):
-      step_value = state.step.get_value() if hasattr(state.step, "get_value") else state.step
-      state = train_state_nnx.to_linen_checkpoint_dict({"model": state.params, "optimizer": {"step": step_value}})
-    else:
+    if isinstance(state, nnx.State):
       state = train_state_nnx.to_checkpoint_dict(state)
 
   if config and getattr(config, "enable_checkpointing", False):

src/maxtext/common/data_loader.py

@@ -23,7 +23,7 @@
     GoodputEvent,
     maybe_record_goodput,
 )
-from maxtext.trainers.diloco import diloco
+from maxtext.utils import diloco_sharding
 from maxtext.utils import elastic_utils
 from maxtext.utils import exceptions
 from maxtext.utils.sharding import get_input_data_sharding
@@ -89,7 +89,7 @@ def load_next_batch(self, *args, **kwargs):
     """Loads the next batch with sharding hint."""
     example_batch = self.load_next_batch_pre_sharding()
     if self.config.enable_diloco:
-      example_batch = diloco.reshape_first_axis_with_diloco(self.config.num_diloco_replicas, example_batch)
+      example_batch = diloco_sharding.reshape_first_axis_with_diloco(self.config.num_diloco_replicas, example_batch)
     return jax.device_put(example_batch, self.input_data_shardings)
 
   def check_example_batch(self):
@@ -171,7 +171,7 @@ def _slice(data):
       output = jax.tree.map(_slice, self.batch_buffer)
     self.rampup_active = rampup_manager.update()
     if self.config.enable_diloco:
-      output = diloco.reshape_first_axis_with_diloco(self.config.num_diloco_replicas, output)
+      output = diloco_sharding.reshape_first_axis_with_diloco(self.config.num_diloco_replicas, output)
     return jax.device_put(output, self.input_data_shardings)
 
 

src/maxtext/configs/base.yml

@@ -943,6 +943,14 @@ dcn_bandwidth_latency: "50ms"
 # The network interface to apply throttling rules to.
 dcn_bandwidth_interface: "eth0"
 
+# Streaming DiLoCo params
+enable_streaming_diloco: false
+num_diloco_fragments: null
+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

@@ -1781,11 +1781,26 @@ class DilocoParams(BaseModel):
 
   enable_diloco: bool = Field(False, description="Enable Diloco parallelism")
   diloco_sync_period: int = Field(36, description="Diloco sync period.")
+
+  @model_validator(mode="after")
+  def validate_streaming_diloco_params(self) -> "DilocoParams":
+    """Validates streaming DiLoCo parameters."""
+    if self.enable_streaming_diloco:
+      if not self.enable_diloco:
+        raise ValueError("enable_diloco must be True when enable_streaming_diloco is True.")
+      if self.num_diloco_fragments is None:
+        raise ValueError("num_diloco_fragments must be specified when enable_streaming_diloco is True.")
+      if self.num_diloco_fragments < 2:
+        raise ValueError(
+            f"num_diloco_fragments ({self.num_diloco_fragments}) must be at least 2 when enable_streaming_diloco "
+            "is True (1 for non-scanned parameters, at least 1 for scanned layers)."
+        )
+    return self
+
   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(
@@ -1794,6 +1809,40 @@ class DilocoParams(BaseModel):
   )
   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 (https://arxiv.org/abs/2501.18512). Streaming DiLoCo partitions"
+          " model parameters into fragments and pipelines cross-island synchronization one fragment per inner step,"
+          " overlapping inter-cluster communication with accelerator computation."
+      ),
+  )
+  num_diloco_fragments: int | None = Field(
+      None,
+      description=(
+          "Total number of fragments to partition the model layers into (including 1 fragment for non-scanned"
+          " parameters). Required when enable_streaming_diloco is True."
+      ),
+  )
+  use_sequential_layers: bool = Field(False, description="Whether to sync layers sequentially (or interleaved).")
+  num_communication_overlapping_steps: NonNegativeInt = Field(
+      0, description="Steps of communication overlap with computation. \\tau from the paper."
+  )
+  communication_overlapping_alpha: float = Field(
+      0.0,
+      ge=0.0,
+      le=1.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."
+      ),
+  )
+
 
 class Optimizer(BaseModel):
   """Configuration for the optimizer and learning rate schedule."""
@@ -3829,6 +3878,20 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
       self.validate_ragged_buffer_factor()
     self.validate_num_moe_emb_chunks()
 
+    if self.enable_diloco and not self.pure_nnx:
+      raise ValueError("enable_diloco=True requires pure_nnx=True (Linen support for DiLoCo has been removed).")
+
+    if self.enable_streaming_diloco:
+      if not self.scan_layers:
+        raise ValueError("enable_streaming_diloco=True requires scan_layers=True.")
+      if self.num_diloco_fragments is not None and self.num_diloco_fragments > 1:
+        num_transformer_fragments = self.num_diloco_fragments - 1
+        if self.num_decoder_layers % num_transformer_fragments != 0:
+          raise ValueError(
+              f"The number of decoder layers ({self.num_decoder_layers}) must be divisible by "
+              f"(num_diloco_fragments - 1) ({num_transformer_fragments}) when enable_streaming_diloco is True."
+          )
+
     # Gemma 4 small (E2B / E4B) uses per-layer KV sharing, which is incompatible with nn.scan.
     if self.model_name in ("gemma4-e2b", "gemma4-e4b") and self.scan_layers:
       raise ValueError(

src/maxtext/input_pipeline/synthetic_data_processing.py

@@ -25,6 +25,7 @@
 
 from maxtext.input_pipeline import multihost_dataloading
 from maxtext.configs import pyconfig
+from maxtext.utils.diloco_sharding import reshape_first_axis_with_diloco
 from maxtext.utils import sharding
 
 
@@ -126,6 +127,8 @@ def raw_generate_synthetic_data(config: pyconfig.HyperParameters, data):
     output["targets"] = tokens[:, 1:]
     output["targets_position"] = positions[:, 1:]
     output["targets_segmentation"] = segmentation
+    if config.enable_diloco:
+      output = reshape_first_axis_with_diloco(config.num_diloco_replicas, output)
     return output