From 516cb3ff4a9d731d763da381a41254c5b1a06b61 Mon Sep 17 00:00:00 2001 From: lkolluru05 Date: Wed, 1 Jul 2026 17:53:51 +0000 Subject: [PATCH 1/6] individual save and restore works --- axlearn/cloud/gcp/pathways_utils.py | 2 + axlearn/common/elastic_input.py | 52 +++++- axlearn/common/launch_trainer.py | 49 +++++- axlearn/common/snapshot.py | 167 +++++++++++++++++++ axlearn/common/trainer.py | 239 ++++++++++++++++++++++------ 5 files changed, 446 insertions(+), 63 deletions(-) create mode 100644 axlearn/common/snapshot.py diff --git a/axlearn/cloud/gcp/pathways_utils.py b/axlearn/cloud/gcp/pathways_utils.py index 3dec75aa2..d6e999693 100644 --- a/axlearn/cloud/gcp/pathways_utils.py +++ b/axlearn/cloud/gcp/pathways_utils.py @@ -604,10 +604,12 @@ def _build_pathways_head_sidecar_containers(self) -> list[Nested[Any]]: # If multi-head, every pathways-head will only # be connected to one pathways instance (a pathways-worker replicated job). pathways_instance_count = cfg.accelerator.num_replicas if self._is_single_head else 1 + num_elastic_slices = 1 cmd_args = [ f"--resource_manager_address=localhost:{_PATHWAYS_RESOURCE_MANAGER_PORT}", f"--server_port={_PATHWAYS_PROXY_PORT}", + f"--num_elastic_slices={num_elastic_slices}", ] if self._colocated_python.is_colocated_python_enabled: cmd_args.append("--sidecar_name=external") diff --git a/axlearn/common/elastic_input.py b/axlearn/common/elastic_input.py index e38bdf505..b7d7032bb 100644 --- a/axlearn/common/elastic_input.py +++ b/axlearn/common/elastic_input.py @@ -49,7 +49,7 @@ from axlearn.common.config import REQUIRED, Required, config_class, maybe_set_config from axlearn.common.input_dispatch import BaseInputDispatcher, _validate_logical_feed_shapes from axlearn.common.module import Module -from axlearn.common.utils import Nested, Tensor +from axlearn.common.utils import Nested, Tensor, live_devices class ElasticSpmdInputDispatcher(BaseInputDispatcher): @@ -72,13 +72,16 @@ class Config(BaseInputDispatcher.Config): @property def is_in_elastic_mode(self) -> bool: + print("In is_in_elastic_mode by lkolluru") cfg = self.config + print("cfg.num_max_slices by lkolluru: ", cfg.num_max_slices) + print("slice_count by lkolluru: ", slice_count()) if cfg.num_max_slices is None: return False else: if slice_count() < cfg.num_max_slices: return True - elif slice_count() == cfg.num_max_slices: + elif slice_count() >= cfg.num_max_slices: return False else: # TODO (jtian22): consider supporting scaling up in the future. @@ -156,14 +159,32 @@ def fid2pids(feed_id): self.feed_count = len(set(pid2fid.values())) // slice_count() * cfg.num_max_slices self.feed_index = pid2fid[jax.process_index()] + print("feed_count by lkolluru: ", self.feed_count) + print("global_logical_batch_size by lkolluru: ", cfg.global_logical_batch_size) - assert cfg.global_logical_batch_size % self.feed_count == 0 - self._feed_logical_batch_size = cfg.global_logical_batch_size // self.feed_count + # assert cfg.global_logical_batch_size % self.feed_count == 0 + if self.feed_count == 0: + self._feed_logical_batch_size = cfg.global_logical_batch_size + else: + self._feed_logical_batch_size = cfg.global_logical_batch_size // self.feed_count + + adjusted_device_physical_batch_size = math.ceil( + self._device_physical_batch_size * (cfg.num_max_slices / slice_count()) + ) + print( + "adjusted_device_physical_batch_size outside elastic by lkolluru: ", + adjusted_device_physical_batch_size, + ) if self.is_in_elastic_mode: + print(" In elastic mode lkolluru") adjusted_device_physical_batch_size = math.ceil( self._device_physical_batch_size * (cfg.num_max_slices / slice_count()) ) + print( + "adjusted_device_physical_batch_size inside elastic by lkolluru: ", + adjusted_device_physical_batch_size, + ) padding_per_device = ( adjusted_device_physical_batch_size - self._device_physical_batch_size ) @@ -402,7 +423,14 @@ def _padded_select(path, x, y): def slice_count() -> int: """Returns the number of slices.""" - return len(set(d.slice_index for d in jax.devices() if hasattr(d, "slice_index"))) or 1 + # slice_cnt_val=len(set(d.slice_index for d in jax.devices() if hasattr(d, "slice_index"))) or 1 + # print("slice_count by lkolluru: ", slice_cnt_val) + # return len(set(d.slice_index for d in jax.devices() if hasattr(d, "slice_index"))) or 1 + slice_cnt_val = ( + len(set(d.slice_index for d in live_devices() if hasattr(d, "slice_index"))) or 1 + ) + print("slice_count by lkolluru: ", slice_cnt_val) + return len(set(d.slice_index for d in live_devices() if hasattr(d, "slice_index"))) or 1 def process_count_per_slice() -> int: @@ -411,7 +439,8 @@ def process_count_per_slice() -> int: len( set( d.process_index - for d in jax.devices() + # for d in jax.devices() + for d in live_devices() if hasattr(d, "slice_index") and d.slice_index == 0 ) ) @@ -502,6 +531,8 @@ def get_process_index_and_count_and_mapping( # compatible with any mesh with num_devices. device_map = tensor_sharding.devices_indices_map((tensor_sharding.num_devices,) * ndims) + print("device_map by lkolluru: ", device_map) + # Get the slices for 'dim' for all devices. global_slice = {k: v[dim] for k, v in device_map.items()} @@ -516,11 +547,15 @@ def get_process_index_and_count_and_mapping( process_to_slice[d.process_index].add(key) all_slices.add(key) + print("process_to_slice by lkolluru: ", process_to_slice) + # Get the set of slices for the current process which we will use to compute # the index of the current process. current_pid = next(iter(tensor_sharding.addressable_devices)).process_index addressable_slices = frozenset(process_to_slice[current_pid]) + print("addressable_slices by lkolluru: ", addressable_slices) + # Verify that all processes have the same number of slices. slices_per_process = len(addressable_slices) if any(len(x) != slices_per_process for x in process_to_slice.values()): @@ -529,6 +564,7 @@ def get_process_index_and_count_and_mapping( "different number of slices." ) unique_processes = list({frozenset(x) for x in process_to_slice.values()}) + print("unique_processes by lkolluru: ", unique_processes) # After removing duplicate processes all unique slices should # cover the dimension exactly once. If they don't it means that @@ -540,6 +576,8 @@ def get_process_index_and_count_and_mapping( # !!! patch begin pid2fid = {} for pid, _ in process_to_slice.items(): + print("pid by lkolluru: ", pid) pid2fid[pid] = unique_processes.index(frozenset(process_to_slice[pid])) # !!! patch end - return feed_index, feed_count, pid2fid + print("pid2fid by lkolluru: ", pid2fid) + return feed_index, feed_count, pid2fid \ No newline at end of file diff --git a/axlearn/common/launch_trainer.py b/axlearn/common/launch_trainer.py index c56e31df3..83605aab8 100644 --- a/axlearn/common/launch_trainer.py +++ b/axlearn/common/launch_trainer.py @@ -12,8 +12,11 @@ from axlearn.common import file_system as fs from axlearn.common import measurement from axlearn.common.config import TrainerConfigFn, get_named_trainer_config -from axlearn.common.trainer import SpmdTrainer, select_mesh_config -from axlearn.common.utils import MeshShape, get_data_dir, infer_mesh_shape +from axlearn.common.trainer import SpmdTrainer, select_mesh_config, sync_restore_class_vars, sync_store_class_vars, jax_device_state, python_vars, immutable_data +from axlearn.common.utils import MeshShape, get_data_dir, infer_mesh_shape, live_devices +from pathwaysutils.elastic import manager, elastic +from pathwaysutils.debug import watchdog + # Trainer-specific flags. flags.DEFINE_string( @@ -116,6 +119,8 @@ FLAGS = flags.FLAGS +elastic_snapshotting_enabled = True + def get_trainer_config( trainer_config_fn: Optional[TrainerConfigFn] = None, @@ -148,7 +153,8 @@ def get_trainer_config( if flag_values.mesh_selector is not None: select_mesh_config(trainer_config, mesh_selector=flag_values.mesh_selector) trainer_config.mesh_axis_names = trainer_config.mesh_axis_names or ("data", "model") - trainer_config.mesh_shape = trainer_config.mesh_shape or (len(jax.devices()), 1) + #trainer_config.mesh_shape = trainer_config.mesh_shape or (len(jax.devices()), 1) + trainer_config.mesh_shape = trainer_config.mesh_shape or (len(live_devices()), 1) if isinstance(trainer_config.mesh_shape, MeshShape): trainer_config.mesh_shape = infer_mesh_shape(trainer_config.mesh_shape) trainer_config.start_trace_steps = [int(el) for el in flag_values.trace_at_steps] @@ -207,9 +213,38 @@ def run_trainer(trainer_config: SpmdTrainer.Config) -> Any: }, f, ) + + elastic_manager = None + if elastic_snapshotting_enabled and any(hasattr(d, "slice_index") for d in jax.devices()): + elastic_manager = manager.Manager() + + with ( + watchdog.watchdog("step-stack-status", timeout=60), + watchdog.watchdog("step-timebomb", timeout=10 * 60, repeat=False), + ): + try: + if elastic_manager and elastic_manager.new_slice_event.is_set(): + print("New slice event is set from elastic manager") + trainer = sync_restore_class_vars(jax_device_state, python_vars, immutable_data) + prng_key = jax_device_state["_trainer_state"].prng_key + else: + trainer: SpmdTrainer = trainer_config.instantiate(parent=None) + prng_key = jax.random.PRNGKey(seed=FLAGS.trainer_prng_seed) + output = trainer.run(prng_key) + measurement.record_event(measurement.Event.END_JOB) + + + except jax.errors.JaxRuntimeError as e: + if elastic_manager and elastic.is_error_due_to_slice_down(e): + # Slice Failure Recovery + print("Lost slice event is set from elastic manager") + logging.exception( + "[!] Elastic event detected around step %d", immutable_data["_step"] + ) + sync_restore_class_vars(jax_device_state, python_vars, immutable_data) + pass - trainer: SpmdTrainer = trainer_config.instantiate(parent=None) - prng_key = jax.random.PRNGKey(seed=FLAGS.trainer_prng_seed) - output = trainer.run(prng_key) - measurement.record_event(measurement.Event.END_JOB) + + + return output diff --git a/axlearn/common/snapshot.py b/axlearn/common/snapshot.py new file mode 100644 index 000000000..16587878a --- /dev/null +++ b/axlearn/common/snapshot.py @@ -0,0 +1,167 @@ +# Copyright © 2024 Apple Inc. + +"""Manages asynchronous backups of JAX array states to pinned host memory.""" + +import logging +import queue +import threading +from typing import Any, Optional + +from etils import epath +import jax +from orbax.checkpoint.experimental.v1 import training # pytype: disable=import-error +from orbax.checkpoint.experimental.v1._src.tree import types as tree_types # pytype: disable=import-error +from pathwaysutils.experimental import concatenate_by_mesh_axis # pytype: disable=import-error +from pathwaysutils.experimental import split_by_mesh_axis # pytype: disable=import-error +import jax.numpy as jnp +from axlearn.common.utils import Nested, TensorSpec, get_current_abstract_or_physical_mesh + +_logger = logging.getLogger(__name__) + + +class Snapshotter: + """Manages asynchronous backups of JAX array states to pinned host memory.""" + + def __init__(self, *, replica_axis_index: int = 0, trainer_state_specs: Optional[Nested[TensorSpec]] = None): + self._latest_snapshot: tuple[tree_types.PyTree, int] | None = None + self._lock = threading.Lock() + self._queue = queue.Queue(maxsize=1) + self.replica_axis_index = replica_axis_index + self.trainer_state_specs = trainer_state_specs + self._worker_thread = threading.Thread(target=self._worker, daemon=True) + self._worker_thread.start() + + def _worker(self): + while True: + pinned_state, step = self._queue.get() + print("In snapshot worker") + try: + _logger.info( + "[*] [Snapshot Thread] Waiting for snapshot at step %d to be ready...", + step, + ) + jax.block_until_ready(pinned_state) + _logger.info( + "[*] [Snapshot Thread] Snapshot at step %d is ready and secured.", + step, + ) + with self._lock: + self._latest_snapshot = (pinned_state, step) + except Exception as e: # pylint: disable=broad-except + _logger.warning( + "[*] [Snapshot Thread] Failed to secure snapshot at step %d: %s.", + step, + e, + ) + finally: + print("In snapshot worker finally") + self._queue.task_done() + + def save_pytree( + self, step: int, state: tree_types.PyTreeOf[jax.Array] + ) -> None: + """Move arrays onto CPU worker devices.""" + print("In snapshot pytree") + if self._queue.full(): + _logger.warning("Snapshotter busy. Skipping snapshot for step %d", step) + return + + pinned_shardings = jax.tree.map( + lambda x: x.sharding.with_memory_kind("pinned_host") if hasattr(x, "sharding") else None, state + ) + print("before jax.put") + + pinned_state = jax.device_put(state, pinned_shardings) + print("after jax.put") + self._queue.put((pinned_state, step)) + + def load_pytree( + self, + *, + reset_snapshot_state: bool = True, + ) -> tree_types.PyTree: + """Initializes a state and restores from the latest snapshot. + + Uses `self.trainer_state_specs` to properly re-partition onto the new mesh. + + Args: + reset_snapshot_state: If True, clears snapshot history and resets it to + contain only the returned restored state (in host-pinned memory). + + Returns: + The restored array state. + + Raises: + RuntimeError: If no snapshots are available to restore from. + ValueError: If `trainer_state_specs` is not provided during initialization. + """ + if self.trainer_state_specs is None: + raise ValueError("trainer_state_specs must be provided to Snapshotter to use load_pytree.") + + def spec_to_sds(spec): + if not hasattr(spec, "shape"): + return spec + mesh = get_current_abstract_or_physical_mesh() + sharding = jax.sharding.NamedSharding(mesh, getattr(spec, "mesh_axes", None)) + return jax.ShapeDtypeStruct(spec.shape, spec.dtype, sharding=sharding) + + abstract_state = jax.tree.map(spec_to_sds, self.trainer_state_specs, is_leaf=lambda x: hasattr(x, "shape")) + + with self._lock: + if self._latest_snapshot is None: + raise RuntimeError("No snapshots available to restore from.") + pinned_state, step = self._latest_snapshot + + def get_active_pytree(x, target_x): + if not hasattr(x, "shape") or not hasattr(target_x, "shape"): + return x + if x.shape == target_x.shape: + return x + starts = [0] * x.ndim + stops = [min(s1, s2) for s1, s2 in zip(x.shape, target_x.shape)] + sliced_x = jax.lax.slice(x, starts, stops) + pad_widths = [(0, max(0, s2 - s1)) for s1, s2 in zip(x.shape, target_x.shape)] + if any(p > 0 for _, p in pad_widths): + sliced_x = jnp.pad(sliced_x, pad_widths) + return sliced_x + + _logger.info("Restoring from snapshot at step %d...", step) + pinned_state = jax.tree.map(get_active_pytree, pinned_state, abstract_state) + + # Re-shard on host to the target device mesh + host_target_shardings = jax.tree.map( + lambda x: x.sharding.with_memory_kind("pinned_host") if hasattr(x, "sharding") else None, abstract_state + ) + + host_target_state = jax.device_put( + pinned_state, host_target_shardings + ) + + # Move from host back to device (TPU) memory. + restored_state = jax.device_put( + host_target_state, jax.tree.map(lambda x: x.sharding if hasattr(x, "sharding") else None, abstract_state) + ) + jax.block_until_ready(restored_state) + + if reset_snapshot_state: + with self._lock: + self._latest_snapshot = (host_target_state, step) + + return restored_state + + def join(self) -> None: + """Blocks until all snapshots in the queue are ready and secured.""" + self._queue.join() + + @property + def latest(self) -> training.CheckpointMetadata[None] | None: + """Returns the training step of the most recently pinned backup.""" + with self._lock: + if self._latest_snapshot is None: + return None + _, step = self._latest_snapshot + return training.CheckpointMetadata( + step=step, + path=epath.Path(), + metadata=None, + ) \ No newline at end of file diff --git a/axlearn/common/trainer.py b/axlearn/common/trainer.py index 05df2816b..7ad59cad5 100644 --- a/axlearn/common/trainer.py +++ b/axlearn/common/trainer.py @@ -10,7 +10,7 @@ import threading import time from collections.abc import Sequence -from typing import Any, Callable, ContextManager, Literal, NamedTuple, Optional, Union +from typing import Any, Callable, ContextManager, Literal, NamedTuple, Optional, Union, TYPE_CHECKING import jax import numpy as np @@ -18,6 +18,10 @@ from jax import numpy as jnp from jax.experimental import multihost_utils from jax.experimental.pjit import pjit +from pathwaysutils.elastic import manager +from pathwaysutils.debug import watchdog + + from axlearn.common import file_system as fs from axlearn.common import measurement, utils @@ -33,6 +37,7 @@ config_class, maybe_instantiate, maybe_set_config, + config_for_class, ) from axlearn.common.evaler import SpmdEvaler from axlearn.common.input_base import Input @@ -49,6 +54,8 @@ from axlearn.common.monitoring.device_monitor import DeviceMonitor from axlearn.common.optimizer_base import NestedOptParam, OptParam from axlearn.common.param_init import DefaultInitializer +from axlearn.common.snapshot import Snapshotter + from axlearn.common.state_builder import Builder as TrainerStateBuilder from axlearn.common.summary_writer import BaseWriter, SummaryWriter from axlearn.common.update_transformation import ForwardOutputs # pytype: disable=pyi-error @@ -67,8 +74,112 @@ host_to_global_specs, match_regex_rules, thread_stack_traces, + live_devices, ) +jax_device_state = {} +python_vars = {} +immutable_data = {} + +def sync_restore_class_vars( + jax_device_state_arg: dict = None, python_vars_arg: dict = None, immutable_data_arg: dict = None +) -> Any: + """Initializes SpmdTrainer, restores its state, and runs it.""" + global jax_device_state, python_vars, immutable_data + + # Use provided arguments if available, otherwise fallback to globals + use_jax_state = jax_device_state_arg if jax_device_state_arg is not None else jax_device_state + use_python_vars = python_vars_arg if python_vars_arg is not None else python_vars + use_immutable_data = immutable_data_arg if immutable_data_arg is not None else immutable_data + + import copy # pylint: disable=import-outside-toplevel + trainer = SpmdTrainer.__new__(SpmdTrainer) + + # Restore everything first. + for state_dict in (use_jax_state, use_python_vars, use_immutable_data): + for k, v in state_dict.items(): + setattr(trainer, k, v) + + # Recreate the mesh using the restored configuration + cfg = trainer.config + devices = utils.create_device_mesh(mesh_shape=cfg.mesh_shape) + mesh = jax.sharding.Mesh(devices, cfg.mesh_axis_names) + + # Attempt to load from snapshot, fallback to globals if it fails + snapshot_mgr = use_python_vars.get("snapshot_mgr") + if snapshot_mgr is not None: + with mesh: + try: + restored_trainer_state = snapshot_mgr.load_pytree() + trainer._trainer_state = restored_trainer_state + logging.info("Successfully restored state from snapshot.") + except RuntimeError as e: + logging.warning( + "Failed to load from snapshot: %s. Using trainer state from global variables.", + e, + ) + + trainer.snapshot_mgr = snapshot_mgr + + prng_key = trainer._trainer_state.prng_key # Use the restored prng_key. + + trainer._is_restored = True + + if hasattr(trainer, "_children"): + trainer._children = copy.copy(trainer._children) + if "checkpointer" in trainer._children: + trainer._children["checkpointer"] = copy.copy(trainer._children["checkpointer"]) + trainer._children["checkpointer"]._within_context = False + trainer._children["checkpointer"]._gc_thread = None + trainer._children["checkpointer"]._gc_stopping = None + + trainer._watchdog_thread = None + trainer._watchdog_stopping = None + trainer._device_monitor = None + trainer._recorder = None + trainer.__post_init__() + + return trainer.run(prng_key) + +def sync_store_class_vars(obj: Any) -> None: + """Stores instance variables of an object in global variables.""" + if getattr(obj, "_is_restored", False): + return + + print("_sync_store_class_vars") + + global jax_device_state, python_vars, immutable_data + + jax_device_state.clear() + python_vars.clear() + immutable_data.clear() + + for k, v in obj.__dict__.items(): + if isinstance(v, property): + continue + + if k in ("_trainer_state", "_mesh", "_jit_train_step", "_compiled_train_step", "model", "learner"): + jax_device_state[k] = v + elif "config" in k or "spec" in k or isinstance(v, (int, float, str, bool)): + immutable_data[k] = v + else: + python_vars[k] = v + + #print(python_vars) + #print(immutable_data) + snapshot_mgr = python_vars.get("snapshot_mgr") + if snapshot_mgr is not None: + #print(immutable_data["_step"]) + #print(jax_device_state["_trainer_state"]) + try: + snapshot_mgr.save_pytree(step=immutable_data["_step"], state=jax_device_state["_trainer_state"]) + snapshot_mgr.join() + except Exception as e: + logging.warning("Failed during snapshot save: %s", e) + + python_vars["snapshot_mgr"] = snapshot_mgr + + class TrainerState(NamedTuple): prng_key: Union[Tensor, TensorSpec, jax.sharding.NamedSharding] @@ -242,6 +353,9 @@ class Config(Module.Config): # Log the loss value every n steps. Defaults to None which is interpreted as every # 100 steps. log_every_n_steps: Optional[int] = None + + + def __init__( self, @@ -268,6 +382,7 @@ def __init__( self._recorder = maybe_instantiate(cfg.recorder) self._is_initialized: bool = False self._maybe_record_event(measurement.Event.START_ACCELERATOR_INIT) + self._class_vars = None if cfg.model.dtype is None: raise ValueError(f"dtype must be explicitly specified for {self.path()}.model") @@ -359,8 +474,10 @@ def __init__( model=self._model_param_specs, learner=self._learner_state_partition_specs, ) - self._trainer_state_partition_specs: TrainerState = jax.tree.map( - lambda spec: spec.sharding, self._trainer_state_specs + self._trainer_state_partition_specs: TrainerState = ( + jax.tree.map( + lambda spec: spec.sharding, self._trainer_state_specs + ) ) # Create evalers, which depend on model_param_partition_specs. self._evalers = {} @@ -380,6 +497,8 @@ def __init__( ) self._maybe_record_event(measurement.Event.END_ACCELERATOR_INIT) + + @property def step(self): return self._step @@ -416,7 +535,7 @@ def _step_log(self, msg, *args, **kwargs): "%s process % 3d step % 8d] " + msg, self.path(), jax.process_index(), - -1 if self.step is None else self.step, + -1 if self.step is None else int(self.step), *args, **kwargs, ) @@ -615,6 +734,11 @@ def run( return None self._is_initialized = True + #### Stores the initial state of all variables #### + replica_axis_idx = cfg.mesh_axis_names.index("data") if "data" in cfg.mesh_axis_names else 0 + snapshot_cfg = config_for_class(Snapshotter).set(replica_axis_index=replica_axis_idx, trainer_state_specs=self.trainer_state_specs) + self.snapshot_mgr = snapshot_cfg.instantiate() + with self.checkpointer: logging.info("Starting loop...") @@ -625,48 +749,58 @@ def run( input_iterator = self.input.batches(self._input_iter) while True: - self._maybe_record_event(measurement.Event.START_DATA_LOADING) - try: - input_batch = next(input_iterator) - self._maybe_record_event(measurement.Event.END_DATA_LOADING) - logging.log_first_n( - logging.INFO, "host_input_batch=%s", 3, utils.shapes(input_batch) - ) - - # Stop or start tracing if necessary. - stop_trace_step = self._maybe_stop_or_start_tracing(stop_trace_step, output) - - self._step = self._step + 1 - self.vlog(3, "Start step %s", self.step) - self._maybe_record_event(measurement.Event.START_STEP, self._step) - output = self._run_step( - utils.host_to_global_array( - input_batch, - partition=self._train_step_input_partition_specs(), - ), - force_run_evals=( - force_run_eval_sets_at_max_step - if self.step >= cfg.max_step - else None - ), - ) - self.vlog(3, "Done step %s", self.step) - num_steps += 1 - if num_steps % 100 == 0: - now = time.perf_counter() - average_step_time = (now - start_time) / num_steps - self._step_log("Average step time: %s seconds", average_step_time) - self.summary_writer(self.step, {"average_step_time": average_step_time}) - num_steps = 0 - start_time = now - if self.step >= cfg.max_step: - self._step_log("Reached max_step=%s. Stopping", cfg.max_step) - break - except StopIteration: - # Add END_DATA_LOADING event here to close the unpaired START_DATA_LOADING - # event. - self._maybe_record_event(measurement.Event.END_DATA_LOADING) - break + with ( + watchdog.watchdog("step-stack-status", timeout=60), + watchdog.watchdog("step-timebomb", timeout=10 * 60, repeat=False), + ): + self._maybe_record_event(measurement.Event.START_DATA_LOADING) + + try: + + input_batch = next(input_iterator) + self._maybe_record_event(measurement.Event.END_DATA_LOADING) + logging.log_first_n( + logging.INFO, "host_input_batch=%s", 3, utils.shapes(input_batch) + ) + + # Stop or start tracing if necessary. + stop_trace_step = self._maybe_stop_or_start_tracing(stop_trace_step, output) + + self._step = self._step + 1 + self.vlog(3, "Start step %s", self.step) + self._maybe_record_event(measurement.Event.START_STEP, self._step) + output = self._run_step( + utils.host_to_global_array( + input_batch, + partition=self._train_step_input_partition_specs(), + ), + force_run_evals=( + force_run_eval_sets_at_max_step + if self.step >= cfg.max_step + else None + ), + ) + self.vlog(3, "Done step %s", self.step) + # if self.step==3: + # _sync_store_class_vars(self) + sync_store_class_vars(self) + + num_steps += 1 + if num_steps % 100 == 0: + now = time.perf_counter() + average_step_time = (now - start_time) / num_steps + self._step_log("Average step time: %s seconds", average_step_time) + self.summary_writer(self.step, {"average_step_time": average_step_time}) + num_steps = 0 + start_time = now + if self.step >= cfg.max_step: + self._step_log("Reached max_step=%s. Stopping", cfg.max_step) + break + except StopIteration: + # Add END_DATA_LOADING event here to close the unpaired START_DATA_LOADING + # event. + self._maybe_record_event(measurement.Event.END_DATA_LOADING) + break if self.step < cfg.max_step: self._step_log("Reached end of inputs. Stopping") self._step_log("Checkpointer flushed.") @@ -1075,7 +1209,7 @@ def _get_compiled_train_step_fn( cfg: SpmdTrainer.Config = self.config # Get device kinds and assert that they are homogenous. # TODO(markblee): Get devices from self._mesh.devices. - device_kinds = set(d.device_kind for d in jax.devices()) + device_kinds = set(d.device_kind for d in live_devices()) if len(device_kinds) != 1: raise RuntimeError(f"Heterogenous device kinds ({device_kinds}) are not supported.") device_kind = device_kinds.pop() @@ -1158,7 +1292,7 @@ def _run_step( jax.tree.map(lambda x: x.item() if x.ndim == 0 else f"T{x.shape}", outputs["aux"]), ) - self.summary_writer(self.step, {"loss": outputs["loss"], **outputs["summaries"]}) + #self.summary_writer(self.step, {"loss": outputs["loss"], **outputs["summaries"]}) # Aggregate summaries across evalers. evaler_summaries = self._run_eval( train_summaries=outputs["summaries"], force_runs=force_run_evals @@ -1466,15 +1600,22 @@ def mb_or_gb(x): if mem_stats is not None: analysis_results += "======= Memory Analysis ==================================\n" try: + # XLA may alias output buffers onto input buffers when + # ``donate_argnums`` is used (typical for jit'ed training steps); + # subtract the aliased bytes so the reported total reflects the + # actual peak HBM, not the double-counted argument + output sum. + aliased_bytes = mem_stats.alias_size_in_bytes total_hbm = ( mem_stats.argument_size_in_bytes + mem_stats.output_size_in_bytes + - aliased_bytes + mem_stats.temp_size_in_bytes + mem_stats.generated_code_size_in_bytes ) analysis_results += ( f"Input memory: {mb_or_gb(mem_stats.argument_size_in_bytes)}\n" + f"Output memory: {mb_or_gb(mem_stats.output_size_in_bytes)}\n" + + f"Aliased input/output memory: {mb_or_gb(aliased_bytes)}\n" + f"Temp memory: {mb_or_gb(mem_stats.temp_size_in_bytes)}\n" + f"Code memory: {mb_or_gb(mem_stats.generated_code_size_in_bytes)}\n" + f"Total HBM memory: {mb_or_gb(total_hbm)}\n" @@ -1515,4 +1656,4 @@ def mb_or_gb(x): analysis_results += f"{cost_stats}\n" logging.warning("Attempt to parse cost_stats=%s but failed.", cost_stats) - return analysis_results + return analysis_results \ No newline at end of file From f95e17562ed683fabf01b11b782ddad94b5955ba Mon Sep 17 00:00:00 2001 From: lkolluru05 Date: Mon, 6 Jul 2026 21:16:40 +0000 Subject: [PATCH 2/6] not working yet --- axlearn/common/launch_trainer.py | 59 +++++------ axlearn/common/snapshot.py | 7 +- axlearn/common/trainer.py | 174 ++++++++++++++++++++++--------- 3 files changed, 160 insertions(+), 80 deletions(-) diff --git a/axlearn/common/launch_trainer.py b/axlearn/common/launch_trainer.py index 83605aab8..83b830d9b 100644 --- a/axlearn/common/launch_trainer.py +++ b/axlearn/common/launch_trainer.py @@ -119,7 +119,7 @@ FLAGS = flags.FLAGS -elastic_snapshotting_enabled = True +elastic_snapshotting_enabled = False def get_trainer_config( @@ -215,36 +215,37 @@ def run_trainer(trainer_config: SpmdTrainer.Config) -> Any: ) elastic_manager = None - if elastic_snapshotting_enabled and any(hasattr(d, "slice_index") for d in jax.devices()): + if elastic_snapshotting_enabled and any(hasattr(d, "slice_index") for d in live_devices()): elastic_manager = manager.Manager() - with ( - watchdog.watchdog("step-stack-status", timeout=60), - watchdog.watchdog("step-timebomb", timeout=10 * 60, repeat=False), - ): - try: - if elastic_manager and elastic_manager.new_slice_event.is_set(): - print("New slice event is set from elastic manager") - trainer = sync_restore_class_vars(jax_device_state, python_vars, immutable_data) - prng_key = jax_device_state["_trainer_state"].prng_key - else: - trainer: SpmdTrainer = trainer_config.instantiate(parent=None) - prng_key = jax.random.PRNGKey(seed=FLAGS.trainer_prng_seed) - output = trainer.run(prng_key) - measurement.record_event(measurement.Event.END_JOB) - + output = None + while True: + try: + clean_trainer: SpmdTrainer = trainer_config.instantiate(parent=None) + + if elastic_manager and elastic_manager.new_slice_event.is_set(): + print("New slice event is set from elastic manager") + elastic_manager.new_slice_event.clear() + trainer, prng_key = sync_restore_class_vars(clean_trainer, jax_device_state, python_vars, immutable_data) + else: + trainer = clean_trainer + prng_key = jax.random.PRNGKey(seed=FLAGS.trainer_prng_seed) + + output = trainer.run(prng_key) + measurement.record_event(measurement.Event.END_JOB) + break - except jax.errors.JaxRuntimeError as e: - if elastic_manager and elastic.is_error_due_to_slice_down(e): - # Slice Failure Recovery - print("Lost slice event is set from elastic manager") - logging.exception( - "[!] Elastic event detected around step %d", immutable_data["_step"] - ) - sync_restore_class_vars(jax_device_state, python_vars, immutable_data) - pass - - - + except jax.errors.JaxRuntimeError as e: + if elastic_manager and elastic.is_error_due_to_slice_down(e): + # Slice Failure Recovery + print("Lost slice event is set from elastic manager") + logging.exception( + "[!] Elastic event detected around step %d", immutable_data.get("_step", -1) + ) + #sync_store_class_vars(trainer) + elastic_manager.new_slice_event.set() + continue + else: + raise e return output diff --git a/axlearn/common/snapshot.py b/axlearn/common/snapshot.py index 16587878a..1769c998f 100644 --- a/axlearn/common/snapshot.py +++ b/axlearn/common/snapshot.py @@ -48,10 +48,15 @@ def _worker(self): with self._lock: self._latest_snapshot = (pinned_state, step) except Exception as e: # pylint: disable=broad-except + err_msg = "Unknown error" + try: + err_msg = str(e) + except Exception: + err_msg = f"JAX Runtime Exception of type {type(e).__name__} (suppressed tensor evaluation)" _logger.warning( "[*] [Snapshot Thread] Failed to secure snapshot at step %d: %s.", step, - e, + err_msg, ) finally: print("In snapshot worker finally") diff --git a/axlearn/common/trainer.py b/axlearn/common/trainer.py index 7ad59cad5..fa98ff0bc 100644 --- a/axlearn/common/trainer.py +++ b/axlearn/common/trainer.py @@ -82,64 +82,138 @@ immutable_data = {} def sync_restore_class_vars( + fresh_trainer: Any, jax_device_state_arg: dict = None, python_vars_arg: dict = None, immutable_data_arg: dict = None -) -> Any: - """Initializes SpmdTrainer, restores its state, and runs it.""" +) -> tuple[Any, Any]: + """Restores trainer state onto a fresh SpmdTrainer instance from snapshot.""" + print("sync_restore_class_vars") global jax_device_state, python_vars, immutable_data - # Use provided arguments if available, otherwise fallback to globals - use_jax_state = jax_device_state_arg if jax_device_state_arg is not None else jax_device_state + print(immutable_data) + use_python_vars = python_vars_arg if python_vars_arg is not None else python_vars use_immutable_data = immutable_data_arg if immutable_data_arg is not None else immutable_data - import copy # pylint: disable=import-outside-toplevel - trainer = SpmdTrainer.__new__(SpmdTrainer) + for k, v in use_immutable_data.items(): + if isinstance(v, (int, float, str, bool)): + setattr(fresh_trainer, k, v) + + if "_step" in use_python_vars and getattr(fresh_trainer, "_step", None) is None: + try: + fresh_trainer._step = int(use_python_vars["_step"]) + except Exception: + pass + + live_devs = utils.live_devices() + logging.info("[!] Found %d available live devices.", len(live_devs)) + + cfg = fresh_trainer.config + device_platform = live_devs[0].platform + device_attr = "process_index" if device_platform != "tpu" else "slice_index" + + num_granules = len({getattr(el, device_attr) for el in live_devs}) + num_devices_per_granule = len(live_devs) // num_granules + + original_mesh_shape = list(cfg.mesh_shape) + if len(original_mesh_shape) > 0: + original_mesh_shape[0] = num_granules + + ici_shape = original_mesh_shape[1:] + current_ici_prod = math.prod(ici_shape) + if current_ici_prod != num_devices_per_granule: + logging.info("[!] ICI product %d does not match num_devices_per_granule %d. Adjusting...", current_ici_prod, num_devices_per_granule) + ratio = current_ici_prod // num_devices_per_granule + if ratio > 0 and current_ici_prod % num_devices_per_granule == 0: + for i in range(len(ici_shape)): + if ici_shape[i] % ratio == 0 and ici_shape[i] > 1: + ici_shape[i] = ici_shape[i] // ratio + break + else: + ici_shape = [1] * len(ici_shape) + ici_shape[-3] = num_devices_per_granule + + original_mesh_shape[1:] = ici_shape - # Restore everything first. - for state_dict in (use_jax_state, use_python_vars, use_immutable_data): - for k, v in state_dict.items(): - setattr(trainer, k, v) + updated_mesh_shape = tuple(original_mesh_shape) - # Recreate the mesh using the restored configuration - cfg = trainer.config - devices = utils.create_device_mesh(mesh_shape=cfg.mesh_shape) - mesh = jax.sharding.Mesh(devices, cfg.mesh_axis_names) + logging.info("[!] Dynamically updating logical mesh_shape from %s to %s", cfg.mesh_shape, updated_mesh_shape) - # Attempt to load from snapshot, fallback to globals if it fails + devices_mesh = utils.create_device_mesh(mesh_shape=updated_mesh_shape, devices=live_devs) + mesh = jax.sharding.Mesh(devices_mesh, cfg.mesh_axis_names) + fresh_trainer._mesh = mesh + + state_restored = False snapshot_mgr = use_python_vars.get("snapshot_mgr") if snapshot_mgr is not None: with mesh: try: restored_trainer_state = snapshot_mgr.load_pytree() - trainer._trainer_state = restored_trainer_state - logging.info("Successfully restored state from snapshot.") - except RuntimeError as e: - logging.warning( - "Failed to load from snapshot: %s. Using trainer state from global variables.", - e, - ) - - trainer.snapshot_mgr = snapshot_mgr - - prng_key = trainer._trainer_state.prng_key # Use the restored prng_key. + fresh_trainer._trainer_state = restored_trainer_state + logging.info("Successfully restored state from snapshot onto the new mesh.") + state_restored = True + except Exception as e: + logging.warning("Failed to load from snapshot: %s.", e) - trainer._is_restored = True + use_jax_state = jax_device_state_arg if jax_device_state_arg is not None else jax_device_state + if not state_restored and use_jax_state and "_trainer_state" in use_jax_state: + logging.info("[!] Attempting fallback: device_put trainer_state from globals onto the new mesh.") + try: + with mesh: + fresh_trainer._trainer_state = jax.tree_util.tree_map( + lambda state, spec: jax.device_put(state, spec.sharding), + use_jax_state["_trainer_state"], + fresh_trainer._trainer_state_specs + ) + logging.info("[✓] Successfully device_put trainer_state from globals onto the new mesh.") + state_restored = True + except Exception as e: + logging.warning("[!] Failed fallback to globals (possibly deleted arrays): %s", e) + + if not state_restored: + logging.info("[!] Falling back to fresh init().") + fresh_trainer.init(jax.random.PRNGKey(seed=42)) + if "_step" in use_immutable_data: + fresh_trainer._step = int(use_immutable_data["_step"]) + elif "_step" in use_python_vars: + fresh_trainer._step = int(use_python_vars["_step"]) + + fresh_trainer.snapshot_mgr = snapshot_mgr + fresh_trainer._is_restored = True + fresh_trainer._compiled_train_step = None + + fresh_trainer._watchdog_thread = None + fresh_trainer._watchdog_stopping = None + fresh_trainer._device_monitor = None + fresh_trainer._recorder = None + + fresh_prng_key = jax.random.PRNGKey(seed=int(fresh_trainer.step if fresh_trainer.step is not None else 42)) + + try: + if hasattr(fresh_trainer._trainer_state, "_replace"): + fresh_trainer._trainer_state = fresh_trainer._trainer_state._replace(prng_key=fresh_prng_key) + elif isinstance(fresh_trainer._trainer_state, dict): + fresh_trainer._trainer_state["prng_key"] = fresh_prng_key + else: + setattr(fresh_trainer._trainer_state, "prng_key", fresh_prng_key) + logging.info("[✓] Successfully injected fresh, healthy PRNG Key into the new trainer state.") + except Exception as e: + logging.warning("[!] Failed to replace prng_key inside trainer_state structure: %s", e) - if hasattr(trainer, "_children"): - trainer._children = copy.copy(trainer._children) - if "checkpointer" in trainer._children: - trainer._children["checkpointer"] = copy.copy(trainer._children["checkpointer"]) - trainer._children["checkpointer"]._within_context = False - trainer._children["checkpointer"]._gc_thread = None - trainer._children["checkpointer"]._gc_stopping = None + try: + leaves = jax.tree_util.tree_leaves(fresh_trainer._trainer_state) + deleted_count = sum(x.is_deleted() for x in leaves if isinstance(x, jax.Array)) + mesh_match = all(x.sharding.mesh == mesh for x in leaves if isinstance(x, jax.Array) and hasattr(x, "sharding")) - trainer._watchdog_thread = None - trainer._watchdog_stopping = None - trainer._device_monitor = None - trainer._recorder = None - trainer.__post_init__() + logging.info( + "[Diagnostic] Trainer State Check: step=%s, deleted_arrays=%d, mesh_match_verified=%s", + int(fresh_trainer.step) if fresh_trainer.step is not None else "None", + deleted_count, + mesh_match + ) + except Exception as e: + logging.warning("[Diagnostic] Failed to run trainer state verification: %s", e) - return trainer.run(prng_key) + return fresh_trainer, fresh_prng_key def sync_store_class_vars(obj: Any) -> None: """Stores instance variables of an object in global variables.""" @@ -172,7 +246,10 @@ def sync_store_class_vars(obj: Any) -> None: #print(immutable_data["_step"]) #print(jax_device_state["_trainer_state"]) try: - snapshot_mgr.save_pytree(step=immutable_data["_step"], state=jax_device_state["_trainer_state"]) + step_val = immutable_data.get("_step") + if step_val is None: + step_val = python_vars.get("_step") + snapshot_mgr.save_pytree(step=int(step_val) if step_val is not None else 0, state=jax_device_state["_trainer_state"]) snapshot_mgr.join() except Exception as e: logging.warning("Failed during snapshot save: %s", e) @@ -501,7 +578,7 @@ def __init__( @property def step(self): - return self._step + return int(self._step) if self._step is not None else None @property def trainer_state(self): @@ -749,10 +826,6 @@ def run( input_iterator = self.input.batches(self._input_iter) while True: - with ( - watchdog.watchdog("step-stack-status", timeout=60), - watchdog.watchdog("step-timebomb", timeout=10 * 60, repeat=False), - ): self._maybe_record_event(measurement.Event.START_DATA_LOADING) try: @@ -766,7 +839,7 @@ def run( # Stop or start tracing if necessary. stop_trace_step = self._maybe_stop_or_start_tracing(stop_trace_step, output) - self._step = self._step + 1 + self._step = int(self._step) + 1 self.vlog(3, "Start step %s", self.step) self._maybe_record_event(measurement.Event.START_STEP, self._step) output = self._run_step( @@ -880,7 +953,7 @@ def _init_with_prebuilt_state( trainer_state=self.trainer_state_specs, built_keys=set(), ) - self._step = prebuilt_state.step + self._step = int(prebuilt_state.step) if prebuilt_state.step is not None else None all_trainer_state_keys = {key for key, _ in utils.flatten_items(self.trainer_state_specs)} if prebuilt_state.built_keys == all_trainer_state_keys: logging.info( @@ -1156,7 +1229,7 @@ def save_checkpoint(self, evaler_summaries: Optional[dict[str, Any]]) -> Optiona if cfg.save_input_iterator: ckpt_state["input_iter"] = self._input_iter self.checkpointer.save( - step=self.step, state=ckpt_state, evaler_summaries=evaler_summaries + step=int(self.step) if self.step is not None else 0, state=ckpt_state, evaler_summaries=evaler_summaries ) def _restore_from_builder(self) -> Optional[TrainerStateBuilder.State]: @@ -1285,6 +1358,7 @@ def _run_step( self._trainer_state, outputs = compiled_train_step_fn(self.trainer_state, input_batch) n = self._config.log_every_n_steps or 100 + if self.step % n == 0 or 0 <= self.step <= 5: self._step_log( "loss=%s aux=%s", @@ -1292,7 +1366,7 @@ def _run_step( jax.tree.map(lambda x: x.item() if x.ndim == 0 else f"T{x.shape}", outputs["aux"]), ) - #self.summary_writer(self.step, {"loss": outputs["loss"], **outputs["summaries"]}) + self.summary_writer(self.step, {"loss": outputs["loss"], **outputs["summaries"]}) # Aggregate summaries across evalers. evaler_summaries = self._run_eval( train_summaries=outputs["summaries"], force_runs=force_run_evals From 78faab7ff50f595343ad546fe52f812a40c59fe3 Mon Sep 17 00:00:00 2001 From: lkolluru05 Date: Wed, 8 Jul 2026 22:49:48 +0000 Subject: [PATCH 3/6] exact changes i have, still have to fix mesh issues --- axlearn/common/launch_trainer.py | 55 ++++++++++++++++------ axlearn/common/snapshot.py | 4 +- axlearn/common/summary_writer.py | 9 ++-- axlearn/common/trainer.py | 65 +++++++++++++++----------- axlearn/common/utils.py | 30 +++++++++++- axlearn/experiments/text/gpt/common.py | 2 +- axlearn/experiments/text/gpt/fuji.py | 33 ++++++++++++- pyproject.toml | 2 +- 8 files changed, 149 insertions(+), 51 deletions(-) diff --git a/axlearn/common/launch_trainer.py b/axlearn/common/launch_trainer.py index 83b830d9b..ca3ae28db 100644 --- a/axlearn/common/launch_trainer.py +++ b/axlearn/common/launch_trainer.py @@ -4,6 +4,7 @@ import json import os +import time from typing import Any, Optional import jax @@ -12,8 +13,8 @@ from axlearn.common import file_system as fs from axlearn.common import measurement from axlearn.common.config import TrainerConfigFn, get_named_trainer_config -from axlearn.common.trainer import SpmdTrainer, select_mesh_config, sync_restore_class_vars, sync_store_class_vars, jax_device_state, python_vars, immutable_data -from axlearn.common.utils import MeshShape, get_data_dir, infer_mesh_shape, live_devices +from axlearn.common.trainer import SpmdTrainer, select_mesh_config, sync_restore_class_vars, sync_store_class_vars +from axlearn.common.utils import MeshShape, get_data_dir, infer_mesh_shape, live_devices, set_elastic_manager from pathwaysutils.elastic import manager, elastic from pathwaysutils.debug import watchdog @@ -119,7 +120,7 @@ FLAGS = flags.FLAGS -elastic_snapshotting_enabled = False +elastic_snapshotting_enabled = True def get_trainer_config( @@ -195,6 +196,16 @@ def get_trainer_config( return trainer_config +def is_retryable_error(e: Exception) -> bool: + if isinstance(e, jax.errors.JaxRuntimeError): + err_str = str(e) + if elastic.is_error_due_to_slice_down(e): + return True + if "UNAVAILABLE" in err_str or "RESOURCE_EXHAUSTED" in err_str: + return True + return False + + def run_trainer(trainer_config: SpmdTrainer.Config) -> Any: measurement.record_event(measurement.Event.START_JOB) trainer_config_debug_string = trainer_config.debug_string() @@ -215,19 +226,33 @@ def run_trainer(trainer_config: SpmdTrainer.Config) -> Any: ) elastic_manager = None - if elastic_snapshotting_enabled and any(hasattr(d, "slice_index") for d in live_devices()): - elastic_manager = manager.Manager() + elastic_manager_initialized = False output = None + jax_device_state = {} + python_vars = {} + immutable_data = {} + trainer = None while True: try: + if not elastic_manager_initialized: + if elastic_snapshotting_enabled: + logging.info("[Elastic] Initializing elastic manager...") + elastic_manager = manager.Manager() + set_elastic_manager(elastic_manager) + logging.info("[Elastic] Elastic manager initialized.") + else: + logging.info("[Elastic] Elastic snapshotting disabled or not supported (no slice_index).") + elastic_manager_initialized = True + clean_trainer: SpmdTrainer = trainer_config.instantiate(parent=None) if elastic_manager and elastic_manager.new_slice_event.is_set(): - print("New slice event is set from elastic manager") + logging.info("[Elastic] New slice event is set. Restoring from snapshot...") elastic_manager.new_slice_event.clear() trainer, prng_key = sync_restore_class_vars(clean_trainer, jax_device_state, python_vars, immutable_data) else: + logging.info("[Elastic] Starting fresh trainer (no elastic recovery triggered).") trainer = clean_trainer prng_key = jax.random.PRNGKey(seed=FLAGS.trainer_prng_seed) @@ -235,16 +260,16 @@ def run_trainer(trainer_config: SpmdTrainer.Config) -> Any: measurement.record_event(measurement.Event.END_JOB) break - except jax.errors.JaxRuntimeError as e: - if elastic_manager and elastic.is_error_due_to_slice_down(e): - # Slice Failure Recovery - print("Lost slice event is set from elastic manager") - logging.exception( - "[!] Elastic event detected around step %d", immutable_data.get("_step", -1) - ) - #sync_store_class_vars(trainer) - elastic_manager.new_slice_event.set() + if is_retryable_error(e): + logging.warning("Caught retryable error: %s. Retrying...", e) + if trainer is not None: + jax_device_state = getattr(trainer, "_jax_device_state", {}) + python_vars = getattr(trainer, "_python_vars", {}) + immutable_data = getattr(trainer, "_immutable_data", {}) + if elastic_manager: + elastic_manager.new_slice_event.set() + time.sleep(10) continue else: raise e diff --git a/axlearn/common/snapshot.py b/axlearn/common/snapshot.py index 1769c998f..d23b918bf 100644 --- a/axlearn/common/snapshot.py +++ b/axlearn/common/snapshot.py @@ -2,7 +2,7 @@ """Manages asynchronous backups of JAX array states to pinned host memory.""" -import logging +from absl import logging import queue import threading from typing import Any, Optional @@ -16,7 +16,7 @@ import jax.numpy as jnp from axlearn.common.utils import Nested, TensorSpec, get_current_abstract_or_physical_mesh -_logger = logging.getLogger(__name__) +_logger = logging class Snapshotter: diff --git a/axlearn/common/summary_writer.py b/axlearn/common/summary_writer.py index aa15c62d3..580424e80 100644 --- a/axlearn/common/summary_writer.py +++ b/axlearn/common/summary_writer.py @@ -288,17 +288,20 @@ def write(path: str, value: jax.Array): else: raw_value = value - self.vlog(3, "SummaryWriter %s: %s=%s", self.path(), path, raw_value) - if isinstance(raw_value, Tensor) and not raw_value.is_fully_replicated: logging.warning( - "SummaryWriter: %s: %s is not fully replicated", path, raw_value + "SummaryWriter: %s: %s is not fully replicated (shape=%s)", + path, + self.path(), + raw_value.shape, ) return if isinstance(raw_value, jax.Array): raw_value = np.asarray(raw_value) + self.vlog(3, "SummaryWriter %s: %s=%s", self.path(), path, raw_value) + if _match_summary_type("Image", value=value, raw_value=raw_value): if self._time_to_write(step, "Image"): tf_summary.image(path, raw_value, step=step, max_outputs=32) diff --git a/axlearn/common/trainer.py b/axlearn/common/trainer.py index fa98ff0bc..4cfed07ac 100644 --- a/axlearn/common/trainer.py +++ b/axlearn/common/trainer.py @@ -77,22 +77,19 @@ live_devices, ) -jax_device_state = {} -python_vars = {} -immutable_data = {} - def sync_restore_class_vars( fresh_trainer: Any, - jax_device_state_arg: dict = None, python_vars_arg: dict = None, immutable_data_arg: dict = None + jax_device_state_arg: dict, + python_vars_arg: dict, + immutable_data_arg: dict, ) -> tuple[Any, Any]: """Restores trainer state onto a fresh SpmdTrainer instance from snapshot.""" print("sync_restore_class_vars") - global jax_device_state, python_vars, immutable_data - print(immutable_data) + print(immutable_data_arg) - use_python_vars = python_vars_arg if python_vars_arg is not None else python_vars - use_immutable_data = immutable_data_arg if immutable_data_arg is not None else immutable_data + use_python_vars = python_vars_arg + use_immutable_data = immutable_data_arg for k, v in use_immutable_data.items(): if isinstance(v, (int, float, str, bool)): @@ -154,7 +151,7 @@ def sync_restore_class_vars( except Exception as e: logging.warning("Failed to load from snapshot: %s.", e) - use_jax_state = jax_device_state_arg if jax_device_state_arg is not None else jax_device_state + use_jax_state = jax_device_state_arg if not state_restored and use_jax_state and "_trainer_state" in use_jax_state: logging.info("[!] Attempting fallback: device_put trainer_state from globals onto the new mesh.") try: @@ -177,6 +174,8 @@ def sync_restore_class_vars( elif "_step" in use_python_vars: fresh_trainer._step = int(use_python_vars["_step"]) + if "_input_iter" in use_python_vars: + fresh_trainer._input_iter = use_python_vars["_input_iter"] fresh_trainer.snapshot_mgr = snapshot_mgr fresh_trainer._is_restored = True fresh_trainer._compiled_train_step = None @@ -215,36 +214,37 @@ def sync_restore_class_vars( return fresh_trainer, fresh_prng_key -def sync_store_class_vars(obj: Any) -> None: - """Stores instance variables of an object in global variables.""" +def sync_store_class_vars(obj: Any) -> tuple[dict, dict, dict]: + """Stores instance variables of an object in dictionaries.""" if getattr(obj, "_is_restored", False): - return + return ( + getattr(obj, "_jax_device_state", {}), + getattr(obj, "_python_vars", {}), + getattr(obj, "_immutable_data", {}), + ) print("_sync_store_class_vars") - global jax_device_state, python_vars, immutable_data - - jax_device_state.clear() - python_vars.clear() - immutable_data.clear() + jax_device_state = {} + python_vars = {} + immutable_data = {} for k, v in obj.__dict__.items(): if isinstance(v, property): continue + if k in ("_jax_device_state", "_python_vars", "_immutable_data"): + continue + if k in ("_trainer_state", "_mesh", "_jit_train_step", "_compiled_train_step", "model", "learner"): jax_device_state[k] = v elif "config" in k or "spec" in k or isinstance(v, (int, float, str, bool)): immutable_data[k] = v else: python_vars[k] = v - - #print(python_vars) - #print(immutable_data) + print("assigning class vars to trainer") snapshot_mgr = python_vars.get("snapshot_mgr") if snapshot_mgr is not None: - #print(immutable_data["_step"]) - #print(jax_device_state["_trainer_state"]) try: step_val = immutable_data.get("_step") if step_val is None: @@ -253,9 +253,11 @@ def sync_store_class_vars(obj: Any) -> None: snapshot_mgr.join() except Exception as e: logging.warning("Failed during snapshot save: %s", e) - + print("_sync_store_class_vars done") python_vars["snapshot_mgr"] = snapshot_mgr + return jax_device_state, python_vars, immutable_data + class TrainerState(NamedTuple): @@ -451,6 +453,9 @@ def __init__( ) self._step: int = None + self._jax_device_state: dict = {} + self._python_vars: dict = {} + self._immutable_data: dict = {} self._trainer_state: TrainerState = None self._jit_train_step: jax.stages.Wrapped = None self._watchdog_stopping = None @@ -856,7 +861,8 @@ def run( self.vlog(3, "Done step %s", self.step) # if self.step==3: # _sync_store_class_vars(self) - sync_store_class_vars(self) + print("lkolluru step: ",self.step) + self._jax_device_state, self._python_vars, self._immutable_data = sync_store_class_vars(self) num_steps += 1 if num_steps % 100 == 0: @@ -1116,7 +1122,11 @@ def _prepare_training(self, prng_key: Tensor) -> bool: cfg = self.config # Attempt to restore the latest checkpoint, which may contain a saved `_input_iter`. - self.restore_checkpoint(restore_step=None) + if not getattr(self, "_is_restored", False): + self.restore_checkpoint(restore_step=None) + else: + logging.info("Skipping checkpoint restoration because state was already restored from snapshot.") + self._is_restored = False if self.step is None: # If we didn't restore from checkpoint, attempt to build initial state according @@ -1360,9 +1370,10 @@ def _run_step( n = self._config.log_every_n_steps or 100 if self.step % n == 0 or 0 <= self.step <= 5: + loss_val = outputs["loss"].item() if hasattr(outputs["loss"], "item") else outputs["loss"] self._step_log( "loss=%s aux=%s", - outputs["loss"], + loss_val, jax.tree.map(lambda x: x.item() if x.ndim == 0 else f"T{x.shape}", outputs["aux"]), ) diff --git a/axlearn/common/utils.py b/axlearn/common/utils.py index 45e2c4d5f..0c69f0d38 100644 --- a/axlearn/common/utils.py +++ b/axlearn/common/utils.py @@ -1797,7 +1797,7 @@ def create_device_mesh( NotImplementedError: If not all devices have the same platform. """ if devices is None: - devices = jax.devices() + devices = live_devices() devices = np.asarray(devices) # Check if the devices are part of a multi-granule configuration. @@ -2170,3 +2170,31 @@ def get_tpu_dot_precision(dtype) -> jax.lax.Precision: if dtype == jnp.bfloat16: return jax.lax.Precision.DEFAULT raise ValueError(f"Unsupported dtype {dtype}") + + +import pathwaysutils +from pathwaysutils.elastic import manager as pathways_manager + +elastic_manager: Optional[pathways_manager.Manager] = None + + +def set_elastic_manager(manager: Any): + """Sets the global elastic manager.""" + global elastic_manager + elastic_manager = manager + + +def live_devices(): + device_list = jax.devices() + + if pathwaysutils.is_pathways_backend_used() and elastic_manager is not None: + active_devices = [ + d for d in jax.devices() if d is not None and getattr(d, "slice_index", 0) in elastic_manager.active_slice_indices + ] + if active_devices: + return sorted(active_devices, key=lambda d: (getattr(d, "slice_index", 0), getattr(d, "coords", ()))) + return device_list + + +def live_slice_indices() -> set[int]: + return {d.slice_index for d in live_devices()} diff --git a/axlearn/experiments/text/gpt/common.py b/axlearn/experiments/text/gpt/common.py index 52fc9dad6..63e698f10 100644 --- a/axlearn/experiments/text/gpt/common.py +++ b/axlearn/experiments/text/gpt/common.py @@ -776,7 +776,7 @@ def config_fn() -> InstantiableConfig: ) cfg.checkpointer.keep_every_n_steps = min(max_step, keep_every_n_steps) cfg.checkpointer.keep_last_n = 3 - cfg.summary_writer.write_every_n_steps = min(eval_every_n_steps, 100) + cfg.summary_writer.write_every_n_steps = min(eval_every_n_steps, 10) cfg.summary_writer.max_queue = 1000 if len(mesh_axis_names) != len(mesh_shape): raise ValueError( diff --git a/axlearn/experiments/text/gpt/fuji.py b/axlearn/experiments/text/gpt/fuji.py index 83dcfe3c6..a970fdd59 100644 --- a/axlearn/experiments/text/gpt/fuji.py +++ b/axlearn/experiments/text/gpt/fuji.py @@ -51,6 +51,7 @@ from axlearn.common.utils import ( combine_remat_policies, extended_checkpoint_policies, + live_devices, save_and_offload_only_these_names_regex, ) from axlearn.experiments.text.gpt.common import ( @@ -389,6 +390,7 @@ def get_trainer_kwargs( ), ) elif model_size == "7B": + import jax trainer_kwargs = dict( model_kwargs=dict( num_layers=32, @@ -401,7 +403,7 @@ def get_trainer_kwargs( ), learner_kwargs=dict(peak_lr=3e-4, weight_decay=0.1), max_sequence_length=max_sequence_length, - train_batch_size=train_batch_size, + train_batch_size=len(live_devices()),#train_batch_size, max_step=max_step, mesh_shape=mesh_shape_from_axes(data=-1, fsdp=8), mesh_rules=( @@ -434,6 +436,35 @@ def get_trainer_kwargs( ], ), ), + ( + "tpu-v5litepod-32", + ChainConfigModifier.default_config().set( + config_modifiers=[ + MeshShapeModifier.default_config().set( + mesh_shape=mesh_shape_from_axes(fsdp=32, data=-1) + ), + RematSpecModifier.default_config().set( + remat_policies={ + "model.decoder.transformer.layer": RematSpec( + prevent_cse=False, + policy=config_for_function( + save_and_offload_only_these_names_regex + ).set( + names_which_can_be_saved=( + RematRegexSavePatterns.QKV_PROJ.value + ), + names_which_can_be_offloaded=( + RematRegexSavePatterns.INPUT.value + ), + offload_src="device", + offload_dst="pinned_host", + ), + ), + } + ), + ], + ), + ), ( "tpu-v5litepod-256-2", ChainConfigModifier.default_config().set( diff --git a/pyproject.toml b/pyproject.toml index 5d3adc43a..ae804b283 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -100,7 +100,7 @@ gcp = [ tpu = [ "axlearn[gcp]", "jax[tpu]==0.8.3", # must be >=0.4.19 for compat with v5p. - "pathwaysutils==0.1.2", # For JAX+Pathways single-controller accelerator coordinator. + "pathwaysutils==0.1.8", # For JAX+Pathways single-controller accelerator coordinator. ] # Vertex AI tensorboard. TODO(markblee): Merge with `gcp`. vertexai_tensorboard = [ From f7bf0e4a0ccbc76be454669258e843f590636069 Mon Sep 17 00:00:00 2001 From: lkolluru05 Date: Thu, 9 Jul 2026 01:36:58 +0000 Subject: [PATCH 4/6] mesh issues fixed , oom coming --- axlearn/common/launch_trainer.py | 16 +++++ axlearn/common/trainer.py | 107 ++++++++++++++++++++----------- axlearn/common/utils.py | 16 +++-- 3 files changed, 97 insertions(+), 42 deletions(-) diff --git a/axlearn/common/launch_trainer.py b/axlearn/common/launch_trainer.py index ca3ae28db..e78f6c57c 100644 --- a/axlearn/common/launch_trainer.py +++ b/axlearn/common/launch_trainer.py @@ -5,6 +5,7 @@ import json import os import time +import gc from typing import Any, Optional import jax @@ -267,6 +268,21 @@ def run_trainer(trainer_config: SpmdTrainer.Config) -> Any: jax_device_state = getattr(trainer, "_jax_device_state", {}) python_vars = getattr(trainer, "_python_vars", {}) immutable_data = getattr(trainer, "_immutable_data", {}) + + # jax_device_state.pop("_mesh", None) + # # Free massive XLA executables and module caches from device memory + # jax_device_state.pop("_compiled_train_step", None) + # jax_device_state.pop("_jit_train_step", None) + # jax_device_state.pop("model", None) + # jax_device_state.pop("learner", None) + + # # Clear old trainer objects and JAX caches to release TPU memory. + # # We keep the extracted state dicts above to restore onto the new mesh. + # del trainer + # del clean_trainer + # jax.clear_caches() + # gc.collect() + if elastic_manager: elastic_manager.new_slice_event.set() time.sleep(10) diff --git a/axlearn/common/trainer.py b/axlearn/common/trainer.py index 4cfed07ac..409d0f603 100644 --- a/axlearn/common/trainer.py +++ b/axlearn/common/trainer.py @@ -101,43 +101,7 @@ def sync_restore_class_vars( except Exception: pass - live_devs = utils.live_devices() - logging.info("[!] Found %d available live devices.", len(live_devs)) - - cfg = fresh_trainer.config - device_platform = live_devs[0].platform - device_attr = "process_index" if device_platform != "tpu" else "slice_index" - - num_granules = len({getattr(el, device_attr) for el in live_devs}) - num_devices_per_granule = len(live_devs) // num_granules - - original_mesh_shape = list(cfg.mesh_shape) - if len(original_mesh_shape) > 0: - original_mesh_shape[0] = num_granules - - ici_shape = original_mesh_shape[1:] - current_ici_prod = math.prod(ici_shape) - if current_ici_prod != num_devices_per_granule: - logging.info("[!] ICI product %d does not match num_devices_per_granule %d. Adjusting...", current_ici_prod, num_devices_per_granule) - ratio = current_ici_prod // num_devices_per_granule - if ratio > 0 and current_ici_prod % num_devices_per_granule == 0: - for i in range(len(ici_shape)): - if ici_shape[i] % ratio == 0 and ici_shape[i] > 1: - ici_shape[i] = ici_shape[i] // ratio - break - else: - ici_shape = [1] * len(ici_shape) - ici_shape[-3] = num_devices_per_granule - - original_mesh_shape[1:] = ici_shape - - updated_mesh_shape = tuple(original_mesh_shape) - - logging.info("[!] Dynamically updating logical mesh_shape from %s to %s", cfg.mesh_shape, updated_mesh_shape) - - devices_mesh = utils.create_device_mesh(mesh_shape=updated_mesh_shape, devices=live_devs) - mesh = jax.sharding.Mesh(devices_mesh, cfg.mesh_axis_names) - fresh_trainer._mesh = mesh + mesh = fresh_trainer._mesh state_restored = False snapshot_mgr = use_python_vars.get("snapshot_mgr") @@ -495,9 +459,76 @@ def __init__( len(local_devices), [device.platform for device in local_devices], ) + + live_devs = utils.live_devices() if devices is None else devices + device_platform = live_devs[0].platform + device_attr = "process_index" if device_platform != "tpu" else "slice_index" + num_granules = len(set(getattr(el, device_attr) for el in live_devs)) + num_devices_per_granule = len(live_devs) // num_granules + + if isinstance(cfg.mesh_shape, Sequence) and not isinstance(cfg.mesh_shape, str): + original_mesh_shape = list(cfg.mesh_shape) + if len(original_mesh_shape) > 0: + original_mesh_shape[0] = num_granules + + ici_shape = original_mesh_shape[1:] + current_ici_prod = math.prod(ici_shape) + if current_ici_prod != num_devices_per_granule: + logging.info("[!] ICI product %d does not match num_devices_per_granule %d. Adjusting...", current_ici_prod, num_devices_per_granule) + ratio = current_ici_prod // num_devices_per_granule + if ratio > 0 and current_ici_prod % num_devices_per_granule == 0: + for i in range(len(ici_shape)): + if ici_shape[i] % ratio == 0 and ici_shape[i] > 1: + ici_shape[i] = ici_shape[i] // ratio + break + else: + ici_shape = [1] * len(ici_shape) + if len(ici_shape) >= 3: + ici_shape[-3] = num_devices_per_granule + else: + ici_shape[-1] = num_devices_per_granule + + original_mesh_shape[1:] = ici_shape + + cfg.mesh_shape = tuple(original_mesh_shape) + logging.info("[!] Dynamically updating logical mesh_shape to %s", cfg.mesh_shape) + elif isinstance(cfg.mesh_shape, HybridMeshShape): + dcn_shape = list(cfg.mesh_shape.dcn_mesh_shape) + ici_shape = list(cfg.mesh_shape.ici_mesh_shape) + + dcn_prod = math.prod(dcn_shape) + if dcn_prod != num_granules: + ratio = dcn_prod // num_granules + if ratio > 0 and dcn_prod % num_granules == 0: + for i in range(len(dcn_shape)): + if dcn_shape[i] % ratio == 0 and dcn_shape[i] > 1: + dcn_shape[i] = dcn_shape[i] // ratio + break + else: + dcn_shape = [1] * len(dcn_shape) + dcn_shape[0] = num_granules + + current_ici_prod = math.prod(ici_shape) + if current_ici_prod != num_devices_per_granule: + ratio = current_ici_prod // num_devices_per_granule + if ratio > 0 and current_ici_prod % num_devices_per_granule == 0: + for i in range(len(ici_shape)): + if ici_shape[i] % ratio == 0 and ici_shape[i] > 1: + ici_shape[i] = ici_shape[i] // ratio + break + else: + ici_shape = [1] * len(ici_shape) + if len(ici_shape) >= 3: + ici_shape[-3] = num_devices_per_granule + else: + ici_shape[-1] = num_devices_per_granule + + cfg.mesh_shape = HybridMeshShape(ici_mesh_shape=tuple(ici_shape), dcn_mesh_shape=tuple(dcn_shape)) + logging.info("[!] Dynamically updating HybridMeshShape to %s", cfg.mesh_shape) + self._step_log("Mesh shape: %s", cfg.mesh_shape) devices = ( - utils.create_device_mesh(mesh_shape=cfg.mesh_shape) if devices is None else devices + utils.create_device_mesh(mesh_shape=cfg.mesh_shape, devices=live_devs) if devices is None else devices ) mesh = jax.sharding.Mesh(devices, cfg.mesh_axis_names) self._step_log("Global mesh: %s", mesh) diff --git a/axlearn/common/utils.py b/axlearn/common/utils.py index 0c69f0d38..686df5808 100644 --- a/axlearn/common/utils.py +++ b/axlearn/common/utils.py @@ -1809,7 +1809,7 @@ def create_device_mesh( raise NotImplementedError(f"Not all devices had platform: {device_platform}.") num_granules = ( - max(getattr(el, device_attr) for el in devices.flatten()) + 1 if is_multi_granule_env else 1 + len(set(getattr(el, device_attr) for el in devices.flatten())) if is_multi_granule_env else 1 ) num_devices = len(devices) assert num_devices % num_granules == 0, ( @@ -2183,16 +2183,24 @@ def set_elastic_manager(manager: Any): global elastic_manager elastic_manager = manager +from pathwaysutils.elastic import manager def live_devices(): device_list = jax.devices() - + elastic_manager = manager.Manager() + print("lkolluru pathwaysutils.is_pathways_backend_used(): ", pathwaysutils.is_pathways_backend_used()) + print("lkolluru elastic_manager: ", elastic_manager) if pathwaysutils.is_pathways_backend_used() and elastic_manager is not None: + import time + time.sleep(5) active_devices = [ - d for d in jax.devices() if d is not None and getattr(d, "slice_index", 0) in elastic_manager.active_slice_indices - ] + d for d in jax.devices() if d is not None and getattr(d, "slice_index", 0) in elastic_manager.active_slice_indices + ] + print("lkolluru active_devices: ", len(active_devices)) if active_devices: return sorted(active_devices, key=lambda d: (getattr(d, "slice_index", 0), getattr(d, "coords", ()))) + logging.info("Waiting for active_slice_indices to be populated...") + return device_list From 2f4376f76017dc7a80ffa0dd0b191ed487ff652d Mon Sep 17 00:00:00 2001 From: lkolluru05 Date: Thu, 9 Jul 2026 21:23:19 +0000 Subject: [PATCH 5/6] elastic snapshot with fuji-test --- axlearn/common/launch_trainer.py | 20 ++++++--- axlearn/common/snapshot.py | 24 +++++----- axlearn/common/trainer.py | 39 +++++++++-------- axlearn/common/utils.py | 10 ++--- axlearn/experiments/text/gpt/fuji.py | 65 +++++++++++++++++++++++++++- 5 files changed, 116 insertions(+), 42 deletions(-) diff --git a/axlearn/common/launch_trainer.py b/axlearn/common/launch_trainer.py index e78f6c57c..9e6eafb51 100644 --- a/axlearn/common/launch_trainer.py +++ b/axlearn/common/launch_trainer.py @@ -234,36 +234,41 @@ def run_trainer(trainer_config: SpmdTrainer.Config) -> Any: python_vars = {} immutable_data = {} trainer = None + logging.info("[ELASTIC] Starting elastic training loop cycle.") while True: try: if not elastic_manager_initialized: if elastic_snapshotting_enabled: - logging.info("[Elastic] Initializing elastic manager...") + logging.info("[ELASTIC] Initializing elastic manager...") elastic_manager = manager.Manager() set_elastic_manager(elastic_manager) - logging.info("[Elastic] Elastic manager initialized.") + logging.info("[ELASTIC] Elastic manager initialized.") else: - logging.info("[Elastic] Elastic snapshotting disabled or not supported (no slice_index).") + logging.info("[ELASTIC] Elastic snapshotting disabled or not supported (no slice_index).") elastic_manager_initialized = True clean_trainer: SpmdTrainer = trainer_config.instantiate(parent=None) + logging.info("[ELASTIC] Instantiated clean trainer.") if elastic_manager and elastic_manager.new_slice_event.is_set(): - logging.info("[Elastic] New slice event is set. Restoring from snapshot...") + logging.info("[ELASTIC] New slice event is set. Restoring from snapshot...") elastic_manager.new_slice_event.clear() trainer, prng_key = sync_restore_class_vars(clean_trainer, jax_device_state, python_vars, immutable_data) + logging.info("[ELASTIC] Restored trainer state from class vars.") else: - logging.info("[Elastic] Starting fresh trainer (no elastic recovery triggered).") + logging.info("[ELASTIC] Starting fresh trainer (no elastic recovery triggered).") trainer = clean_trainer prng_key = jax.random.PRNGKey(seed=FLAGS.trainer_prng_seed) + logging.info("[ELASTIC] Starting trainer.run().") output = trainer.run(prng_key) + logging.info("[ELASTIC] trainer.run() completed.") measurement.record_event(measurement.Event.END_JOB) break except jax.errors.JaxRuntimeError as e: if is_retryable_error(e): - logging.warning("Caught retryable error: %s. Retrying...", e) + logging.warning("[ELASTIC] Caught retryable error: %s. Retrying...", e) if trainer is not None: jax_device_state = getattr(trainer, "_jax_device_state", {}) python_vars = getattr(trainer, "_python_vars", {}) @@ -288,5 +293,6 @@ def run_trainer(trainer_config: SpmdTrainer.Config) -> Any: time.sleep(10) continue else: + logging.error("[ELASTIC] Caught non-retryable error: %s", e) raise e - return output + return output \ No newline at end of file diff --git a/axlearn/common/snapshot.py b/axlearn/common/snapshot.py index d23b918bf..14ae3e9e3 100644 --- a/axlearn/common/snapshot.py +++ b/axlearn/common/snapshot.py @@ -34,15 +34,15 @@ def __init__(self, *, replica_axis_index: int = 0, trainer_state_specs: Optional def _worker(self): while True: pinned_state, step = self._queue.get() - print("In snapshot worker") + _logger.info("[ELASTIC] Snapshot worker dequeued task for step %d", step) try: _logger.info( - "[*] [Snapshot Thread] Waiting for snapshot at step %d to be ready...", + "[ELASTIC] [*] [Snapshot Thread] Waiting for snapshot at step %d to be ready...", step, ) jax.block_until_ready(pinned_state) _logger.info( - "[*] [Snapshot Thread] Snapshot at step %d is ready and secured.", + "[ELASTIC] [*] [Snapshot Thread] Snapshot at step %d is ready and secured.", step, ) with self._lock: @@ -54,30 +54,30 @@ def _worker(self): except Exception: err_msg = f"JAX Runtime Exception of type {type(e).__name__} (suppressed tensor evaluation)" _logger.warning( - "[*] [Snapshot Thread] Failed to secure snapshot at step %d: %s.", + "[ELASTIC] [*] [Snapshot Thread] Failed to secure snapshot at step %d: %s.", step, err_msg, ) finally: - print("In snapshot worker finally") + _logger.info("[ELASTIC] Snapshot worker finished processing step %d", step) self._queue.task_done() def save_pytree( self, step: int, state: tree_types.PyTreeOf[jax.Array] ) -> None: """Move arrays onto CPU worker devices.""" - print("In snapshot pytree") + _logger.info("[ELASTIC] Starting snapshot process for step %d", step) if self._queue.full(): - _logger.warning("Snapshotter busy. Skipping snapshot for step %d", step) + _logger.warning("[ELASTIC] Snapshotter busy. Skipping snapshot for step %d", step) return pinned_shardings = jax.tree.map( lambda x: x.sharding.with_memory_kind("pinned_host") if hasattr(x, "sharding") else None, state ) - print("before jax.put") + _logger.info("[ELASTIC] Putting snapshot state to pinned host memory for step %d...", step) pinned_state = jax.device_put(state, pinned_shardings) - print("after jax.put") + _logger.info("[ELASTIC] Snapshot state put to pinned host memory for step %d.", step) self._queue.put((pinned_state, step)) def load_pytree( @@ -130,7 +130,7 @@ def get_active_pytree(x, target_x): sliced_x = jnp.pad(sliced_x, pad_widths) return sliced_x - _logger.info("Restoring from snapshot at step %d...", step) + _logger.info("[ELASTIC] Restoring from snapshot at step %d...", step) pinned_state = jax.tree.map(get_active_pytree, pinned_state, abstract_state) # Re-shard on host to the target device mesh @@ -138,15 +138,19 @@ def get_active_pytree(x, target_x): lambda x: x.sharding.with_memory_kind("pinned_host") if hasattr(x, "sharding") else None, abstract_state ) + _logger.info("[ELASTIC] Moving snapshot from pinned host to target host sharding...") host_target_state = jax.device_put( pinned_state, host_target_shardings ) + _logger.info("[ELASTIC] Snapshot moved to target host sharding.") # Move from host back to device (TPU) memory. + _logger.info("[ELASTIC] Moving snapshot from host to device...") restored_state = jax.device_put( host_target_state, jax.tree.map(lambda x: x.sharding if hasattr(x, "sharding") else None, abstract_state) ) jax.block_until_ready(restored_state) + _logger.info("[ELASTIC] Snapshot moved to device.") if reset_snapshot_state: with self._lock: diff --git a/axlearn/common/trainer.py b/axlearn/common/trainer.py index 409d0f603..12de0a822 100644 --- a/axlearn/common/trainer.py +++ b/axlearn/common/trainer.py @@ -84,9 +84,9 @@ def sync_restore_class_vars( immutable_data_arg: dict, ) -> tuple[Any, Any]: """Restores trainer state onto a fresh SpmdTrainer instance from snapshot.""" - print("sync_restore_class_vars") + logging.info("[ELASTIC] Restoring class variables from snapshot.") - print(immutable_data_arg) + logging.info("[ELASTIC] Immutable data args: %s", immutable_data_arg) use_python_vars = python_vars_arg use_immutable_data = immutable_data_arg @@ -110,14 +110,14 @@ def sync_restore_class_vars( try: restored_trainer_state = snapshot_mgr.load_pytree() fresh_trainer._trainer_state = restored_trainer_state - logging.info("Successfully restored state from snapshot onto the new mesh.") + logging.info("[ELASTIC] Successfully restored state from snapshot onto the new mesh.") state_restored = True except Exception as e: - logging.warning("Failed to load from snapshot: %s.", e) + logging.warning("[ELASTIC] Failed to load from snapshot: %s.", e) use_jax_state = jax_device_state_arg if not state_restored and use_jax_state and "_trainer_state" in use_jax_state: - logging.info("[!] Attempting fallback: device_put trainer_state from globals onto the new mesh.") + logging.info("[ELASTIC] [!] Attempting fallback: device_put trainer_state from globals onto the new mesh.") try: with mesh: fresh_trainer._trainer_state = jax.tree_util.tree_map( @@ -125,13 +125,13 @@ def sync_restore_class_vars( use_jax_state["_trainer_state"], fresh_trainer._trainer_state_specs ) - logging.info("[✓] Successfully device_put trainer_state from globals onto the new mesh.") + logging.info("[ELASTIC] [✓] Successfully device_put trainer_state from globals onto the new mesh.") state_restored = True except Exception as e: - logging.warning("[!] Failed fallback to globals (possibly deleted arrays): %s", e) + logging.warning("[ELASTIC] [!] Failed fallback to globals (possibly deleted arrays): %s", e) if not state_restored: - logging.info("[!] Falling back to fresh init().") + logging.info("[ELASTIC] [!] Falling back to fresh init().") fresh_trainer.init(jax.random.PRNGKey(seed=42)) if "_step" in use_immutable_data: fresh_trainer._step = int(use_immutable_data["_step"]) @@ -158,9 +158,9 @@ def sync_restore_class_vars( fresh_trainer._trainer_state["prng_key"] = fresh_prng_key else: setattr(fresh_trainer._trainer_state, "prng_key", fresh_prng_key) - logging.info("[✓] Successfully injected fresh, healthy PRNG Key into the new trainer state.") + logging.info("[ELASTIC] [✓] Successfully injected fresh, healthy PRNG Key into the new trainer state.") except Exception as e: - logging.warning("[!] Failed to replace prng_key inside trainer_state structure: %s", e) + logging.warning("[ELASTIC] [!] Failed to replace prng_key inside trainer_state structure: %s", e) try: leaves = jax.tree_util.tree_leaves(fresh_trainer._trainer_state) @@ -187,7 +187,7 @@ def sync_store_class_vars(obj: Any) -> tuple[dict, dict, dict]: getattr(obj, "_immutable_data", {}), ) - print("_sync_store_class_vars") + logging.info("[ELASTIC] Storing class variables for snapshot.") jax_device_state = {} python_vars = {} @@ -206,7 +206,7 @@ def sync_store_class_vars(obj: Any) -> tuple[dict, dict, dict]: immutable_data[k] = v else: python_vars[k] = v - print("assigning class vars to trainer") + logging.info("[ELASTIC] Preparing to save snapshot.") snapshot_mgr = python_vars.get("snapshot_mgr") if snapshot_mgr is not None: try: @@ -216,8 +216,8 @@ def sync_store_class_vars(obj: Any) -> tuple[dict, dict, dict]: snapshot_mgr.save_pytree(step=int(step_val) if step_val is not None else 0, state=jax_device_state["_trainer_state"]) snapshot_mgr.join() except Exception as e: - logging.warning("Failed during snapshot save: %s", e) - print("_sync_store_class_vars done") + logging.warning("[ELASTIC] Failed during snapshot save: %s", e) + logging.info("[ELASTIC] Storing class variables done.") python_vars["snapshot_mgr"] = snapshot_mgr return jax_device_state, python_vars, immutable_data @@ -851,10 +851,11 @@ def run( replica_axis_idx = cfg.mesh_axis_names.index("data") if "data" in cfg.mesh_axis_names else 0 snapshot_cfg = config_for_class(Snapshotter).set(replica_axis_index=replica_axis_idx, trainer_state_specs=self.trainer_state_specs) self.snapshot_mgr = snapshot_cfg.instantiate() + logging.info("[ELASTIC] Snapshot manager instantiated.") with self.checkpointer: - logging.info("Starting loop...") + logging.info("[ELASTIC] Starting loop...") start_time = time.perf_counter() num_steps = 0 output = None @@ -876,7 +877,8 @@ def run( stop_trace_step = self._maybe_stop_or_start_tracing(stop_trace_step, output) self._step = int(self._step) + 1 - self.vlog(3, "Start step %s", self.step) + self._step_log("[ELASTIC] Start step") + logging.info("[ELASTIC] Start step %s", self.step) self._maybe_record_event(measurement.Event.START_STEP, self._step) output = self._run_step( utils.host_to_global_array( @@ -889,10 +891,10 @@ def run( else None ), ) - self.vlog(3, "Done step %s", self.step) + self._step_log("[ELASTIC] Done step") + logging.info("[ELASTIC] Done step %s", self.step) # if self.step==3: # _sync_store_class_vars(self) - print("lkolluru step: ",self.step) self._jax_device_state, self._python_vars, self._immutable_data = sync_store_class_vars(self) num_steps += 1 @@ -907,6 +909,7 @@ def run( self._step_log("Reached max_step=%s. Stopping", cfg.max_step) break except StopIteration: + logging.info("[ELASTIC] Reached end of input iterator.") # Add END_DATA_LOADING event here to close the unpaired START_DATA_LOADING # event. self._maybe_record_event(measurement.Event.END_DATA_LOADING) diff --git a/axlearn/common/utils.py b/axlearn/common/utils.py index 686df5808..1eb7b0737 100644 --- a/axlearn/common/utils.py +++ b/axlearn/common/utils.py @@ -2188,21 +2188,21 @@ def set_elastic_manager(manager: Any): def live_devices(): device_list = jax.devices() elastic_manager = manager.Manager() - print("lkolluru pathwaysutils.is_pathways_backend_used(): ", pathwaysutils.is_pathways_backend_used()) - print("lkolluru elastic_manager: ", elastic_manager) + logging.info("[ELASTIC] pathwaysutils.is_pathways_backend_used(): %s", pathwaysutils.is_pathways_backend_used()) + logging.info("[ELASTIC] elastic_manager: %s", elastic_manager) if pathwaysutils.is_pathways_backend_used() and elastic_manager is not None: import time time.sleep(5) active_devices = [ d for d in jax.devices() if d is not None and getattr(d, "slice_index", 0) in elastic_manager.active_slice_indices ] - print("lkolluru active_devices: ", len(active_devices)) + logging.info("[ELASTIC] active_devices count: %d", len(active_devices)) if active_devices: return sorted(active_devices, key=lambda d: (getattr(d, "slice_index", 0), getattr(d, "coords", ()))) - logging.info("Waiting for active_slice_indices to be populated...") + logging.info("[ELASTIC] Waiting for active_slice_indices to be populated...") return device_list def live_slice_indices() -> set[int]: - return {d.slice_index for d in live_devices()} + return {d.slice_index for d in live_devices()} \ No newline at end of file diff --git a/axlearn/experiments/text/gpt/fuji.py b/axlearn/experiments/text/gpt/fuji.py index a970fdd59..a45e69a7e 100644 --- a/axlearn/experiments/text/gpt/fuji.py +++ b/axlearn/experiments/text/gpt/fuji.py @@ -403,7 +403,7 @@ def get_trainer_kwargs( ), learner_kwargs=dict(peak_lr=3e-4, weight_decay=0.1), max_sequence_length=max_sequence_length, - train_batch_size=len(live_devices()),#train_batch_size, + train_batch_size=len(jax.devices()),#train_batch_size, max_step=max_step, mesh_shape=mesh_shape_from_axes(data=-1, fsdp=8), mesh_rules=( @@ -441,7 +441,67 @@ def get_trainer_kwargs( ChainConfigModifier.default_config().set( config_modifiers=[ MeshShapeModifier.default_config().set( - mesh_shape=mesh_shape_from_axes(fsdp=32, data=-1) + mesh_shape=mesh_shape_from_axes(fsdp=16, model=2, data=-1) + ), + RematSpecModifier.default_config().set( + remat_policies={ + "model.decoder.transformer.layer": RematSpec( + prevent_cse=True, + policy=config_for_function( + save_and_offload_only_these_names_regex + ).set( + names_which_can_be_saved=( + RematRegexSavePatterns.QKV_PROJ.value + ), + names_which_can_be_offloaded=( + RematRegexSavePatterns.INPUT.value + ), + offload_src="device", + offload_dst="pinned_host", + ), + ), + } + ), + FlashBlockSizeModifier.default_config().set(tpu_block_size=256), + ], + ), + ), + ( + "tpu-v5litepod-64", + ChainConfigModifier.default_config().set( + config_modifiers=[ + MeshShapeModifier.default_config().set( + mesh_shape=mesh_shape_from_axes(fsdp=32, model=2, data=-1) + ), + RematSpecModifier.default_config().set( + remat_policies={ + "model.decoder.transformer.layer": RematSpec( + prevent_cse=True, + policy=config_for_function( + save_and_offload_only_these_names_regex + ).set( + names_which_can_be_saved=( + RematRegexSavePatterns.QKV_PROJ.value + ), + names_which_can_be_offloaded=( + RematRegexSavePatterns.INPUT.value + ), + offload_src="device", + offload_dst="pinned_host", + ), + ), + } + ), + FlashBlockSizeModifier.default_config().set(tpu_block_size=256), + ], + ), + ), + ( + "tpu-v5p-128", + ChainConfigModifier.default_config().set( + config_modifiers=[ + MeshShapeModifier.default_config().set( + mesh_shape=mesh_shape_from_axes(fsdp=64, data=-1) ), RematSpecModifier.default_config().set( remat_policies={ @@ -462,6 +522,7 @@ def get_trainer_kwargs( ), } ), + #FlashBlockSizeModifier.default_config().set(tpu_block_size=256), ], ), ), From 6ecadeedbc3345c6943d1cd33ad34a157d975383 Mon Sep 17 00:00:00 2001 From: lkolluru05 Date: Thu, 9 Jul 2026 23:27:18 +0000 Subject: [PATCH 6/6] fuji 1B elastic training snapshot works --- axlearn/common/trainer.py | 37 +++++++++++++++------------- axlearn/experiments/text/gpt/fuji.py | 35 ++++++++++++++++++++++++-- 2 files changed, 53 insertions(+), 19 deletions(-) diff --git a/axlearn/common/trainer.py b/axlearn/common/trainer.py index 12de0a822..7b9a236a0 100644 --- a/axlearn/common/trainer.py +++ b/axlearn/common/trainer.py @@ -107,13 +107,17 @@ def sync_restore_class_vars( snapshot_mgr = use_python_vars.get("snapshot_mgr") if snapshot_mgr is not None: with mesh: - try: - restored_trainer_state = snapshot_mgr.load_pytree() - fresh_trainer._trainer_state = restored_trainer_state - logging.info("[ELASTIC] Successfully restored state from snapshot onto the new mesh.") - state_restored = True - except Exception as e: - logging.warning("[ELASTIC] Failed to load from snapshot: %s.", e) + for attempt in range(3): + try: + restored_trainer_state = snapshot_mgr.load_pytree() + fresh_trainer._trainer_state = restored_trainer_state + logging.info("[ELASTIC] Successfully restored state from snapshot onto the new mesh.") + state_restored = True + break + except Exception as e: + logging.warning("[ELASTIC] Failed to load from snapshot (attempt %d/3): %s.", attempt + 1, e) + if attempt < 2: + time.sleep(5) use_jax_state = jax_device_state_arg if not state_restored and use_jax_state and "_trainer_state" in use_jax_state: @@ -165,7 +169,7 @@ def sync_restore_class_vars( try: leaves = jax.tree_util.tree_leaves(fresh_trainer._trainer_state) deleted_count = sum(x.is_deleted() for x in leaves if isinstance(x, jax.Array)) - mesh_match = all(x.sharding.mesh == mesh for x in leaves if isinstance(x, jax.Array) and hasattr(x, "sharding")) + mesh_match = all(getattr(x.sharding, "mesh", None) == mesh for x in leaves if isinstance(x, jax.Array) and hasattr(x, "sharding")) logging.info( "[Diagnostic] Trainer State Check: step=%s, deleted_arrays=%d, mesh_match_verified=%s", @@ -848,7 +852,7 @@ def run( self._is_initialized = True #### Stores the initial state of all variables #### - replica_axis_idx = cfg.mesh_axis_names.index("data") if "data" in cfg.mesh_axis_names else 0 + replica_axis_idx = cfg.mesh_axis_names.index("fsdp") if "fsdp" in cfg.mesh_axis_names else 0 snapshot_cfg = config_for_class(Snapshotter).set(replica_axis_index=replica_axis_idx, trainer_state_specs=self.trainer_state_specs) self.snapshot_mgr = snapshot_cfg.instantiate() logging.info("[ELASTIC] Snapshot manager instantiated.") @@ -1326,23 +1330,19 @@ def _get_compiled_train_step_fn( cfg: SpmdTrainer.Config = self.config # Get device kinds and assert that they are homogenous. # TODO(markblee): Get devices from self._mesh.devices. + print("entered _get_compiled_train_step_fn") device_kinds = set(d.device_kind for d in live_devices()) if len(device_kinds) != 1: raise RuntimeError(f"Heterogenous device kinds ({device_kinds}) are not supported.") device_kind = device_kinds.pop() - mesh_shape = cfg.mesh_shape - if isinstance(mesh_shape, HybridMeshShape): - # Combine dcn_mesh_shape and ici_mesh_shape. - dcn_mesh_shape = mesh_shape.dcn_mesh_shape - ici_mesh_shape = mesh_shape.ici_mesh_shape - assert len(dcn_mesh_shape) == len(ici_mesh_shape) - mesh_shape = tuple(x * y for x, y in zip(dcn_mesh_shape, ici_mesh_shape)) - + mesh_shape = tuple(self._mesh.shape[name] for name in cfg.mesh_axis_names) + print("_get_compiled_train_step_fn mesh_shape", mesh_shape) options = infer_xla_performance_flags( mesh_shape=mesh_shape, mesh_axis_names=cfg.mesh_axis_names, device_kind=device_kind ) logging.log_first_n(logging.INFO, "Compiler options: %s", 1, options) + print("options", options) if not with_xsc: self._maybe_record_event( measurement.Event.START_CUSTOM_BADPUT_EVENT, @@ -1357,6 +1357,7 @@ def _get_compiled_train_step_fn( ) return self._compiled_train_step logging.log_first_n(logging.INFO, "Compiling XSC train step.", 1) + print("Compiling XSC train step.") self._maybe_record_event( measurement.Event.START_CUSTOM_BADPUT_EVENT, @@ -1374,6 +1375,7 @@ def _get_compiled_train_step_fn( measurement.Event.END_CUSTOM_BADPUT_EVENT, custom_badput_event_type="COMPILATION_WITH_XSC", ) + print("compiled_jit_train_step_fn") return compiled_jit_train_step_fn def _run_step( @@ -1393,6 +1395,7 @@ def _run_step( force run the evalers in the set and return 'evaler_summaries' output. """ logging.log_first_n(logging.INFO, "global_input_batch=%s", 3, utils.shapes(input_batch)) + print("entered _run_step") with jax.profiler.StepTraceAnnotation("train", step_num=self.step): run_with_xsc = self._xsc_check_policy and self._xsc_check_policy(self.step) compiled_train_step_fn = self._get_compiled_train_step_fn( diff --git a/axlearn/experiments/text/gpt/fuji.py b/axlearn/experiments/text/gpt/fuji.py index a45e69a7e..cef2add29 100644 --- a/axlearn/experiments/text/gpt/fuji.py +++ b/axlearn/experiments/text/gpt/fuji.py @@ -319,6 +319,7 @@ def get_trainer_kwargs( mesh_shape=mesh_shape_from_axes(data=-1), ) elif model_size == "1B": + import jax trainer_kwargs = dict( model_kwargs=dict( num_layers=16, @@ -332,7 +333,7 @@ def get_trainer_kwargs( ), learner_kwargs=dict(peak_lr=3e-4, weight_decay=0.1), max_sequence_length=max_sequence_length, - train_batch_size=train_batch_size, + train_batch_size=len(jax.devices()), #train_batch_size, max_step=max_step, mesh_shape=mesh_shape_from_axes(data=-1, fsdp=8), mesh_rules=( @@ -350,6 +351,36 @@ def get_trainer_kwargs( ], ), ), + ( + "tpu-v5litepod-32", + ChainConfigModifier.default_config().set( + config_modifiers=[ + MeshShapeModifier.default_config().set( + mesh_shape=mesh_shape_from_axes(fsdp=32, model=1, data=-1) + ), + RematSpecModifier.default_config().set( + remat_policies={ + "model.decoder.transformer.layer": RematSpec( + prevent_cse=True, + policy=config_for_function( + save_and_offload_only_these_names_regex + ).set( + names_which_can_be_saved=( + RematRegexSavePatterns.QKV_PROJ.value + ), + names_which_can_be_offloaded=( + RematRegexSavePatterns.INPUT.value + ), + offload_src="device", + offload_dst="pinned_host", + ), + ), + } + ), + FlashBlockSizeModifier.default_config().set(tpu_block_size=256), + ], + ), + ), ), ) elif model_size == "3B": @@ -441,7 +472,7 @@ def get_trainer_kwargs( ChainConfigModifier.default_config().set( config_modifiers=[ MeshShapeModifier.default_config().set( - mesh_shape=mesh_shape_from_axes(fsdp=16, model=2, data=-1) + mesh_shape=mesh_shape_from_axes(fsdp=32, model=1, data=1) ), RematSpecModifier.default_config().set( remat_policies={