From 1e259fb65c58704c8b5048166bfb2ad2e9352e46 Mon Sep 17 00:00:00 2001 From: Chris Zuo Date: Sat, 27 Jun 2026 00:04:35 +0000 Subject: [PATCH] [DiLoCo] Streaming diloco with vmap (SPMD Version) 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 Resolve checkpointing issues --- src/maxtext/common/checkpointing.py | 86 ++--- src/maxtext/common/data_loader.py | 11 +- src/maxtext/common/train_state_nnx.py | 23 ++ src/maxtext/configs/base.yml | 8 + src/maxtext/configs/types.py | 60 +++- .../synthetic_data_processing.py | 3 + src/maxtext/trainers/diloco/diloco.py | 309 +++++++----------- .../trainers/diloco/scripts/run_gemma4_moe.sh | 135 ++++++++ .../trainers/diloco/scripts/run_moe.sh | 142 ++++++++ .../run_olmo_qwen3_30b_streaming_diloco.sh | 150 +++++++++ .../scripts/run_spmd_streaming_diloco.sh | 129 ++++++++ src/maxtext/trainers/diloco/utils/__init__.py | 51 +++ .../trainers/diloco/utils/spmd/__init__.py | 55 ++++ .../diloco/utils/spmd/checkpoint_utils.py | 224 +++++++++++++ .../diloco/utils/spmd/fragment_utils.py | 258 +++++++++++++++ .../trainers/diloco/utils/spmd/state_utils.py | 227 +++++++++++++ src/maxtext/trainers/pre_train/train.py | 9 +- src/maxtext/utils/maxtext_utils.py | 31 +- src/maxtext/utils/train_utils.py | 25 +- tests/integration/diloco_test.py | 105 ++++++ 20 files changed, 1796 insertions(+), 245 deletions(-) create mode 100755 src/maxtext/trainers/diloco/scripts/run_gemma4_moe.sh create mode 100755 src/maxtext/trainers/diloco/scripts/run_moe.sh create mode 100755 src/maxtext/trainers/diloco/scripts/run_olmo_qwen3_30b_streaming_diloco.sh create mode 100644 src/maxtext/trainers/diloco/scripts/run_spmd_streaming_diloco.sh create mode 100644 src/maxtext/trainers/diloco/utils/__init__.py create mode 100644 src/maxtext/trainers/diloco/utils/spmd/__init__.py create mode 100644 src/maxtext/trainers/diloco/utils/spmd/checkpoint_utils.py create mode 100644 src/maxtext/trainers/diloco/utils/spmd/fragment_utils.py create mode 100644 src/maxtext/trainers/diloco/utils/spmd/state_utils.py diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py index 5632f8a3a6..5010e997be 100644 --- a/src/maxtext/common/checkpointing.py +++ b/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.spmd import checkpoint_utils as diloco_checkpoint_utils from maxtext.utils import elastic_utils from maxtext.utils import exceptions from maxtext.utils import gcs_utils @@ -149,26 +151,8 @@ def _raise_on_weight_mismatch(want, have, config=None): def _linen_items_to_nnx(restored_linen, abstract_nnx_state): - """Reshapes a restored Linen-layout `items` dict into an NNX state. - - The inverse of `to_checkpoint_dict`, over the same `split_for_checkpoint` partition. The Linen - weights + optimizer fill `linen_state`; the `nnx_aux` state (rngs/dropout, batch stats, custom - variables) fills `aux`; the two are recombined with `nnx.merge_state`. The split copies, so the - caller's abstract is untouched. Leaves the checkpoint didn't carry -- including the caches it - never stores -- stay unmaterialized `ShapeDtypeStruct`s; the caller fills them from a fresh init. - """ - linen_state, aux_state, ephemeral = train_state_nnx.split_for_checkpoint(abstract_nnx_state) - weights = train_state_nnx.from_linen_checkpoint_dict(restored_linen) - if "model" in weights: - nnx.replace_by_pure_dict(linen_state, {"model": weights["model"]}) - if "optimizer" in weights: - nnx.replace_by_pure_dict(linen_state, {"optimizer": weights["optimizer"]}) - - nnx_aux = restored_linen.get("nnx_aux") - if nnx_aux: - nnx.replace_by_pure_dict(aux_state, nnx_aux) - - return nnx.merge_state(linen_state, aux_state, ephemeral) + """Reshapes a restored Linen-layout `items` dict into an NNX state.""" + return train_state_nnx.linen_items_to_nnx(restored_linen, abstract_nnx_state) def _load_linen_checkpoint_into_nnx( @@ -185,7 +169,25 @@ def _load_linen_checkpoint_into_nnx( `_linen_items_to_nnx`. rngs/dropout/batch stats come from `items/nnx_aux` when 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}") + p = epath.Path(path) + if (p / "items").exists(): + p = p / "items" + max_logging.log(f"Restoring Linen-layout checkpoint into NNX state at {p}") + if config and getattr(config, "enable_diloco", False): + diloco_abstract = diloco_checkpoint_utils.to_diloco_checkpoint_dict(abstract_nnx_state, config=config) + ckptr = ocp.Checkpointer( + ocp.PyTreeCheckpointHandler( + restore_concurrent_gb=checkpoint_storage_concurrent_gb, + save_concurrent_gb=checkpoint_storage_concurrent_gb, + use_ocdbt=use_ocdbt, + use_zarr3=use_zarr3, + ) + ) + restore_args = ocp.checkpoint_utils.construct_restore_args(diloco_abstract) + restored = ocp.args.PyTreeRestore(item=diloco_abstract, restore_args=restore_args, partial_restore=True) + restored = ckptr.restore(p, args=restored) + return diloco_checkpoint_utils.from_diloco_checkpoint_dict(restored, abstract_nnx_state, 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) @@ -199,7 +201,7 @@ def _load_linen_checkpoint_into_nnx( ) restore_args = ocp.checkpoint_utils.construct_restore_args(linen_abstract) restored = ocp.args.PyTreeRestore(item=linen_abstract, restore_args=restore_args, partial_restore=True) - restored = ckptr.restore(epath.Path(path), args=restored) + restored = ckptr.restore(p, args=restored) return _restored_linen_to_nnx(restored, abstract_nnx_state, config=config) @@ -288,7 +290,7 @@ def _load_full_state_from_path( if enable_orbax_v1: if source_checkpoint_layout == "orbax": # pure_nnx saves in the Linen on-disk layout; reshape it back into the NNX state. - if isinstance(abstract_unboxed_pre_state, nnx.State): + if isinstance(abstract_unboxed_pre_state, (nnx.State, train_state_nnx.TrainStateNNX)): return _load_linen_checkpoint_into_nnx( path, abstract_unboxed_pre_state, @@ -317,14 +319,16 @@ def combine_sharding(sds, shardings): state = conversion_fn(pre_transformed_state) # The conversion fn returns MaxText's on-disk (Linen) layout, which is what pure_nnx reads, # so NNX needs the same reshape as every other restore. An NNX state passes through. - if isinstance(abstract_unboxed_pre_state, nnx.State) and not isinstance(state, nnx.State): + if isinstance(abstract_unboxed_pre_state, (nnx.State, train_state_nnx.TrainStateNNX)) and not isinstance( + state, nnx.State + ): state = _restored_linen_to_nnx(state, abstract_unboxed_pre_state, config=maxtext_config) return state else: raise ocp_v1.errors.InvalidLayoutError(f"Unknown checkpoint layout: {source_checkpoint_layout}") else: # pure_nnx saves in the Linen on-disk layout; reshape it back into the NNX state. - if isinstance(abstract_unboxed_pre_state, nnx.State): + if isinstance(abstract_unboxed_pre_state, (nnx.State, train_state_nnx.TrainStateNNX)): return _load_linen_checkpoint_into_nnx( path, abstract_unboxed_pre_state, @@ -672,7 +676,10 @@ def load_params_from_path( ): """Load decode params from checkpoint at specified path.""" assert load_parameters_from_path, "load_parameters_from_path is not defined." - max_logging.log(f"restoring params from {load_parameters_from_path}") + p = epath.Path(load_parameters_from_path) + if (p / "items").exists(): + p = p / "items" + max_logging.log(f"restoring params from {p}") # On disk the weights live at `params/params/...`: an outer key naming the item, and Flax's # `params` collection inside it. A Linen TrainState.params is that collection; an NNX params @@ -681,7 +688,7 @@ def load_params_from_path( want = abstract_unboxed_params.to_pure_dict() if is_nnx else abstract_unboxed_params # Determine the restore key based on the leaf directory name to support native and custom SFT - restore_key = os.path.basename(load_parameters_from_path) + restore_key = os.path.basename(str(p)) if restore_key not in ("model_params", "model"): restore_key = "params" @@ -705,12 +712,11 @@ def load_params_from_path( # Rather than pass the entire abstract state, which could unnecessarily restore opt_state and such and waste # memory, we instead specify here that we are just restoring the params field of the checkpoint # (which itself may be a dictionary containing a key named 'params' or 'model'). - restore_args = ocp.checkpoint_utils.construct_restore_args(params_collection) + restore_target = {restore_key: params_collection} + restore_args = ocp.checkpoint_utils.construct_restore_args(restore_target) restored = ckptr.restore( - epath.Path(load_parameters_from_path), - item={restore_key: params_collection}, - transforms={}, - restore_args={restore_key: restore_args}, + p, + args=ocp.args.PyTreeRestore(item=restore_target, restore_args=restore_args, partial_restore=True), ) restored_collection = restored[restore_key] @@ -848,7 +854,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 +911,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): diff --git a/src/maxtext/common/data_loader.py b/src/maxtext/common/data_loader.py index 2bd04c527b..040b01fb78 100644 --- a/src/maxtext/common/data_loader.py +++ b/src/maxtext/common/data_loader.py @@ -23,7 +23,7 @@ GoodputEvent, maybe_record_goodput, ) -from maxtext.trainers.diloco import diloco +from maxtext.trainers.diloco import utils as diloco_utils from maxtext.utils import elastic_utils from maxtext.utils import exceptions from maxtext.utils.sharding import get_input_data_sharding @@ -89,8 +89,11 @@ 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) - return jax.device_put(example_batch, self.input_data_shardings) + example_batch = diloco_utils.reshape_first_axis_with_diloco(self.config.num_diloco_replicas, example_batch) + sharded_batch = jax.device_put(example_batch, self.input_data_shardings) + if self.config.reuse_example_batch: + self.last_batch = sharded_batch + return sharded_batch def check_example_batch(self): if self.config.max_checkify: @@ -171,7 +174,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_utils.reshape_first_axis_with_diloco(self.config.num_diloco_replicas, output) return jax.device_put(output, self.input_data_shardings) diff --git a/src/maxtext/common/train_state_nnx.py b/src/maxtext/common/train_state_nnx.py index 45dc386576..6528339cc9 100644 --- a/src/maxtext/common/train_state_nnx.py +++ b/src/maxtext/common/train_state_nnx.py @@ -245,3 +245,26 @@ def to_checkpoint_dict(state: nnx.State | nnx.Module): if aux: linen_dict["nnx_aux"] = aux return linen_dict + + +def linen_items_to_nnx(restored_linen: dict[str, Any], abstract_nnx_state: nnx.State | TrainStateNNX) -> nnx.State: + """Reshapes a restored Linen-layout `items` dict into an NNX state. + + The inverse of `to_checkpoint_dict`, over the same `split_for_checkpoint` partition. The Linen + weights + optimizer fill `linen_state`; the `nnx_aux` state (rngs/dropout, batch stats, custom + variables) fills `aux`; the two are recombined with `nnx.merge_state`. The split copies, so the + caller's abstract is untouched. Leaves the checkpoint didn't carry -- including the caches it + never stores -- stay unmaterialized `ShapeDtypeStruct`s; the caller fills them from a fresh init. + """ + linen_state, aux_state, ephemeral = split_for_checkpoint(abstract_nnx_state) + weights = from_linen_checkpoint_dict(restored_linen) + if "model" in weights: + nnx.replace_by_pure_dict(linen_state, {"model": weights["model"]}) + if "optimizer" in weights: + nnx.replace_by_pure_dict(linen_state, {"optimizer": weights["optimizer"]}) + + nnx_aux = restored_linen.get("nnx_aux") + if nnx_aux: + nnx.replace_by_pure_dict(aux_state, nnx_aux) + + return nnx.merge_state(linen_state, aux_state, ephemeral) diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index db0b1a68a7..b733a0886d 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -911,6 +911,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 diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 6416d0ae9f..0482edc297 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1680,11 +1680,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( @@ -1693,6 +1708,33 @@ 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.") + 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.""" @@ -3529,6 +3571,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( diff --git a/src/maxtext/input_pipeline/synthetic_data_processing.py b/src/maxtext/input_pipeline/synthetic_data_processing.py index d10816435c..b6b970bc8e 100644 --- a/src/maxtext/input_pipeline/synthetic_data_processing.py +++ b/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.trainers.diloco import utils as diloco_utils from maxtext.utils import sharding @@ -78,6 +79,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 = diloco_utils.reshape_first_axis_with_diloco(config.num_diloco_replicas, output) return output diff --git a/src/maxtext/trainers/diloco/diloco.py b/src/maxtext/trainers/diloco/diloco.py index a91d052c69..9e8d3008a6 100644 --- a/src/maxtext/trainers/diloco/diloco.py +++ b/src/maxtext/trainers/diloco/diloco.py @@ -1,16 +1,16 @@ -# Copyright 2025 Google LLC +# Copyright 2025 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 +# 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 +# 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. +# 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. """An implementation of Distributed Low-Communication (DiLoCo) training. @@ -22,7 +22,6 @@ https://arxiv.org/abs/2501.18512 """ -from collections.abc import Sequence from typing import Any, Callable import drjax @@ -31,10 +30,11 @@ import jax import jax.numpy as jnp from jaxtyping import Array, Int32, Key, PyTree, UInt32 -from maxtext.common.train_state_nnx import TrainStateNNX from maxtext.configs import pyconfig +from maxtext.trainers.diloco import utils as diloco_utils import optax + Batch = Any Params = PyTree Metrics = PyTree @@ -49,11 +49,11 @@ class DiLoCoTrainState(struct.PyTreeNode): Attributes: inner_state: A `flax.training.train_state.TrainState` of the state for each - step of the inner optimization. All arrays are expected to have a leading + step of the inner optimization. All arrays are expected to have a leading dimension with size of the number of diloco replicas so that training steps can be mapped over this dimension. - params: A PyTree of the global model weights. These will mimic a - sub-PyTree in `inner_state`, which rank-1 shape. + params: A PyTree of the global model weights. These will mimic a sub-PyTree + in `inner_state`, which rank-1 shape. outer_opt_state: The state for the outer Nesterov momentum optimizer. step: The step counter of the training process. """ @@ -64,59 +64,6 @@ class DiLoCoTrainState(struct.PyTreeNode): step: Step -def add_diloco_to_sharding(pytree): - """ - Recursively traverses a PyTree and prepends 'diloco' to the PartitionSpec - of any NamedSharding object that doesn't have an empty PartitionSpec. - """ - - def map_fn(leaf): - if isinstance(leaf, jax.sharding.NamedSharding): - new_spec = jax.sharding.PartitionSpec("diloco", *leaf.spec) - return jax.sharding.NamedSharding(mesh=leaf.mesh, spec=new_spec) - return leaf - - return jax.tree_util.tree_map(map_fn, pytree) - - -def reshape_first_axis_with_diloco(num_diloco_replicas: int, pytree: PyTree) -> PyTree: - """Reshapes the first dimension of each array in the PyTree to include a DiLoCo axis. - - This function takes a a batch of data represented as a PyTree - and reshapes the leading dimension of each array within it. The purpose is - to introduce a new 'diloco' axis, which is used for distributing data - across DiLoCo replicas. - - Args: - num_diloco_replicas: The number of DiLoCo replicas. This determines the - size of the new leading dimension. - pytree: The input PyTree, where each array is expected to have a batch - dimension as its first axis. - - Returns: - A new PyTree with the same structure as the input, but with each array's - first dimension reshaped to `(num_diloco_replicas, original_batch_dim // num_diloco_replicas, ...)`. - The sharding specification is also updated to include the 'diloco' axis. - """ - - def extend_pspec(pspec: jax.sharding.PartitionSpec | Sequence[str | Sequence[str]] = ()) -> jax.sharding.PartitionSpec: - if tuple(*pspec)[0] == "diloco": - # pull out diloco axis if already present - return jax.sharding.PartitionSpec("diloco", (*pspec[0][1:],), (*pspec[1:],)) - return jax.sharding.PartitionSpec("diloco", *pspec) - - def reshape_for_diloco(arr): - batch_dim, *example_shape = arr.shape - diloco_shape = (num_diloco_replicas, batch_dim // num_diloco_replicas, *example_shape) - if hasattr(arr, "sharding"): - s = arr.sharding - s = jax.sharding.NamedSharding(mesh=s.mesh, spec=extend_pspec(s.spec)) - return jax.lax.with_sharding_constraint(jnp.reshape(arr, shape=diloco_shape), s) - return jnp.reshape(arr, shape=diloco_shape) - - return jax.tree.map(reshape_for_diloco, pytree) - - def build_abstract_diloco_state( config: "pyconfig.HyperParameters", abstract_state: PyTree, @@ -153,16 +100,11 @@ def add_diloco_dim(x): momentum=config.diloco_outer_momentum, nesterov=True, ) - # For NNX, model params (Param variables only) live under abstract_state.model; - # for Linen under abstract_state.params. - if config.pure_nnx: - _, model_params, _ = nnx.split(abstract_state.model, nnx.Param, ...) - model_params = model_params.to_pure_dict() # pyrefly: ignore[missing-attribute] - _, model_params_sharding, _ = nnx.split(state_mesh_shardings.model, nnx.Param, ...) - model_params_sharding = model_params_sharding.to_pure_dict() # pyrefly: ignore[missing-attribute] - else: - model_params = abstract_state.params - model_params_sharding = state_mesh_shardings.params + # Model params (Param variables only) live under abstract_state.model. + _, model_params, _ = nnx.split(abstract_state.model, nnx.Param, ...) + model_params = model_params.to_pure_dict() # pyrefly: ignore[missing-attribute] + _, model_params_sharding, _ = nnx.split(state_mesh_shardings.model, nnx.Param, ...) + model_params_sharding = model_params_sharding.to_pure_dict() # pyrefly: ignore[missing-attribute] outer_opt_state = jax.eval_shape(outer_optimizer.init, model_params) # Create abstract step @@ -177,7 +119,7 @@ def add_diloco_dim(x): ) # Build shardings - inner_state_shardings = add_diloco_to_sharding(state_mesh_shardings) + inner_state_shardings = diloco_utils.add_diloco_to_sharding(state_mesh_shardings) # Sharding for outer_opt_state. For SGD with momentum, it is (TraceState(trace=...), EmptyState()) # We shard the momentum trace the same way as the parameters. outer_opt_state_sharding = ( @@ -215,17 +157,12 @@ def init_diloco_state() -> tuple[DiLoCoTrainState, PyTree]: # mesh automatically when jax.set_mesh is used. inner_state = drjax.broadcast(state, mesh=mesh) # Outer state retains a single copy of the model parameters and optimizer state. - # For NNX, model params (Param variables only) live under state.model; - # for Linen under state.params. - if config.pure_nnx: - _, outer_params, _ = nnx.split(state.model, nnx.Param, ...) - outer_params = outer_params.to_pure_dict() # pyrefly: ignore[missing-attribute] - else: - outer_params = state.params + # Model params (Param variables only) live under state.model. + _, outer_params, _ = nnx.split(state.model, nnx.Param, ...) + outer_params = outer_params.to_pure_dict() # pyrefly: ignore[missing-attribute] outer_opt_state = outer_optimizer.init(outer_params) outer_opt_state_sharding = jax.tree_util.tree_map(lambda x: x.sharding, outer_opt_state) - # For NNX, the step counter lives at state.optimizer.step; for Linen at state.step. - step = state.optimizer.step if config.pure_nnx else state.step + step = state.optimizer.step return ( DiLoCoTrainState(inner_state=inner_state, params=outer_params, outer_opt_state=outer_opt_state, step=step), outer_opt_state_sharding, @@ -234,119 +171,123 @@ def init_diloco_state() -> tuple[DiLoCoTrainState, PyTree]: return init_diloco_state() -def build_diloco_train_step( +def build_vanilla_diloco_train_step( config: pyconfig.HyperParameters, train_step: Callable[[Any, Batch, PRNGKey], tuple[Any, Metrics]], mesh: jax.sharding.Mesh | None = None, ) -> Callable[[DiLoCoTrainState, Batch, PRNGKey], tuple[DiLoCoTrainState, Metrics]]: - """Convert a local state and train step into DiLoCo-compatible versions. - - This is an implementation of the original (non-streaming) DiLoCo algorithm - which syncs all model parameters across the replicas every - `config.diloco_sync_period` steps, treating the difference accumulated over - non-sync steps as a pseudo gradient and applying SGD with Nesterov momentum on - the "global" model. - - Args: - config: The config used to set up training. - train_step: A local train step. This will be executed independently within - each replica. - """ + """Convert a local state and train step into vanilla DiLoCo train step.""" outer_optimizer = optax.sgd( config.diloco_outer_lr, momentum=config.diloco_outer_momentum, nesterov=True, ) - def synchronize(state): - # Calculate the delta between the current replica's state and the global - # state (since last synchronization). - broadcast_outer_params = drjax.broadcast(state.params, mesh=mesh) - # For NNX, model Param vars live under inner_state.model; for Linen under inner_state.params. - if config.pure_nnx: - _, inner_model_params, _ = nnx.split(state.inner_state.model, nnx.Param, ...) - inner_model_params = inner_model_params.to_pure_dict() # pyrefly: ignore[missing-attribute] - else: - inner_model_params = state.inner_state.params - model_delta = jax.tree.map(lambda x, y: y - x, inner_model_params, broadcast_outer_params) - # Treat the average delta as the outer optimizer's gradient and apply to - # the global (outer) model params. - averaged_pseudo_grad = drjax.reduce_mean(model_delta) - updates, new_opt_state = outer_optimizer.update(averaged_pseudo_grad, state.outer_opt_state, state.params) - new_outer_params = optax.apply_updates(state.params, updates) - # Replace inner model params with the new global model params. - # NOTE: inner optimizer state is retained despite the change in parameters, - # see section 6.1 in https://arxiv.org/pdf/2311.08105. - if config.pure_nnx: - # For NNX: merge new Param vars back with the non-Param model vars (e.g. RNG state). - def replace_nnx_model_params(s, new_params): - s_model = s["model"] if hasattr(s, "keys") else s.model - s_opt = s["optimizer"] if hasattr(s, "keys") else s.optimizer - - graphdef, _, non_param_state = nnx.split(s_model, nnx.Param, ...) - new_model = nnx.merge(graphdef, new_params, non_param_state) - - 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 hasattr(s, "keys"): - # Replace "model" leaves by path, keeping s's treedef. Picking by position - # (leaves[N:]) breaks if a key sorts before "model"; reconstructing via - # type(s)({...}) breaks the lax.cond match — nnx.State recursive-wraps. - leaves_with_paths, treedef = jax.tree_util.tree_flatten_with_path(s) - new_model_iter = iter(jax.tree_util.tree_leaves(new_model)) - - def _is_model_leaf(path): - if not path: - return False - k = path[0] - return getattr(k, "key", None) == "model" or getattr(k, "name", None) == "model" - - new_leaves = [next(new_model_iter) if _is_model_leaf(p) else leaf for p, leaf in leaves_with_paths] - return jax.tree_util.tree_unflatten(treedef, new_leaves) - else: - return TrainStateNNX(new_model, s_opt) - - new_inner_state = drjax.map_fn( - lambda s: replace_nnx_model_params(s, new_outer_params), - state.inner_state, - mesh=mesh, - ) - else: - new_inner_state = drjax.map_fn(lambda s: s.replace(params=new_outer_params), state.inner_state, mesh=mesh) - return state.replace( - params=new_outer_params, - outer_opt_state=new_opt_state, - inner_state=new_inner_state, - ) - - def typed_reduce_mean(in_tree): - total = drjax.reduce_sum(in_tree) - avg = jax.tree.map(lambda x: (x / config.num_diloco_replicas).astype(x.dtype), total) - return avg - @drjax.program(placements={"diloco": config.num_diloco_replicas}) - def diloco_train_step(state, batch, prng): - # Broadcast the RNG across replicas. - broadcast_rng = drjax.broadcast(prng, mesh=mesh) - inner_state, metrics = drjax.map_fn(train_step, (state.inner_state, batch, broadcast_rng), mesh=mesh) - avg_metrics = typed_reduce_mean(metrics) - # For NNX, the step counter lives at inner_state.optimizer.step; for Linen at inner_state.step. - new_step = inner_state.optimizer.step[0] if config.pure_nnx else inner_state.step[0] + def vanilla_diloco_train_step(state: DiLoCoTrainState, batch: Batch, prng: PRNGKey): + keys = jax.random.split(prng, config.num_diloco_replicas) if prng is not None else None + inner_state, metrics = drjax.map_fn(train_step, (state.inner_state, batch, keys), mesh=mesh) + default_metrics = diloco_utils.extract_per_island_metrics(metrics, config.num_diloco_replicas) + new_step = inner_state.optimizer.step[0] state = state.replace( inner_state=inner_state, step=new_step, ) - # Either synchronize the model, or no-op, depending on whether the current - # step falls on the synchronization period. + state = jax.lax.cond( new_step % config.diloco_sync_period == 0, - synchronize, - lambda x: x, # no-op + lambda s: diloco_utils.synchronize_full_state(s, outer_optimizer, mesh=mesh), + lambda x: x, state, ) - return state, avg_metrics + return state, default_metrics - return diloco_train_step + return vanilla_diloco_train_step + + +def build_streaming_diloco_train_step( + config: pyconfig.HyperParameters, + train_step: Callable[[Any, Batch, PRNGKey], tuple[Any, Metrics]], + mesh: jax.sharding.Mesh | None = None, +) -> Callable[[DiLoCoTrainState, Batch, PRNGKey], tuple[DiLoCoTrainState, Metrics]]: + """Convert a local state and train step into streaming DiLoCo train step.""" + outer_optimizer = optax.sgd( + config.diloco_outer_lr, + momentum=config.diloco_outer_momentum, + nesterov=True, + ) + num_fragments = config.num_diloco_fragments + steps_between_syncs, period = diloco_utils.get_streaming_schedule(config) + delay_v = config.num_communication_overlapping_steps + alpha = config.communication_overlapping_alpha + + @drjax.program(placements={"diloco": config.num_diloco_replicas}) + def streaming_diloco_train_step(state: DiLoCoTrainState, batch: Batch, prng: PRNGKey): + keys = jax.random.split(prng, config.num_diloco_replicas) if prng is not None else None + inner_state, metrics = drjax.map_fn(train_step, (state.inner_state, batch, keys), mesh=mesh) + default_metrics = diloco_utils.extract_per_island_metrics(metrics, config.num_diloco_replicas) + new_step = inner_state.optimizer.step[0] + state = state.replace( + inner_state=inner_state, + step=new_step, + ) + + manipulator = diloco_utils.FragmentedTreeManipulator.create(state.params, config) + + # Step 1: Run the synchronization logic if we hit a sync step + is_sync_step = (new_step > 0) & (new_step % steps_between_syncs == 0) + + def do_sync(s): + frag_idx = (new_step % period) // steps_between_syncs + return jax.lax.switch( + frag_idx, + [ + lambda s_arg, idx=i: diloco_utils.synchronize_fragment_state( + s_arg, manipulator, idx, outer_optimizer, mesh=mesh + ) + for i in range(num_fragments) + ], + s, + ) + + state = jax.lax.cond(is_sync_step, do_sync, lambda s: s, state) + + # Step 2: Apply the synced parameters (with delay V) + is_apply_step = (new_step - delay_v > 0) & ((new_step - delay_v) % steps_between_syncs == 0) + + def do_apply(s): + frag_idx = ((new_step - delay_v) % period) // steps_between_syncs + return jax.lax.switch( + frag_idx, + [ + lambda s_arg, idx=i: diloco_utils.apply_fragment_to_inner_state( + s_arg, manipulator, idx, alpha=alpha, mesh=mesh + ) + for i in range(num_fragments) + ], + s, + ) + + state = jax.lax.cond(is_apply_step, do_apply, lambda s: s, state) + + return state, default_metrics + + return streaming_diloco_train_step + + +def build_diloco_train_step( + config: pyconfig.HyperParameters, + train_step: Callable[[Any, Batch, PRNGKey], tuple[Any, Metrics]], + mesh: jax.sharding.Mesh | None = None, +) -> Callable[[DiLoCoTrainState, Batch, PRNGKey], tuple[DiLoCoTrainState, Metrics]]: + """Convert a local state and train step into DiLoCo-compatible versions. + + Args: + config: The config used to set up training. + train_step: A local train step. This will be executed independently within + each replica. + mesh: The mesh for sharding. + """ + if config.enable_streaming_diloco: + return build_streaming_diloco_train_step(config, train_step, mesh=mesh) + return build_vanilla_diloco_train_step(config, train_step, mesh=mesh) diff --git a/src/maxtext/trainers/diloco/scripts/run_gemma4_moe.sh b/src/maxtext/trainers/diloco/scripts/run_gemma4_moe.sh new file mode 100755 index 0000000000..c9077c08c2 --- /dev/null +++ b/src/maxtext/trainers/diloco/scripts/run_gemma4_moe.sh @@ -0,0 +1,135 @@ +#!/bin/bash +set -e + +# Configuration for Gemma4-26B MoE (Streaming) DiLoCo training job logging diagnostic metrics +# Setup: 2-Slice v5p-256 GKE cluster (64 nodes / 256 TPU chips) +# Model: gemma4-26b (MoE: 128 routed experts + 1 shared expert) + +CLUSTER="${CLUSTER:-mlperf-v5p}" +PROJECT="${PROJECT:-cloud-tpu-multipod-dev}" +ZONE="${ZONE:-europe-west4}" + +NUM_SLICES="${NUM_SLICES:-2}" +DEVICE_TYPE="${DEVICE_TYPE:-v5p-256}" + +RUNNAME="${RUNNAME:-dlco-gemma4-moe}" +DOCKER_IMAGE_BASE="${DOCKER_IMAGE_BASE:-gcr.io/tpu-prod-env-multipod/maxtext_jax_stable:2026-07-17}" +MY_IMAGE="gcr.io/${PROJECT}/$(whoami)-runner:${RUNNAME}" + +BASE_OUTPUT_DIRECTORY="${BASE_OUTPUT_DIRECTORY:-gs://chriszuo-maxtext-logs}" +DATASET_PATH="${DATASET_PATH:-gs://chriszuo-maxtext-datasets}" + +# DiLoCo Hyperparameters (Optimal parameters from MoE sweep) +DILOCO_SYNC_PERIOD="${DILOCO_SYNC_PERIOD:-31}" +DILOCO_OUTER_LR="${DILOCO_OUTER_LR:-0.3}" +DILOCO_OUTER_MOMENTUM="${DILOCO_OUTER_MOMENTUM:-0.9}" +ENABLE_STREAMING_DILOCO="${ENABLE_STREAMING_DILOCO:-true}" +DILOCO_NUM_FRAGMENTS="${DILOCO_NUM_FRAGMENTS:-31}" +DILOCO_USE_SEQUENTIAL_LAYERS="${DILOCO_USE_SEQUENTIAL_LAYERS:-false}" + +# Gemma4-26B Model Parameters +MODEL_NAME="${MODEL_NAME:-gemma4-26b}" +PER_DEVICE_BATCH_SIZE="${PER_DEVICE_BATCH_SIZE:-8}" +MAX_TARGET_LENGTH="${MAX_TARGET_LENGTH:-2048}" +STEPS="${STEPS:-1024}" + +# Learning Rate & Optimizer Parameters (Identical to pre-training runs) +LEARNING_RATE="${LEARNING_RATE:-1.0e-4}" +WARMUP_STEPS="${WARMUP_STEPS:-2000}" +LR_SCHEDULE_STEPS="${LR_SCHEDULE_STEPS:-18596}" +COSINE_FINAL_FRAC="${COSINE_FINAL_FRAC:-0.1}" +WARMUP_FRAC=$(python3 -c "print(${WARMUP_STEPS}/${LR_SCHEDULE_STEPS})") +ADAM_B1="${ADAM_B1:-0.9}" +ADAM_B2="${ADAM_B2:-0.95}" +ADAM_EPS="${ADAM_EPS:-1e-8}" +ADAM_WD="${ADAM_WD:-0.1}" +GRAD_CLIP="${GRAD_CLIP:-1.0}" + +# Z-Loss & MoE Load Balancing Parameters +Z_LOSS="${Z_LOSS:-1.0e-5}" +LOAD_BALANCE_LOSS_WEIGHT="${LOAD_BALANCE_LOSS_WEIGHT:-0.001}" + +LIBTPU_INIT_ARGS=" \ + --xla_tpu_scoped_vmem_limit_kib=65472 \ + --xla_tpu_bf16_emission_mode=NATIVE_EMISSION \ + --xla_tpu_enable_sparse_core_reduce_scatter_v2=true \ + --xla_tpu_enable_sparse_core_collective_offload_all_gather=true \ + --xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true \ + --xla_tpu_enable_all_gather_offload_tracing=true \ + --xla_tpu_use_tc_device_shape_on_sc=True \ + --xla_sc_disable_megacore_partitioning=True \ + --xla_tpu_enable_async_collective_fusion_fuse_all_gather=false \ + --xla_enable_async_all_gather=true \ + --xla_tpu_prefer_async_allgather_to_allreduce=true \ + --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true \ + --xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true \ + --xla_tpu_enable_sparse_core_collective_offload_3d_all_gather=true \ + --xla_tpu_use_single_sparse_core_for_all_gather_offload=true \ + --xla_tpu_enable_concurrent_sparse_core_offloading=true \ + --xla_tpu_aggressive_opt_barrier_removal=true \ + --xla_tpu_enable_offloading_gather_to_sparsecore=true \ + --xla_tpu_sparse_core_all_gather_latency_multiplier=1 \ + --xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3 \ + --xla_tpu_enable_sparse_core_collective_aggregator=true \ + --xla_tpu_enable_latency_hiding_layer_scheduler=true \ + --xla_tpu_scheduler_percent_shared_memory_limit=150 \ + --xla_tpu_enable_layer_scheduler_for_dependent_collectives=true \ + --xla_tpu_enable_sparse_core_collective_offload_nd_reduce_scatter=true \ + --xla_tpu_pcie_bandwidth_multiplier=0.03 \ + --xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true \ + --xla_tpu_enable_multi_compute_overlap_in_layer_scheduler=false \ + --xla_tpu_enable_3d_reduce_scatter_decomposer=false " + +CMD="export PYTHONPATH=/app/src:\$PYTHONPATH && unset XLA_FLAGS && export LIBTPU_INIT_ARGS=\"${LIBTPU_INIT_ARGS}\" && cd /app/src/ && python3 maxtext/trainers/pre_train/train.py \ + maxtext/configs/base.yml \ + run_name=${RUNNAME} \ + save_config_to_gcs=true \ + base_output_directory=${BASE_OUTPUT_DIRECTORY} \ + dataset_path=${DATASET_PATH} \ + dataset_name='c4/en:3.0.1' \ + eval_dataset_name='c4/en:3.0.1' \ + model_name=${MODEL_NAME} \ + tokenizer_path=maxtext/assets/tokenizers/tokenizer.gemma3 \ + per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \ + max_target_length=${MAX_TARGET_LENGTH} \ + learning_rate=${LEARNING_RATE} \ + learning_rate_schedule_steps=${LR_SCHEDULE_STEPS} \ + learning_rate_final_fraction=${COSINE_FINAL_FRAC} \ + warmup_steps_fraction=${WARMUP_FRAC} \ + adam_b1=${ADAM_B1} \ + adam_b2=${ADAM_B2} \ + adam_eps=${ADAM_EPS} \ + adam_weight_decay=${ADAM_WD} \ + gradient_clipping_threshold=${GRAD_CLIP} \ + z_loss_multiplier=${Z_LOSS} \ + load_balance_loss_weight=${LOAD_BALANCE_LOSS_WEIGHT} \ + enable_diloco=true \ + enable_streaming_diloco=${ENABLE_STREAMING_DILOCO} \ + pure_nnx=true \ + dcn_diloco_parallelism=${NUM_SLICES} \ + diloco_sync_period=${DILOCO_SYNC_PERIOD} \ + diloco_outer_lr=${DILOCO_OUTER_LR} \ + diloco_outer_momentum=${DILOCO_OUTER_MOMENTUM} \ + num_diloco_fragments=${DILOCO_NUM_FRAGMENTS} \ + use_sequential_layers=${DILOCO_USE_SEQUENTIAL_LAYERS} \ + steps=${STEPS}" + +docker build -t "${MY_IMAGE}" -f - . << 'EOF' +FROM gcr.io/tpu-prod-env-multipod/maxtext_jax_stable:2026-07-17 +WORKDIR /app +COPY . . +RUN find /app -name "*.pyc" -delete && find /app -name "__pycache__" -type d -exec rm -rf {} + 2>/dev/null || true +EOF + +echo "Pushing docker image: ${MY_IMAGE}" +docker push "${MY_IMAGE}" + +echo "Submitting Gemma4-26B MoE Streaming DiLoCo training workload using XPK..." +xpk workload create --workload "${RUNNAME}" \ + --docker-image "${MY_IMAGE}" \ + --command "${CMD}" \ + --num-slices="${NUM_SLICES}" \ + --enable-debug-logs \ + --cluster "${CLUSTER}" --tpu-type "${DEVICE_TYPE}" --project "${PROJECT}" --zone "${ZONE}" + +echo "Gemma4-26B MoE Streaming DiLoCo workload submission complete!" diff --git a/src/maxtext/trainers/diloco/scripts/run_moe.sh b/src/maxtext/trainers/diloco/scripts/run_moe.sh new file mode 100755 index 0000000000..aac8d17c57 --- /dev/null +++ b/src/maxtext/trainers/diloco/scripts/run_moe.sh @@ -0,0 +1,142 @@ +#!/bin/bash +set -e + +# Configuration for MoE (Streaming) DiLoCo training job logging all 5 diagnostic metrics: +# 1. Inter-Replica Router Distance (d_router / diloco/inter_replica_router_distance) +# 2. Top-K Token Routing Overlap (J_route / diloco/topk_token_routing_overlap) +# 3. Jensen-Shannon Routing Divergence Index (RDI / diloco/js_routing_divergence_index) +# 4. Post-Sync Loss Spike Severity (Delta L_sync / diloco/post_sync_loss_spike_severity) +# 5. Post-Sync Expert Utilization Entropy (EUE / diloco/post_sync_expert_utilization_entropy) + +CLUSTER="${CLUSTER:-mlperf-v5p}" +PROJECT="${PROJECT:-cloud-tpu-multipod-dev}" +ZONE="${ZONE:-europe-west4-b}" + +NUM_SLICES="${NUM_SLICES:-2}" +DEVICE_TYPE="${DEVICE_TYPE:-v5p-128}" + +RUNNAME="${RUNNAME:-dlco-moe-$(date +%H%M)}" +DOCKER_IMAGE_BASE="${DOCKER_IMAGE_BASE:-gcr.io/tpu-prod-env-multipod/maxtext_jax_stable:latest}" +MY_IMAGE="gcr.io/${PROJECT}/$(whoami)-runner:${RUNNAME}" + +BASE_OUTPUT_DIRECTORY="${BASE_OUTPUT_DIRECTORY:-gs://chriszuo-maxtext-logs}" +DATASET_PATH="${DATASET_PATH:-gs://chriszuo-maxtext-datasets}" + +# DiLoCo Hyperparameters +DILOCO_SYNC_PERIOD="${DILOCO_SYNC_PERIOD:-49}" +DILOCO_OUTER_LR="${DILOCO_OUTER_LR:-0.1}" +DILOCO_OUTER_MOMENTUM="${DILOCO_OUTER_MOMENTUM:-0.9}" +ENABLE_STREAMING_DILOCO="${ENABLE_STREAMING_DILOCO:-true}" +DILOCO_NUM_FRAGMENTS="${DILOCO_NUM_FRAGMENTS:-49}" +DILOCO_USE_SEQUENTIAL_LAYERS="${DILOCO_USE_SEQUENTIAL_LAYERS:-false}" + +# MoE Model Parameters +MODEL_NAME="${MODEL_NAME:-qwen3-30b-a3b}" +PER_DEVICE_BATCH_SIZE="${PER_DEVICE_BATCH_SIZE:-8}" +MAX_TARGET_LENGTH="${MAX_TARGET_LENGTH:-2048}" +STEPS="${STEPS:-100}" + +# Learning Rate & Optimizer Parameters +LEARNING_RATE="${LEARNING_RATE:-1.0e-4}" +WARMUP_STEPS="${WARMUP_STEPS:-2000}" +LR_SCHEDULE_STEPS="${LR_SCHEDULE_STEPS:-18596}" +COSINE_FINAL_FRAC="${COSINE_FINAL_FRAC:-0.1}" +WARMUP_FRAC=$(python3 -c "print(${WARMUP_STEPS}/${LR_SCHEDULE_STEPS})") +ADAM_B1="${ADAM_B1:-0.9}" +ADAM_B2="${ADAM_B2:-0.95}" +ADAM_EPS="${ADAM_EPS:-1e-8}" +ADAM_WD="${ADAM_WD:-0.1}" +GRAD_CLIP="${GRAD_CLIP:-1.0}" + +# Z-Loss & MoE Load Balancing Parameters +Z_LOSS="${Z_LOSS:-1.0e-5}" +LOAD_BALANCE_LOSS_WEIGHT="${LOAD_BALANCE_LOSS_WEIGHT:-0.01}" + +LIBTPU_INIT_ARGS=" \ + --xla_tpu_scoped_vmem_limit_kib=65472 \ + --xla_tpu_bf16_emission_mode=NATIVE_EMISSION \ + --xla_tpu_enable_sparse_core_reduce_scatter_v2=true \ + --xla_tpu_enable_sparse_core_collective_offload_all_gather=true \ + --xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true \ + --xla_tpu_enable_all_gather_offload_tracing=true \ + --xla_tpu_use_tc_device_shape_on_sc=True \ + --xla_sc_disable_megacore_partitioning=True \ + --xla_tpu_enable_async_collective_fusion_fuse_all_gather=false \ + --xla_enable_async_all_gather=true \ + --xla_tpu_prefer_async_allgather_to_allreduce=true \ + --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true \ + --xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true \ + --xla_tpu_enable_sparse_core_collective_offload_3d_all_gather=true \ + --xla_tpu_use_single_sparse_core_for_all_gather_offload=true \ + --xla_tpu_enable_concurrent_sparse_core_offloading=true \ + --xla_tpu_aggressive_opt_barrier_removal=true \ + --xla_tpu_enable_offloading_gather_to_sparsecore=true \ + --xla_tpu_sparse_core_all_gather_latency_multiplier=1 \ + --xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3 \ + --xla_tpu_enable_sparse_core_collective_aggregator=true \ + --xla_tpu_enable_latency_hiding_layer_scheduler=true \ + --xla_tpu_scheduler_percent_shared_memory_limit=150 \ + --xla_tpu_enable_layer_scheduler_for_dependent_collectives=true \ + --xla_tpu_enable_sparse_core_collective_offload_nd_reduce_scatter=true \ + --xla_tpu_pcie_bandwidth_multiplier=0.03 \ + --xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true \ + --xla_tpu_enable_multi_compute_overlap_in_layer_scheduler=false \ + --xla_tpu_enable_3d_reduce_scatter_decomposer=false " + +CMD="export PYTHONPATH=/app/src:\$PYTHONPATH && unset XLA_FLAGS && export LIBTPU_INIT_ARGS=\"${LIBTPU_INIT_ARGS}\" && cd /app/src/ && python3 maxtext/trainers/pre_train/train.py \ + maxtext/configs/base.yml \ + run_name=${RUNNAME} \ + save_config_to_gcs=true \ + base_output_directory=${BASE_OUTPUT_DIRECTORY} \ + dataset_path=${DATASET_PATH} \ + dataset_name='c4/en:3.0.1' \ + eval_dataset_name='c4/en:3.0.1' \ + model_name=${MODEL_NAME} \ + tokenizer_type=huggingface \ + tokenizer_path=maxtext/assets/tokenizers/qwen3-tokenizer \ + per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \ + max_target_length=${MAX_TARGET_LENGTH} \ + learning_rate=${LEARNING_RATE} \ + learning_rate_schedule_steps=${LR_SCHEDULE_STEPS} \ + learning_rate_final_fraction=${COSINE_FINAL_FRAC} \ + warmup_steps_fraction=${WARMUP_FRAC} \ + adam_b1=${ADAM_B1} \ + adam_b2=${ADAM_B2} \ + adam_eps=${ADAM_EPS} \ + adam_weight_decay=${ADAM_WD} \ + gradient_clipping_threshold=${GRAD_CLIP} \ + z_loss_multiplier=${Z_LOSS} \ + load_balance_loss_weight=${LOAD_BALANCE_LOSS_WEIGHT} \ + enable_diloco=true \ + enable_streaming_diloco=${ENABLE_STREAMING_DILOCO} \ + pure_nnx=true \ + dcn_diloco_parallelism=${NUM_SLICES} \ + diloco_sync_period=${DILOCO_SYNC_PERIOD} \ + diloco_outer_lr=${DILOCO_OUTER_LR} \ + diloco_outer_momentum=${DILOCO_OUTER_MOMENTUM} \ + num_diloco_fragments=${DILOCO_NUM_FRAGMENTS} \ + use_sequential_layers=${DILOCO_USE_SEQUENTIAL_LAYERS} \ + jax_distributed_initialization_timeout=1200 \ + steps=${STEPS}" + +docker build -t "${MY_IMAGE}" -f - . << EOF +FROM ${DOCKER_IMAGE_BASE} +WORKDIR /app +COPY . . +RUN find /app -name "*.pyc" -delete && find /app -name "__pycache__" -type d -exec rm -rf {} + 2>/dev/null || true +EOF + +echo "Pushing docker image: ${MY_IMAGE}" +docker push "${MY_IMAGE}" + +echo "Submitting MoE DiLoCo training workload using XPK..." +xpk workload create --workload "${RUNNAME}" \ + --docker-image "${MY_IMAGE}" \ + --command "${CMD}" \ + --num-slices="${NUM_SLICES}" \ + --priority "${PRIORITY:-medium}" \ + --reservation "${RESERVATION:-cloudtpu-20240716121201-595617744}" \ + --enable-debug-logs \ + --cluster "${CLUSTER}" --tpu-type "${DEVICE_TYPE}" --project "${PROJECT}" --zone "${ZONE}" + +echo "MoE DiLoCo workload submission complete!" diff --git a/src/maxtext/trainers/diloco/scripts/run_olmo_qwen3_30b_streaming_diloco.sh b/src/maxtext/trainers/diloco/scripts/run_olmo_qwen3_30b_streaming_diloco.sh new file mode 100755 index 0000000000..6f8873b523 --- /dev/null +++ b/src/maxtext/trainers/diloco/scripts/run_olmo_qwen3_30b_streaming_diloco.sh @@ -0,0 +1,150 @@ +#!/bin/bash +# +# Script to train Qwen3-30B-A3B (MoE) on OLMo dataset with Streaming DiLoCo across TPU slices. +# References: +# - src/maxtext/trainers/diloco/scripts/run_moe.sh (Qwen3-30B MoE & DiLoCo hyperparameters, SparseCore flags) +# - MyStuff/scripts/pretrain/run_olmo_streaming_diloco_v5p-128.sh (OLMo grain dataset pipeline & gcsfuse setup) +# + +set -euo pipefail + +# -------------------------- Cluster & Topology Settings -------------------------- +: "${XPK_PROJECT:=cloud-tpu-multipod-dev}" +: "${XPK_ZONE:=europe-west4-b}" +: "${XPK_CLUSTER:=mlperf-v5p}" +: "${XPK_RESERVATION:=cloudtpu-20240716121201-595617744}" +: "${XPK_DEVICE_TYPE:=v5p-128}" +: "${XPK_NUM_SLICES:=2}" +: "${XPK_PRIORITY:=medium}" +: "${XPK_MAX_RESTARTS:=50}" + +# -------------------------- Docker & Storage -------------------------- +: "${XPK_DOCKER_IMAGE:=gcr.io/tpu-prod-env-multipod/maxtext_jax_stable:2026-07-17}" +: "${BASE_OUTPUT_DIRECTORY:=gs://chriszuo-maxtext-logs}" + +# -------------------------- OLMo Data & Tokenizer -------------------------- +: "${OLMO_INDEX_PATH:=/tmp/olmo-data/olmo/indices/olmo_index_seq8192.json}" +: "${OLMO_GCS_BASE:=gs://chriszuo-maxtext-datasets}" +: "${OLMO_LOCAL_MOUNT:=/tmp/olmo-data}" + +# HuggingFace Token (autodetect if ~/.hf_token.sh exists) +if [ -z "${HF_TOKEN:-}" ] && [ -f "${HOME}/.hf_token.sh" ]; then + # shellcheck disable=SC1090 + source "${HOME}/.hf_token.sh" +fi +: "${HF_TOKEN:=}" + +# -------------------------- MoE Model Configuration -------------------------- +: "${MODEL_NAME:=qwen3-30b-a3b}" +: "${TOKENIZER_TYPE:=huggingface}" +: "${TOKENIZER_PATH:=maxtext/assets/tokenizers/qwen3-tokenizer}" +: "${MAX_TARGET_LENGTH:=8192}" + +# -------------------------- DiLoCo Hyperparameters -------------------------- +: "${ENABLE_STREAMING_DILOCO:=true}" +: "${DILOCO_SYNC_PERIOD:=49}" +: "${DILOCO_OUTER_LR:=0.1}" +: "${DILOCO_OUTER_MOMENTUM:=0.9}" +: "${DILOCO_NUM_FRAGMENTS:=49}" +: "${DILOCO_USE_SEQUENTIAL_LAYERS:=false}" + +# -------------------------- Training & Optimizer Hyperparameters -------------------------- +: "${RUN_NAME:=jzuo-qwen3-olmo-dlco-$(date +%m%d%H%M)}" +: "${WORKLOAD_NAME:=qw3-olmo-$(date +%m%d%H%M)}" +: "${WARMUP_STEPS:=2000}" +: "${TARGET_GLOBAL_BATCH:=512}" + +# Calculate total steps for 1 full epoch of OLMo dataset (24,357,482 sequences / global batch 512 = 47,573 steps, ~199.5B tokens) +: "${TOTAL_INSTANCES:=24357482}" +TOTAL_DATASET_STEPS=$(( TOTAL_INSTANCES / TARGET_GLOBAL_BATCH )) +: "${LR_SCHEDULE_STEPS:=${TOTAL_DATASET_STEPS}}" +: "${STEPS:=${LR_SCHEDULE_STEPS}}" +: "${LEARNING_RATE:=1.0e-4}" +: "${COSINE_FINAL_FRAC:=0.1}" +: "${ADAM_B1:=0.9}" +: "${ADAM_B2:=0.95}" +: "${ADAM_EPS:=1e-8}" +: "${ADAM_WD:=0.1}" +: "${GRAD_CLIP:=1.0}" + +# MoE Z-loss & Load Balancing +: "${Z_LOSS:=1.0e-5}" +: "${LOAD_BALANCE_LOSS_WEIGHT:=0.001}" +: "${FLOAT32_GATE_LOGITS:=true}" +: "${DATA_SEED:=42}" +: "${LOAD_FULL_STATE_PATH:=}" +: "${LOAD_PARAMETERS_PATH:=}" + + +# Determine batch size per device +DEVICE_NUM=$(echo "${XPK_DEVICE_TYPE}" | cut -d'-' -f2) +TOTAL_DEVICES=$(( DEVICE_NUM * XPK_NUM_SLICES )) +: "${PER_DEVICE_BATCH_SIZE:=$(( TARGET_GLOBAL_BATCH / TOTAL_DEVICES ))}" + +# Resolve container-side index path +if [[ "${OLMO_INDEX_PATH}" = /* ]]; then + OLMO_INDEX_PATH_IN_CONTAINER="${OLMO_INDEX_PATH}" +else + OLMO_INDEX_PATH_IN_CONTAINER="/deps/${OLMO_INDEX_PATH}" +fi + +WARMUP_FRAC=$(python3 -c "print(${WARMUP_STEPS}/${LR_SCHEDULE_STEPS})") + +# -------------------------- LibTPU & SparseCore Flags -------------------------- +LIBTPU_INIT_ARGS=" \ + --xla_tpu_scoped_vmem_limit_kib=65472 \ + --xla_tpu_bf16_emission_mode=NATIVE_EMISSION \ + --xla_tpu_enable_sparse_core_reduce_scatter_v2=true \ + --xla_tpu_enable_sparse_core_collective_offload_all_gather=true \ + --xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true \ + --xla_tpu_enable_all_gather_offload_tracing=true \ + --xla_tpu_use_tc_device_shape_on_sc=True \ + --xla_sc_disable_megacore_partitioning=True \ + --xla_tpu_enable_async_collective_fusion_fuse_all_gather=false \ + --xla_enable_async_all_gather=true \ + --xla_tpu_prefer_async_allgather_to_allreduce=true \ + --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true \ + --xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true \ + --xla_tpu_enable_sparse_core_collective_offload_3d_all_gather=true \ + --xla_tpu_use_single_sparse_core_for_all_gather_offload=true \ + --xla_tpu_enable_concurrent_sparse_core_offloading=true \ + --xla_tpu_aggressive_opt_barrier_removal=true \ + --xla_tpu_enable_offloading_gather_to_sparsecore=true \ + --xla_tpu_sparse_core_all_gather_latency_multiplier=1 \ + --xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3 \ + --xla_tpu_enable_sparse_core_collective_aggregator=true \ + --xla_tpu_enable_latency_hiding_layer_scheduler=true \ + --xla_tpu_scheduler_percent_shared_memory_limit=150 \ + --xla_tpu_enable_layer_scheduler_for_dependent_collectives=true \ + --xla_tpu_enable_sparse_core_collective_offload_nd_reduce_scatter=true \ + --xla_tpu_pcie_bandwidth_multiplier=0.03 \ + --xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true \ + --xla_tpu_enable_multi_compute_overlap_in_layer_scheduler=false \ + --xla_tpu_enable_3d_reduce_scatter_decomposer=false " + +## -------------------------- Command Execution -------------------------- +CMD="set -euo pipefail; mkdir -p /deps/src/src/dependencies/scripts && ( [ -f /deps/src/src/dependencies/scripts/setup_gcsfuse.sh ] || ln -sf /deps/src/dependencies/scripts/setup_gcsfuse.sh /deps/src/src/dependencies/scripts/setup_gcsfuse.sh ); mkdir -p /deps/src/src/maxtext && ( [ -d /deps/src/src/maxtext/configs ] || ln -sf /deps/src/maxtext/configs /deps/src/src/maxtext/configs ); if [[ \"\${MOUNT_GCSFUSE:-1}\" = \"1\" ]]; then bucket=\"\${GCS_BASE#gs://}\"; bucket=\"\${bucket%/}\"; bash /deps/src/dependencies/scripts/setup_gcsfuse.sh DATASET_GCS_BUCKET=\"\${bucket}\" MOUNT_PATH=\"\${LOCAL_MOUNT}\"; fi; export PYTHONPATH=/app/src:\${PYTHONPATH:-}; cd /app/src && unset XLA_FLAGS; export LIBTPU_INIT_ARGS=\"${LIBTPU_INIT_ARGS}\"; LOAD_PARAM_ARG=\"\"; if [ -n \"\${LOAD_FULL_STATE_PATH}\" ]; then LOAD_PARAM_ARG=\"load_full_state_path=\${LOAD_FULL_STATE_PATH}\"; elif [ -n \"\${LOAD_PARAMETERS_PATH}\" ]; then LOAD_PARAM_ARG=\"load_parameters_path=\${LOAD_PARAMETERS_PATH}\"; fi; python3 -m maxtext.trainers.pre_train.train maxtext/configs/base.yml run_name=\"\${RUN_NAME}\" save_config_to_gcs=true base_output_directory=\"\${OUTPUT_DIR}\" \${LOAD_PARAM_ARG} dataset_type=olmo_grain olmo_index_path=\"\${INDEX_PATH}\" olmo_path_remap_from=\"\${GCS_BASE}\" olmo_path_remap_to=\"\${LOCAL_MOUNT}\" olmo_apply_ngram_filter=True data_shuffle_seed=\"\${DATA_SEED}\" model_name=\"\${MODEL_NAME}\" tokenizer_type=\"\${TOKENIZER_TYPE}\" tokenizer_path=\"\${TOKENIZER_PATH}\" per_device_batch_size=\"\${PER_DEVICE_BATCH_SIZE}\" max_target_length=\"\${MAX_TARGET_LENGTH}\" learning_rate=\"\${LEARNING_RATE}\" learning_rate_schedule_steps=\"\${LR_SCHEDULE_STEPS}\" learning_rate_final_fraction=\"\${COSINE_FINAL_FRAC}\" warmup_steps_fraction=\"\${WARMUP_FRAC}\" adam_b1=\"\${ADAM_B1}\" adam_b2=\"\${ADAM_B2}\" adam_eps=\"\${ADAM_EPS}\" adam_weight_decay=\"\${ADAM_WD}\" gradient_clipping_threshold=\"\${GRAD_CLIP}\" z_loss_multiplier=\"\${Z_LOSS}\" load_balance_loss_weight=\"\${LOAD_BALANCE_LOSS_WEIGHT}\" float32_gate_logits=\"\${FLOAT32_GATE_LOGITS}\" enable_diloco=true enable_streaming_diloco=\"\${ENABLE_STREAMING_DILOCO}\" pure_nnx=true dcn_diloco_parallelism=\"\${XPK_NUM_SLICES}\" diloco_sync_period=\"\${DILOCO_SYNC_PERIOD}\" diloco_outer_lr=\"\${DILOCO_OUTER_LR}\" diloco_outer_momentum=\"\${DILOCO_OUTER_MOMENTUM}\" num_diloco_fragments=\"\${DILOCO_NUM_FRAGMENTS}\" use_sequential_layers=\"\${DILOCO_USE_SEQUENTIAL_LAYERS}\" jax_distributed_initialization_timeout=1200 steps=\"\${STEPS}\"" + +echo "Submitting Qwen3-30B-A3B + OLMo DiLoCo training workload using XPK..." +echo "Workload Name: ${WORKLOAD_NAME}" +echo "Run Name: ${RUN_NAME}" +echo "Model: ${MODEL_NAME}" +echo "Dataset: olmo_grain (${OLMO_INDEX_PATH_IN_CONTAINER})" +echo "Topology: ${XPK_DEVICE_TYPE} x ${XPK_NUM_SLICES} slices" + +xpk workload create \ + --cluster "${XPK_CLUSTER}" \ + --workload "${WORKLOAD_NAME}" \ + --project "${XPK_PROJECT}" \ + --zone "${XPK_ZONE}" \ + --tpu-type "${XPK_DEVICE_TYPE}" \ + --num-slices "${XPK_NUM_SLICES}" \ + --priority "${XPK_PRIORITY}" \ + --max-restarts "${XPK_MAX_RESTARTS}" \ + --reservation "${XPK_RESERVATION}" \ + --base-docker-image "${XPK_DOCKER_IMAGE}" \ + --script-dir "$(pwd)" \ + --command "export HF_TOKEN='${HF_TOKEN}'; export INDEX_PATH='${OLMO_INDEX_PATH_IN_CONTAINER}'; export GCS_BASE='${OLMO_GCS_BASE}'; export LOCAL_MOUNT='${OLMO_LOCAL_MOUNT}'; export OUTPUT_DIR='${BASE_OUTPUT_DIRECTORY}'; export RUN_NAME='${RUN_NAME}'; export LOAD_FULL_STATE_PATH='${LOAD_FULL_STATE_PATH}'; export LOAD_PARAMETERS_PATH='${LOAD_PARAMETERS_PATH}'; export MODEL_NAME='${MODEL_NAME}'; export TOKENIZER_TYPE='${TOKENIZER_TYPE}'; export TOKENIZER_PATH='${TOKENIZER_PATH}'; export MAX_TARGET_LENGTH='${MAX_TARGET_LENGTH}'; export STEPS='${STEPS}'; export WARMUP_STEPS='${WARMUP_STEPS}'; export LR_SCHEDULE_STEPS='${LR_SCHEDULE_STEPS}'; export LEARNING_RATE='${LEARNING_RATE}'; export COSINE_FINAL_FRAC='${COSINE_FINAL_FRAC}'; export ADAM_B1='${ADAM_B1}'; export ADAM_B2='${ADAM_B2}'; export ADAM_EPS='${ADAM_EPS}'; export ADAM_WD='${ADAM_WD}'; export GRAD_CLIP='${GRAD_CLIP}'; export Z_LOSS='${Z_LOSS}'; export LOAD_BALANCE_LOSS_WEIGHT='${LOAD_BALANCE_LOSS_WEIGHT}'; export FLOAT32_GATE_LOGITS='${FLOAT32_GATE_LOGITS}'; export ENABLE_STREAMING_DILOCO='${ENABLE_STREAMING_DILOCO}'; export DILOCO_SYNC_PERIOD='${DILOCO_SYNC_PERIOD}'; export DILOCO_OUTER_LR='${DILOCO_OUTER_LR}'; export DILOCO_OUTER_MOMENTUM='${DILOCO_OUTER_MOMENTUM}'; export DILOCO_NUM_FRAGMENTS='${DILOCO_NUM_FRAGMENTS}'; export DILOCO_USE_SEQUENTIAL_LAYERS='${DILOCO_USE_SEQUENTIAL_LAYERS}'; export DATA_SEED='${DATA_SEED}'; export MOUNT_GCSFUSE=1; export XPK_NUM_SLICES='${XPK_NUM_SLICES}'; export PER_DEVICE_BATCH_SIZE='${PER_DEVICE_BATCH_SIZE}'; export WARMUP_FRAC='${WARMUP_FRAC}'; ${CMD}" + + +echo "Qwen3-30B-A3B + OLMo DiLoCo workload submission complete!" diff --git a/src/maxtext/trainers/diloco/scripts/run_spmd_streaming_diloco.sh b/src/maxtext/trainers/diloco/scripts/run_spmd_streaming_diloco.sh new file mode 100644 index 0000000000..de3bcf53be --- /dev/null +++ b/src/maxtext/trainers/diloco/scripts/run_spmd_streaming_diloco.sh @@ -0,0 +1,129 @@ +#!/bin/bash +set -e + +# cluster +# cluster +CLUSTER="${CLUSTER:-mlperf-v5p}" +PROJECT="${PROJECT:-cloud-tpu-multipod-dev}" +ZONE="${ZONE:-europe-west4-b}" + +# specify resource +NUM_SLICES="${NUM_SLICES:-2}" +DEVICE_TYPE="${DEVICE_TYPE:-v5p-128}" + +# command +RUNNAME="${RUNNAME:-spmd-dlco-$(date +%H%M)}" +XPK_WORKLOAD="${XPK_WORKLOAD:-$RUNNAME}" +DOCKER_IMAGE_BASE="${DOCKER_IMAGE_BASE:-gcr.io/tpu-prod-env-multipod/maxtext_jax_stable:latest}" +MY_IMAGE="gcr.io/${PROJECT}/$(whoami)-runner:${XPK_WORKLOAD}" + +BASE_OUTPUT_DIRECTORY="${BASE_OUTPUT_DIRECTORY:-gs://chriszuo-maxtext-logs}" +DATASET_PATH="${DATASET_PATH:-gs://chriszuo-maxtext-datasets}" +DILOCO_SYNC_PERIOD="${DILOCO_SYNC_PERIOD:-37}" +DILOCO_OUTER_LR="${DILOCO_OUTER_LR:-0.1}" +DILOCO_OUTER_MOMENTUM="${DILOCO_OUTER_MOMENTUM:-0.9}" +DILOCO_NUM_FRAGMENTS="${DILOCO_NUM_FRAGMENTS:-37}" # 36 decoder layers + 1 embedding fragment = 37 +DILOCO_USE_SEQUENTIAL_LAYERS="${DILOCO_USE_SEQUENTIAL_LAYERS:-false}" +DILOCO_NUM_COMM_OVERLAP_STEPS="${DILOCO_NUM_COMM_OVERLAP_STEPS:-2}" +DILOCO_COMM_OVERLAP_ALPHA="${DILOCO_COMM_OVERLAP_ALPHA:-0.0}" +MODEL_NAME="${MODEL_NAME:-qwen3-8b}" +PER_DEVICE_BATCH_SIZE="${PER_DEVICE_BATCH_SIZE:-8}" +MAX_TARGET_LENGTH="${MAX_TARGET_LENGTH:-2048}" +STEPS="${STEPS:-100}" + +XLA_FLAGS=" \ + --xla_tpu_scoped_vmem_limit_kib=65536 \ + --xla_tpu_bf16_emission_mode=NATIVE_EMISSION \ + --xla_tpu_enable_sparse_core_reduce_scatter_v2=true \ + --xla_tpu_enable_sparse_core_collective_offload_all_gather=true \ + --xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true \ + --xla_tpu_enable_all_gather_offload_tracing=true \ + --xla_tpu_use_tc_device_shape_on_sc=True \ + --xla_sc_disable_megacore_partitioning=True \ + --xla_tpu_enable_async_collective_fusion_fuse_all_gather=false \ + --xla_enable_async_all_gather=true \ + --xla_tpu_prefer_async_allgather_to_allreduce=true \ + --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true \ + --xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true \ + --xla_tpu_enable_sparse_core_collective_offload_3d_all_gather=true \ + --xla_tpu_use_single_sparse_core_for_all_gather_offload=true \ + --xla_tpu_enable_concurrent_sparse_core_offloading=true \ + --xla_tpu_aggressive_opt_barrier_removal=true \ + --xla_tpu_enable_offloading_gather_to_sparsecore=true \ + --xla_tpu_sparse_core_all_gather_latency_multiplier=1 \ + --xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3 \ + --xla_tpu_enable_sparse_core_collective_aggregator=true \ + --xla_tpu_enable_latency_hiding_layer_scheduler=true \ + --xla_tpu_scheduler_percent_shared_memory_limit=150 \ + --xla_tpu_enable_layer_scheduler_for_dependent_collectives=true \ + --xla_tpu_enable_sparse_core_collective_offload_nd_reduce_scatter=true \ + --xla_tpu_pcie_bandwidth_multiplier=0.03 \ + --xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true \ + --xla_tpu_enable_multi_compute_overlap_in_layer_scheduler=false \ + --xla_tpu_enable_3d_reduce_scatter_decomposer=false " + +CMD="export PYTHONPATH=/app/src:\$PYTHONPATH && cd /app/src/ && python3 maxtext/trainers/pre_train/train.py \ + maxtext/configs/base.yml \ + run_name=${RUNNAME} \ + save_config_to_gcs=true \ + base_output_directory=${BASE_OUTPUT_DIRECTORY} \ + dataset_path=${DATASET_PATH} \ + dataset_name='c4/en:3.0.1' \ + eval_dataset_name='c4/en:3.0.1' \ + model_name=${MODEL_NAME} \ + tokenizer_type=huggingface \ + tokenizer_path=maxtext/assets/tokenizers/qwen3-tokenizer \ + per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \ + max_target_length=${MAX_TARGET_LENGTH} \ + enable_diloco=true \ + enable_streaming_diloco=true \ + pure_nnx=true \ + num_diloco_fragments=${DILOCO_NUM_FRAGMENTS} \ + use_sequential_layers=${DILOCO_USE_SEQUENTIAL_LAYERS} \ + num_communication_overlapping_steps=${DILOCO_NUM_COMM_OVERLAP_STEPS} \ + communication_overlapping_alpha=${DILOCO_COMM_OVERLAP_ALPHA} \ + dcn_diloco_parallelism=${NUM_SLICES} \ + diloco_sync_period=${DILOCO_SYNC_PERIOD} \ + diloco_outer_lr=${DILOCO_OUTER_LR} \ + diloco_outer_momentum=${DILOCO_OUTER_MOMENTUM} \ + profiler=xplane \ + skip_first_n_steps_for_profiler=5 \ + profiler_steps=5 \ + upload_all_profiler_results=true \ + jax_distributed_initialization_timeout=1200 \ + steps=${STEPS}" + +# 1. Build and push the docker image manually containing your local changes +echo "Building docker image containing local changes..." +docker build -t "${MY_IMAGE}" -f - . </dev/null || true +EOF + +echo "Pushing image ${MY_IMAGE}..." +docker push "${MY_IMAGE}" + +# 2. Create the workload directly using xpk +echo "Creating workload: ${XPK_WORKLOAD}" +XPK_ARGS=( + --workload "${XPK_WORKLOAD}" + --docker-image "${MY_IMAGE}" + --command "${CMD}" + --num-slices "${NUM_SLICES}" + --priority "${PRIORITY:-medium}" + --enable-debug-logs + --cluster "${CLUSTER}" + --tpu-type "${DEVICE_TYPE}" + --project "${PROJECT}" + --zone "${ZONE}" +) + +if [ -n "${RESERVATION}" ] && [ "${RESERVATION}" != "NONE" ]; then + XPK_ARGS+=(--reservation "${RESERVATION}") +elif [ -z "${RESERVATION+x}" ]; then + XPK_ARGS+=(--reservation "cloudtpu-20240716121201-595617744") +fi + +xpk workload create "${XPK_ARGS[@]}" diff --git a/src/maxtext/trainers/diloco/utils/__init__.py b/src/maxtext/trainers/diloco/utils/__init__.py new file mode 100644 index 0000000000..a8afd9f7f1 --- /dev/null +++ b/src/maxtext/trainers/diloco/utils/__init__.py @@ -0,0 +1,51 @@ +# Copyright 2023-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. + +"""DiLoCo utilities package.""" + +from maxtext.trainers.diloco.utils.spmd import ( + FragmentedTreeManipulator, + add_diloco_to_sharding, + apply_fragment_to_inner_state, + extract_per_island_metrics, + extract_replica_0, + from_diloco_checkpoint_dict, + get_streaming_schedule, + is_diloco_checkpoint, + replace_nnx_model_params, + replace_nnx_model_params_frag, + reshape_first_axis_with_diloco, + setup_diloco_initial_state, + synchronize_fragment_state, + synchronize_full_state, + to_diloco_checkpoint_dict, +) + +__all__ = [ + "FragmentedTreeManipulator", + "apply_fragment_to_inner_state", + "get_streaming_schedule", + "replace_nnx_model_params_frag", + "synchronize_fragment_state", + "add_diloco_to_sharding", + "extract_per_island_metrics", + "extract_replica_0", + "replace_nnx_model_params", + "reshape_first_axis_with_diloco", + "setup_diloco_initial_state", + "synchronize_full_state", + "from_diloco_checkpoint_dict", + "is_diloco_checkpoint", + "to_diloco_checkpoint_dict", +] diff --git a/src/maxtext/trainers/diloco/utils/spmd/__init__.py b/src/maxtext/trainers/diloco/utils/spmd/__init__.py new file mode 100644 index 0000000000..9e22c01844 --- /dev/null +++ b/src/maxtext/trainers/diloco/utils/spmd/__init__.py @@ -0,0 +1,55 @@ +# Copyright 2023-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. + +"""SPMD utilities for DiLoCo and Streaming DiLoCo.""" + +from maxtext.trainers.diloco.utils.spmd.checkpoint_utils import ( + from_diloco_checkpoint_dict, + is_diloco_checkpoint, + to_diloco_checkpoint_dict, +) +from maxtext.trainers.diloco.utils.spmd.fragment_utils import ( + FragmentedTreeManipulator, + apply_fragment_to_inner_state, + get_streaming_schedule, + replace_nnx_model_params_frag, + synchronize_fragment_state, +) +from maxtext.trainers.diloco.utils.spmd.state_utils import ( + add_diloco_to_sharding, + extract_per_island_metrics, + extract_replica_0, + replace_nnx_model_params, + reshape_first_axis_with_diloco, + setup_diloco_initial_state, + synchronize_full_state, +) + +__all__ = [ + "FragmentedTreeManipulator", + "apply_fragment_to_inner_state", + "get_streaming_schedule", + "replace_nnx_model_params_frag", + "synchronize_fragment_state", + "add_diloco_to_sharding", + "extract_per_island_metrics", + "extract_replica_0", + "replace_nnx_model_params", + "reshape_first_axis_with_diloco", + "setup_diloco_initial_state", + "synchronize_full_state", + "from_diloco_checkpoint_dict", + "is_diloco_checkpoint", + "to_diloco_checkpoint_dict", +] diff --git a/src/maxtext/trainers/diloco/utils/spmd/checkpoint_utils.py b/src/maxtext/trainers/diloco/utils/spmd/checkpoint_utils.py new file mode 100644 index 0000000000..2c9d572f23 --- /dev/null +++ b/src/maxtext/trainers/diloco/utils/spmd/checkpoint_utils.py @@ -0,0 +1,224 @@ +# Copyright 2023-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. + +"""Checkpointing and restoration utilities for SPMD DiLoCo in MaxText.""" + +from typing import Any +from flax import nnx +import jax +import jax.numpy as jnp +from maxtext.common import train_state_nnx +from maxtext.trainers.diloco import diloco +from maxtext.trainers.diloco.utils.spmd import state_utils +import optax + + +def is_diloco_checkpoint(restored_dict: Any) -> bool: + """Checks if a restored checkpoint dictionary contains a full multi-replica DiLoCo state.""" + if not isinstance(restored_dict, dict): + return False + # A full DiLoCo checkpoint contains 'inner_state' and 'outer_opt_state' + return "inner_state" in restored_dict and "outer_opt_state" in restored_dict + + +def to_diloco_checkpoint_dict(state: Any, config: Any = None) -> dict[str, Any]: + """Packages a DiLoCoTrainState or abstract state into an Orbax-serializable checkpoint dictionary. + + Saves / Serializes: + 1. inner_state: per-replica model weights + Adam first/second moment buffers (m, v) + step + 2. params: outer global model weights + 3. outer_opt_state: outer SGD Nesterov momentum state + 4. step: global step counter + + Args: + state: The DiLoCoTrainState instance or single-replica state (for abstract restore targets). + config: MaxText configuration object. + + Returns: + A dictionary formatted for Orbax saving / restoring. + """ + num_replicas = getattr(config, "dcn_diloco_parallelism", getattr(config, "num_diloco_replicas", 2)) + + if isinstance(state, diloco.DiLoCoTrainState): + inner_state = state.inner_state + params = state.params + outer_opt_state = state.outer_opt_state + step = state.step + elif hasattr(state, "inner_state"): + inner_state = state.inner_state + params = getattr(state, "params", None) + outer_opt_state = getattr(state, "outer_opt_state", None) + step = getattr(state, "step", jnp.array(0, dtype=jnp.int32)) + else: + # Single-replica state (e.g. abstract TrainStateNNX passed during checkpoint restoration) + def _add_diloco_dim(leaf): + if hasattr(leaf, "shape") and hasattr(leaf, "dtype"): + new_shape = (num_replicas, *leaf.shape) + sharding = getattr(leaf, "sharding", None) + if isinstance(sharding, jax.sharding.NamedSharding) and "diloco" in sharding.mesh.axis_names: + new_spec = jax.sharding.PartitionSpec("diloco", *sharding.spec) + sharding = jax.sharding.NamedSharding(mesh=sharding.mesh, spec=new_spec) + if isinstance(leaf, jax.ShapeDtypeStruct): + return jax.ShapeDtypeStruct(new_shape, leaf.dtype, sharding=sharding) + return jnp.broadcast_to(leaf, new_shape) + return leaf + + inner_state = jax.tree_util.tree_map(_add_diloco_dim, state) + + if hasattr(state, "model"): + _, params, _ = nnx.split(state.model, nnx.Param, ...) + params = params.to_pure_dict() if hasattr(params, "to_pure_dict") else params + elif hasattr(state, "params"): + params = state.params + else: + params = state + + outer_optimizer = optax.sgd( + getattr(config, "diloco_outer_lr", 0.1), + momentum=getattr(config, "diloco_outer_momentum", 0.9), + nesterov=True, + ) + outer_opt_state = outer_optimizer.init(params) + step = getattr(getattr(state, "optimizer", None), "step", jnp.array(0, dtype=jnp.int32)) + + # 1. Inner state: convert per-replica NNX state to Linen checkpoint layout + if isinstance(inner_state, (nnx.State, train_state_nnx.TrainStateNNX)): + inner_state_dict = train_state_nnx.to_checkpoint_dict(inner_state) + elif hasattr(inner_state, "to_pure_dict"): + inner_state_dict = inner_state.to_pure_dict() + elif isinstance(inner_state, dict): + inner_state_dict = inner_state + else: + inner_state_dict = inner_state + + # 2. Outer params: outer global model parameters + if hasattr(params, "to_pure_dict"): + params_dict = params.to_pure_dict() + elif isinstance(params, (nnx.State, nnx.Module)): + params_dict = nnx.state(params).to_pure_dict() + elif isinstance(params, dict): + params_dict = params + else: + params_dict = params + + # 3. Outer optimizer state: outer SGD momentum trace + # 4. Global step + step_val = step.get_value() if hasattr(step, "get_value") else step + + return { + "inner_state": inner_state_dict, + "params": params_dict, + "outer_opt_state": outer_opt_state, + "step": step_val, + } + + +def from_diloco_checkpoint_dict( + restored_dict: dict[str, Any], + abstract_diloco_state: Any, + config: Any = None, + mesh: jax.sharding.Mesh | None = None, +) -> Any: + """Restores a DiLoCoTrainState from an Orbax checkpoint dictionary. + + Supports: + - Full DiLoCo checkpoints (restores inner_state, outer_opt_state, params, step). + - Legacy / params-only checkpoints (restores params & step, broadcasts to inner_state, + and initializes fresh optimizer states). + + Args: + restored_dict: The dictionary loaded by Orbax. + abstract_diloco_state: Abstract DiLoCoTrainState or TrainStateNNX with expected shapes. + config: MaxText configuration object. + mesh: Device mesh. + + Returns: + A concrete DiLoCoTrainState instance. + """ + num_replicas = getattr(config, "dcn_diloco_parallelism", getattr(config, "num_diloco_replicas", 2)) + + def _add_diloco_dim(leaf): + if hasattr(leaf, "shape") and hasattr(leaf, "dtype"): + new_shape = (num_replicas, *leaf.shape) + sharding = getattr(leaf, "sharding", None) + if isinstance(sharding, jax.sharding.NamedSharding) and "diloco" in sharding.mesh.axis_names: + new_spec = jax.sharding.PartitionSpec("diloco", *sharding.spec) + sharding = jax.sharding.NamedSharding(mesh=sharding.mesh, spec=new_spec) + if isinstance(leaf, jax.ShapeDtypeStruct): + return jax.ShapeDtypeStruct(new_shape, leaf.dtype, sharding=sharding) + return jnp.broadcast_to(leaf, new_shape) + return leaf + + if is_diloco_checkpoint(restored_dict): + # Full multi-replica DiLoCo checkpoint restoration + inner_dict = restored_dict["inner_state"] + if isinstance(abstract_diloco_state, diloco.DiLoCoTrainState): + abstract_inner = abstract_diloco_state.inner_state + else: + abstract_inner = jax.tree_util.tree_map(_add_diloco_dim, abstract_diloco_state) + + if isinstance(abstract_inner, nnx.Module): + abstract_inner = nnx.state(abstract_inner) + + if isinstance(abstract_inner, (nnx.State, train_state_nnx.TrainStateNNX)): + inner_state = train_state_nnx.linen_items_to_nnx(inner_dict, abstract_inner) + else: + inner_state = inner_dict + + params = restored_dict["params"] + outer_opt_state = restored_dict["outer_opt_state"] + step = restored_dict["step"] + + return diloco.DiLoCoTrainState( + inner_state=inner_state, + params=params, + outer_opt_state=outer_opt_state, + step=step, + ) + + # Legacy checkpoint fallback (only params / model and step present) + raw_params = None + if "params" in restored_dict: + raw_params = restored_dict["params"] + if isinstance(raw_params, dict) and "params" in raw_params: + raw_params = raw_params["params"] + elif "model" in restored_dict: + raw_params = restored_dict["model"] + else: + raw_params = restored_dict + + step = restored_dict.get("step", jnp.array(0, dtype=jnp.int32)) + + broadcasted_model_params = jax.tree_util.tree_map(_add_diloco_dim, raw_params) + + if isinstance(abstract_diloco_state, diloco.DiLoCoTrainState): + abstract_inner = abstract_diloco_state.inner_state + else: + abstract_inner = jax.tree_util.tree_map(_add_diloco_dim, abstract_diloco_state) + + inner_state = state_utils.replace_nnx_model_params(abstract_inner, broadcasted_model_params) + + outer_optimizer = optax.sgd( + getattr(config, "diloco_outer_lr", 0.1), + momentum=getattr(config, "diloco_outer_momentum", 0.9), + nesterov=True, + ) + outer_opt_state = outer_optimizer.init(raw_params) + + return diloco.DiLoCoTrainState( + inner_state=inner_state, + params=raw_params, + outer_opt_state=outer_opt_state, + step=step, + ) diff --git a/src/maxtext/trainers/diloco/utils/spmd/fragment_utils.py b/src/maxtext/trainers/diloco/utils/spmd/fragment_utils.py new file mode 100644 index 0000000000..7cab8ab71d --- /dev/null +++ b/src/maxtext/trainers/diloco/utils/spmd/fragment_utils.py @@ -0,0 +1,258 @@ +# Copyright 2023-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. + +"""Streaming DiLoCo fragment manipulation and synchronization utilities for SPMD.""" + +import re +from typing import Any + +import drjax +from flax import nnx +import jax +import jax.numpy as jnp +from maxtext.common.train_state_nnx import TrainStateNNX +import optax + + +class FragmentedTreeManipulator: + """For Streaming DiLoCo: Partitions and manipulates fragments of a JAX PyTree, supporting scanned layers.""" + + def __init__( + self, + keypath_to_is_scanned: dict[str, bool], + fragment_to_layer_indices: dict[int, jax.Array], + num_fragments: int, + ): + self.keypath_to_is_scanned = keypath_to_is_scanned + self.fragment_to_layer_indices = fragment_to_layer_indices + self.num_fragments = num_fragments + + @classmethod + def create(cls, params_tree, config): + """Creates a FragmentedTreeManipulator from the parameters PyTree and configuration.""" + kvs, _ = jax.tree_util.tree_flatten_with_path(params_tree) + + num_layers = config.num_decoder_layers + num_fragments = config.num_diloco_fragments + num_transformer_fragments = num_fragments - 1 + + if num_transformer_fragments <= 0: + raise ValueError( + f"num_diloco_fragments ({num_fragments}) must be at least 2 (1 for non-scanned parameters, at least 1 for" + " scanned layers)." + ) + if num_layers % num_transformer_fragments != 0: + raise ValueError( + f"num_decoder_layers ({num_layers}) must be divisible by " + f"num_diloco_fragments - 1 ({num_transformer_fragments}) for now." + ) + + num_synced = num_layers // num_transformer_fragments + use_sequential = config.use_sequential_layers + + # Pre-compute layer indices for each fragment 1 ... num_transformer_fragments + fragment_to_layer_indices = {} + for i in range(1, num_fragments): + sync_id = i - 1 + if use_sequential: + indices = list(range(sync_id * num_synced, (sync_id + 1) * num_synced)) + else: + indices = list(range(sync_id, num_layers, num_transformer_fragments)) + fragment_to_layer_indices[i] = jnp.array(indices) + + # Regex to identify scanned layer parameters + 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 + + return cls(keypath_to_is_scanned, fragment_to_layer_indices, num_fragments) + + def get_flat_fragment(self, tree, fragment_idx: int, has_replica_dim: bool = False) -> dict[str, Any]: + """Extracts a flat dictionary containing parameters for the specified fragment index.""" + kvs, _ = jax.tree_util.tree_flatten_with_path(tree) + flat_frag = {} + for k, v in kvs: + keystr = jax.tree_util.keystr(k) + is_scanned = self.keypath_to_is_scanned.get(keystr, False) + if fragment_idx == 0: + if not is_scanned: + flat_frag[keystr] = v + else: + if is_scanned: + indices = self.fragment_to_layer_indices[fragment_idx] + if isinstance(v, jax.ShapeDtypeStruct): + new_shape = (v.shape[0], len(indices), *v.shape[2:]) if has_replica_dim else (len(indices), *v.shape[1:]) + flat_frag[keystr] = jax.ShapeDtypeStruct(new_shape, v.dtype) + elif has_replica_dim: + flat_frag[keystr] = v[:, indices] # Slice second dimension (layer axis) + else: + flat_frag[keystr] = v[indices] # Slice first dimension (layer axis) + return flat_frag + + def apply_flat_fragment( + self, + tree, + fragment_idx: int, + flat_fragment: dict[str, Any], + has_replica_dim: bool = False, + ): + """Merges a flat fragment dictionary back into the full parameters PyTree structure.""" + kvs, treedef = jax.tree_util.tree_flatten_with_path(tree) + new_kvs = [] + for k, v in kvs: + keystr = jax.tree_util.keystr(k) + is_scanned = self.keypath_to_is_scanned.get(keystr, False) + if fragment_idx == 0: + if not is_scanned: + new_kvs.append(flat_fragment[keystr]) + else: + new_kvs.append(v) + else: + 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) + return jax.tree_util.tree_unflatten(treedef, new_kvs) + + +def get_streaming_schedule(config) -> tuple[int, int]: + """Computes steps_between_syncs and synchronization period for streaming DiLoCo.""" + num_fragments = config.num_diloco_fragments + steps_between_syncs = int(round(config.diloco_sync_period / num_fragments)) + steps_between_syncs = max(1, steps_between_syncs) + period = num_fragments * steps_between_syncs + return steps_between_syncs, period + + +def replace_nnx_model_params_frag( + s, + manipulator: FragmentedTreeManipulator, + frag_idx: int, + outer_frag_replica: dict[str, Any], + alpha: float = 0.0, +): + """Replaces a single parameter fragment in an NNX TrainState with optional alpha interpolation.""" + s_model = s["model"] if hasattr(s, "keys") else s.model + graphdef, full_params, non_param_state = nnx.split(s_model, nnx.Param, ...) + full_params_dict = full_params.to_pure_dict() + if alpha > 0.0: + inner_frag = manipulator.get_flat_fragment(full_params_dict, frag_idx, has_replica_dim=False) + merged_frag = jax.tree.map(lambda i, o: alpha * i + (1 - alpha) * o, inner_frag, outer_frag_replica) + else: + merged_frag = outer_frag_replica + + new_full_params = manipulator.apply_flat_fragment(full_params_dict, frag_idx, merged_frag, has_replica_dim=False) + new_model = nnx.merge(graphdef, new_full_params, non_param_state) + + 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 hasattr(s, "keys"): + leaves_with_paths, treedef = jax.tree_util.tree_flatten_with_path(s) + new_model_iter = iter(jax.tree_util.tree_leaves(new_model)) + + def _is_model_leaf(path): + if not path: + return False + k = path[0] + return getattr(k, "key", None) == "model" or getattr(k, "name", None) == "model" + + new_leaves = [next(new_model_iter) if _is_model_leaf(p) else leaf for p, leaf in leaves_with_paths] + return jax.tree_util.tree_unflatten(treedef, new_leaves) + else: + s_opt = s["optimizer"] if hasattr(s, "keys") else s.optimizer + return TrainStateNNX(new_model, s_opt) + + +def synchronize_fragment_state( + state, + manipulator: FragmentedTreeManipulator, + frag_idx: int, + outer_optimizer: optax.GradientTransformation, + mesh: jax.sharding.Mesh | None = None, +): + """Synchronizes a single parameter fragment across DiLoCo replicas in streaming DiLoCo.""" + # 1. Extract global and local parameters for the fragment + outer_params_frag = manipulator.get_flat_fragment(state.params, frag_idx, has_replica_dim=False) + inner_model_params = nnx.filter_state(state.inner_state.model, nnx.Param).to_pure_dict() + inner_params_frag = manipulator.get_flat_fragment(inner_model_params, frag_idx, has_replica_dim=True) + + # 2. Compute the pseudo-gradient: outer - inner + broadcast_outer_frag = drjax.broadcast(outer_params_frag, mesh=mesh) + unreduced_grads = jax.tree.map(lambda x, y: x - y, broadcast_outer_frag, inner_params_frag) + + # 3. Average gradients across replicas + averaged_pseudo_grad = drjax.reduce_mean(unreduced_grads) + + # 4. Extract outer optimizer state for this fragment (TraceState is (trace, EmptyState)) + trace_frag = manipulator.get_flat_fragment(state.outer_opt_state[0].trace, frag_idx, has_replica_dim=False) + opt_state_frag = (optax.TraceState(trace=trace_frag), optax.EmptyState()) + + # 5. Run outer optimizer on the fragment + updates_frag, new_opt_state_frag = outer_optimizer.update( + averaged_pseudo_grad, opt_state_frag, params=outer_params_frag + ) + new_outer_params_frag = optax.apply_updates(outer_params_frag, updates_frag) + + # 6. Re-merge updated params and optimizer states back to full PyTree + new_params = manipulator.apply_flat_fragment(state.params, frag_idx, new_outer_params_frag, has_replica_dim=False) + new_trace = manipulator.apply_flat_fragment( + state.outer_opt_state[0].trace, frag_idx, new_opt_state_frag[0].trace, has_replica_dim=False + ) + new_outer_opt_state = (optax.TraceState(trace=new_trace), state.outer_opt_state[1]) + + return state.replace( + params=new_params, + outer_opt_state=new_outer_opt_state, + ) + + +def apply_fragment_to_inner_state( + state, + manipulator: FragmentedTreeManipulator, + frag_idx: int, + alpha: float = 0.0, + mesh: jax.sharding.Mesh | None = None, +): + """Broadcasts synced outer parameter fragment and updates inner state across replicas.""" + outer_params_frag = manipulator.get_flat_fragment(state.params, frag_idx, has_replica_dim=False) + broadcast_outer_frag = drjax.broadcast(outer_params_frag, mesh=mesh) + + new_inner_state = drjax.map_fn( + lambda s, frag: replace_nnx_model_params_frag(s, manipulator, frag_idx, frag, alpha=alpha), + (state.inner_state, broadcast_outer_frag), + mesh=mesh, + ) + return state.replace(inner_state=new_inner_state) diff --git a/src/maxtext/trainers/diloco/utils/spmd/state_utils.py b/src/maxtext/trainers/diloco/utils/spmd/state_utils.py new file mode 100644 index 0000000000..af08cbf74d --- /dev/null +++ b/src/maxtext/trainers/diloco/utils/spmd/state_utils.py @@ -0,0 +1,227 @@ +# Copyright 2023-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. + +"""Core DiLoCo state initialization, sharding, and synchronization utilities for SPMD.""" + +from collections.abc import Sequence +from typing import Any + +import drjax +from flax import nnx +import jax +import jax.numpy as jnp +from jaxtyping import PyTree +from maxtext.common.train_state_nnx import TrainStateNNX +from maxtext.trainers.diloco import diloco +import optax + + +def add_diloco_to_sharding(pytree: PyTree) -> PyTree: + """Recursively traverses a PyTree and prepends 'diloco' to the PartitionSpec + of any NamedSharding object that contains 'diloco' in its mesh axis names. + """ + + def map_fn(leaf): + if isinstance(leaf, jax.sharding.NamedSharding): + if "diloco" not in leaf.mesh.axis_names: + return leaf + new_spec = jax.sharding.PartitionSpec("diloco", *leaf.spec) + return jax.sharding.NamedSharding(mesh=leaf.mesh, spec=new_spec) + return leaf + + return jax.tree_util.tree_map(map_fn, pytree) + + +def reshape_first_axis_with_diloco(num_diloco_replicas: int, pytree: PyTree) -> PyTree: + """Reshapes the first dimension of each array in the PyTree to include a DiLoCo axis.""" + + def extend_pspec( + pspec: jax.sharding.PartitionSpec | Sequence[str | Sequence[str]] = (), + ) -> jax.sharding.PartitionSpec: + if pspec and isinstance(pspec[0], (tuple, list)) and len(pspec[0]) > 0 and pspec[0][0] == "diloco": + remaining = tuple(pspec[0][1:]) + if len(remaining) == 1: + return jax.sharding.PartitionSpec("diloco", remaining[0], *pspec[1:]) + elif len(remaining) > 1: + return jax.sharding.PartitionSpec("diloco", remaining, *pspec[1:]) + else: + return jax.sharding.PartitionSpec("diloco", *pspec[1:]) + return jax.sharding.PartitionSpec("diloco", *pspec) + + def reshape_for_diloco(arr): + if not hasattr(arr, "shape"): + return arr + if ( + hasattr(arr, "ndim") + and arr.ndim >= 3 + and arr.shape[0] == num_diloco_replicas + and hasattr(arr, "sharding") + and isinstance(arr.sharding, jax.sharding.NamedSharding) + and arr.sharding.spec + and arr.sharding.spec[0] == "diloco" + and isinstance(arr.sharding.spec[0], str) + ): + return arr + batch_dim, *example_shape = arr.shape + if batch_dim % num_diloco_replicas != 0: + raise ValueError(f"Batch dimension {batch_dim} is not divisible by num_diloco_replicas {num_diloco_replicas}.") + diloco_shape = (num_diloco_replicas, batch_dim // num_diloco_replicas, *example_shape) + if hasattr(arr, "sharding") and arr.sharding is not None: + s = arr.sharding + s = jax.sharding.NamedSharding(mesh=s.mesh, spec=extend_pspec(s.spec)) + return jax.lax.with_sharding_constraint(jnp.reshape(arr, shape=diloco_shape), s) + return jnp.reshape(arr, shape=diloco_shape) + + return jax.tree.map(reshape_for_diloco, pytree) + + +def extract_replica_0(metrics: PyTree) -> PyTree: + """Extracts metrics from replica 0 across DiLoCo islands.""" + + def select_first_replica(x): + if not hasattr(x, "shape") or len(x.shape) == 0: + return x + r = x.shape[0] + mask = (jnp.arange(r) == 0).reshape((r,) + (1,) * (x.ndim - 1)) + return drjax.reduce_sum(x * mask) + + return jax.tree.map(select_first_replica, metrics) + + +def extract_per_island_metrics(metrics: PyTree, num_diloco_replicas: int) -> PyTree: + """Extracts replica 0 metrics and appends per-island loss metrics.""" + default_metrics = extract_replica_0(metrics) + if isinstance(metrics, dict) and "scalar" in metrics and "learning/loss" in metrics["scalar"]: + for i in range(num_diloco_replicas): + mask_i = jnp.arange(num_diloco_replicas) == i + loss_i = drjax.reduce_sum(metrics["scalar"]["learning/loss"] * mask_i) + default_metrics["scalar"][f"learning/loss_island_{i}"] = loss_i + return default_metrics + + +def replace_nnx_model_params(s, new_params): + """Replaces model parameters in an NNX TrainState or dictionary structure.""" + s_model = s["model"] if hasattr(s, "keys") else s.model + s_opt = s["optimizer"] if hasattr(s, "keys") else s.optimizer + + graphdef, _, non_param_state = nnx.split(s_model, nnx.Param, ...) + new_model = nnx.merge(graphdef, new_params, non_param_state) + + 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 hasattr(s, "keys"): + leaves_with_paths, treedef = jax.tree_util.tree_flatten_with_path(s) + new_model_iter = iter(jax.tree_util.tree_leaves(new_model)) + + def _is_model_leaf(path): + if not path: + return False + k = path[0] + return getattr(k, "key", None) == "model" or getattr(k, "name", None) == "model" + + new_leaves = [next(new_model_iter) if _is_model_leaf(p) else leaf for p, leaf in leaves_with_paths] + return jax.tree_util.tree_unflatten(treedef, new_leaves) + else: + return TrainStateNNX(new_model, s_opt) + + +def synchronize_full_state( + state, + outer_optimizer: optax.GradientTransformation, + mesh: jax.sharding.Mesh | None = None, +): + """Synchronizes all parameters across DiLoCo replicas for vanilla DiLoCo.""" + broadcast_outer_params = drjax.broadcast(state.params, mesh=mesh) + _, inner_model_params, _ = nnx.split(state.inner_state.model, nnx.Param, ...) + inner_model_params = inner_model_params.to_pure_dict() + + model_delta = jax.tree.map(lambda x, y: y - x, inner_model_params, broadcast_outer_params) + averaged_pseudo_grad = drjax.reduce_mean(model_delta) + updates, new_opt_state = outer_optimizer.update(averaged_pseudo_grad, state.outer_opt_state, state.params) + new_outer_params = optax.apply_updates(state.params, updates) + + new_inner_state = drjax.map_fn( + lambda s: replace_nnx_model_params(s, new_outer_params), + state.inner_state, + mesh=mesh, + ) + return state.replace( + params=new_outer_params, + outer_opt_state=new_opt_state, + inner_state=new_inner_state, + ) + + +def setup_diloco_initial_state( + state: Any, + config: Any, + mesh: jax.sharding.Mesh, + state_mesh_shardings: PyTree, + restored: Any = None, +) -> Any: + """Builds the full DiLoCoTrainState from the restored or freshly initialized single-replica state.""" + if isinstance(state, diloco.DiLoCoTrainState): + return state + + # 1. Compute per-replica shardings with 'diloco' axis prepended + inner_state_shardings = add_diloco_to_sharding(state_mesh_shardings) + + # 2. Extract concrete outer model parameters from state.model (or state.params) + if hasattr(state, "model"): + _, outer_params, _ = nnx.split(state.model, nnx.Param, ...) + outer_params = outer_params.to_pure_dict() if hasattr(outer_params, "to_pure_dict") else outer_params + else: + outer_params = getattr(state, "params", state) + + # 3. Broadcast single-replica state to multi-replica inner_state across the diloco axis + def _broadcast_to_replicas(leaf, sharding): + if hasattr(leaf, "shape"): + target_shape = (config.num_diloco_replicas, *leaf.shape) + if isinstance(leaf, jax.ShapeDtypeStruct): + sharding_arg = sharding if isinstance(sharding, jax.sharding.NamedSharding) else None + return jax.ShapeDtypeStruct(target_shape, leaf.dtype, sharding=sharding_arg) + if isinstance(sharding, jax.sharding.NamedSharding): + return jax.jit( + lambda x: jnp.broadcast_to(x, target_shape), + out_shardings=sharding, + )(leaf) + return jnp.broadcast_to(leaf, target_shape) + return leaf + + inner_state = jax.tree_util.tree_map( + _broadcast_to_replicas, + state, + inner_state_shardings, + ) + + # 4. Initialize outer optimizer state with outer SGD momentum + outer_optimizer = optax.sgd( + config.diloco_outer_lr, + momentum=config.diloco_outer_momentum, + nesterov=True, + ) + outer_opt_state = outer_optimizer.init(outer_params) + + # 5. Extract global step + step = getattr(getattr(state, "optimizer", None), "step", jnp.array(0, dtype=jnp.int32)) + + return diloco.DiLoCoTrainState( + inner_state=inner_state, + params=outer_params, + outer_opt_state=outer_opt_state, + step=step, + ) diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py index 138bfe1ec7..96f168572f 100644 --- a/src/maxtext/trainers/pre_train/train.py +++ b/src/maxtext/trainers/pre_train/train.py @@ -77,8 +77,9 @@ def get_first_step(model, state): if isinstance(model, nn.Module): return int(state.step) - if hasattr(state, "inner_state"): # DiLoCoTrainState (NNX DiLoCo): step is the optimizer step var - return int(state.step.get_value()) + if hasattr(state, "inner_state"): # DiLoCoTrainState (NNX DiLoCo) + step_val = state.step.get_value() if hasattr(state.step, "get_value") else state.step + return int(step_val) return int(state.optimizer.step.get_value()) @@ -802,7 +803,7 @@ def train_loop(config, recorder, state=None): start_step = get_first_step(model, state) # this is the start_step for training train_utils.validate_completed_steps(start_step, config.steps) - if isinstance(model, nn.Module): + if config.enable_diloco or isinstance(model, nn.Module): jit_model = model elif config.enable_diloco: # state is the DiLoCoTrainState; `model` is already the TrainStateNNX graphdef the inner step needs. @@ -852,7 +853,7 @@ def train_loop(config, recorder, state=None): metric_logger_instance = metric_logger.MetricLogger(config=config, learning_rate_schedule=learning_rate_schedule) # 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): setup_params = state.params elif config.enable_diloco: setup_params = state.params # DiLoCoTrainState.params: the outer (global) params diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py index 769e86b4bd..284a9eab44 100644 --- a/src/maxtext/utils/maxtext_utils.py +++ b/src/maxtext/utils/maxtext_utils.py @@ -39,6 +39,7 @@ from maxtext.configs import pyconfig from maxtext.configs import types from maxtext.multimodal import processor as mm_processor +from maxtext.trainers.diloco import utils as diloco_utils from maxtext.utils import elastic_utils from maxtext.utils import gcs_utils from maxtext.utils import max_logging @@ -1690,6 +1691,8 @@ def setup_initial_state( config, mesh, init_state_fn, is_training ) + abstract_restore_state = unboxed_abstract_state + # Initialization with nn_partitioning.axis_rules(config.logical_axis_rules): restored, raw_params = checkpointing.load_state_if_possible( @@ -1698,7 +1701,7 @@ def setup_initial_state( config.load_parameters_path, config.load_full_state_path, config.checkpoint_storage_concurrent_gb, - unboxed_abstract_state, + abstract_restore_state, config.enable_single_replica_ckpt_restoring, config.dataset_type, use_ocdbt=config.checkpoint_storage_use_ocdbt, @@ -1814,21 +1817,33 @@ def _merge_restored_overlay(ckpt_node, init_node): )() if raw_params: # If we loaded a partial state, we need to merge it. sparsity_enabled = config.weight_sparsity_n and config.weight_sparsity_m + target_params = raw_params if sparsity_enabled: - # Sparsity-init keeps freshly initialized params for any leaf still - # represented as an abstract ShapeDtypeStruct in raw_params (i.e. not - # actually restored), and uses the restored value otherwise. + def _merge_params(p_raw, p_init): if isinstance(p_raw, jax.ShapeDtypeStruct): return p_init return p_raw - merged_params = jax.tree_util.tree_map(_merge_params, raw_params, state.params) - state = state.replace(params=merged_params) - else: - state = state.replace(params=raw_params) + target_params = jax.tree_util.tree_map(_merge_params, raw_params, state.params) + + if hasattr(state, "keys") and "model" in state: + nnx.update(state["model"], target_params) + elif hasattr(state, "model"): + nnx.update(state.model, target_params) + elif hasattr(state, "replace"): + state = state.replace(params=target_params) + if not config.pure_nnx: state = max_utils.unbox_logicallypartioned(state) + if config.enable_diloco: + state = diloco_utils.setup_diloco_initial_state( + state=state, + config=config, + mesh=mesh, + state_mesh_shardings=state_mesh_shardings, + restored=restored, + ) return state, state_mesh_annotations, state_mesh_shardings, data_iterator, was_restored diff --git a/src/maxtext/utils/train_utils.py b/src/maxtext/utils/train_utils.py index 69e4e5fbaf..d73cea00be 100644 --- a/src/maxtext/utils/train_utils.py +++ b/src/maxtext/utils/train_utils.py @@ -17,6 +17,7 @@ import subprocess import jax +import optax import functools import orbax.checkpoint.pathways as ocp_pathways from functools import partial @@ -32,6 +33,7 @@ from maxtext.common.goodput import GoodputEvent, maybe_record_goodput from maxtext.optimizers import optimizers from maxtext.trainers.diloco import diloco +from maxtext.trainers.diloco import utils as diloco_utils from maxtext.utils import lora_utils from maxtext.utils import max_logging from maxtext.utils import max_utils @@ -329,6 +331,15 @@ def create_train_state_fn(): # logical_axis_rules (e.g. concat_embed on the MTP kernel). Tracing shapes # without a mesh skips sharding resolution, so it avoids the crash. state_graphdef = nnx.graphdef(nnx.eval_shape(init_state_fn)) + + if isinstance(state, diloco.DiLoCoTrainState): + state_params = state.params + if hasattr(state_mesh_shardings, "model"): + _, state_mesh_shardings_params, _ = nnx.split(state_mesh_shardings.model, nnx.Param, ...) + else: + state_mesh_shardings_params = state_mesh_shardings.params + elif config.pure_nnx: + with nn_partitioning.axis_rules(config.logical_axis_rules): _, state_params, _ = nnx.split(state.model, nnx.Param, ...) _, state_mesh_shardings_params, _ = nnx.split(state_mesh_shardings.model, nnx.Param, ...) else: @@ -337,17 +348,23 @@ def create_train_state_fn(): if config.enable_diloco: with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules): - state, outer_opt_state_sharding = diloco.build_diloco_state(config, lambda: state, mesh=mesh) + if not isinstance(state, diloco.DiLoCoTrainState): + state, outer_opt_state_sharding = diloco.build_diloco_state(config, lambda: state, mesh=mesh) + else: + outer_opt_state_sharding = ( + optax.TraceState(trace=state_mesh_shardings_params), + optax.EmptyState(), + ) # create state_mesh_shardings for the DilocoState - step_mesh = state_mesh_shardings.optimizer.step.mesh if config.pure_nnx else state_mesh_shardings.step.mesh - inner_state_shardings = diloco.add_diloco_to_sharding(state_mesh_shardings) + step_mesh = state_mesh_shardings.optimizer.step.mesh + inner_state_shardings = diloco_utils.add_diloco_to_sharding(state_mesh_shardings) state_mesh_shardings = diloco.DiLoCoTrainState( inner_state_shardings, # Match the outer params' pure-dict structure (build_diloco_state stores # outer_params via to_pure_dict), so the sharding tree matches the state tree. state_mesh_shardings_params.to_pure_dict() # pyrefly: ignore[missing-attribute] - if config.pure_nnx + if isinstance(state_mesh_shardings_params, nnx.State) else state_mesh_shardings_params, outer_opt_state_sharding, jax.sharding.NamedSharding( # pyrefly: ignore[bad-argument-type] diff --git a/tests/integration/diloco_test.py b/tests/integration/diloco_test.py index bfd54d21c1..bff359f1ff 100644 --- a/tests/integration/diloco_test.py +++ b/tests/integration/diloco_test.py @@ -28,6 +28,7 @@ from maxtext.common.train_state_nnx import TrainStateNNX from maxtext.configs.pyconfig import initialize_pydantic from maxtext.trainers.diloco import diloco +from maxtext.trainers.diloco import utils as diloco_utils from maxtext.trainers.pre_train.train_compile import main as train_compile_main from tests.utils.test_helpers import get_test_config_path import numpy as np @@ -392,3 +393,107 @@ def test_diloco_two_slices(self): "head_dim=4", ) ) + + @pytest.mark.cpu_only + @pytest.mark.tpu_backend + def test_streaming_diloco_two_slices(self): + temp_dir = gettempdir() + compiled_trainstep_file = os.path.join(temp_dir, "test_compiled_streaming_diloco.pickle") + train_compile_main( + ( + None, + get_test_config_path(), + f"compiled_trainstep_file={compiled_trainstep_file}", + "compile_topology=tpu7x-8", + "compile_topology_num_slices=2", + "ici_fsdp_parallelism=-1", + "dcn_diloco_parallelism=2", + "enable_diloco=true", + "enable_streaming_diloco=true", + "num_diloco_fragments=2", + "model_name=gemma2-2b", + "override_model_config=True", + "base_emb_dim=32", + "base_num_decoder_layers=2", + "base_mlp_dim=64", + "base_num_query_heads=1", + "base_num_kv_heads=1", + "head_dim=4", + ) + ) + + def test_fragmented_tree_manipulator_scanned_filter(self): + """Tests that parameters matching regex but lacking leading layer dim are NOT marked scanned.""" + num_layers = 4 + config = initialize_pydantic( + [ + "", + get_test_config_path(), + "enable_diloco=true", + "enable_streaming_diloco=true", + "num_diloco_fragments=3", + f"base_num_decoder_layers={num_layers}", + ] + ) + # Scanned param has leading dim = num_layers; non-scanned param matches regex name but lacks leading layer dim + params_tree = { + "decoder": { + "layers": jnp.ones((num_layers, 16, 16)), + "layers_outside_pipeline": jnp.ones((16, 16)), # Lacks leading layer dim = 4 + } + } + manipulator = diloco_utils.FragmentedTreeManipulator.create(params_tree, config) + # Check that layers_outside_pipeline is NOT treated as scanned + scanned_map = manipulator.keypath_to_is_scanned + for keystr, is_scanned in scanned_map.items(): + if "layers_outside_pipeline" in keystr: + self.assertFalse(is_scanned) + elif "decoder/layers" in keystr: + self.assertTrue(is_scanned) + + def test_streaming_diloco_requires_scan_layers(self): + """Tests that enable_streaming_diloco=True raises ValueError if scan_layers=False.""" + with self.assertRaises(ValueError) as ctx: + initialize_pydantic( + [ + "", + get_test_config_path(), + "enable_diloco=true", + "enable_streaming_diloco=true", + "num_diloco_fragments=2", + "scan_layers=false", + ] + ) + self.assertIn("enable_streaming_diloco=True requires scan_layers=True", str(ctx.exception)) + + def test_apply_flat_fragment_shapedtypestruct(self): + """Tests that FragmentedTreeManipulator handles ShapeDtypeStruct leaves during abstract tracing.""" + num_layers = 2 + config = initialize_pydantic( + [ + "", + get_test_config_path(), + "enable_diloco=true", + "enable_streaming_diloco=true", + "num_diloco_fragments=2", + f"base_num_decoder_layers={num_layers}", + ] + ) + abstract_tree = {"decoder": {"layers": jax.ShapeDtypeStruct((num_layers, 8), jnp.float32)}} + manipulator = diloco_utils.FragmentedTreeManipulator.create(abstract_tree, config) + frag = manipulator.get_flat_fragment(abstract_tree, fragment_idx=1) + res = manipulator.apply_flat_fragment(abstract_tree, fragment_idx=1, flat_fragment=frag) + self.assertIsInstance(res["decoder"]["layers"], jax.ShapeDtypeStruct) + + def test_diloco_requires_pure_nnx(self): + """Tests that enable_diloco=True raises ValueError if pure_nnx=False.""" + with self.assertRaises(ValueError) as ctx: + initialize_pydantic( + [ + "", + get_test_config_path(), + "enable_diloco=true", + "pure_nnx=false", + ] + ) + self.assertIn("enable_diloco=True requires pure_nnx=True", str(ctx.exception))