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..9e6eafb51 100644 --- a/axlearn/common/launch_trainer.py +++ b/axlearn/common/launch_trainer.py @@ -4,6 +4,8 @@ import json import os +import time +import gc from typing import Any, Optional import jax @@ -12,8 +14,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 +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 + # Trainer-specific flags. flags.DEFINE_string( @@ -116,6 +121,8 @@ FLAGS = flags.FLAGS +elastic_snapshotting_enabled = True + def get_trainer_config( trainer_config_fn: Optional[TrainerConfigFn] = None, @@ -148,7 +155,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] @@ -189,6 +197,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() @@ -207,9 +225,74 @@ def run_trainer(trainer_config: SpmdTrainer.Config) -> Any: }, f, ) + + elastic_manager = None + elastic_manager_initialized = False + + output = None + jax_device_state = {} + 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...") + 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) + 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...") + 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).") + 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("[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", {}) + immutable_data = getattr(trainer, "_immutable_data", {}) - 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 + # 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) + continue + else: + logging.error("[ELASTIC] Caught non-retryable error: %s", e) + raise e + return output \ No newline at end of file diff --git a/axlearn/common/snapshot.py b/axlearn/common/snapshot.py new file mode 100644 index 000000000..14ae3e9e3 --- /dev/null +++ b/axlearn/common/snapshot.py @@ -0,0 +1,176 @@ +# Copyright © 2024 Apple Inc. + +"""Manages asynchronous backups of JAX array states to pinned host memory.""" + +from absl 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 + + +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() + _logger.info("[ELASTIC] Snapshot worker dequeued task for step %d", step) + try: + _logger.info( + "[ELASTIC] [*] [Snapshot Thread] Waiting for snapshot at step %d to be ready...", + step, + ) + jax.block_until_ready(pinned_state) + _logger.info( + "[ELASTIC] [*] [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 + 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( + "[ELASTIC] [*] [Snapshot Thread] Failed to secure snapshot at step %d: %s.", + step, + err_msg, + ) + 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.""" + _logger.info("[ELASTIC] Starting snapshot process for step %d", step) + if self._queue.full(): + _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 + ) + _logger.info("[ELASTIC] Putting snapshot state to pinned host memory for step %d...", step) + + pinned_state = jax.device_put(state, pinned_shardings) + _logger.info("[ELASTIC] Snapshot state put to pinned host memory for step %d.", step) + 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("[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 + host_target_shardings = jax.tree.map( + 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: + 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/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 05df2816b..7b9a236a0 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,159 @@ host_to_global_specs, match_regex_rules, thread_stack_traces, + live_devices, ) +def sync_restore_class_vars( + fresh_trainer: Any, + 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.""" + logging.info("[ELASTIC] Restoring class variables from snapshot.") + + logging.info("[ELASTIC] Immutable data args: %s", immutable_data_arg) + + 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)): + 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 + + mesh = fresh_trainer._mesh + + state_restored = False + snapshot_mgr = use_python_vars.get("snapshot_mgr") + if snapshot_mgr is not None: + with mesh: + 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: + 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( + lambda state, spec: jax.device_put(state, spec.sharding), + use_jax_state["_trainer_state"], + fresh_trainer._trainer_state_specs + ) + logging.info("[ELASTIC] [✓] Successfully device_put trainer_state from globals onto the new mesh.") + state_restored = True + except Exception as e: + logging.warning("[ELASTIC] [!] Failed fallback to globals (possibly deleted arrays): %s", e) + + if not state_restored: + 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"]) + 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 + + 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("[ELASTIC] [✓] Successfully injected fresh, healthy PRNG Key into the new trainer state.") + except Exception as 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) + deleted_count = sum(x.is_deleted() for x in leaves if isinstance(x, jax.Array)) + 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", + 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 fresh_trainer, fresh_prng_key + +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 ( + getattr(obj, "_jax_device_state", {}), + getattr(obj, "_python_vars", {}), + getattr(obj, "_immutable_data", {}), + ) + + logging.info("[ELASTIC] Storing class variables for snapshot.") + + 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 + logging.info("[ELASTIC] Preparing to save snapshot.") + snapshot_mgr = python_vars.get("snapshot_mgr") + if snapshot_mgr is not None: + try: + 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("[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 + + class TrainerState(NamedTuple): prng_key: Union[Tensor, TensorSpec, jax.sharding.NamedSharding] @@ -242,6 +400,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, @@ -260,6 +421,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 @@ -268,6 +432,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") @@ -298,9 +463,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) @@ -359,8 +591,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,9 +614,11 @@ def __init__( ) self._maybe_record_event(measurement.Event.END_ACCELERATOR_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): @@ -416,7 +652,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,9 +851,15 @@ def run( return None self._is_initialized = True + #### Stores the initial state of all variables #### + 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.") + with self.checkpointer: - logging.info("Starting loop...") + logging.info("[ELASTIC] Starting loop...") start_time = time.perf_counter() num_steps = 0 output = None @@ -625,48 +867,57 @@ 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 + 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 = int(self._step) + 1 + 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( + 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._step_log("[ELASTIC] Done step") + logging.info("[ELASTIC] Done step %s", self.step) + # if self.step==3: + # _sync_store_class_vars(self) + self._jax_device_state, self._python_vars, self._immutable_data = 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: + 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) + break if self.step < cfg.max_step: self._step_log("Reached end of inputs. Stopping") self._step_log("Checkpointer flushed.") @@ -746,7 +997,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( @@ -909,7 +1160,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 @@ -1022,7 +1277,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]: @@ -1075,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. - device_kinds = set(d.device_kind for d in jax.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, @@ -1106,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, @@ -1123,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( @@ -1142,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( @@ -1151,10 +1405,12 @@ 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: + 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"]), ) @@ -1466,15 +1722,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 +1778,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 diff --git a/axlearn/common/utils.py b/axlearn/common/utils.py index 45e2c4d5f..1eb7b0737 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. @@ -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, ( @@ -2170,3 +2170,39 @@ 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 + +from pathwaysutils.elastic import manager + +def live_devices(): + device_list = jax.devices() + elastic_manager = manager.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 + ] + 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("[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()} \ No newline at end of file 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..cef2add29 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 ( @@ -318,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, @@ -331,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=( @@ -349,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": @@ -389,6 +421,7 @@ def get_trainer_kwargs( ), ) elif model_size == "7B": + import jax trainer_kwargs = dict( model_kwargs=dict( num_layers=32, @@ -401,7 +434,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=( @@ -434,6 +467,96 @@ 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), + ], + ), + ), + ( + "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={ + "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", + ), + ), + } + ), + #FlashBlockSizeModifier.default_config().set(tpu_block_size=256), + ], + ), + ), ( "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 = [