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))