[DiLoCo] Streaming diloco with vmap (SPMD Version) - #4759
Conversation
|
🤖 Hi @Dr-Left, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces "Streaming DiLoCo" parallelism to MaxText, enabling communication-computation overlap by partitioning model layers into fragments. Key additions include configuration parameters, validation logic, the FragmentedTreeManipulator utility for PyTree partitioning, and several training run scripts. The review feedback highlights several critical areas for improvement: resolving unreachable branches in get_first_step, enhancing the robustness of scanned parameter identification to avoid false positives, handling jax.ShapeDtypeStruct during abstract tracing, replacing fragile type name comparisons with idiomatic isinstance checks, and enforcing strict validation constraints on the new configuration parameters.
| scanned_regex = re.compile(r"/(?:layers|blocks|moe_layers|dense_layers|layers_outside_pipeline)(?:/|$)") | ||
| keypath_to_is_scanned = {} | ||
|
|
||
| for keypath, _ in kvs: | ||
| parts = [] | ||
| for k in keypath: | ||
| parts.append(str(k.key) if hasattr(k, "key") else (str(k.idx) if hasattr(k, "idx") else str(k))) | ||
| serialized_path = "/" + "/".join(parts) | ||
| keypath_to_is_scanned[jax.tree_util.keystr(keypath)] = bool(scanned_regex.search(serialized_path)) |
There was a problem hiding this comment.
Identifying scanned parameters solely based on a regex match on the keypath can lead to false positives. For example, non-scanned parameters (such as layer-specific biases, layernorm weights, or parameters in layers_outside_pipeline) might match the regex but will not have the extra leading layer dimension. If these are treated as scanned, get_flat_fragment and apply_flat_fragment will attempt to slice them along the first dimension, leading to shape mismatches or silent correctness bugs. To prevent this, we should also verify that the parameter has a leading dimension matching num_layers.
| scanned_regex = re.compile(r"/(?:layers|blocks|moe_layers|dense_layers|layers_outside_pipeline)(?:/|$)") | |
| keypath_to_is_scanned = {} | |
| for keypath, _ in kvs: | |
| parts = [] | |
| for k in keypath: | |
| parts.append(str(k.key) if hasattr(k, "key") else (str(k.idx) if hasattr(k, "idx") else str(k))) | |
| serialized_path = "/" + "/".join(parts) | |
| keypath_to_is_scanned[jax.tree_util.keystr(keypath)] = bool(scanned_regex.search(serialized_path)) | |
| scanned_regex = re.compile(r"/(?:layers|blocks|moe_layers|dense_layers|layers_outside_pipeline)(?:/|$)") | |
| keypath_to_is_scanned = {} | |
| for keypath, v in kvs: | |
| parts = [] | |
| for k in keypath: | |
| parts.append(str(k.key) if hasattr(k, "key") else (str(k.idx) if hasattr(k, "idx") else str(k))) | |
| serialized_path = "/" + "/".join(parts) | |
| is_scanned = bool(scanned_regex.search(serialized_path)) and hasattr(v, "shape") and len(v.shape) > 0 and v.shape[0] == num_layers | |
| keypath_to_is_scanned[jax.tree_util.keystr(keypath)] = is_scanned |
| if self.enable_streaming_diloco and 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." | ||
| ) |
There was a problem hiding this comment.
Streaming DiLoCo (enable_streaming_diloco=True) relies on FragmentedTreeManipulator to slice stacked parameters along the layer dimension. This assumption only holds when scan_layers is enabled (scan_layers=True). If scan_layers is disabled (e.g., for models like Gemma 4 small which use per-layer KV sharing), the layer parameters are not stacked, and slicing them will result in silent correctness bugs or shape mismatches. We should add a validation check to ensure scan_layers is enabled when streaming DiLoCo is active.
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."
)| if hasattr(state, "optimizer") and hasattr(state.optimizer, "step"): | ||
| return int(state.optimizer.step.get_value()) | ||
| elif hasattr(state, "step"): | ||
| if hasattr(state.step, "get_value"): | ||
| return int(state.step.get_value()) | ||
| return int(state.step) | ||
| if hasattr(state, "inner_state"): # DiLoCoTrainState (NNX DiLoCo): step is the optimizer step var | ||
| return int(state.step.get_value()) | ||
| return int(state.optimizer.step.get_value()) | ||
| return int(state.step.get_value()) if hasattr(state.step, "get_value") else int(state.step) | ||
| elif isinstance(model, nn.Module): | ||
| return int(state.step) | ||
| return 0 |
There was a problem hiding this comment.
The get_first_step function contains unreachable code. Specifically, because DiLoCoTrainState has a step attribute, any state that is a DiLoCoTrainState will be caught by the elif hasattr(state, "step"): branch on line 81. As a result, the subsequent if hasattr(state, "inner_state"): check on line 85 is dead code and will never be executed. Similarly, the elif isinstance(model, nn.Module): check on line 87 is also unreachable. We can simplify this function to be much cleaner, more readable, and free of dead code.
if hasattr(state, "optimizer") and hasattr(state.optimizer, "step"):
return int(state.optimizer.step.get_value())
if hasattr(state, "step"):
return int(state.step.get_value()) if hasattr(state.step, "get_value") else int(state.step)
return 0| if is_scanned: | ||
| indices = self.fragment_to_layer_indices[fragment_idx] | ||
| if has_replica_dim: | ||
| new_v = v.at[:, indices].set(flat_fragment[keystr]) | ||
| else: | ||
| new_v = v.at[indices].set(flat_fragment[keystr]) | ||
| new_kvs.append(new_v) | ||
| else: | ||
| new_kvs.append(v) |
There was a problem hiding this comment.
During abstract tracing or evaluation (e.g., when running jax.eval_shape or generating HLO dumps), the parameters v in the PyTree may be represented as jax.ShapeDtypeStruct objects rather than actual JAX arrays or tracers. Since ShapeDtypeStruct does not support .at indexing or .set(), calling apply_flat_fragment in these contexts will raise an AttributeError. To ensure robust defensive programming, we should check if v is a ShapeDtypeStruct and return it directly, as the shape and dtype of the parameter remain unchanged.
| if is_scanned: | |
| indices = self.fragment_to_layer_indices[fragment_idx] | |
| if has_replica_dim: | |
| new_v = v.at[:, indices].set(flat_fragment[keystr]) | |
| else: | |
| new_v = v.at[indices].set(flat_fragment[keystr]) | |
| new_kvs.append(new_v) | |
| else: | |
| new_kvs.append(v) | |
| if is_scanned: | |
| indices = self.fragment_to_layer_indices[fragment_idx] | |
| if isinstance(v, jax.ShapeDtypeStruct): | |
| new_v = v | |
| elif has_replica_dim: | |
| new_v = v.at[:, indices].set(flat_fragment[keystr]) | |
| else: | |
| new_v = v.at[indices].set(flat_fragment[keystr]) | |
| new_kvs.append(new_v) | |
| else: | |
| new_kvs.append(v) |
| if type(s_model).__name__ == "State": | ||
| new_model = nnx.state(new_model) | ||
| elif isinstance(s_model, dict): | ||
| new_model = nnx.to_pure_dict(new_model) |
There was a problem hiding this comment.
Checking the type of s_model by comparing type(s_model).__name__ == "State" is fragile and non-idiomatic. Since nnx.State is imported and available, we should use isinstance(s_model, nnx.State) instead, which is safer and more robust.
| if type(s_model).__name__ == "State": | |
| new_model = nnx.state(new_model) | |
| elif isinstance(s_model, dict): | |
| new_model = nnx.to_pure_dict(new_model) | |
| if isinstance(s_model, nnx.State): | |
| new_model = nnx.state(new_model) | |
| elif isinstance(s_model, dict): | |
| new_model = nnx.to_pure_dict(new_model) |
| if arr.ndim >= 3 and arr.shape[0] == num_diloco_replicas: | ||
| return arr |
There was a problem hiding this comment.
The check arr.ndim >= 3 is used as a heuristic to avoid double-reshaping already processed arrays. However, this heuristic can fail for 1D or 2D arrays (e.g., 1D metadata or labels) that have already been reshaped to 2D, causing them to be reshaped again incorrectly. A more robust and precise way to check if an array has already been reshaped for DiLoCo is to inspect its sharding and verify if diloco is already the first axis in its PartitionSpec.
| if arr.ndim >= 3 and arr.shape[0] == num_diloco_replicas: | |
| return arr | |
| if hasattr(arr, "sharding") and isinstance(arr.sharding, jax.sharding.NamedSharding): | |
| if arr.sharding.spec and arr.sharding.spec[0] == "diloco": | |
| return arr |
| 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." | ||
| ), |
There was a problem hiding this comment.
The communication_overlapping_alpha parameter is described as an interpolation factor between 0.0 and 1.0, but there are no constraints enforced on its range in the Pydantic model. We should add ge=0.0 and le=1.0 to the Field definition to prevent users from specifying invalid out-of-bounds values.
| 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." | |
| ), | |
| 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." | |
| ), | |
| ) |
| num_communication_overlapping_steps: int = Field( | ||
| 0, description="Steps of communication overlap with computation. \\tau from the paper." | ||
| ) |
There was a problem hiding this comment.
The num_communication_overlapping_steps parameter represents a count of steps and must be non-negative. We should use NonNegativeInt instead of int to enforce this constraint at the configuration level.
| num_communication_overlapping_steps: int = Field( | |
| 0, description="Steps of communication overlap with computation. \\tau from the paper." | |
| ) | |
| num_communication_overlapping_steps: NonNegativeInt = Field( | |
| 0, description="Steps of communication overlap with computation. \\tau from the paper." | |
| ) |
| return state.replace(inner_state=new_inner_state) | ||
|
|
||
| # Step 1: Run the synchronization logic if we hit a sync step | ||
| is_sync_step = jax.lax.bitwise_and(new_step > 0, new_step % steps_between_syncs_plus_1 == 0) |
There was a problem hiding this comment.
Using jax.lax.bitwise_and on boolean conditions is less readable and non-idiomatic in JAX. We should use the standard bitwise & operator or jnp.logical_and instead, which are fully supported on JAX tracers and much more readable.
| is_sync_step = jax.lax.bitwise_and(new_step > 0, new_step % steps_between_syncs_plus_1 == 0) | |
| is_sync_step = (new_step > 0) & (new_step % steps_between_syncs_plus_1 == 0) |
| if latest_step(checkpoint_manager) == actual_step or ( | ||
| checkpoint_manager is not None and actual_step in checkpoint_manager.all_steps() | ||
| ): |
There was a problem hiding this comment.
The check checkpoint_manager is not None is redundant because checkpoint_manager is already guaranteed to be not None by the early return check at the beginning of the maybe_save_checkpoint function (line 824). We can simplify this condition to improve readability.
if latest_step(checkpoint_manager) == actual_step or actual_step in checkpoint_manager.all_steps():![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
This PR introduces substantial enhancements to the DiLoCo training pipeline in MaxText, specifically introducing SPMD Streaming DiLoCo training via jax.vmap across replicas, experimental MoE configurations, and auto-resume robustness improvements. While the architecture and streaming logic are highly performant and well-structured, several critical correctness bugs in the NNX/pure-NNX codepaths would prevent successful training or resume workloads.
🔍 General Feedback
- Great Test Coverage & Real-world Validation: Excellent inclusion of integration tests for streaming DiLoCo on CPU/TPU backends and highly thorough validation across TPU slices.
- NNX State vs. Dictionary Mismatches: A recurring pattern in the changes involves handling model states as dictionaries versus specific class types (like
TrainStateNNXornnx.State), which introduces runtime type and JAX PyTree structure mismatches. - Dead Config & Parameter Handling: Some MoE configurations are defined and validated but not yet utilized in the runtime code; cleaning up or completing these features will prevent future maintainability issues.
| if hasattr(state, "inner_state"): # DiLoCoTrainState (NNX DiLoCo): step is the optimizer step var | ||
| return int(state.step.get_value()) | ||
| return int(state.optimizer.step.get_value()) | ||
| return int(state.step.get_value()) if hasattr(state.step, "get_value") else int(state.step) | ||
| elif isinstance(model, nn.Module): | ||
| return int(state.step) | ||
| return 0 |
There was a problem hiding this comment.
Because DiLoCoTrainState defines a step attribute, the preceding elif hasattr(state, "step"): check evaluates to True and returns immediately. This prevents the execution from ever reaching this inner_state check.
Since state.step already correctly returns the initial step, this block can be safely removed or refactored.
| if hasattr(state, "inner_state"): # DiLoCoTrainState (NNX DiLoCo): step is the optimizer step var | |
| return int(state.step.get_value()) | |
| return int(state.optimizer.step.get_value()) | |
| return int(state.step.get_value()) if hasattr(state.step, "get_value") else int(state.step) | |
| elif isinstance(model, nn.Module): | |
| return int(state.step) | |
| return 0 | |
| elif hasattr(state, "step"): | |
| if hasattr(state.step, "get_value"): | |
| return int(state.step.get_value()) | |
| return int(state.step) | |
| elif isinstance(model, nn.Module): | |
| return int(state.step) | |
| return 0 |
| new_model = nnx.to_pure_dict(new_model) | ||
|
|
||
| result = type(s)({}) | ||
| result["model"] = new_model |
There was a problem hiding this comment.
This will fail catastrophically in two ways when pure_nnx=True:
type(s)isTrainStateNNX, and its constructor__init__does not have default arguments for all parameters (e.g.optimizerhas no default), so callingtype(s)({})will raise aTypeError: __init__() missing 1 required positional argument: 'optimizer'.- Even if instantiation succeeded,
TrainStateNNXis annnx.Moduleand does not support item assignment (i.e.,__setitem__is not defined), soresult["model"] = new_modelwill raise aTypeError.
Additionally, as noted in the non-streaming replace_nnx_model_params helper, reconstructing via type(s)({...}) on dictionary/State wrapper types breaks lax.cond matching.
| new_model = nnx.to_pure_dict(new_model) | |
| result = type(s)({}) | |
| result["model"] = new_model | |
| if hasattr(s, "keys"): | |
| # If s is a dict or State wrapper, reconstruct keeping the original type structure | |
| result = type(s)({}) | |
| result["model"] = new_model | |
| result["optimizer"] = s["optimizer"] | |
| return result | |
| else: | |
| # If s is a TrainStateNNX instance, return a new instance via standard constructor | |
| return type(s)(new_model, s.optimizer) |
| -1, | ||
| description=( | ||
| "Period (in steps) for synchronizing MoE router parameters across replicas." | ||
| " Must divide diloco_sync_period when > 0. Defaults to -1." | ||
| ), | ||
| ) | ||
|
|
||
| @model_validator(mode="after") |
There was a problem hiding this comment.
However, there is no corresponding implementation or usage of moe_router_syncing_period in the actual DiLoCo synchronization loop/logic inside diloco.py.
If MoE router parameter synchronization is required at a different frequency than the standard outer syncing period, it should be implemented in diloco.py; otherwise, this config parameter and the unused exclude_router argument should be removed to avoid dead code and confusion.
| if latest_step(checkpoint_manager) == actual_step or ( | ||
| checkpoint_manager is not None and actual_step in checkpoint_manager.all_steps() | ||
| ): |
There was a problem hiding this comment.
If checkpoint_manager is indeed None, calling latest_step(checkpoint_manager) will immediately crash with an AttributeError (since latest_step tries to call checkpoint_manager.latest_step()) before checking checkpoint_manager is not None in the second half of the statement.
For safety, the checkpoint_manager is not None check should guard the entire expression.
| if latest_step(checkpoint_manager) == actual_step or ( | |
| checkpoint_manager is not None and actual_step in checkpoint_manager.all_steps() | |
| ): | |
| if checkpoint_manager is not None and ( | |
| latest_step(checkpoint_manager) == actual_step or actual_step in checkpoint_manager.all_steps() | |
| ): |
|
|
||
| if isinstance(state, diloco.DiLoCoTrainState): | ||
| state_params = state.params | ||
| state_mesh_shardings_params = ( | ||
| state_mesh_shardings.model.filter(nnx.Param) if config.pure_nnx else state_mesh_shardings.params |
There was a problem hiding this comment.
However, during normal initialization (when state is not yet a DiLoCoTrainState), build_diloco_state converts the model parameters to a pure dictionary using outer_params.to_pure_dict(), and constructs the outer optimizer sharding trace as a pure dictionary.
Using an nnx.State instead of a pure dict for state_mesh_shardings_params creates a PyTree structure and sharding mismatch (since dict and nnx.State have different JAX PyTree definitions/treedefs), causing JAX to raise a sharding-tree mismatch error on checkpoint restoration/resumption.
| if isinstance(state, diloco.DiLoCoTrainState): | |
| state_params = state.params | |
| state_mesh_shardings_params = ( | |
| state_mesh_shardings.model.filter(nnx.Param) if config.pure_nnx else state_mesh_shardings.params | |
| if isinstance(state, diloco.DiLoCoTrainState): | |
| state_params = state.params | |
| state_mesh_shardings_params = ( | |
| state_mesh_shardings.model.filter(nnx.Param).to_pure_dict() if config.pure_nnx else state_mesh_shardings.params | |
| ) |
167332d to
b6dc5ec
Compare
working vmap streaming add harness for MoE Experiments. Fix SPMD DiLoCo compatibility for NNX state and synthetic dataset (rank-3 sharding) Update num_diloco_fragments to represent total fragments and validate required streaming DiLoCo arguments [DiLoCo]Use split PRNG key and log losses separatedly on tb Add exception check for num_decoder_layers divisibility by num_diloco_fragments - 1 in streaming DiLoCo Apply fixes to fit for the current main head (checkpoint filtering and prng null key More robust checkpointing checking logic for DiLoCo Update spmd diloco script Remove Linen support and clean code for PR Refactor DiLoCo to have separate diloco_utils.py
b6dc5ec to
66caa65
Compare
|
🤖 Hi @Dr-Left, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
|
🤖 Hi @khatwanimohit, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
There was a problem hiding this comment.
This Pull Request introduces support for SPMD Streaming DiLoCo training using JAX's vmap and contains key improvements to Orbax checkpointing robustness, MoE harnesses, and NNX state compatibility. While the streaming DiLoCo logic and optimization sweeps are highly comprehensive, a critical mismatch in the abstract checkpoint restoration logic was introduced which completely breaks checkpoint restoration for DiLoCo under pure NNX.
🔍 General Feedback
- Restoration Mismatch: Re-building the
abstract_restore_stateas aDiLoCoTrainStatewhen restoring checkpoint introduces an incompatibility with the saved Linen-style checkpoint layout. Reverting tounboxed_abstract_state(i.e.TrainStateNNXlayout) correctly allows checkpoint loading and subsequent initialization of the replica model state wrapper. - Logical Redundancies: Several unreachable
elif config.enable_diloco:branches exist in the pre-training loop due to the outerifcondition already covering the same logic, which should be simplified. - Robust Integration Testing: The additions to
diloco_test.pyare robust and provide excellent coverage of configuration, shape evaluation, and SPMD streaming compilation.
| # Build abstract DiLoCo state structure so checkpoint restoration expects the full | ||
| # DiLoCoTrainState layout (outer optimizer + replica state) when DiLoCo is enabled. | ||
| abstract_restore_state = unboxed_abstract_state | ||
| if config.enable_diloco: | ||
| abstract_restore_state, _, _ = diloco.build_abstract_diloco_state( | ||
| config, unboxed_abstract_state, state_mesh_shardings, mesh | ||
| ) |
There was a problem hiding this comment.
When checkpoints are saved, save_checkpoint manually extracts only the global model parameters and step value in the Linen on-disk layout (using train_state_nnx.to_linen_checkpoint_dict).
By setting abstract_restore_state to DiLoCoTrainState during restoration, the Orbax restorer expects the full DiLoCoTrainState (with nested inner_state and outer_opt_state fields), which do not exist in the checkpoint. Consequently:
- Since
is_nnxevaluates toFalse, the code bypasses_load_linen_checkpoint_into_nnx. - Orbax attempts to restore a
DiLoCoTrainStatefrom the on-disk Linen dict. Because of the key/nesting mismatch ("params"is nested as"params/params"on-disk but is expected directly as"params"in the target), Orbax silently fails to restore the model weights (due topartial_restore=True). - During the overlay merge in
maxtext_utils.py,_has_shape_dtype_structreturnsTrue(since parameters were not restored and remain asShapeDtypeStructs). This causes the restorer to silently throw away the restored state and revert to a freshly initialized state. - Finally, training resumes with random/freshly-initialized weights, completely defeating the auto-resume and checkpointing functionality!
Reverting abstract_restore_state back to unboxed_abstract_state fixes the issue. train_utils.py will then correctly construct DiLoCoTrainState around the successfully restored TrainStateNNX state as designed.
| # Build abstract DiLoCo state structure so checkpoint restoration expects the full | |
| # DiLoCoTrainState layout (outer optimizer + replica state) when DiLoCo is enabled. | |
| abstract_restore_state = unboxed_abstract_state | |
| if config.enable_diloco: | |
| abstract_restore_state, _, _ = diloco.build_abstract_diloco_state( | |
| config, unboxed_abstract_state, state_mesh_shardings, mesh | |
| ) | |
| abstract_restore_state = unboxed_abstract_state |
| train_utils.validate_completed_steps(start_step, config.steps) | ||
|
|
||
| if isinstance(model, nn.Module): | ||
| if config.enable_diloco or isinstance(model, nn.Module): |
There was a problem hiding this comment.
|
|
||
| # Write train config params, num model params, and XLA flags to tensorboard | ||
| if isinstance(model, nn.Module): | ||
| if config.enable_diloco or isinstance(model, nn.Module): |
There was a problem hiding this comment.
2 participants
Description
This PR adds full support for SPMD Streaming DiLoCo training using
jax.vmap, introduces experimental harnesses for MoE models, fixes NNX state/synthetic dataset compatibility, and improves Orbax checkpointing robustness for auto-resume continuation workloads.Key Changes & Motivation:
SPMD Streaming DiLoCo with
vmap:jax.vmapacross DiLoCo replicas insrc/maxtext/trainers/diloco/diloco.py.num_diloco_fragmentsto represent total fragments and added validation checks to ensurenum_decoder_layersis evenly divisible by(num_diloco_fragments - 1).Robust Checkpointing Logic (
src/maxtext/common/checkpointing.py):actual_step in checkpoint_manager.all_steps()check inmaybe_save_checkpoint().ValueError: FAILED_PRECONDITION: Configuration mismatch on uuid) when auto-resuming from existing step directories (e.g., step 0 or prior run checkpoints).MoE Harness & Scripts:
run_gemma4_moe.sh,run_moe.sh,run_olmo_qwen3_30b_streaming_diloco.sh,run_spmd_streaming_diloco.sh).PRNG Key & TensorBoard Metrics:
Tests
End-to-End TPU Multi-Slice Verification
Tested SPMD Streaming DiLoCo training for
qwen3-8bacross 2xv5p-8TPU slices on themlperf-v5pcluster.Initial 100-Step Baseline Run (
spmd-v5p8-0705):30.5 GiBweights) to GCS atgs://chriszuo-maxtext-logs/smoketest/spmd-v5p8-0705/checkpoints/99.200-Step Checkpoint Auto-Resume Run (
spmd-v5p8-0705b):RUNNAME="spmd-v5p8-0705"andSTEPS=200.restoring from this run's directory step 99) anddrjax.broadcastparameter synchronization across slices.loss: 6.761), completed all 200 total steps (steps 0–199), and committed final step 199 checkpoint (loss: 6.328).0/1 Completed) with 0 error matches.Test Links & Paths
gs://chriszuo-maxtext-logs/smoketest/spmd-v5p8-0705/checkpoints/Checklist
gemini-reviewlabel.