diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index ecc349d9be..1e57fe6adc 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -896,6 +896,15 @@ dcn_bandwidth_latency: "50ms" # The network interface to apply throttling rules to. dcn_bandwidth_interface: "eth0" +# Streaming DiLoCo params +enable_streaming_diloco: false +enable_non_spmd_diloco: false +num_diloco_fragments: 16 +use_sequential_layers: false +num_communication_overlapping_steps: 0 +communication_overlapping_alpha: 0.0 + + # You may disable clipping by setting gradient_clipping_threshold to zero. gradient_clipping_threshold: 1.0 diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 757e04e515..d4c5073794 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1521,7 +1521,7 @@ class DilocoParams(BaseModel): diloco_outer_lr: float = Field(0.3, description="learning rate for outer optimizer.") diloco_outer_momentum: float = Field(0.9, description="momentum for outer optimizer.") dcn_bandwidth_limit: str = Field( - "", description="Programmatic DCN egress bandwidth limit (e.g., '28gbit'). Empty means no limit." + "", description="Programmatic DCN egress bandwidth limit per VM (e.g., '28gbit'). Empty means no limit." ) dcn_bandwidth_burst: str = Field("10mb", description="Burst size for Token Bucket Filter (TBF) traffic shaping.") dcn_bandwidth_latency: str = Field( @@ -1529,6 +1529,37 @@ 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.") + enable_non_spmd_diloco: bool = Field(False, description="Enable non-SPMD, multi-threaded streaming DiLoCo.") + num_diloco_fragments: int = Field(16, description="Number of fragments to partition the model layers into.") + use_sequential_layers: bool = Field(False, description="Whether to sync layers sequentially (or interleaved).") + num_communication_overlapping_steps: int = Field( + 0, description="Steps of communication overlap with computation. \\tau from the paper." + ) + communication_overlapping_alpha: float = Field( + 0.0, + description=( + "Interpolation factor between local and global parameters. alpha=1" + " means no communication between islands, alpha=0 means discards any" + " updates done in the inner optimizer in the first" + " `num_communication_overlapping_steps` steps. alpha=0.5 does a" + " uniform average between the local fragment parameters and the" + " globally shared one." + ), + ) + + @model_validator(mode="after") + def validate_overlap_steps(self) -> "DilocoParams": + if self.enable_streaming_diloco: + if self.num_communication_overlapping_steps >= self.diloco_sync_period: + raise ValueError( + "num_communication_overlapping_steps must be strictly less than" + f" diloco_sync_period. Got {self.num_communication_overlapping_steps=}" + f" and {self.diloco_sync_period=}" + ) + return self + class Optimizer(BaseModel): """Configuration for the optimizer and learning rate schedule.""" @@ -3025,6 +3056,8 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de "Colocated python data input is only supported with Pathways (single" " controller) enabled (`enable_single_controller=True`)." ) + if self.enable_non_spmd_diloco and self.colocated_python_data_input: + raise ValueError("Non-SPMD DiLoCo does not support colocated python data input.") if self.grain_use_elastic_iterator and self.grain_file_type != "arrayrecord": raise ValueError( "`grain_use_elastic_iterator=True` only supports `grain_file_type=arrayrecord`. " diff --git a/src/maxtext/input_pipeline/grain_data_processing.py b/src/maxtext/input_pipeline/grain_data_processing.py index 114b49abac..063f3b7a28 100644 --- a/src/maxtext/input_pipeline/grain_data_processing.py +++ b/src/maxtext/input_pipeline/grain_data_processing.py @@ -422,6 +422,8 @@ def make_grain_train_iterator( config: ml_collections.ConfigDict, global_mesh, process_indices, + learner_idx: int = 0, + num_learners: int = 1, ): """Load, preprocess dataset and return iterators""" assert ( @@ -478,9 +480,11 @@ def make_grain_train_iterator( if 0 < config.expansion_factor_real_data < 1: num_dataloader_to_restore = int(1 / config.expansion_factor_real_data) train_dataloader_list = [] - dataloading_host_count = len(process_indices) * num_dataloader_to_restore + dataloading_host_count = len(process_indices) * num_dataloader_to_restore * num_learners for i in range(num_dataloader_to_restore): - dataloading_host_index = len(process_indices) * i + process_indices.index(jax.process_index()) + dataloading_host_index = ( + len(process_indices) * i + process_indices.index(jax.process_index()) + ) * num_learners + learner_idx train_ds = get_ds_fn(dataloading_host_index=dataloading_host_index, dataloading_host_count=dataloading_host_count) train_dataloader = preprocessing_fn(dataset=train_ds) train_dataloader_list.append(train_dataloader) @@ -490,8 +494,8 @@ def make_grain_train_iterator( ] # Default non-colocated, non-expansion path - shard_index = process_indices.index(jax.process_index()) - shard_count = len(process_indices) + shard_index = process_indices.index(jax.process_index()) * num_learners + learner_idx + shard_count = len(process_indices) * num_learners train_ds = get_ds_fn( dataloading_host_index=shard_index, dataloading_host_count=shard_count, @@ -523,6 +527,8 @@ def make_grain_eval_iterator( config: ml_collections.ConfigDict, global_mesh, process_indices, + learner_idx: int = 0, + num_learners: int = 1, ): """Load, preprocess dataset and return iterators""" assert ( @@ -556,8 +562,8 @@ def make_grain_eval_iterator( if not config.colocated_python_data_input: eval_ds = get_ds_fn( - dataloading_host_index=process_indices.index(jax.process_index()), - dataloading_host_count=len(process_indices), + dataloading_host_index=process_indices.index(jax.process_index()) * num_learners + learner_idx, + dataloading_host_count=len(process_indices) * num_learners, ) eval_dataloader = preprocessing_fn(dataset=eval_ds) return multihost_dataloading.MultiHostDataLoadIterator( diff --git a/src/maxtext/input_pipeline/input_pipeline_interface.py b/src/maxtext/input_pipeline/input_pipeline_interface.py index 4a0b37532f..f6a2f02676 100644 --- a/src/maxtext/input_pipeline/input_pipeline_interface.py +++ b/src/maxtext/input_pipeline/input_pipeline_interface.py @@ -14,6 +14,7 @@ """Input pipeline""" import functools +import inspect import jax from jax.sharding import PartitionSpec as P @@ -46,20 +47,33 @@ def get_process_loading_real_data( return list(process_loading_real_data) -def create_process_specific_iterator(config: pyconfig.HyperParameters, mesh, process_indices, input_iterator): +def create_process_specific_iterator( + config: pyconfig.HyperParameters, + mesh, + process_indices, + input_iterator, + learner_idx: int = 0, + num_learners: int = 1, +): """ If the current process's index is among the `process_indices`, a real data iterator is created. Otherwise, a placeholder iterator is returned. """ if jax.process_index() in process_indices: - iterator_fn = functools.partial(input_iterator, config, mesh, process_indices) + sig = inspect.signature(input_iterator) + kwargs = {} + if "learner_idx" in sig.parameters: + kwargs["learner_idx"] = learner_idx + if "num_learners" in sig.parameters: + kwargs["num_learners"] = num_learners + iterator_fn = functools.partial(input_iterator, config, mesh, process_indices, **kwargs) output_iterator = iterator_fn() else: output_iterator = PlaceHolderDataIterator(config, mesh) return output_iterator -def create_data_iterator(config: pyconfig.HyperParameters, mesh): +def create_data_iterator(config: pyconfig.HyperParameters, mesh, learner_idx: int = 0, num_learners: int = 1): """Create train and eval data iterators given configs and mesh.""" # Return synthetic dataset if selected @@ -99,7 +113,9 @@ def create_data_iterator(config: pyconfig.HyperParameters, mesh): config.max_target_length, mesh, ) - output_train_iterator = create_process_specific_iterator(config, mesh, process_indices_train, train_iterator) + output_train_iterator = create_process_specific_iterator( + config, mesh, process_indices_train, train_iterator, learner_idx=learner_idx, num_learners=num_learners + ) if config.expansion_factor_real_data > 1: # assert number of hosts loading real data assert len(process_indices_train) == jax.process_count() // config.expansion_factor_real_data @@ -116,5 +132,7 @@ def create_data_iterator(config: pyconfig.HyperParameters, mesh): if config.expansion_factor_real_data > 1: assert len(process_indices_eval) == jax.process_count() // config.expansion_factor_real_data - output_eval_iterator = create_process_specific_iterator(config, mesh, process_indices_eval, eval_iterator) + output_eval_iterator = create_process_specific_iterator( + config, mesh, process_indices_eval, eval_iterator, learner_idx=learner_idx, num_learners=num_learners + ) return output_train_iterator, output_eval_iterator diff --git a/src/maxtext/input_pipeline/tfds_data_processing.py b/src/maxtext/input_pipeline/tfds_data_processing.py index 9296b1e00d..b79c564db2 100644 --- a/src/maxtext/input_pipeline/tfds_data_processing.py +++ b/src/maxtext/input_pipeline/tfds_data_processing.py @@ -185,6 +185,8 @@ def make_tfds_train_iterator( config: ml_collections.ConfigDict, global_mesh, process_indices_train, + learner_idx: int = 0, + num_learners: int = 1, ): """load dataset, preprocess and return iterators""" assert ( @@ -200,8 +202,8 @@ def make_tfds_train_iterator( } if not config.colocated_python_data_input: train_ds = get_datasets( - dataloading_host_index=process_indices_train.index(jax.process_index()), - dataloading_host_count=len(process_indices_train), + dataloading_host_index=process_indices_train.index(jax.process_index()) * num_learners + learner_idx, + dataloading_host_count=len(process_indices_train) * num_learners, **get_datasets_kwargs, ) train_dataloader = preprocessing_pipeline( @@ -254,6 +256,8 @@ def make_tfds_eval_iterator( config: ml_collections.ConfigDict, global_mesh, process_indices_eval, + learner_idx: int = 0, + num_learners: int = 1, ): """load eval dataset, preprocess and return iterators""" assert ( @@ -266,8 +270,8 @@ def make_tfds_eval_iterator( data_split=config.eval_split, shuffle_files=False, shuffle_seed=config.data_shuffle_seed, - dataloading_host_index=process_indices_eval.index(jax.process_index()), - dataloading_host_count=len(process_indices_eval), + dataloading_host_index=process_indices_eval.index(jax.process_index()) * num_learners + learner_idx, + dataloading_host_count=len(process_indices_eval) * num_learners, ) eval_dataloader = preprocessing_pipeline( dataset=eval_ds, diff --git a/src/maxtext/trainers/diloco/decomposed_transport.py b/src/maxtext/trainers/diloco/decomposed_transport.py new file mode 100644 index 0000000000..a59a9afdd8 --- /dev/null +++ b/src/maxtext/trainers/diloco/decomposed_transport.py @@ -0,0 +1,148 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Thread-safe transport layer for non-SPMD multi-threaded DiLoCo.""" + +import queue +from typing import Any +from concurrent.futures import ThreadPoolExecutor +import traceback +import jax +import jax.numpy as jnp +from jax.experimental import colocated_python +from maxtext.utils import max_logging + + + +class ThreadedTransportManager: + """Manages in-memory communication between learner threads and the syncer thread.""" + + def __init__(self, num_learners: int): + self.num_learners = num_learners + + # Thread-safe FIFO queues for each learner. + # Stores tuples of (step, fragment_id, data). + self._learner_to_syncer_queues = [queue.Queue() for _ in range(num_learners)] + self._syncer_to_learner_queues = [queue.Queue() for _ in range(num_learners)] + + # Local buffers for out-of-order storage. + # Since only the receiving thread (syncer for learner_to_syncer, and respective + # learner for syncer_to_learner) accesses these buffers, we do not need locks. + self._learner_to_syncer_buffers = [{} for _ in range(num_learners)] + self._syncer_to_learner_buffers = [{} for _ in range(num_learners)] + + def send_to_syncer(self, learner_idx: int, step: int, fragment_id: int, data: Any): + """Learner sends data to the syncer.""" + self._learner_to_syncer_queues[learner_idx].put((step, fragment_id, data)) + + def recv_from_learner(self, learner_idx: int, step: int, fragment_id: int) -> Any: + """Syncer receives data from a specific learner. Blocks if not available.""" + key = (step, fragment_id) + buffer = self._learner_to_syncer_buffers[learner_idx] + if key in buffer: + return buffer.pop(key) + + while True: + rec_step, rec_frag, data = self._learner_to_syncer_queues[learner_idx].get() + if rec_step == step and rec_frag == fragment_id: + return data + buffer[(rec_step, rec_frag)] = data + + def send_to_learner(self, learner_idx: int, step: int, fragment_id: int, data: Any): + """Syncer sends data to a specific learner.""" + self._syncer_to_learner_queues[learner_idx].put((step, fragment_id, data)) + + def recv_from_syncer(self, learner_idx: int, step: int, fragment_id: int) -> Any: + """Learner receives data from the syncer. Blocks if not available.""" + key = (step, fragment_id) + buffer = self._syncer_to_learner_buffers[learner_idx] + if key in buffer: + return buffer.pop(key) + + while True: + rec_step, rec_frag, data = self._syncer_to_learner_queues[learner_idx].get() + if rec_step == step and rec_frag == fragment_id: + return data + buffer[(rec_step, rec_frag)] = data + + +class LearnerTransport: + """Wrapper for learner threads to communicate with the syncer.""" + + def __init__( + self, + manager: ThreadedTransportManager, + learner_idx: int, + local_cpu_mesh: jax.sharding.Mesh, + ): + self.manager = manager + self.learner_idx = learner_idx + self.local_cpu_mesh = local_cpu_mesh + self._executor = ThreadPoolExecutor(max_workers=1) + + def send_to_syncer_async(self, step: int, fragment_id: int, data: Any): + """Asynchronously offloads TPU data to local CPU mesh and sends to syncer.""" + # 1. Asynchronously offload to CPU colocated mesh (non-blocking on main thread) + cpu_sharding = jax.tree_util.tree_map( + lambda s: jax.sharding.NamedSharding(self.local_cpu_mesh, s.spec), + jax.tree_util.tree_map(lambda x: x.sharding, data), + ) + frag_cpu = jax.device_put(data, cpu_sharding) + + # 2. Block and send in the background executor thread + def _send(): + try: + max_logging.log(f"Learner {self.learner_idx}: async send starting for step {step} frag {fragment_id}") + jax.block_until_ready(frag_cpu) + max_logging.log(f"Learner {self.learner_idx}: async send block_until_ready done") + self.manager.send_to_syncer(self.learner_idx, step, fragment_id, frag_cpu) + max_logging.log(f"Learner {self.learner_idx}: async send sent to syncer") + except Exception as e: + max_logging.error(f"Learner {self.learner_idx}: async send failed: {e}") + max_logging.error(traceback.format_exc()) + raise e + + self._executor.submit(_send) + + def send_to_syncer(self, step: int, fragment_id: int, data: Any): + """Synchronously offloads TPU data to local CPU mesh and sends to syncer.""" + cpu_sharding = jax.tree_util.tree_map( + lambda s: jax.sharding.NamedSharding(self.local_cpu_mesh, s.spec), + jax.tree_util.tree_map(lambda x: x.sharding, data), + ) + frag_cpu = jax.device_put(data, cpu_sharding) + jax.block_until_ready(frag_cpu) + self.manager.send_to_syncer(self.learner_idx, step, fragment_id, frag_cpu) + + def recv_from_syncer(self, step: int, fragment_id: int) -> Any: + return self.manager.recv_from_syncer(self.learner_idx, step, fragment_id) + + def close(self): + """Shutdown background thread executor.""" + self._executor.shutdown(wait=True) + + +class SyncerTransport: + """Wrapper for the syncer thread to communicate with learners.""" + + def __init__(self, manager: ThreadedTransportManager): + self.manager = manager + + def send_to_learner(self, learner_idx: int, step: int, fragment_id: int, data: Any): + self.manager.send_to_learner(learner_idx, step, fragment_id, data) + + def recv_from_learner(self, learner_idx: int, step: int, fragment_id: int) -> Any: + return self.manager.recv_from_learner(learner_idx, step, fragment_id) + + diff --git a/src/maxtext/trainers/diloco/diloco.py b/src/maxtext/trainers/diloco/diloco.py index ef650b872e..2f78f00bd6 100644 --- a/src/maxtext/trainers/diloco/diloco.py +++ b/src/maxtext/trainers/diloco/diloco.py @@ -28,13 +28,13 @@ import drjax from flax import nnx from flax import struct -from flax.training import train_state import jax import jax.numpy as jnp from jaxtyping import Array, Int32, Key, PyTree, UInt32 import optax from maxtext.configs import pyconfig +from maxtext.trainers.diloco.fragmenter import FragmentedTreeManipulator Batch = Any Params = PyTree @@ -229,6 +229,17 @@ def init_diloco_state() -> tuple[DiLoCoTrainState, PyTree]: return init_diloco_state() +def extract_replica_0(metrics): + 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 build_diloco_train_step( config: pyconfig.HyperParameters, train_step: Callable[[Any, Batch, PRNGKey], tuple[Any, Metrics]], @@ -296,31 +307,148 @@ def replace_nnx_model_params(s, new_params): 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) + # Split the RNG key across replicas. + keys = jax.random.split(prng, config.num_diloco_replicas) + inner_state, metrics = drjax.map_fn(train_step, (state.inner_state, batch, keys), mesh=mesh) + default_metrics = extract_replica_0(metrics) + # TODO Remove this all-reduce in the future to avoid cross-DCN comm completely + # TODO require to modify the metric_logger.py, as we need separate loggers for separate islands. + # not critical in streaming diloco as the communication happens almost everystep so + # this will not cause additional stall + if isinstance(metrics, dict) and "scalar" in metrics and "learning/loss" in metrics["scalar"]: + for i in range(config.num_diloco_replicas): + mask_i = jnp.arange(config.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 # 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] 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 - state, - ) - return state, avg_metrics + + if config.enable_streaming_diloco: + manipulator = FragmentedTreeManipulator.create(state.params, config) + + num_transformer_fragments = config.num_diloco_fragments + num_fragments = 1 + num_transformer_fragments + + steps_between_syncs_plus_1 = int(round(config.diloco_sync_period / num_fragments)) + steps_between_syncs_plus_1 = max(1, steps_between_syncs_plus_1) + period = num_fragments * steps_between_syncs_plus_1 + + def synchronize_fragment(state, idx): + # 1. Extract global and local parameters for the fragment + outer_params_frag = manipulator.get_flat_fragment(state.params, idx, has_replica_dim=False) + inner_model_params = ( + nnx.filter_state(state.inner_state.model, nnx.Param) if config.pure_nnx else state.inner_state.params + ) + inner_params_frag = manipulator.get_flat_fragment(inner_model_params, 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, 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, idx, new_outer_params_frag, has_replica_dim=False) + new_trace = manipulator.apply_flat_fragment( + state.outer_opt_state[0].trace, 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(state, idx): + # Get synced global params fragment + outer_params_frag = manipulator.get_flat_fragment(state.params, idx, has_replica_dim=False) + + # Broadcast to replicas + broadcast_outer_frag = drjax.broadcast(outer_params_frag, mesh=mesh) + + # Interpolation functions per replica + def replace_nnx_model_params_frag(s, outer_frag_replica): + full_params = nnx.filter_state(s.model, nnx.Param) + if config.communication_overlapping_alpha > 0.0: + inner_frag = manipulator.get_flat_fragment(full_params, idx, has_replica_dim=False) + alpha = config.communication_overlapping_alpha + 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, idx, merged_frag, has_replica_dim=False) + non_param_model = nnx.filter_state(s.model, nnx.Not(nnx.Param)) + new_model = nnx.merge_state(non_param_model, new_full_params) + + result = type(s)({}) + result["model"] = new_model + result["optimizer"] = s["optimizer"] + return result + + def replace_linen_model_params_frag(s, outer_frag_replica): + if config.communication_overlapping_alpha > 0.0: + inner_frag = manipulator.get_flat_fragment(s.params, idx, has_replica_dim=False) + alpha = config.communication_overlapping_alpha + 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_params = manipulator.apply_flat_fragment(s.params, idx, merged_frag, has_replica_dim=False) + return s.replace(params=new_params) + + # Apply to replica inner states + replace_fn = replace_nnx_model_params_frag if config.pure_nnx else replace_linen_model_params_frag + new_inner_state = drjax.map_fn(replace_fn, (state.inner_state, broadcast_outer_frag), mesh=mesh) + return state.replace(inner_state=new_inner_state) + + # Step 1: Run the synchronization logic if we hit a sync step + is_sync_step = jax.lax.bitwise_and(new_step > 0, new_step % steps_between_syncs_plus_1 == 0) + + def do_sync(s): + frag_idx = (new_step % period) // steps_between_syncs_plus_1 + return jax.lax.switch( + frag_idx, [lambda state_arg, idx=i: synchronize_fragment(state_arg, idx) 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) + V = config.num_communication_overlapping_steps + is_apply_step = jax.lax.bitwise_and(new_step - V > 0, (new_step - V) % steps_between_syncs_plus_1 == 0) + + def do_apply(s): + frag_idx = ((new_step - V) % period) // steps_between_syncs_plus_1 + return jax.lax.switch( + frag_idx, [lambda state_arg, idx=i: apply_fragment(state_arg, idx) for i in range(num_fragments)], s + ) + + state = jax.lax.cond(is_apply_step, do_apply, lambda s: s, state) + + else: + # 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 + state, + ) + return state, default_metrics return diloco_train_step diff --git a/src/maxtext/trainers/diloco/fragmenter.py b/src/maxtext/trainers/diloco/fragmenter.py new file mode 100644 index 0000000000..9e10d7a2c2 --- /dev/null +++ b/src/maxtext/trainers/diloco/fragmenter.py @@ -0,0 +1,116 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""FragmentedTreeManipulator for parameter tree slicing.""" + +import re +from typing import Any +import jax +import jax.numpy as jnp + + +class FragmentedTreeManipulator: + """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, + param_scan_axis: int = 0, + ): + self.keypath_to_is_scanned = keypath_to_is_scanned + self.fragment_to_layer_indices = fragment_to_layer_indices + self.num_fragments = num_fragments + self.param_scan_axis = param_scan_axis + + @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_transformer_fragments = config.num_diloco_fragments + + assert num_layers % num_transformer_fragments == 0, ( + f"num_decoder_layers ({num_layers}) must be divisible by " + f"num_diloco_fragments ({num_transformer_fragments}) for now." + ) + + num_synced = num_layers // num_transformer_fragments + use_sequential = config.use_sequential_layers + num_fragments = 1 + num_transformer_fragments + + # 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, _ in kvs: + parts = [] + for k in keypath: + parts.append(str(k.key) if hasattr(k, "key") else (str(k.idx) if hasattr(k, "idx") else str(k))) + serialized_path = "/" + "/".join(parts) + keypath_to_is_scanned[jax.tree_util.keystr(keypath)] = bool(scanned_regex.search(serialized_path)) + + return cls(keypath_to_is_scanned, fragment_to_layer_indices, num_fragments, config.param_scan_axis) + + 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] + axis = self.param_scan_axis + 1 if has_replica_dim else self.param_scan_axis + flat_frag[keystr] = jnp.take(v, indices, axis=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] + axis = self.param_scan_axis + 1 if has_replica_dim else self.param_scan_axis + idx_tuple = tuple(slice(None) if i != axis else indices for i in range(v.ndim)) + new_v = v.at[idx_tuple].set(flat_fragment[keystr]) + new_kvs.append(new_v) + else: + new_kvs.append(v) + return jax.tree_util.tree_unflatten(treedef, new_kvs) diff --git a/src/maxtext/trainers/diloco/threaded_diloco.py b/src/maxtext/trainers/diloco/threaded_diloco.py new file mode 100644 index 0000000000..c81cb08711 --- /dev/null +++ b/src/maxtext/trainers/diloco/threaded_diloco.py @@ -0,0 +1,732 @@ +# 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 +# +# 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. + +# pylint: disable=protected-access +"""Non-SPMD, multi-threaded streaming DiLoCo implementation with single client Pathways.""" + +import copy +import datetime +import functools +import threading +import traceback +from typing import Any +from concurrent.futures import ThreadPoolExecutor + +from flax import linen as nn, nnx, struct +from flax.linen import partitioning as nn_partitioning +import jax +import jax.numpy as jnp +from jax.experimental import colocated_python +import numpy as np +import optax + +from maxtext.common import checkpointing, profiler, metric_logger +from maxtext.common.goodput import maybe_record_goodput, GoodputEvent +from maxtext.trainers.diloco.decomposed_transport import ThreadedTransportManager, LearnerTransport, SyncerTransport +from maxtext.trainers.diloco.fragmenter import FragmentedTreeManipulator +from maxtext.utils import exceptions +from maxtext.utils import max_logging +from maxtext.utils import maxtext_utils +from maxtext.utils import model_creation_utils +from maxtext.utils import sharding +from maxtext.utils import train_utils +from maxtext.utils.mesh_utils import partition_mesh_by_diloco_axis, stack_across_meshes_pytree, unstack_across_meshes_pytree + + +@jax.jit +def mix_frags(i_frag, o_frag, alpha): + return jax.tree_util.tree_map(lambda x, y: alpha * x + (1 - alpha) * y, i_frag, o_frag) + + +def cpu_clone(x): + return jnp.copy(x) + + +class SyncerState(struct.PyTreeNode): + params: Any + opt_state: optax.OptState + step: int + + +def get_first_step(model, state): + if isinstance(model, nn.Module): + return int(state.step) + return int(state.optimizer.step.get_value()) + + +def make_learner_config(config, learner_idx, num_learners): + """Creates a modified deep copy of the global configuration for a specific learner.""" + learner_config = copy.deepcopy(config) + + # Remove 'diloco' from mesh_axes + mesh_axes = list(learner_config.mesh_axes) + if "diloco" in mesh_axes: + mesh_axes.remove("diloco") + learner_config._flat_config["mesh_axes"] = mesh_axes + + # Adjust logical_axis_rules to remove 'diloco' + new_logical_axis_rules = [] + for logical_axis, physical_axes in learner_config.logical_axis_rules: + if isinstance(physical_axes, str): + if physical_axes == "diloco": + continue + elif isinstance(physical_axes, (list, tuple)): + physical_axes = [ax for ax in physical_axes if ax != "diloco"] + new_logical_axis_rules.append((logical_axis, physical_axes)) + learner_config._flat_config["logical_axis_rules"] = new_logical_axis_rules + + # Enable local data loading for each learner + learner_config._flat_config["enable_local_data_loading"] = True + learner_config._flat_config["learner_idx"] = learner_idx + learner_config._flat_config["num_learners"] = num_learners + + # Disable SPMD diloco for learners + learner_config._flat_config["enable_streaming_diloco"] = False + learner_config._flat_config["enable_diloco"] = False + + return learner_config + + +def get_abstract_syncer_state(config, local_cpu_mesh): + """Computes abstract state shapes and types for the syncer's parameters and optimizer.""" + + if config.pure_nnx: + with nn_partitioning.axis_rules(config.logical_axis_rules): + _, abstract_model = model_creation_utils.create_nnx_abstract_model(config, local_cpu_mesh) + abstract_params = nnx.state(abstract_model, nnx.Param) + outer_optimizer = optax.sgd( + learning_rate=config.diloco_outer_lr, + momentum=config.diloco_outer_momentum, + nesterov=True, + ) + + @jax.jit + def init_opt(p): + return outer_optimizer.init(p) + + abstract_opt_state = init_opt.eval_shape(abstract_params) + else: + model = model_creation_utils.from_config(config, mesh=local_cpu_mesh) + abstract_vars = maxtext_utils.get_abstract_param(model, config) + abstract_params = abstract_vars["params"] + + params_logical_annotations = nn.get_partition_spec(abstract_params) + params_mesh_shardings = nn.logical_to_mesh_sharding( + params_logical_annotations, local_cpu_mesh, config.logical_axis_rules + ) + + @jax.jit + def dummy_init(): + return abstract_params + + abstract_params = jax.jit(dummy_init, out_shardings=params_mesh_shardings).eval_shape() + + outer_optimizer = optax.sgd( + learning_rate=config.diloco_outer_lr, + momentum=config.diloco_outer_momentum, + nesterov=True, + ) + + @jax.jit + def init_opt(p): + return outer_optimizer.init(p) + + opt_state_shardings = (optax.TraceState(trace=params_mesh_shardings), optax.EmptyState()) + abstract_opt_state = jax.jit(init_opt, out_shardings=opt_state_shardings).eval_shape(abstract_params) + + return abstract_params, abstract_opt_state + + +def _run_learner_loop( + learner_idx, config, submesh, local_cpu_mesh, transport, recorder, train_step, eval_step, init_lock +): + """Runs the main training and communication loop for a single learner replica.""" + max_logging.log(f"Learner {learner_idx}: Starting loop") + learner_config = make_learner_config(config, learner_idx, config.num_diloco_replicas) + learner_config._flat_config["run_name"] = config.run_name + f"_learner_{learner_idx}" + + with jax.set_mesh(submesh), submesh, nn_partitioning.axis_rules(learner_config.logical_axis_rules): + learner_config._flat_config["checkpoint_dir"] = config.checkpoint_dir + f"/learner_{learner_idx}" + + with init_lock: + ( + init_rng, + checkpoint_manager, + state_mesh_shardings, + model, + mesh, + learning_rate_schedule, + data_iterator, + data_loader, + rampup_manager, + eval_data_iterator, + state, + ) = train_utils.setup_train_loop(learner_config, recorder, mesh=submesh) + + params_shardings, state_mesh_shardings = sharding.maybe_update_params_sharding_with_opt( + learner_config, state_mesh_shardings + ) + flat_params_shardings, _ = jax.tree_util.tree_flatten_with_path(params_shardings) + flat_params_shardings = {jax.tree_util.keystr(p): leaf for p, leaf in flat_params_shardings} + + if isinstance(model, nn.Module): + jit_model = model + else: + jit_model, state = nnx.split(state) + + p_train_step, p_eval_step = train_utils.jit_train_and_eval_step( + learner_config, + jit_model, + mesh, + state, + state_mesh_shardings, + train_step, + eval_step, + eval_data_iterator, + params_shardings, + ) + + start_step = get_first_step(model, state) + + # Synchronized Initialization / Resume + try: + if start_step == 0: + if learner_idx == 0: + params = nnx.state(state.model, nnx.Param) if learner_config.pure_nnx else state.params + max_logging.log(f"Learner {learner_idx}: sending init params") + transport.send_to_syncer(step=0, fragment_id=-1, data=params) + max_logging.log(f"Learner {learner_idx}: waiting for init params") + initial_params = transport.recv_from_syncer(step=0, fragment_id=-1) + max_logging.log(f"Learner {learner_idx}: received init params") + else: + max_logging.log(f"Learner {learner_idx}: waiting for init params") + initial_params = transport.recv_from_syncer(step=0, fragment_id=-1) + max_logging.log(f"Learner {learner_idx}: received init params") + + tpu_param_sharding = jax.tree_util.tree_map( + lambda s: jax.sharding.NamedSharding(submesh, s.spec), params_shardings + ) + initial_params_tpu = jax.device_put(initial_params, tpu_param_sharding) + if learner_config.pure_nnx: + non_param_model = nnx.filter_state(state.model, nnx.Not(nnx.Param)) + new_model = nnx.merge_state(non_param_model, initial_params_tpu) + new_state = type(state)({}) + new_state["model"] = new_model + new_state["optimizer"] = state["optimizer"] + state = new_state + else: + state = state.replace(params=initial_params_tpu) + else: + global_params = transport.recv_from_syncer(step=start_step, fragment_id=-1) + tpu_param_sharding = jax.tree_util.tree_map( + lambda s: jax.sharding.NamedSharding(submesh, s.spec), params_shardings + ) + global_params_tpu = jax.device_put(global_params, tpu_param_sharding) + if learner_config.pure_nnx: + non_param_model = nnx.filter_state(state.model, nnx.Not(nnx.Param)) + new_model = nnx.merge_state(non_param_model, global_params_tpu) + new_state = type(state)({}) + new_state["model"] = new_model + new_state["optimizer"] = state["optimizer"] + state = new_state + else: + state = state.replace(params=global_params_tpu) + except Exception as e: + max_logging.error(f"Learner {learner_idx} crashed in init: {e}") + max_logging.error(traceback.format_exc()) + raise e + + params_template = nnx.state(state.model, nnx.Param) if learner_config.pure_nnx else state.params + manipulator = FragmentedTreeManipulator.create(params_template, learner_config) + num_fragments = manipulator.num_fragments + + tau = learner_config.num_communication_overlapping_steps + alpha = learner_config.communication_overlapping_alpha + + steps_between_syncs_plus_1 = int(round(learner_config.diloco_sync_period / num_fragments)) + steps_between_syncs_plus_1 = max(1, steps_between_syncs_plus_1) + period = num_fragments * steps_between_syncs_plus_1 + + prof = profiler.Profiler(learner_config, offset_step=start_step) + metric_logger_instance = metric_logger.MetricLogger( + config=learner_config, learning_rate_schedule=learning_rate_schedule + ) + metric_logger_instance.write_setup_info_to_tensorboard(params_template) + + # Pre-compile the mix function for each fragment to avoid concurrent compilation crashes + with init_lock: + for f_idx in range(num_fragments): + dummy_frag = manipulator.get_flat_fragment(params_template, f_idx) + _ = mix_frags(dummy_frag, dummy_frag, alpha) + + try: + last_step_completion = datetime.datetime.now() + for step in range(start_step, learner_config.steps): + max_logging.log(f"Learner {learner_idx}: Step {step} starting") + prof.maybe_activate_profiler(step, state) + + with jax.profiler.StepTraceAnnotation("train", step_num=step): + example_batch = data_loader.load_next_batch(rampup_manager=rampup_manager) + if isinstance(model, nn.Module): + step_rng_args = (jax.jit(jax.random.fold_in)(init_rng, step),) + else: + step_rng_args = () + + with maybe_record_goodput(recorder, GoodputEvent.STEP, step): + with jax.set_mesh(mesh), nn_partitioning.axis_rules(learner_config.logical_axis_rules): + if learner_config.shard_optimizer_over_data and isinstance(model, nn.Module): + state = sharding.maybe_shard_with_name(state, state_mesh_shardings, learner_config.shard_mode) + # DEBUG + state_leaves, _ = jax.tree_util.tree_flatten(state) + batch_leaves, _ = jax.tree_util.tree_flatten(example_batch) + print( + f"DEBUG_DEVICES: Learner {learner_idx} step {step} before train_step state leaves devices: {[set(x.devices()) for x in state_leaves]}", + flush=True, + ) + print( + f"DEBUG_DEVICES: Learner {learner_idx} step {step} before train_step batch leaves devices: {[set(x.devices()) for x in batch_leaves]}", + flush=True, + ) + state, metrics = p_train_step(state, example_batch, *step_rng_args) + # Force block to catch async errors immediately + for leaf in jax.tree_util.tree_flatten((state, metrics))[0]: + if hasattr(leaf, "block_until_ready"): + leaf.block_until_ready() + + max_logging.log(f"Learner {learner_idx}: Step {step} finished") + step_time_delta = datetime.datetime.now() - last_step_completion + + completed_step = step + 1 + + if completed_step > 0 and completed_step % steps_between_syncs_plus_1 == 0: + frag_idx = (completed_step % period) // steps_between_syncs_plus_1 + params = nnx.state(state.model, nnx.Param) if learner_config.pure_nnx else state.params + frag_data = manipulator.get_flat_fragment(params, frag_idx) + transport.send_to_syncer_async(completed_step, frag_idx, frag_data) + + if completed_step - tau > 0 and (completed_step - tau) % steps_between_syncs_plus_1 == 0: + frag_idx = ((completed_step - tau) % period) // steps_between_syncs_plus_1 + received_frag = transport.recv_from_syncer(completed_step - tau, frag_idx) + + tpu_frag_sharding = { + k: jax.sharding.NamedSharding(submesh, flat_params_shardings[k].spec) for k in received_frag.keys() + } + received_frag_tpu = jax.device_put(received_frag, tpu_frag_sharding) + + params = nnx.state(state.model, nnx.Param) if learner_config.pure_nnx else state.params + inner_frag = manipulator.get_flat_fragment(params, frag_idx) + + # DEBUG + print( + f"DEBUG_DEVICES: Learner {learner_idx} step {step} mix_frags check start. Keys in fragment: {list(inner_frag.keys())}", + flush=True, + ) + for k in inner_frag.keys(): + l_leaf = inner_frag[k] + s_leaf = received_frag_tpu[k] + l_devs = set(l_leaf.devices()) + s_devs = set(s_leaf.devices()) + print(f"DEBUG_DEVICES: key {k}: shape={l_leaf.shape}") + print(f"DEBUG_DEVICES: inner_frag sharding={l_leaf.sharding}, devices={l_devs}") + print(f"DEBUG_DEVICES: received_frag_tpu sharding={s_leaf.sharding}, devices={s_devs}") + if l_devs != s_devs: + print(f"DEBUG_DEVICES: MISMATCH for key {k}:", flush=True) + print(f"DEBUG_DEVICES: inner_frag devices: {l_devs}", flush=True) + print(f"DEBUG_DEVICES: received_frag_tpu devices: {s_devs}", flush=True) + print(f"DEBUG_DEVICES: received_frag devices: {set(received_frag[k].devices())}", flush=True) + print(f"DEBUG_DEVICES: Learner {learner_idx} step {step} mix_frags check end", flush=True) + + mixed_frag = mix_frags(inner_frag, received_frag_tpu, alpha) + new_params = manipulator.apply_flat_fragment(params, frag_idx, mixed_frag) + # DEBUG + new_params_leaves, _ = jax.tree_util.tree_flatten(new_params) + print( + f"DEBUG_DEVICES: Learner {learner_idx} step {step} new_params leaves devices: {[set(x.devices()) for x in new_params_leaves]}", + flush=True, + ) + + if learner_config.pure_nnx: + non_param_model = nnx.filter_state(state.model, nnx.Not(nnx.Param)) + new_model = nnx.merge_state(non_param_model, new_params) + new_state = type(state)({}) + new_state["model"] = new_model + new_state["optimizer"] = state["optimizer"] + state = new_state + else: + state = state.replace(params=new_params) + + checkpointing.maybe_save_checkpoint(checkpoint_manager, state, learner_config, data_iterator, step) + + metric_logger_instance.buffer_and_write_metrics(metrics, step, step_time_delta) + + eval_step_count = None + if learner_config.eval_interval > 0 and step > start_step and (step + 1) % learner_config.eval_interval == 0: + assert eval_data_iterator + eval_data_iterator.reset() + metric_logger_instance.reset_eval_metrics() + max_logging.log(f"Learner {learner_idx}: Starting eval after train step {step}") + + eval_step_count = 0 + last_eval_step_completion = datetime.datetime.now() + for eval_batch in eval_data_iterator: + eval_batch = jax.device_put(eval_batch, sharding.get_input_data_sharding(learner_config, mesh)) + if learner_config.eval_steps > 0 and eval_step_count >= learner_config.eval_steps: + break + with jax.set_mesh(mesh), nn_partitioning.axis_rules(learner_config.logical_axis_rules): + eval_metrics = p_eval_step(state, eval_batch, *step_rng_args) + eval_step_time_delta = datetime.datetime.now() - last_eval_step_completion + last_eval_step_completion = datetime.datetime.now() + metric_logger_instance.buffer_and_write_metrics( + eval_metrics, eval_step_count, step_time_delta=eval_step_time_delta, is_training=False + ) + eval_step_count += 1 + + prof.maybe_deactivate_profiler(step, state) + last_step_completion = datetime.datetime.now() + + if learner_config.save_checkpoint_on_completion: + checkpointing.maybe_save_checkpoint(checkpoint_manager, state, learner_config, data_iterator) + if checkpoint_manager is not None: + checkpoint_manager.wait_until_finished() + + except exceptions.StopTraining as e: + prof.deactivate() + max_logging.log(f"Learner {learner_idx} training stopped: {str(e)}") + finally: + metric_logger_instance.flush_metrics_and_cleanup() + transport.close() + + +def learner_loop(learner_idx, config, submesh, local_cpu_mesh, transport, recorder, train_step, eval_step, init_lock): + """Wrapper to run the learner loop and handle/log top-level exceptions.""" + try: + _run_learner_loop(learner_idx, config, submesh, local_cpu_mesh, transport, recorder, train_step, eval_step, init_lock) + except Exception as e: + max_logging.error(f"Learner {learner_idx} crashed: {e}") + max_logging.error(traceback.format_exc()) + raise e + + +def syncer_loop( + config, global_cpu_mesh, cpu_submeshes, transport, recorder, abstract_params=None, abstract_opt_state=None +): + """Wrapper to run the syncer loop and handle/log top-level exceptions.""" + try: + _run_syncer_loop(config, global_cpu_mesh, cpu_submeshes, transport, recorder, abstract_params, abstract_opt_state) + except Exception as e: + max_logging.error(f"Syncer crashed: {e}") + max_logging.error(traceback.format_exc()) + raise e + + +def make_step_fns(global_cpu_mesh, flat_params_shardings, frag_keys, outer_optimizer, num_learners): + """Creates JIT-compiled functions for computing gradients and applying outer steps.""" + global_sharding_tree = { + k: jax.sharding.NamedSharding(global_cpu_mesh, flat_params_shardings[k].spec) for k in frag_keys + } + stacked_sharding_tree = { + k: jax.sharding.NamedSharding(global_cpu_mesh, jax.sharding.PartitionSpec("diloco", *flat_params_shardings[k].spec)) + for k in frag_keys + } + opt_state_sharding_tree = (optax.TraceState(trace=global_sharding_tree), optax.EmptyState()) + + @functools.partial( + jax.jit, + in_shardings=(global_sharding_tree, stacked_sharding_tree), + out_shardings=global_sharding_tree, + ) + def compute_grad(o_frag, stacked_i_frag): + averaged_i_frag = jax.tree_util.tree_map(lambda x: jnp.mean(x, axis=0), stacked_i_frag) + return jax.tree_util.tree_map(lambda x, y: x - y, o_frag, averaged_i_frag) + + @functools.partial( + jax.jit, + in_shardings=(global_sharding_tree, opt_state_sharding_tree, global_sharding_tree), + out_shardings=(global_sharding_tree, opt_state_sharding_tree), + ) + def apply_outer_step(g_frag, o_state_frag, p_frag): + updates_frag, new_o_state_frag = outer_optimizer.update(g_frag, o_state_frag, params=p_frag) + new_p_frag = optax.apply_updates(p_frag, updates_frag) + return new_p_frag, new_o_state_frag + + @functools.partial( + jax.jit, + in_shardings=(global_sharding_tree,), + out_shardings=stacked_sharding_tree, + ) + def broadcast_frag(frag): + return jax.tree_util.tree_map( + lambda x: jnp.broadcast_to(jnp.expand_dims(x, axis=0), (num_learners,) + x.shape), + frag + ) + + return compute_grad, apply_outer_step, broadcast_frag + + +def _run_syncer_loop( + config, global_cpu_mesh, cpu_submeshes, transport, recorder, abstract_params=None, abstract_opt_state=None +): + """Runs the main syncer loop that coordinates parameter averaging and outer optimization.""" + max_logging.log("Syncer: Starting loop") + + num_learners = config.num_diloco_replicas + + if abstract_params is None or abstract_opt_state is None: + abstract_params, abstract_opt_state = get_abstract_syncer_state(config, global_cpu_mesh) + abstract_step = jax.ShapeDtypeStruct( + (), jnp.int32, sharding=jax.sharding.NamedSharding(global_cpu_mesh, jax.sharding.PartitionSpec()) + ) + abstract_syncer_state = SyncerState(params=abstract_params, opt_state=abstract_opt_state, step=abstract_step) + + # Init(1): Loading checkpoints + logger = checkpointing.setup_checkpoint_logger(config) + checkpoint_manager = checkpointing.create_orbax_checkpoint_manager( + config.checkpoint_dir, + config.enable_checkpointing, + config.async_checkpointing, + config.checkpoint_period, + config.dataset_type, + logger, + config.checkpoint_storage_use_ocdbt, + config.checkpoint_storage_use_zarr3, + config.enable_continuous_checkpointing, + config.max_num_checkpoints_to_keep, + config.checkpoint_storage_concurrent_gb, + config.enable_single_controller, + config.colocated_python_checkpointing, + config.enable_single_replica_ckpt_restoring, + config.enable_autocheckpoint, + config.checkpoint_todelete_subdir, + config.checkpoint_todelete_full_path, + ) + + restored_state, _ = checkpointing.load_state_if_possible( + checkpoint_manager=checkpoint_manager, + data_iterator=None, + load_parameters_from_path="", + load_full_state_from_path="", + checkpoint_storage_concurrent_gb=config.checkpoint_storage_concurrent_gb, + abstract_unboxed_pre_state=abstract_syncer_state, + enable_single_replica_ckpt_restoring=config.enable_single_replica_ckpt_restoring, + dataset_type=config.dataset_type, + use_ocdbt=config.checkpoint_storage_use_ocdbt, + use_zarr3=config.checkpoint_storage_use_zarr3, + ) + + # Get abstract shardings for params and opt_state + params_shardings = jax.tree_util.tree_map(lambda x: x.sharding, abstract_params) + flat_params_shardings = { + jax.tree_util.keystr(k): v for k, v in jax.tree_util.tree_flatten_with_path(params_shardings)[0] + } + + if restored_state is None: # (1,a) No checkpoint found, start from scratch + max_logging.log("Syncer: waiting for init params from Learner 0") + initial_params_l0 = transport.recv_from_learner(learner_idx=0, step=0, fragment_id=-1) + max_logging.log("Syncer: received init params from Learner 0") + with jax.set_mesh(global_cpu_mesh): + global_params = jax.device_put(initial_params_l0, params_shardings) + outer_optimizer = optax.sgd( + learning_rate=config.diloco_outer_lr, + momentum=config.diloco_outer_momentum, + nesterov=True, + ) + outer_opt_state = outer_optimizer.init(global_params) + + syncer_state = SyncerState(params=global_params, opt_state=outer_opt_state, step=0) + start_step = 0 + + else: # loading checkpoints successfully + syncer_state = restored_state["items"] + start_step = int(syncer_state.step) + max_logging.log(f"Syncer restored from step {start_step}") + + # Init (2): broadcast initial params to all learners + stacked_params_sharding_tree = jax.tree_util.tree_map( + lambda s: jax.sharding.NamedSharding(global_cpu_mesh, jax.sharding.PartitionSpec("diloco", *s.spec)), + params_shardings + ) + + @functools.partial( + jax.jit, + in_shardings=(params_shardings,), + out_shardings=stacked_params_sharding_tree, + ) + def broadcast_params(params): + return jax.tree_util.tree_map( + lambda x: jnp.broadcast_to(jnp.expand_dims(x, axis=0), (num_learners,) + x.shape), + params + ) + + with jax.set_mesh(global_cpu_mesh): + stacked_params = broadcast_params(syncer_state.params) + + local_params_list = unstack_across_meshes_pytree(stacked_params, cpu_submeshes, "diloco") + + for i in range(num_learners): + max_logging.log(f"Syncer: sending params to Learner {i} at step {start_step}") + transport.send_to_learner(learner_idx=i, step=start_step, fragment_id=-1, data=local_params_list[i]) + max_logging.log(f"Syncer: sent params to Learner {i} at step {start_step}") + + manipulator = FragmentedTreeManipulator.create(syncer_state.params, config) + num_fragments = manipulator.num_fragments + + steps_between_syncs_plus_1 = int(round(config.diloco_sync_period / num_fragments)) + steps_between_syncs_plus_1 = max(1, steps_between_syncs_plus_1) + period = num_fragments * steps_between_syncs_plus_1 + + outer_optimizer = optax.sgd( + learning_rate=config.diloco_outer_lr, + momentum=config.diloco_outer_momentum, + nesterov=True, + ) + + # steps that syncing is happening + sync_steps = [step for step in range(start_step + 1, config.steps + 1) if step % steps_between_syncs_plus_1 == 0] + + # Pre-build JIT step functions for each fragment index to avoid late-binding cell-var-from-loop warning + step_fns_by_frag = {} + with jax.set_mesh(global_cpu_mesh): + for f_idx in range(num_fragments): + frag_dict = manipulator.get_flat_fragment(syncer_state.params, f_idx) + step_fns_by_frag[f_idx] = make_step_fns( + global_cpu_mesh, flat_params_shardings, frag_dict, outer_optimizer, num_learners + ) + + # Start main syncer loop + for step in sync_steps: # e.g. 50, 100, 150... if sync_period=50 + max_logging.log(f"Syncer: Step {step} sync starting") + frag_idx = (step % period) // steps_between_syncs_plus_1 + + learner_frags = [] + + # receive the fragment of the current step from each learner. + for i in range(num_learners): + frag_i = transport.recv_from_learner(learner_idx=i, step=step, fragment_id=frag_idx) + learner_frags.append(frag_i) + max_logging.log(f"Syncer: received all fragments for step {step}") + + stacked_inner_frag = stack_across_meshes_pytree(learner_frags, global_cpu_mesh, "diloco") + max_logging.log(f"Syncer: Step {step} stacking done") + + with jax.set_mesh(global_cpu_mesh): + outer_params_frag = manipulator.get_flat_fragment(syncer_state.params, frag_idx) + + compute_grad, apply_outer_step, broadcast_frag = step_fns_by_frag[frag_idx] + + pseudo_grad_frag = compute_grad(outer_params_frag, stacked_inner_frag) + + trace_frag = manipulator.get_flat_fragment(syncer_state.opt_state[0].trace, frag_idx) + opt_state_frag = (optax.TraceState(trace=trace_frag), optax.EmptyState()) + + new_outer_params_frag, new_opt_state_frag = apply_outer_step(pseudo_grad_frag, opt_state_frag, outer_params_frag) + + new_outer_params_frag_for_broadcast = jax.tree_util.tree_map(cpu_clone, new_outer_params_frag) + + # Debug prints + first_key = list(new_outer_params_frag_for_broadcast.keys())[0] + max_logging.log(f"DEBUG_SHARD: first_key={first_key}") + val_global = new_outer_params_frag_for_broadcast[first_key] + max_logging.log(f"DEBUG_SHARD: global val sharding={val_global.sharding}") + for i, shard in enumerate(val_global.addressable_shards): + max_logging.log(f"DEBUG_SHARD: global shard {i}: device={shard.device}") + + # Broadcast the fragment to all learners + stacked_new_outer_params_frag = broadcast_frag(new_outer_params_frag_for_broadcast) + + new_params = manipulator.apply_flat_fragment(syncer_state.params, frag_idx, new_outer_params_frag) + new_trace = manipulator.apply_flat_fragment(syncer_state.opt_state[0].trace, frag_idx, new_opt_state_frag[0].trace) + new_opt_state = (optax.TraceState(trace=new_trace), syncer_state.opt_state[1]) + + syncer_state = syncer_state.replace(params=new_params, opt_state=new_opt_state, step=step) + max_logging.log(f"Syncer: Step {step} outer step applied") + + local_frags = unstack_across_meshes_pytree(stacked_new_outer_params_frag, cpu_submeshes, "diloco") + + for i in range(num_learners): + transport.send_to_learner(learner_idx=i, step=step, fragment_id=frag_idx, data=local_frags[i]) + + checkpointing.maybe_save_checkpoint( + checkpoint_manager=checkpoint_manager, state=syncer_state, config=config, data_iterator=None, step=step + ) + max_logging.log(f"Syncer: Step {step} sync finished") + + if checkpoint_manager is not None: + checkpoint_manager.wait_until_finished() + + +def run_threaded_diloco(config, recorder, train_step, eval_step): + """Orchestrator for multi-threaded DiLoCo.""" + max_logging.log("Starting run_threaded_diloco") + num_learners = config.num_diloco_replicas + + max_logging.log("Creating global mesh") + global_mesh = maxtext_utils.get_mesh_from_config(config) + max_logging.log("Partitioning global mesh") + tpu_submeshes = partition_mesh_by_diloco_axis(global_mesh, num_learners) + cpu_submeshes = [colocated_python.colocated_cpu_devices(submesh) for submesh in tpu_submeshes] + global_cpu_mesh = colocated_python.colocated_cpu_devices(global_mesh) + + transport_manager = ThreadedTransportManager(num_learners) + + # Get abstract syncer state first, on main thread, before spawning learner threads. + max_logging.log("Getting abstract syncer state") + abstract_params, abstract_opt_state = get_abstract_syncer_state(config, global_cpu_mesh) + max_logging.log("Got abstract syncer state") + + init_lock = threading.Lock() + + max_logging.log("Spawning learner threads") + with ThreadPoolExecutor(max_workers=num_learners) as executor: + futures = [] + for i in range(num_learners): + # Determine if this learner should run on this process + learner_devices = tpu_submeshes[i].devices.flat + should_run = any(d in jax.local_devices() for d in learner_devices) + + if should_run: + learner_transport = LearnerTransport(transport_manager, i, cpu_submeshes[i]) + + futures.append( + executor.submit( + learner_loop, + i, + config, + tpu_submeshes[i], # each learner will only see its local TPU submesh + cpu_submeshes[i], + learner_transport, + recorder, + train_step, + eval_step, + init_lock=init_lock, + ) + ) + else: + max_logging.log(f"Learner {i} is remote, not spawning thread") + + syncer_transport = SyncerTransport(transport_manager) + max_logging.log("Starting syncer loop") + syncer_loop( + config, + global_cpu_mesh, + cpu_submeshes, + syncer_transport, + recorder, + abstract_params=abstract_params, + abstract_opt_state=abstract_opt_state, + ) + + max_logging.log("Waiting for learner threads to finish") + for f in futures: + f.result() + max_logging.log("Finished run_threaded_diloco") diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py index 74c41e8e90..48c18c0093 100644 --- a/src/maxtext/trainers/pre_train/train.py +++ b/src/maxtext/trainers/pre_train/train.py @@ -39,6 +39,7 @@ from flax import linen as nn, nnx from flax.linen import partitioning as nn_partitioning from flax.nnx import variablelib +from maxtext.trainers.diloco.threaded_diloco import run_threaded_diloco from maxtext.configs import pyconfig from maxtext.utils.globals import EPS @@ -792,7 +793,10 @@ def initialize(argv: Sequence[str]) -> tuple[pyconfig.HyperParameters, Any]: def run(config, recorder): """Run the job given hyperparameters and utilities.""" with (max_utils.maybe_get_transformer_engine_context(config),): - train_loop(config, recorder) + if config.enable_non_spmd_diloco: + run_threaded_diloco(config, recorder, train_step, eval_step) + else: + train_loop(config, recorder) def get_train_func(config, recorder, argv): diff --git a/src/maxtext/utils/globals.py b/src/maxtext/utils/globals.py index 4ad0f1d9c4..88541a6251 100644 --- a/src/maxtext/utils/globals.py +++ b/src/maxtext/utils/globals.py @@ -24,7 +24,7 @@ MAXTEXT_REPO_ROOT = os.environ.get( "MAXTEXT_REPO_ROOT", r - if os.path.isdir( + if os.path.exists( os.path.join(r := os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))), ".git") ) else MAXTEXT_PKG_DIR, diff --git a/src/maxtext/utils/mesh_utils.py b/src/maxtext/utils/mesh_utils.py new file mode 100644 index 0000000000..84890944d5 --- /dev/null +++ b/src/maxtext/utils/mesh_utils.py @@ -0,0 +1,111 @@ +# 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 +# +# 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. + +"""Mesh utilities for DiLoCo stack/unstack operations across submeshes.""" + +from typing import Any +import jax +import jax.numpy as jnp +import numpy as np + +from pathwaysutils.experimental.concatenate_by_mesh_axis import concatenate_by_mesh_axis +from pathwaysutils.experimental.split_by_mesh_axis import split_by_mesh_axis + + +def partition_mesh_by_diloco_axis( + global_mesh: jax.sharding.Mesh, num_replicas: int, diloco_axis_name: str = "diloco" +) -> list[jax.sharding.Mesh]: + """Slices a global mesh along the diloco axis into multiple submeshes. Won't keep diloco dim.""" + if diloco_axis_name not in global_mesh.axis_names: + raise ValueError(f"Axis {diloco_axis_name} not found in mesh axis names: {global_mesh.axis_names}") + + diloco_axis_index = global_mesh.axis_names.index(diloco_axis_name) + diloco_axis_size = global_mesh.shape[diloco_axis_name] + + if diloco_axis_size != num_replicas: + raise ValueError(f"Diloco axis size ({diloco_axis_size}) must match num_replicas ({num_replicas})") + + devices = global_mesh.devices + submeshes = [] + axis_names = list(global_mesh.axis_names) + axis_names.remove(diloco_axis_name) + + for i in range(num_replicas): + sub_devices = np.take(devices, i, axis=diloco_axis_index) + submesh = jax.sharding.Mesh(sub_devices, axis_names) + submeshes.append(submesh) + + return submeshes + + +def _expand_array_dims_with_mesh( + x: jax.Array, + axis_name: str, +) -> jax.Array: + """Expands array dimensions by introducing a new dim-1 at index 0 and expanding its mesh.""" + # TODO: verify the correctness of this function. compare to the g3 implementation. + # 1. Expand the array shape natively to prepending a size-1 dimension + expanded_x = jnp.expand_dims(x, axis=0) + + # 2. Get original sharding + sharding = x.sharding + assert isinstance(sharding, jax.sharding.NamedSharding) + submesh = sharding.mesh + + # 3. Build expanded mesh (with axis_name of size 1 at the beginning) + expanded_devices = np.expand_dims(np.array(submesh.devices), axis=0) + expanded_mesh = jax.sharding.Mesh(expanded_devices, axis_names=(axis_name,) + submesh.axis_names) + + # 4. Build expanded NamedSharding and reshard (metadata-only operation in JAX) + expanded_sharding = jax.sharding.NamedSharding( + expanded_mesh, jax.sharding.PartitionSpec(axis_name, *sharding.spec), memory_kind=sharding.memory_kind + ) + + return jax.device_put(expanded_x, expanded_sharding) + + +def stack_across_meshes_pytree(trees: list[Any], global_mesh: jax.sharding.Mesh, axis_name: str) -> Any: + """Stacks a list of PyTrees across submeshes into a single global PyTree.""" + # 1. Expand dimensions of all arrays in all PyTrees manually + expanded_trees = [] + for tree in trees: + exp_tree = jax.tree_util.tree_map(lambda x: _expand_array_dims_with_mesh(x, axis_name), tree) + expanded_trees.append(exp_tree) + + # 2. Concatenate along the mesh axis using pathwaysutils + return concatenate_by_mesh_axis(expanded_trees, mesh_axis=axis_name) + + +def unstack_across_meshes_pytree( + global_tree: Any, + submeshes: list[jax.sharding.Mesh], + axis_name: str, +) -> list[Any]: + """Unstacks/splits a global PyTree into a list of submesh-local PyTrees.""" + num_replicas = len(submeshes) + # 1. Split the global PyTree along the mesh axis using pathwaysutils + split_trees = split_by_mesh_axis(global_tree, mesh_axis=axis_name) + + # 2. Squeeze out the leading size-1 axis and reshard each split tree to its submesh + def squeeze_and_reshard_tree(tree, submesh): + def squeeze_leaf(x): + # Squeeze leading size-1 dimension + squeezed = x[0] + # Reshard to target submesh + submesh_sharding = jax.sharding.NamedSharding(submesh, squeezed.sharding.spec) + return jax.device_put(squeezed, submesh_sharding) + + return jax.tree_util.tree_map(squeeze_leaf, tree) + + return [squeeze_and_reshard_tree(split_trees[i], submeshes[i]) for i in range(num_replicas)] diff --git a/src/maxtext/utils/train_utils.py b/src/maxtext/utils/train_utils.py index bdbc24136a..b03e5fc094 100644 --- a/src/maxtext/utils/train_utils.py +++ b/src/maxtext/utils/train_utils.py @@ -189,7 +189,7 @@ def jit_train_and_eval_step( return p_train_step, p_eval_step -def setup_train_loop(config, recorder, devices=None): +def setup_train_loop(config, recorder, devices=None, mesh=None): """Set up prerequisites for the training loop - checkpoint_manager, PRNG keys, Mesh, Model and optimizer. @@ -215,7 +215,8 @@ def setup_train_loop(config, recorder, devices=None): with maybe_record_goodput(recorder, GoodputEvent.TPU_INIT): is_training = True init_rng = jax.random.PRNGKey(config.init_weights_seed) - mesh = maxtext_utils.get_mesh_from_config(config, devices) + if mesh is None: + mesh = maxtext_utils.get_mesh_from_config(config, devices) context_parallel_size = mesh.shape.get(config.context_sharding, 1) if config.pure_nnx: # Create abstract NNX model. @@ -241,7 +242,9 @@ def create_train_state_fn(): validate_completed_steps(checkpoint_step + 1, config.steps) with maybe_record_goodput(recorder, GoodputEvent.TRAINING_PREPARATION): - data_iterator, eval_data_iterator = create_data_iterator(config, mesh) + learner_idx = getattr(config, "learner_idx", 0) + num_learners = getattr(config, "num_learners", 1) + data_iterator, eval_data_iterator = create_data_iterator(config, mesh, learner_idx, num_learners) rampup_manager = create_rampup_manager(config, checkpoint_manager) # Validate context parallelism with packing configuration context_parallel_strategy = config.context_parallel_strategy.lower() diff --git a/tests/integration/diloco_test.py b/tests/integration/diloco_test.py index 68633ee436..616b862f6d 100644 --- a/tests/integration/diloco_test.py +++ b/tests/integration/diloco_test.py @@ -31,6 +31,7 @@ from maxtext.configs.pyconfig import initialize_pydantic from maxtext.trainers.pre_train.train_compile import main as train_compile_main +from maxtext.trainers.pre_train.train import main as train_main from maxtext.trainers.diloco import diloco from tests.utils.test_helpers import get_test_config_path @@ -191,7 +192,7 @@ def loss_fn(params, batch): # = 0.81 diloco_test_state, loss = diloco_train_step(diloco_test_state, (inputs, labels), jax.random.key(seed=42)) chex.assert_equal(diloco_test_state.step, 2.0) - chex.assert_trees_all_close(loss, 0.65) + chex.assert_trees_all_close(loss, 0.49) # Assert no updates to the global model yet (no synchronization) chex.assert_trees_all_equal(diloco_test_state.params, initial_test_state.params) @@ -225,7 +226,7 @@ def loss_fn(params, batch): # based outer optimizer. diloco_test_state, loss = diloco_train_step(diloco_test_state, (inputs, labels), jax.random.key(seed=42)) chex.assert_equal(diloco_test_state.step, 3.0) - chex.assert_trees_all_close(loss, 0.4481) + chex.assert_trees_all_close(loss, 0.2401) # Assert that inner and outer parameters are all equal now that # synchronization has happened. chex.assert_trees_all_equal( @@ -264,7 +265,7 @@ def loss_fn(params, batch): step_three_outer_params = diloco_test_state.params diloco_test_state, loss = diloco_train_step(diloco_test_state, (inputs, labels), jax.random.key(seed=42)) chex.assert_equal(diloco_test_state.step, 4.0) - chex.assert_trees_all_close(loss, 0.574244) + chex.assert_trees_all_close(loss, 0.20754369) # Assert no updates to the global model since previous step (no # synchronization). chex.assert_trees_all_equal(diloco_test_state.params, step_three_outer_params) @@ -320,3 +321,75 @@ 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", + ) + ) + + @pytest.mark.cpu_only + def test_threaded_diloco_minimal_run(self): + """Runs a minimal training run with threaded DiLoCo on CPU.""" + devices = jax.devices() + num_replicas = 2 + if len(devices) < num_replicas: + self.skipTest(f"Test requires {num_replicas} devices, but only {len(devices)} are available.") + + temp_dir = gettempdir() + base_output_directory = os.path.join(temp_dir, "test_threaded_diloco") + run_name = "test_threaded_diloco_run" + + argv = [ + "", + get_test_config_path(), + f"base_output_directory={base_output_directory}", + f"run_name={run_name}", + "steps=4", + "dataset_type=synthetic", + "enable_checkpointing=False", + "enable_goodput_recording=False", + "enable_non_spmd_diloco=True", + "pure_nnx=True", + "num_diloco_fragments=2", + "num_diloco_replicas=2", + "diloco_sync_period=2", + "num_communication_overlapping_steps=1", + "communication_overlapping_alpha=0.5", + "ici_diloco_parallelism=2", + "ici_fsdp_parallelism=2", + "ici_tensor_parallelism=1", + "ici_pipeline_parallelism=1", + "base_emb_dim=16", + "base_num_query_heads=1", + "base_num_kv_heads=1", + "base_mlp_dim=16", + "base_num_decoder_layers=2", + "head_dim=4", + "max_target_length=16", + "vocab_size=32", + ] + + train_main(argv) diff --git a/tests/unit/decomposed_transport_test.py b/tests/unit/decomposed_transport_test.py new file mode 100644 index 0000000000..1a57856cc0 --- /dev/null +++ b/tests/unit/decomposed_transport_test.py @@ -0,0 +1,79 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for decomposed_transport.py.""" + +import threading +import time +from maxtext.trainers.diloco.decomposed_transport import ThreadedTransportManager, LearnerTransport, SyncerTransport + + +def test_threaded_transport_basic(): + manager = ThreadedTransportManager(num_learners=2) + learner0 = LearnerTransport(manager, 0) + learner1 = LearnerTransport(manager, 1) + syncer = SyncerTransport(manager) + + # Test Learner -> Syncer + learner0.send_to_syncer(step=0, fragment_id=0, data="L0_S0_F0") + learner1.send_to_syncer(step=0, fragment_id=0, data="L1_S0_F0") + + assert syncer.recv_from_learner(0, step=0, fragment_id=0) == "L0_S0_F0" + assert syncer.recv_from_learner(1, step=0, fragment_id=0) == "L1_S0_F0" + + # Test Syncer -> Learner + syncer.send_to_learner(0, step=0, fragment_id=0, data="S_L0_S0_F0") + syncer.send_to_learner(1, step=0, fragment_id=0, data="S_L1_S0_F0") + + assert learner0.recv_from_syncer(step=0, fragment_id=0) == "S_L0_S0_F0" + assert learner1.recv_from_syncer(step=0, fragment_id=0) == "S_L1_S0_F0" + + +def test_threaded_transport_blocking(): + manager = ThreadedTransportManager(num_learners=1) + learner = LearnerTransport(manager, 0) + syncer = SyncerTransport(manager) + + received_data = [] + + def reader_thread(): + data = learner.recv_from_syncer(step=5, fragment_id=0) + received_data.append(data) + + t = threading.Thread(target=reader_thread) + t.start() + + # Sleep to ensure reader thread is blocking + time.sleep(0.1) + assert not received_data + + # Send data + syncer.send_to_learner(0, step=5, fragment_id=0, data="late_data") + t.join(timeout=2) + + assert received_data == ["late_data"] + + +def test_threaded_transport_out_of_order(): + manager = ThreadedTransportManager(num_learners=1) + learner = LearnerTransport(manager, 0) + syncer = SyncerTransport(manager) + + # Send step 2 then step 1 + learner.send_to_syncer(step=2, fragment_id=0, data="step2") + learner.send_to_syncer(step=1, fragment_id=0, data="step1") + + # Receive step 1 first + assert syncer.recv_from_learner(0, step=1, fragment_id=0) == "step1" + assert syncer.recv_from_learner(0, step=2, fragment_id=0) == "step2" diff --git a/tests/unit/diloco_data_loader_test.py b/tests/unit/diloco_data_loader_test.py new file mode 100644 index 0000000000..13f27ddb74 --- /dev/null +++ b/tests/unit/diloco_data_loader_test.py @@ -0,0 +1,84 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for DiLoCo data loader sharding.""" + +import sys +import unittest +from unittest.mock import patch +import jax +from jax.sharding import Mesh +import numpy as np +from maxtext.configs import pyconfig +from maxtext.input_pipeline.grain_data_processing import make_grain_train_iterator +from maxtext.input_pipeline.tfds_data_processing import make_tfds_train_iterator + + +class DilocoDataLoaderTest(unittest.TestCase): + + def setUp(self): + super().setUp() + self.config = pyconfig.initialize( + [sys.argv[0], "src/maxtext/configs/base.yml"], + run_name="test", + global_batch_size_to_load=8, + max_target_length=128, + colocated_python_data_input=False, + expansion_factor_real_data=1.0, + ) + + devices = jax.devices() + self.mesh = Mesh(np.array(devices[:1]), ["data"]) + self.process_indices = [0] + + @patch("maxtext.input_pipeline.grain_data_processing.get_datasets") + @patch("maxtext.input_pipeline.grain_data_processing._get_pipeline_fn") + @patch("maxtext.input_pipeline.multihost_dataloading.MultiHostDataLoadIterator") + def test_grain_loader_sharding_diloco(self, mock_iterator, mock_get_pipeline_fn, mock_get_datasets): + mock_get_datasets.return_value = unittest.mock.MagicMock() + mock_get_pipeline_fn.return_value = lambda **kwargs: (lambda dataset: dataset) + + make_grain_train_iterator(self.config, self.mesh, self.process_indices, learner_idx=0, num_learners=2) + kwargs = mock_get_datasets.call_args[1] + self.assertEqual(kwargs["dataloading_host_index"], 0) + self.assertEqual(kwargs["dataloading_host_count"], 2) + + mock_get_datasets.reset_mock() + + make_grain_train_iterator(self.config, self.mesh, self.process_indices, learner_idx=1, num_learners=2) + kwargs = mock_get_datasets.call_args[1] + self.assertEqual(kwargs["dataloading_host_index"], 1) + self.assertEqual(kwargs["dataloading_host_count"], 2) + + @patch("maxtext.input_pipeline.tfds_data_processing.get_datasets") + @patch("maxtext.input_pipeline.tfds_data_processing.preprocessing_pipeline") + @patch("maxtext.input_pipeline.multihost_dataloading.MultiHostDataLoadIterator") + def test_tfds_loader_sharding_diloco(self, mock_iterator, mock_preprocess, mock_get_datasets): + mock_get_datasets.return_value = unittest.mock.MagicMock() + + make_tfds_train_iterator(self.config, self.mesh, self.process_indices, learner_idx=0, num_learners=2) + kwargs = mock_get_datasets.call_args[1] + self.assertEqual(kwargs["dataloading_host_index"], 0) + self.assertEqual(kwargs["dataloading_host_count"], 2) + + mock_get_datasets.reset_mock() + + make_tfds_train_iterator(self.config, self.mesh, self.process_indices, learner_idx=1, num_learners=2) + kwargs = mock_get_datasets.call_args[1] + self.assertEqual(kwargs["dataloading_host_index"], 1) + self.assertEqual(kwargs["dataloading_host_count"], 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/sharding_test.py b/tests/unit/sharding_test.py index 390c865b73..33d825f4ed 100644 --- a/tests/unit/sharding_test.py +++ b/tests/unit/sharding_test.py @@ -22,6 +22,7 @@ import pytest from jax.sharding import PartitionSpec +from maxtext.utils.sharding import partition_mesh_by_diloco_axis from jax.sharding import Mesh from jax.experimental import mesh_utils from jax.lax import with_sharding_constraint @@ -198,3 +199,53 @@ def test_fsdp_sharding(): TFLOPs_per_device = parameters * 6 * BATCH_SIZE / 10**12 / len(jax.devices()) time = simple_timeit(lambda: jax.block_until_ready(jit_func(presharded_X, presharded_layers))) print(f"time is {time} seconds, TFLOP is {TFLOPs_per_device}, TFLOP/s is {TFLOPs_per_device/time}", flush=True) + + +def test_partition_mesh_by_diloco_axis(): + + devices = jax.devices() + if len(devices) < 2: + pytest.skip("partition_mesh_by_diloco_axis test requires at least 2 devices") + + num_diloco = 2 + num_data = len(devices) // num_diloco + + if num_data == 0: + pytest.skip("Not enough devices for 2-way diloco partitioning") + + # Test with diloco as first axis + devices_grid = np.array(devices[: num_diloco * num_data]).reshape((num_diloco, num_data)) + global_mesh = Mesh(devices_grid, ["diloco", "data"]) + + submeshes = partition_mesh_by_diloco_axis(global_mesh, num_diloco) + + assert len(submeshes) == num_diloco + for i, submesh in enumerate(submeshes): + assert submesh.axis_names == ("data",) + assert submesh.shape == {"data": num_data} + expected_devices = devices_grid[i] + np.testing.assert_array_equal(submesh.devices, expected_devices) + + # Test with diloco as second axis + devices_grid = np.array(devices[: num_diloco * num_data]).reshape((num_data, num_diloco)) + global_mesh = Mesh(devices_grid, ["data", "diloco"]) + + submeshes = partition_mesh_by_diloco_axis(global_mesh, num_diloco) + + assert len(submeshes) == num_diloco + for i, submesh in enumerate(submeshes): + assert submesh.axis_names == ("data",) + assert submesh.shape == {"data": num_data} + expected_devices = devices_grid[:, i] + np.testing.assert_array_equal(submesh.devices, expected_devices) + + # Error case: axis not found + global_mesh_no_diloco = Mesh(np.array(devices[:num_data]).reshape((num_data,)), ["data"]) + with pytest.raises(ValueError, match="Axis diloco not found"): + partition_mesh_by_diloco_axis(global_mesh_no_diloco, num_diloco) + + # Error case: size mismatch + devices_grid = np.array(devices[: num_diloco * num_data]).reshape((num_diloco, num_data)) + global_mesh = Mesh(devices_grid, ["diloco", "data"]) + with pytest.raises(ValueError, match="Diloco axis size .* must match num_replicas"): + partition_mesh_by_diloco_axis(global_mesh, num_diloco + 1) diff --git a/tests/unit/threaded_diloco_test.py b/tests/unit/threaded_diloco_test.py new file mode 100644 index 0000000000..1b72fbd75a --- /dev/null +++ b/tests/unit/threaded_diloco_test.py @@ -0,0 +1,110 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for threaded DiLoCo components.""" + +import sys +import unittest +import threading +import time +from maxtext.configs import pyconfig +from maxtext.trainers.diloco.threaded_diloco import make_learner_config +from maxtext.trainers.diloco.decomposed_transport import ThreadedTransportManager + +class ThreadedDilocoUnitTest(unittest.TestCase): + + def setUp(self): + super().setUp() + # Need to add src to path if not already there, but maxtext imports usually assume src is in path. + # We will initialize config with base.yml + self.config = pyconfig.initialize( + [sys.argv[0], "src/maxtext/configs/base.yml"], + run_name="test", + enable_diloco=True, + enable_streaming_diloco=True, + num_diloco_replicas=2, + ) + + def test_make_learner_config(self): + learner_config = make_learner_config(self.config, learner_idx=1, num_learners=2) + + # Check that diloco is removed from mesh_axes + self.assertNotIn("diloco", learner_config.mesh_axes) + + # Check logical_axis_rules + for _, physical_axes in learner_config.logical_axis_rules: + if isinstance(physical_axes, str): + self.assertNotEqual(physical_axes, "diloco") + elif isinstance(physical_axes, (list, tuple)): + self.assertNotIn("diloco", physical_axes) + + # Check other flags + self.assertTrue(learner_config.enable_local_data_loading) + self.assertEqual(learner_config.learner_idx, 1) + self.assertEqual(learner_config.num_learners, 2) + self.assertFalse(learner_config.enable_streaming_diloco) + self.assertFalse(learner_config.enable_diloco) + + def test_transport_manager_basic(self): + manager = ThreadedTransportManager(num_learners=2) + + # Test learner to syncer + manager.send_to_syncer(learner_idx=0, step=1, fragment_id=1, data="l0_s1_f1") + manager.send_to_syncer(learner_idx=1, step=1, fragment_id=1, data="l1_s1_f1") + + self.assertEqual(manager.recv_from_learner(learner_idx=0, step=1, fragment_id=1), "l0_s1_f1") + self.assertEqual(manager.recv_from_learner(learner_idx=1, step=1, fragment_id=1), "l1_s1_f1") + + # Test syncer to learner + manager.send_to_learner(learner_idx=0, step=1, fragment_id=1, data="s_l0_s1_f1") + manager.send_to_learner(learner_idx=1, step=1, fragment_id=1, data="s_l1_s1_f1") + + self.assertEqual(manager.recv_from_syncer(learner_idx=0, step=1, fragment_id=1), "s_l0_s1_f1") + self.assertEqual(manager.recv_from_syncer(learner_idx=1, step=1, fragment_id=1), "s_l1_s1_f1") + + def test_transport_manager_out_of_order(self): + manager = ThreadedTransportManager(num_learners=1) + + # Send out of order + manager.send_to_syncer(learner_idx=0, step=2, fragment_id=1, data="step2") + manager.send_to_syncer(learner_idx=0, step=1, fragment_id=1, data="step1") + + # Receive in order + self.assertEqual(manager.recv_from_learner(learner_idx=0, step=1, fragment_id=1), "step1") + self.assertEqual(manager.recv_from_learner(learner_idx=0, step=2, fragment_id=1), "step2") + + def test_transport_manager_blocking(self): + manager = ThreadedTransportManager(num_learners=1) + results = {} + + def worker(): + results['data'] = manager.recv_from_learner(learner_idx=0, step=1, fragment_id=1) + + t = threading.Thread(target=worker) + t.start() + + # Sleep to ensure worker is blocked + time.sleep(0.1) + self.assertTrue(t.is_alive()) + self.assertNotIn('data', results) + + # Send data + manager.send_to_syncer(learner_idx=0, step=1, fragment_id=1, data="blocked_data") + t.join(timeout=1.0) + + self.assertFalse(t.is_alive()) + self.assertEqual(results['data'], "blocked_data") + +if __name__ == "__main__": + unittest.main()