From e249bb800fc08f543de6ae7f423f4a53370d59a2 Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Mon, 27 Jul 2026 01:47:00 -0400 Subject: [PATCH 1/5] perf(trainer): jit + NamedSharding replaces shard_map, with real FSDP The training step ran under shard_map with in_specs=(P(), P(), P('data'), P('data')) - arg 0 fully replicated, so parameters and optimizer state sat on every device. That was pure DDP; there was no FSDP to speak of, and model size was capped by one device's memory. Replace it with jax.jit + NamedSharding over a two-axis ('data','fsdp') mesh and let GSPMD derive the collectives: - Parameter layout comes from a shape-based heuristic applied to every leaf of the abstract state, so optimizer moments and the EMA copy inherit their params' spec through tx.init. No model file declares partitioning. - The hand-written pmean over gradients is gone. The loss is already a mean over the batch-sharded axis, so its gradient carries the all-reduce. - The per-device fold_in RNG hack is gone with the device-index argument; threefry is partitionable, so one key per step is correct. - donate_argnums now donates the train state instead of the batch, which required best_state to stop aliasing a live state - it is a host-side numpy tree now, which is what get_best_state() already handed out. Batches are assembled with jax.make_array_from_process_local_data instead of the manual split/device_put in form_global_array, and the host-to-device transfer moved into a two-deep prefetch thread so it overlaps compute. The dead DataLoaderWithMesh prefetch path is deleted rather than left as a second way to do the same thing. Checkpoints now save and restore sharded arrays in place via orbax restore_args, instead of gathering the whole state onto the host with get_np_tree and then immediately blocking on wait_until_finished. The on-disk layout is unchanged, so existing pretrained checkpoints and the inference loader keep working. Numerical parity is the gate: single-device, 8-device and 2-way-FSDP runs are asserted to agree over 20 steps. Co-Authored-By: Claude Fable 5 --- flaxdiff/data/dataloaders.py | 114 +---- flaxdiff/trainer/general_diffusion_trainer.py | 108 ++--- flaxdiff/trainer/simple_trainer.py | 436 +++++++++--------- flaxdiff/utils.py | 212 +++++++-- tests/conftest.py | 5 + tests/test_parallelism.py | 255 ++++++++++ 6 files changed, 720 insertions(+), 410 deletions(-) create mode 100644 tests/test_parallelism.py diff --git a/flaxdiff/data/dataloaders.py b/flaxdiff/data/dataloaders.py index a388bec..b3ce58a 100644 --- a/flaxdiff/data/dataloaders.py +++ b/flaxdiff/data/dataloaders.py @@ -4,84 +4,13 @@ import numpy as np import jax import cv2 # Added missing import -from flaxdiff.utils import convert_to_global_tree, AutoTextTokenizer +from flaxdiff.utils import AutoTextTokenizer from .dataset_map import datasetMap, onlineDatasetMap, mediaDatasetMap import traceback from .online_loader import OnlineStreamingDataLoader -import queue -from jax.sharding import Mesh -import threading from functools import partial -def batch_mesh_map(mesh): - """Create an augmenter that maps batches to a mesh.""" - class augmenters(pygrain.MapTransform): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - def map(self, batch) -> Dict[str, jnp.array]: - return convert_to_global_tree(mesh, batch) - return augmenters - - -class DataLoaderWithMesh: - """A wrapper for data loaders that distributes data to a JAX mesh. - - This class wraps any iterable dataset and maps the data to a JAX mesh. - It runs a background thread that fetches data from the loader and - distributes it to the mesh. - """ - - def __init__(self, dataloader, mesh, buffer_size=20): - """Initialize a DataLoaderWithMesh. - - Args: - dataloader: The data loader to wrap. - mesh: The JAX mesh to distribute data to. - buffer_size: Size of the prefetch buffer. - """ - self.dataloader = dataloader - self.mesh = mesh - self.buffer_size = buffer_size - self.tmp_queue = queue.Queue(buffer_size) - self.loader_thread = None - self._start_loader_thread() - - def _start_loader_thread(self): - """Start the background thread for data loading.""" - def batch_loader(): - try: - for batch in self.dataloader: - try: - self.tmp_queue.put(convert_to_global_tree(self.mesh, batch)) - except Exception as e: - print("Error processing batch", e) - traceback.print_exc() - except Exception as e: - print("Error in batch loader thread", e) - traceback.print_exc() - - self.loader_thread = threading.Thread(target=batch_loader, daemon=True) - self.loader_thread.start() - - def __iter__(self): - return self - - def __next__(self): - try: - return self.tmp_queue.get(timeout=60) # Add timeout to prevent hanging - except queue.Empty: - if not self.loader_thread.is_alive(): - raise StopIteration("Loader thread died") - raise queue.Empty("Timed out waiting for batch") - - def __del__(self): - # Clean up resources - if hasattr(self, 'loader_thread') and self.loader_thread is not None: - self.loader_thread.join(timeout=1) - - def generate_collate_fn(media_type="image"): """Generate a collate function based on media type. @@ -407,9 +336,7 @@ def get_dataset_online( default_split="train", ) - def get_trainset(mesh: Mesh = None): - if mesh is not None: - return DataLoaderWithMesh(dataloader, mesh, buffer_size=worker_buffer_size) + def get_trainset(): return dataloader return { @@ -439,7 +366,6 @@ def get_media_dataset_grain( seed: int = 0, dataset_source: str = None, media_type: Optional[str] = None, # Will be auto-detected if None - mesh: Optional[Mesh] = None, additional_transform_kwargs: Dict[str, Any] = None, ): """Get a grain dataset loader for any media type (image or video). @@ -459,7 +385,6 @@ def get_media_dataset_grain( seed: Random seed. dataset_source: Source path for the dataset. media_type: Type of media ("image" or "video"). Auto-detected if None. - mesh: Optional JAX mesh for distributed training. additional_transform_kwargs: Additional arguments for the transform. Returns: @@ -506,25 +431,12 @@ def get_media_dataset_grain( shard_options=pygrain.ShardByJaxProcess(), ) - def get_trainset(mesh_override: Optional[Mesh] = None): - """Get a training dataset iterator. - - Args: - mesh_override: Optional mesh to override the default. - - Returns: - A dataset iterator. - """ - current_mesh = mesh_override or mesh - + def get_trainset(): + """Get a training dataset iterator.""" transformations = [ augmenter(), pygrain.Batch(local_batch_size, drop_remainder=True), ] - - # # Add mesh mapping if needed - # if current_mesh is not None: - # transformations.append(batch_mesh_map(current_mesh)()) loader = pygrain.DataLoader( data_source=data_source, @@ -556,7 +468,6 @@ def get_media_dataset_online( worker_buffer_size: int = 20, dataset_sources: List[str] = None, media_type: str = "image", # Default to image for online datasets - mesh: Optional[Mesh] = None, timeout: int = 15, retries: int = 3, min_media_scale: int = 128, @@ -572,7 +483,6 @@ def get_media_dataset_online( worker_buffer_size: Size of the worker buffer. dataset_sources: Custom dataset sources if data_name is "custom". media_type: Type of media ("image" or "video"). - mesh: Optional JAX mesh for distributed training. timeout: Timeout for dataset operations. retries: Number of retries for dataset operations. min_media_scale: Minimum scale for media items. @@ -616,20 +526,8 @@ def get_media_dataset_online( dataloader = OnlineStreamingDataLoader(sources, **dataloader_kwargs) - def get_trainset(mesh_override: Optional[Mesh] = None): - """Get a training dataset iterator. - - Args: - mesh_override: Optional mesh to override the default. - - Returns: - A dataset iterator. - """ - current_mesh = mesh_override or mesh - - if current_mesh is not None: - return DataLoaderWithMesh(dataloader, current_mesh, buffer_size=worker_buffer_size) - + def get_trainset(): + """Get a training dataset iterator.""" return dataloader return { diff --git a/flaxdiff/trainer/general_diffusion_trainer.py b/flaxdiff/trainer/general_diffusion_trainer.py index b651d9d..a0a7f13 100644 --- a/flaxdiff/trainer/general_diffusion_trainer.py +++ b/flaxdiff/trainer/general_diffusion_trainer.py @@ -7,9 +7,6 @@ import jax.numpy as jnp import optax import functools -from jax.sharding import Mesh, PartitionSpec as P -from jax.experimental.shard_map import shard_map - from ..schedulers import NoiseScheduler, get_coeff_shapes_tuple from ..predictors import DiffusionPredictionTransform, EpsilonPredictionTransform from ..samplers.common import DiffusionSampler @@ -18,7 +15,9 @@ from flaxdiff.utils import RandomMarkovState, serialize_model, get_latest_checkpoint from flaxdiff.inputs import ConditioningEncoder, ConditionalInputConfig, DiffusionInputConfig -from .simple_trainer import SimpleTrainer, SimpleTrainState, Metrics, convert_to_global_tree +from flaxdiff.utils import shard_batch + +from .simple_trainer import SimpleTrainer, SimpleTrainState, Metrics from flaxdiff.models.autoencoder.autoencoder import AutoEncoder from flax.training import dynamic_scale as dynamic_scale_lib @@ -200,23 +199,23 @@ def generate_states( rngs: jax.random.PRNGKey, model: nn.Module = None, use_dynamic_scale: bool = False - ) -> Tuple[TrainState, TrainState]: + ) -> TrainState: print("Generating states for DiffusionTrainer") - rngs, subkey = jax.random.split(rngs) - input_vars = self.get_input_ones() - params = model.init(subkey, **input_vars) + def init_fn(): + next_rngs, subkey = jax.random.split(rngs) + params = model.init(subkey, **self.get_input_ones()) + return TrainState.create( + apply_fn=model.apply, + params=params, + ema_params=params, + tx=optimizer, + rngs=next_rngs, + metrics=Metrics.empty(), + dynamic_scale=dynamic_scale_lib.DynamicScale() if use_dynamic_scale else None, + ) - state = TrainState.create( - apply_fn=model.apply, - params=params, - ema_params=params, - tx=optimizer, - rngs=rngs, - metrics=Metrics.empty(), - dynamic_scale = dynamic_scale_lib.DynamicScale() if use_dynamic_scale else None - ) - return state, state + return self._build_state(init_fn) def fit(self, data, training_steps_per_epoch, epochs, val_steps_per_epoch=8, sampler_class: Type[DiffusionSampler]=DDIMSampler, sampling_noise_schedule: NoiseScheduler=None): local_batch_size = data['local_batch_size'] @@ -263,13 +262,13 @@ def process_conditioning(batch, uncond_mask): ) # Main training step function - optimized for JIT compilation and sharding - def train_step(train_state: TrainState, rng_state: RandomMarkovState, batch, local_device_index): - """Training step optimized for distributed execution.""" - # Random key handling - rng_state, key_fold = rng_state.get_random_key() - folded_key = jax.random.fold_in(key_fold, local_device_index.reshape()) - local_rng_state = RandomMarkovState(folded_key) - + def train_step(train_state: TrainState, rng_state: RandomMarkovState, batch): + """Training step over the global batch; GSPMD partitions it.""" + # One key per step: threefry is partitionable, so every device draws + # its own slice of the same stream without folding in a device index. + rng_state, step_key = rng_state.get_random_key() + local_rng_state = RandomMarkovState(step_key) + # Extract and normalize data (works for both images and videos) data = batch[sample_data_key] local_batch_size = data.shape[0] @@ -324,52 +323,40 @@ def model_loss(params): return jnp.mean(weighted_loss) - # Compute gradients and apply updates + # Compute gradients and apply updates. The loss is a mean over the + # batch-sharded axis, so its gradient carries the cross-device + # all-reduce on its own - no hand-written pmean. if train_state.dynamic_scale is not None: # Mixed precision training with dynamic scale - grad_fn = train_state.dynamic_scale.value_and_grad(model_loss, axis_name="data") - dynamic_scale, is_finite, loss, grads = grad_fn(train_state.params) - + grad_fn = train_state.dynamic_scale.value_and_grad(model_loss) + dynamic_scale, grads_finite, loss, grads = grad_fn(train_state.params) + train_state = train_state.replace(dynamic_scale=dynamic_scale) new_state = train_state.apply_gradients(grads=grads) - + # Handle NaN/Inf gradients - select_fn = functools.partial(jnp.where, is_finite) + select_fn = functools.partial(jnp.where, grads_finite) new_state = new_state.replace( opt_state=jax.tree.map(select_fn, new_state.opt_state, train_state.opt_state), params=jax.tree.map(select_fn, new_state.params, train_state.params) ) else: - # Standard gradient computation grad_fn = jax.value_and_grad(model_loss) loss, grads = grad_fn(train_state.params) - - if distributed_training: - grads = jax.lax.pmean(grads, axis_name="data") - new_state = train_state.apply_gradients(grads=grads) - + # Apply EMA update new_state = new_state.apply_ema(self.ema_decay) - - # Average loss across devices if distributed - if distributed_training: - loss = jax.lax.pmean(loss, axis_name="data") - - return new_state, loss, rng_state - - # Apply sharding for distributed training - if distributed_training: - train_step = shard_map( - train_step, - mesh=self.mesh, - in_specs=(P(), P(), P('data'), P('data')), - out_specs=(P(), P(), P()), - ) - - # Apply JIT compilation - train_step = jax.jit(train_step, donate_argnums=(2)) - return train_step + + return new_state, loss, rng_state, jnp.isfinite(loss) + + replicated = self.replicated + return jax.jit( + train_step, + in_shardings=(self.state_sharding, replicated, self.batch_sharding), + out_shardings=(self.state_sharding, replicated, replicated, replicated), + donate_argnums=(0,), + ) def _define_validation_step(self, sampler_class: Type[DiffusionSampler]=DDIMSampler, sampling_noise_schedule: NoiseScheduler=None): """ @@ -465,9 +452,7 @@ def validation_loop( if val_ds is None: batch = None else: - batch = next(val_ds) - if self.distributed_training and global_device_count > 1: - batch = convert_to_global_tree(self.mesh, batch) + batch = shard_batch(self.batch_sharding, next(val_ds)) # Generate samples samples = generate_samples( val_state, @@ -723,8 +708,11 @@ def __compare_run_against_best__(self, top_k=2, metric="train/best_loss", from_s def save(self, epoch=0, step=0, state=None, rngstate=None): super().save(epoch=epoch, step=step, state=state, rngstate=rngstate) - + if self.wandb is not None: + # Uploading reads the checkpoint back off disk, so the async write + # has to have landed first. + self.wait_for_checkpoints() checkpoint = get_latest_checkpoint(self.checkpoint_path()) try: is_good, is_best = self.__compare_run_against_best__(top_k=5, metric=self.best_tracker_metric, from_sweeps=hasattr(self, "wandb_sweep")) diff --git a/flaxdiff/trainer/simple_trainer.py b/flaxdiff/trainer/simple_trainer.py index e9ff541..8b57104 100644 --- a/flaxdiff/trainer/simple_trainer.py +++ b/flaxdiff/trainer/simple_trainer.py @@ -1,29 +1,23 @@ -import orbax.checkpoint import tqdm from flax import linen as nn import jax -from typing import Callable -from dataclasses import field import jax.numpy as jnp import numpy as np -from functools import partial from clu import metrics from flax.training import train_state # Useful dataclass to keep train state import optax from flax import struct # Flax dataclasses -import flax import time import os -import orbax -from flax.training import orbax_utils -from jax.sharding import Mesh, PartitionSpec as P -from jax.experimental import mesh_utils -from jax.experimental.shard_map import shard_map -from orbax.checkpoint.utils import fully_replicated_host_local_array_to_global_array +import orbax.checkpoint as ocp +from jax.sharding import NamedSharding, PartitionSpec as P from termcolor import colored -from typing import Dict, Callable, Sequence, Any, Union, Tuple -from flax.training.dynamic_scale import DynamicScale -from flaxdiff.utils import RandomMarkovState, convert_to_global_tree +from typing import Dict, Callable, Any, Tuple +from flaxdiff.utils import ( + DEFAULT_MIN_SHARD_SIZE, DevicePrefetchIterator, RandomMarkovState, batch_sharding, + build_mesh, enable_compilation_cache, model_flops_utilization, shard_batch, + state_sharding_tree, step_flops, +) from flax.training import dynamic_scale as dynamic_scale_lib from dataclasses import dataclass import shutil @@ -49,80 +43,10 @@ class SimpleTrainState(train_state.TrainState): metrics: Metrics dynamic_scale: dynamic_scale_lib.DynamicScale -def move_contents_to_subdir(target_dir, new_subdir_name): - # --- 1. Validate Target Directory --- - if not os.path.isdir(target_dir): - print(f"Error: Target directory '{target_dir}' not found or is not a directory.") - return - # --- 2. Define Paths --- - # Construct the full path for the new subdirectory - new_subdir_path = os.path.join(target_dir, new_subdir_name) - # --- 3. Create New Subdirectory --- - try: - # Create the subdirectory. - # exist_ok=True prevents an error if the directory already exists. - os.makedirs(new_subdir_path, exist_ok=True) - print(f"Subdirectory '{new_subdir_path}' created or already exists.") - except OSError as e: - print(f"Error creating subdirectory '{new_subdir_path}': {e}") - return # Stop execution if subdirectory creation fails - # --- 4. List Contents of Target Directory --- - try: - items_to_move = os.listdir(target_dir) - except OSError as e: - print(f"Error listing contents of '{target_dir}': {e}") - return # Stop if we can't list directory contents - # --- 5. Move Items --- - print(f"Moving items from '{target_dir}' to '{new_subdir_path}'...") - moved_count = 0 - error_count = 0 - for item_name in items_to_move: - # Construct the full path of the item in the target directory - source_path = os.path.join(target_dir, item_name) - # IMPORTANT: Skip the newly created subdirectory itself! - if source_path == new_subdir_path: - continue - # Construct the destination path inside the new subdirectory - destination_path = os.path.join(new_subdir_path, item_name) - # Move the item - try: - shutil.move(source_path, destination_path) - # print(f" Moved: '{item_name}'") # Uncomment for verbose output - moved_count += 1 - except Exception as e: - print(f" Error moving '{item_name}': {e}") - error_count += 1 - print(f"\nOperation complete.") - print(f" Successfully moved: {moved_count} item(s).") - if error_count > 0: - print(f" Errors encountered: {error_count} item(s).") - -def load_from_checkpoint( - checkpoint_dir: str, -): - try: - checkpointer = orbax.checkpoint.PyTreeCheckpointer() - options = orbax.checkpoint.CheckpointManagerOptions(create=False) - # Convert checkpoint_dir to absolute path - checkpoint_dir = os.path.abspath(checkpoint_dir) - manager = orbax.checkpoint.CheckpointManager(checkpoint_dir, checkpointer, options) - ckpt = manager.restore(checkpoint_dir) - # Extract as above - state, best_state = None, None - if 'state' in ckpt: - state = ckpt['state'] - if 'best_state' in ckpt: - best_state = ckpt['best_state'] - print(f"Loaded checkpoint from local dir {checkpoint_dir}") - return state, best_state - except Exception as e: - print(f"Warning: Failed to load checkpoint from local dir: {e}") - return None, None - @dataclass class SimpleTrainer: state: SimpleTrainState - best_state: SimpleTrainState + best_state: Any best_loss: float model: nn.Module ema_decay: float = 0.999 @@ -143,21 +67,40 @@ def __init__(self, use_dynamic_scale: bool = False, max_checkpoints_to_keep: int = 2, train_start_step_override: int = None, + fsdp_size: int = 1, + fsdp_min_param_size: int = DEFAULT_MIN_SHARD_SIZE, + compilation_cache_dir: str = None, + profile_steps: int = 0, + log_every: int = 100, + max_bad_loss_steps: int = 5, ): - if distributed_training is None or distributed_training is True: - # Auto-detect if we are running on multiple devices - distributed_training = jax.device_count() > 1 - self.mesh = jax.sharding.Mesh(jax.devices(), 'data') - else: - self.mesh = None + if compilation_cache_dir: + enable_compilation_cache(compilation_cache_dir) + # One code path for every topology: the mesh spans all devices unless + # the caller explicitly opts out, and a 1x1 mesh behaves exactly like + # the old single-device path. + if distributed_training is None: + distributed_training = jax.device_count() > 1 self.distributed_training = distributed_training + devices = jax.devices() if distributed_training else jax.devices()[:1] + self.mesh = build_mesh(fsdp_size, devices=devices) + self.batch_sharding = batch_sharding(self.mesh) + self.replicated = NamedSharding(self.mesh, P()) + self.fsdp_min_param_size = fsdp_min_param_size + self.model = model self.name = name self.loss_fn = loss_fn self.input_shapes = input_shapes self.checkpoint_base_path = checkpoint_base_path - + self.profile_steps = profile_steps + self.log_every = log_every + self.max_bad_loss_steps = max_bad_loss_steps + # Measured from the compiled step the first time it runs. + self.flops_per_step = None + self.global_batch_size = 0 + load_directly_from_dir = False self.wandb = None @@ -202,25 +145,26 @@ def __init__(self, print(f"Running sweep {self.wandb_sweep.id} with id {self.wandb.sweep_id}") # checkpointer = orbax.checkpoint.PyTreeCheckpointer() - async_checkpointer = orbax.checkpoint.AsyncCheckpointer(orbax.checkpoint.PyTreeCheckpointHandler(), timeout_secs=60) - - options = orbax.checkpoint.CheckpointManagerOptions( - max_to_keep=max_checkpoints_to_keep, create=True) - self.checkpointer = orbax.checkpoint.CheckpointManager( - self.checkpoint_path(), async_checkpointer, options) + options = ocp.CheckpointManagerOptions( + max_to_keep=max_checkpoints_to_keep, create=True, + enable_async_checkpointing=True) + self.checkpointer = ocp.CheckpointManager(self.checkpoint_path(), options=options) self.rngstate = RandomMarkovState(rngs) self.rngstate, subkey = self.rngstate.get_random_key() - if train_state == None: - state, best_state = self.generate_states( - optimizer, subkey, model, use_dynamic_scale - ) - self.init_state(state, best_state) + self.best_loss = 1e9 + if train_state is None: + self.state = self.generate_states(optimizer, subkey, model, use_dynamic_scale) else: self.state = train_state - self.best_state = train_state - self.best_loss = 1e9 + self.state_sharding = jax.tree.map(lambda x: x.sharding, train_state) + # Host-side copy: aliasing a live train state here would pin its buffers + # and block donating them to the training step. + self.best_state = self.get_np_tree(self.state) + # Position of the data iterator, carried through checkpoints so a resume + # continues mid-epoch instead of replaying from the top. + self.dataset_state = None self.latest_step = 0 if load_from_checkpoint is not None: @@ -233,47 +177,47 @@ def __init__(self, def get_input_ones(self): return {k: jnp.ones((1, *v)) for k, v in self.input_shapes.items()} + def _build_state(self, init_fn) -> SimpleTrainState: + """Materialise a train state directly into its sharded layout. + + The sharding is derived from the abstract state, so optimizer moments + and EMA copies inherit their params' layout through tx.init without any + model or optimizer having to declare partitioning. + """ + self.state_sharding = state_sharding_tree( + self.mesh, jax.eval_shape(init_fn), self.fsdp_min_param_size) + return jax.jit(init_fn, out_shardings=self.state_sharding)() + def generate_states( self, optimizer: optax.GradientTransformation, rngs: jax.random.PRNGKey, model: nn.Module = None, use_dynamic_scale: bool = False - ) -> Tuple[SimpleTrainState, SimpleTrainState]: + ) -> SimpleTrainState: print("Generating states for SimpleTrainer") - rngs, subkey = jax.random.split(rngs) - input_vars = self.get_input_ones() - params = model.init(subkey, **input_vars) - - state = SimpleTrainState.create( - apply_fn=model.apply, - params=params, - tx=optimizer, - metrics=Metrics.empty(), - dynamic_scale = dynamic_scale_lib.DynamicScale() if use_dynamic_scale else None - ) - return state, state - - def init_state( - self, - state: SimpleTrainState, - best_state: SimpleTrainState, - ): - self.best_loss = 1e9 + def init_fn(): + _, subkey = jax.random.split(rngs) + return SimpleTrainState.create( + apply_fn=model.apply, + params=model.init(subkey, **self.get_input_ones()), + tx=optimizer, + metrics=Metrics.empty(), + dynamic_scale=dynamic_scale_lib.DynamicScale() if use_dynamic_scale else None, + ) - self.state = state - self.best_state = best_state + return self._build_state(init_fn) def get_state(self): return self.get_np_tree(self.state) def get_best_state(self): - return self.get_np_tree(self.best_state) - + return self.best_state + def get_rngstate(self): return self.get_np_tree(self.rngstate) - + def get_np_tree(self, pytree): return jax.tree_util.tree_map(lambda x : np.array(x), pytree) @@ -285,40 +229,51 @@ def checkpoint_path(self): os.makedirs(path) return path - def load(self, checkpoint_path, checkpoint_step=None, load_directly_from_dir=False): - checkpointer = orbax.checkpoint.PyTreeCheckpointer() - options = orbax.checkpoint.CheckpointManagerOptions( - max_to_keep=4, create=False) - checkpointer = orbax.checkpoint.CheckpointManager( - checkpoint_path, checkpointer, options) - - if checkpoint_step is None: - step = checkpointer.latest_step() - else: - step = checkpoint_step - - print("Loading model from checkpoint at step ", step) - loaded_checkpoint_path = os.path.join( - checkpoint_path if checkpoint_path else self.checkpoint_path(), - f"{step}") - self.loaded_checkpoint_path = loaded_checkpoint_path - - # Restore against the freshly-initialized states as a template so orbax - # rebuilds the exact pytree types - optimizer state and step included. - # Restoring untyped used to silently discard opt_state and reset the - # step counter (and with it the lr schedule) on every resume. + def _checkpoint_template(self): + """Restore template plus the per-leaf args that place arrays on the mesh. + + Shapes and types come from the freshly built state, so a checkpoint + written on one mesh restores onto whatever mesh this run is using. + Restoring untyped used to silently discard opt_state and reset the step + counter (and with it the lr schedule) on every resume. + """ + abstract_state = jax.eval_shape(lambda: self.state) template = { 'rngs': self.get_rngstate(), - 'state': self.get_state(), - 'best_state': self.get_best_state(), + 'state': abstract_state, + 'best_state': self.best_state, 'best_loss': np.array(self.best_loss), 'epoch': 0, } - ckpt = checkpointer.restore(step, items=template) if not load_directly_from_dir else checkpointer.restore(checkpoint_path, items=template) - + if self.dataset_state is not None: + template['dataset_state'] = self.dataset_state + restore_args = jax.tree.map(lambda _: ocp.RestoreArgs(), template) + # Only the train state is placed onto the mesh; everything else is + # bookkeeping that belongs on the host. + restore_args['state'] = jax.tree.map( + lambda s: ocp.ArrayRestoreArgs(sharding=s), self.state_sharding) + return template, restore_args + + def load(self, checkpoint_path, checkpoint_step=None, load_directly_from_dir=False): + manager = ocp.CheckpointManager( + checkpoint_path, options=ocp.CheckpointManagerOptions(max_to_keep=4, create=False)) + + step = manager.latest_step() if checkpoint_step is None else checkpoint_step + print("Loading model from checkpoint at step ", step) + self.loaded_checkpoint_path = os.path.join( + checkpoint_path if checkpoint_path else self.checkpoint_path(), f"{step}") + + template, restore_args = self._checkpoint_template() + # A checkpoint written before iterator state was tracked simply has no + # such key, so ask for it only when the run can already produce one. + target = checkpoint_path if load_directly_from_dir else step + ckpt = manager.restore( + target, args=ocp.args.PyTreeRestore(item=template, restore_args=restore_args)) + self.state = ckpt['state'] self.best_state = ckpt['best_state'] self.rngstate = ckpt['rngs'] + self.dataset_state = ckpt.get('dataset_state') self.best_loss = float(ckpt['best_loss']) if self.best_loss == 0: # It cant be zero as that must have been some problem @@ -328,25 +283,22 @@ def load(self, checkpoint_path, checkpoint_step=None, load_directly_from_dir=Fal def save(self, epoch=0, step=0, state=None, rngstate=None): print(f"Saving model at epoch {epoch} step {step}") + # Sharded arrays go straight to orbax: gathering them onto the host + # first would serialise the whole state through one process and undo + # the point of an async checkpointer. + ckpt = { + 'rngs': self.get_rngstate() if rngstate is None else self.get_np_tree(rngstate), + 'state': self.state if state is None else state, + 'best_state': self.best_state, + 'best_loss': np.array(self.best_loss), + 'epoch': epoch, + } + if self.dataset_state is not None: + ckpt['dataset_state'] = self.dataset_state try: - ckpt = { - # 'model': self.model, - 'rngs': self.get_rngstate() if rngstate is None else self.get_np_tree(rngstate), - 'state': self.get_state() if state is None else self.get_np_tree(state), - 'best_state': self.get_best_state(), - 'best_loss': np.array(self.best_loss), - 'epoch': epoch, - } - try: - save_args = orbax_utils.save_args_from_target(ckpt) - self.checkpointer.save(step, ckpt, save_kwargs={ - 'save_args': save_args}, force=True) - self.checkpointer.wait_until_finished() - pass - except Exception as e: - print("Error saving checkpoint", e) + self.checkpointer.save(step, args=ocp.args.PyTreeSave(ckpt), force=True) except Exception as e: - print("Error saving checkpoint outer", e) + print("Error saving checkpoint", e) def _define_train_step(self, **kwargs): raise NotImplementedError("Subclasses must define their train step") @@ -386,9 +338,7 @@ def validation_loop( if val_ds is None: batch = None else: - batch = next(val_ds) - if self.distributed_training and global_device_count > 1: - batch = convert_to_global_tree(self.mesh, batch) + batch = shard_batch(self.batch_sharding, next(val_ds)) if i == 0: print(f"Evaluation started for process index {process_index}") metrics = val_step_fn(val_state, batch) @@ -415,83 +365,128 @@ def train_loop( save_every:int=None, val_every=None, ): - global_device_count = jax.device_count() process_index = jax.process_index() - if self.distributed_training: - global_device_indexes = jnp.arange(global_device_count) - else: - global_device_indexes = 0 - + log_every = self.log_every + epoch_loss = 0 - bad_loss_steps = 0 current_epoch = current_step // train_steps_per_epoch - + + # Both counters live on device so the loop never blocks on a result. + # `worst_bad_run` remembers the longest streak of non-finite losses seen + # since the last host check, which is what decides whether to stop. + bad_run = jnp.zeros((), jnp.int32) + worst_bad_run = jnp.zeros((), jnp.int32) + if process_index == 0: pbar = tqdm.tqdm(total=train_steps_per_epoch, desc=f'\t\tEpoch {current_epoch}', ncols=100, unit='step') else: pbar = None - + + last_log_time = time.time() + steps_since_log = 0 + for i in range(train_steps_per_epoch): batch = next(train_ds) - # if i == 0: - # print(f"First batch loaded at step {current_step}") - - if self.distributed_training and global_device_count > 1: - # # Convert the local device batches to a unified global jax.Array - batch = convert_to_global_tree(self.mesh, batch) - train_state, loss, rng_state = train_step_fn(train_state, rng_state, batch, global_device_indexes) + if i == 0 and self.profile_steps: + jax.profiler.start_trace(self.profile_path()) + + train_state, loss, rng_state, is_finite = train_step_fn(train_state, rng_state, batch) + # No stale alias may outlive the step: its buffers were donated. + self.state, self.rngstate = train_state, rng_state + self.dataset_state = getattr(train_ds, 'source_state', None) + + bad_run = jnp.where(is_finite, 0, bad_run + 1) + worst_bad_run = jnp.maximum(worst_bad_run, bad_run) if i == 0: print(f"Training started for process index {process_index} at step {current_step}") - - if self.distributed_training: - # loss = jax.experimental.multihost_utils.process_allgather(loss) - loss = jnp.mean(loss) # Just to make sure its a scaler value - - if not jnp.isfinite(loss): - # No silent recovery: a diverged run must fail loudly, not be - # papered over with a stale best_state and a cosmetic loss value - print(colored(f"Non-finite loss at step {current_step}: {loss}", 'red')) - bad_loss_steps += 1 - if bad_loss_steps >= 5: - raise RuntimeError( - f"Loss has been non-finite for {bad_loss_steps} consecutive steps, stopping" - ) - else: - bad_loss_steps = 0 + if self.flops_per_step is None: + self.flops_per_step = step_flops( + train_step_fn, train_state, rng_state, batch) epoch_loss += loss current_step += 1 - if i % 100 == 0: + steps_since_log += 1 + + if self.profile_steps and i + 1 == self.profile_steps: + loss.block_until_ready() + jax.profiler.stop_trace() + print(f"Wrote profile for {self.profile_steps} steps to {self.profile_path()}") + + if i % log_every == 0: + self._check_finite(worst_bad_run, current_step) + worst_bad_run = jnp.zeros((), jnp.int32) if pbar is not None: + # The one place per interval where waiting on the device is + # justified: the numbers below are meaningless without it. + loss.block_until_ready() + now = time.time() + elapsed = now - last_log_time pbar.set_postfix(loss=f'{loss:.4f}') - pbar.update(100) + pbar.update(log_every) if self.wandb is not None: self.wandb.log({ - "train/step" : current_step, + "train/step": current_step, "train/loss": loss, + **self._throughput_metrics(elapsed, steps_since_log), }, step=current_step) + last_log_time, steps_since_log = now, 0 # Save the model every few steps if save_every and i % save_every == 0 and i > 0: print(f"Saving model after {save_every} step {current_step}") - print(f"Devices: {len(jax.devices())}") # To sync the devices self.save(current_epoch, current_step, train_state, rng_state) print(f"Saving done by process index {process_index}") print(colored(f"Epoch done on index {process_index} => {current_epoch} Loss: {epoch_loss/train_steps_per_epoch}", 'green')) + self._check_finite(worst_bad_run, current_step) if pbar is not None: pbar.close() return epoch_loss, current_step, train_state, rng_state + def profile_path(self): + return os.path.join(self.checkpoint_path(), 'profile') + + def _check_finite(self, worst_bad_run, current_step): + """Fail a diverged run loudly rather than papering over it. + + Deferred to the logging cadence so the step loop never synchronises; + detection is late by at most that many steps, never missed. + """ + streak = int(worst_bad_run) + if streak >= self.max_bad_loss_steps: + raise RuntimeError( + f"Loss has been non-finite for {streak} consecutive steps " + f"ending near step {current_step}, stopping") + if streak: + print(colored(f"Non-finite loss for {streak} step(s) before {current_step}", 'red')) + + def _throughput_metrics(self, elapsed: float, steps: int) -> Dict[str, float]: + if elapsed <= 0 or steps <= 0: + return {} + step_time = elapsed / steps + metrics = { + "train/step_time_ms": step_time * 1000, + "train/samples_per_sec": self.global_batch_size / step_time, + } + mfu = model_flops_utilization(self.flops_per_step, step_time, + self.mesh.devices.size) + if mfu is not None: + metrics["train/mfu"] = mfu + return metrics + def fit(self, data, train_steps_per_epoch, epochs, train_step_args={}, val_steps_per_epoch=5, validation_step_args={}): - train_ds = iter(data['train']()) + local_batch_size = data.get('local_batch_size', 0) + self.global_batch_size = data.get( + 'global_batch_size', local_batch_size * jax.process_count()) + train_ds = DevicePrefetchIterator( + data['train'](), self.batch_sharding, source_state=self.dataset_state) val_ds = data.get('val', data.get('test', None)) train_step = self._define_train_step(**train_step_args) val_step = self._define_validation_step(**validation_step_args) train_state = self.state rng_state = self.rngstate process_index = jax.process_index() - + if val_steps_per_epoch > 0: # We should first run a validation step to make sure the model is working print(f"Validation run for sanity check for process index {process_index}") @@ -543,7 +538,7 @@ def fit(self, data, train_steps_per_epoch, epochs, train_step_args={}, val_steps avg_loss = epoch_loss / train_steps_per_epoch if avg_loss < self.best_loss: self.best_loss = avg_loss - self.best_state = train_state + self.best_state = self.get_np_tree(train_state) self.save(current_epoch, current_step) if process_index == 0: @@ -558,5 +553,14 @@ def fit(self, data, train_steps_per_epoch, epochs, train_step_args={}, val_steps print(colored(f"\n\tEpoch {current_epoch} completed. Avg Loss: {avg_loss}, Time: {total_time:.2f}s, Best Loss: {self.best_loss}", 'green')) - self.save(epochs)# + self.save(epochs) + self.wait_for_checkpoints() return self.state + + def wait_for_checkpoints(self): + """Block until pending async checkpoint writes have landed on disk. + + Saving is async so it stays off the training loop's critical path; + anything that reads the checkpoint back has to call this first. + """ + self.checkpointer.wait_until_finished() diff --git a/flaxdiff/utils.py b/flaxdiff/utils.py index 3017607..c8b54e3 100644 --- a/flaxdiff/utils.py +++ b/flaxdiff/utils.py @@ -2,11 +2,12 @@ import jax.numpy as jnp import flax.struct as struct import flax.linen as nn -from typing import Any -from functools import partial +from typing import Any, Iterator, Optional import numpy as np import os -from jax.sharding import Mesh, PartitionSpec as P +import queue +import threading +from jax.sharding import AxisType, Mesh, NamedSharding, PartitionSpec as P from flaxdiff.inputs import TextEncoder, CLIPTextEncoder # Setup mappings for dtype, precision, and activation @@ -94,29 +95,188 @@ def denormalize_images(images, target_type=jnp.uint8, source_range=(-1, 1), targ return images -def _build_global_shape_and_sharding( - local_shape: tuple[int, ...], global_mesh: Mesh -) -> tuple[tuple[int, ...], jax.sharding.NamedSharding]: - sharding = jax.sharding.NamedSharding(global_mesh, P(global_mesh.axis_names)) - global_shape = (jax.process_count() * local_shape[0],) + local_shape[1:] - return global_shape, sharding - - -def form_global_array(path, array: np.ndarray, global_mesh: Mesh) -> jax.Array: - """Put local sharded array into local devices""" - global_shape, sharding = _build_global_shape_and_sharding(np.shape(array), global_mesh) - try: - local_device_arrays = np.split(array, len(global_mesh.local_devices), axis=0) - except ValueError as array_split_error: - raise ValueError( - f"Unable to put to devices shape {array.shape} with " - f"local device count {len(global_mesh.local_devices)} " - ) from array_split_error - local_device_buffers = jax.device_put(local_device_arrays, global_mesh.local_devices) - return jax.make_array_from_single_device_arrays(global_shape, sharding, local_device_buffers) - -def convert_to_global_tree(global_mesh, pytree): - return jax.tree_util.tree_map_with_path(partial(form_global_array, global_mesh=global_mesh), pytree) +# --------------------------------------------------------------------------- +# Sharding +# --------------------------------------------------------------------------- + +DATA_AXIS = 'data' +FSDP_AXIS = 'fsdp' + +# Batches are split across every device, whichever axis it sits on; only +# parameters distinguish the two axes. +BATCH_SPEC = P((DATA_AXIS, FSDP_AXIS)) + +# Below this many elements a parameter costs more in collectives than it saves +# in memory, so it stays replicated. +DEFAULT_MIN_SHARD_SIZE = 2 ** 16 + + +def build_mesh(fsdp_size: int = 1, devices: Optional[list] = None) -> Mesh: + """Two-axis device mesh: parameters shard over 'fsdp', replicate over 'data'. + + fsdp_size=1 degenerates to plain data parallelism, so the same code path + serves both without a flag. Axes are Auto so GSPMD infers the collectives + rather than us writing them by hand. + """ + devices = list(devices) if devices is not None else jax.devices() + if fsdp_size < 1 or len(devices) % fsdp_size: + raise ValueError( + f"fsdp_size {fsdp_size} must be a positive divisor of device count {len(devices)}") + return jax.make_mesh( + (len(devices) // fsdp_size, fsdp_size), + (DATA_AXIS, FSDP_AXIS), + devices=devices, + axis_types=(AxisType.Auto, AxisType.Auto), + ) + + +def parameter_spec(shape: tuple, fsdp_size: int, min_shard_size: int) -> P: + """Shard the largest evenly-divisible axis over 'fsdp', else replicate. + + Applied to every leaf of the train state, not just params: optimizer moments + and EMA copies have the same shapes as the params they track, so they pick + up the same spec without anyone having to describe the optimizer's layout. + """ + if fsdp_size == 1 or int(np.prod(shape, dtype=np.int64)) < min_shard_size: + return P() + for axis in sorted(range(len(shape)), key=lambda i: -shape[i]): + if shape[axis] % fsdp_size == 0: + return P(*([None] * axis), FSDP_AXIS) + return P() + + +def state_sharding_tree( + mesh: Mesh, abstract_state, min_shard_size: int = DEFAULT_MIN_SHARD_SIZE +): + """Map a train state of ShapeDtypeStructs to its NamedSharding tree.""" + fsdp_size = mesh.shape[FSDP_AXIS] + return jax.tree.map( + lambda x: NamedSharding(mesh, parameter_spec(x.shape, fsdp_size, min_shard_size)), + abstract_state, + ) + + +def batch_sharding(mesh: Mesh) -> NamedSharding: + return NamedSharding(mesh, BATCH_SPEC) + + +def shard_batch(sharding: NamedSharding, batch): + """Assemble this process's slice of each array into a globally sharded one.""" + return jax.tree.map( + lambda x: jax.make_array_from_process_local_data(sharding, np.asarray(x)), batch) + + +class DevicePrefetchIterator: + """Runs the host-to-device batch transfer a few batches ahead of the loop. + + Without this the transfer sits on the critical path between steps, because + the loop only starts moving batch N+1 after step N has been dispatched. + """ + + def __init__(self, iterator: Iterator, sharding: NamedSharding, depth: int = 2, + source_state=None): + self._iterator = iter(iterator) + self._sharding = sharding + self._queue = queue.Queue(maxsize=depth) + self._terminal: Optional[BaseException] = None + self._checkpointable = hasattr(self._iterator, 'get_state') + if source_state is not None: + if not self._checkpointable: + raise TypeError( + f"{type(self._iterator).__name__} cannot resume from a saved position") + self._iterator.set_state(source_state) + # Position of the source iterator as of the batch most recently handed + # out, so a checkpoint resumes at the next unseen batch rather than at + # whatever the prefetch thread has already raced ahead to. + self.source_state = source_state + self._thread = threading.Thread(target=self._prefetch, daemon=True) + self._thread.start() + + def _prefetch(self): + try: + while True: + batch = next(self._iterator) + state = self._iterator.get_state() if self._checkpointable else None + self._queue.put((shard_batch(self._sharding, batch), state)) + except StopIteration: + self._queue.put(StopIteration()) + except BaseException as error: # surfaced on the consumer's thread + self._queue.put(error) + + def __iter__(self): + return self + + def __next__(self): + if self._terminal is not None: + raise self._terminal + item = self._queue.get() + if isinstance(item, BaseException): + self._terminal = item + raise item + batch, self.source_state = item + return batch + +# --------------------------------------------------------------------------- +# Throughput accounting +# --------------------------------------------------------------------------- + +# Dense bf16 peak per chip, from the vendors' own spec sheets. Only used to turn +# measured FLOPs into a utilisation percentage; unknown hardware just skips MFU. +PEAK_FLOPS_PER_DEVICE = { + 'TPU v2': 45e12, + 'TPU v3': 123e12, + 'TPU v4': 275e12, + 'TPU v5 lite': 197e12, + 'TPU v5e': 197e12, + 'TPU v5': 459e12, + 'TPU v5p': 459e12, + 'TPU v6 lite': 918e12, + 'TPU v6e': 918e12, + 'NVIDIA A100': 312e12, + 'NVIDIA H100': 989e12, + 'NVIDIA H200': 989e12, +} + + +def step_flops(jitted, *args, **kwargs) -> Optional[float]: + """FLOPs for one call of a jitted function, straight from the compiler. + + Measured rather than derived from a hand-written parameter-count formula, so + it stays honest across architectures, remat and gradient accumulation. + """ + analysis = jitted.lower(*args, **kwargs).compile().cost_analysis() + if isinstance(analysis, (list, tuple)): + analysis = analysis[0] if analysis else None + if not analysis or 'flops' not in analysis: + return None + return float(analysis['flops']) + + +def model_flops_utilization( + flops_per_step: Optional[float], step_time: float, device_count: int +) -> Optional[float]: + """Fraction of the cluster's peak FLOPs the training step actually achieved.""" + if not flops_per_step or step_time <= 0: + return None + peak = PEAK_FLOPS_PER_DEVICE.get(jax.devices()[0].device_kind) + if peak is None: + return None + return flops_per_step / step_time / (peak * device_count) + + +def enable_compilation_cache(path: str): + """Persist compiled executables so restarts skip XLA compilation. + + The dominant cost of a restart-heavy TPU workflow, where every run otherwise + recompiles the same step function from scratch. + """ + os.makedirs(path, exist_ok=True) + jax.config.update('jax_compilation_cache_dir', path) + # Defaults skip small/fast compilations; a training step is neither, and + # caching everything keeps startup predictable. + jax.config.update('jax_persistent_cache_min_entry_size_bytes', -1) + jax.config.update('jax_persistent_cache_min_compile_time_secs', 0.0) + class AutoTextTokenizer: def __init__(self, tensor_type="pt", modelname="openai/clip-vit-large-patch14"): diff --git a/tests/conftest.py b/tests/conftest.py index d9c41fd..42f06e9 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,6 +2,11 @@ # Tests must run identically on any machine, CPU is enough os.environ.setdefault("JAX_PLATFORMS", "cpu") +# Enough simulated devices to exercise a 4x2 data/fsdp mesh. Must be set before +# jax initialises its backend. +os.environ["XLA_FLAGS"] = ( + os.environ.get("XLA_FLAGS", "") + " --xla_force_host_platform_device_count=8" +).strip() import jax import jax.numpy as jnp diff --git a/tests/test_parallelism.py b/tests/test_parallelism.py new file mode 100644 index 0000000..b007140 --- /dev/null +++ b/tests/test_parallelism.py @@ -0,0 +1,255 @@ +"""Sharding, FSDP and data-pipeline tests on a simulated 8-device CPU mesh. + +The parity tests are the safety net for the shard_map -> jit + NamedSharding +migration: a partitioned run has to produce the same numbers as a single-device +one, otherwise the collectives GSPMD derived are not the ones we meant. +""" + +import jax +import jax.numpy as jnp +import numpy as np +import optax +import pytest +from jax.sharding import PartitionSpec as P + +from flaxdiff.inputs import DiffusionInputConfig +from flaxdiff.models.simple_dit import SimpleDiT +from flaxdiff.predictors import get_diffusion_preset +from flaxdiff.trainer import GeneralDiffusionTrainer +from flaxdiff.utils import ( + DevicePrefetchIterator, batch_sharding, build_mesh, parameter_spec, shard_batch, +) + +RES = 8 +BATCH = 8 +# The test model's parameters are far below the production shard threshold, so +# lower it or "FSDP on" would silently mean "everything replicated". +TINY = 256 + + +def make_trainer(tmp_path, name, distributed_training, fsdp_size=1, **kwargs): + train_schedule, _, transform = get_diffusion_preset("edm") + return GeneralDiffusionTrainer( + model=SimpleDiT(patch_size=4, emb_features=32, num_layers=1, num_heads=2, mlp_ratio=1), + optimizer=optax.adam(1e-3), + noise_schedule=train_schedule, + model_output_transform=transform, + input_config=DiffusionInputConfig( + sample_data_key="image", sample_data_shape=(RES, RES, 3), conditions=[]), + rngs=jax.random.PRNGKey(0), + name=name, + wandb_config=None, + distributed_training=distributed_training, + fsdp_size=fsdp_size, + checkpoint_base_path=str(tmp_path), + **kwargs, + ) + + +def batches(): + rng = np.random.default_rng(0) + images = rng.integers(0, 256, size=(BATCH, RES, RES, 3)).astype(np.float32) + while True: + yield {"image": images} + + +def run_losses(trainer, steps): + """Per-step losses from the real compiled training step.""" + train_step = trainer._define_train_step(batch_size=BATCH) + source = DevicePrefetchIterator(batches(), trainer.batch_sharding) + state, rng = trainer.state, trainer.rngstate + losses = [] + for _ in range(steps): + state, loss, rng, is_finite = train_step(state, rng, next(source)) + assert bool(is_finite) + losses.append(float(loss)) + return losses + + +# -------------------------------------------------------------------------- +# Sharding heuristic +# -------------------------------------------------------------------------- + +def test_parameter_spec_replicates_without_fsdp(): + assert parameter_spec((1024, 1024), fsdp_size=1, min_shard_size=16) == P() + + +def test_parameter_spec_replicates_small_params(): + assert parameter_spec((8, 8), fsdp_size=2, min_shard_size=2 ** 16) == P() + + +def test_parameter_spec_shards_largest_divisible_axis(): + assert parameter_spec((64, 1024), fsdp_size=2, min_shard_size=16) == P(None, 'fsdp') + assert parameter_spec((1024, 64), fsdp_size=2, min_shard_size=16) == P('fsdp') + + +def test_parameter_spec_falls_back_to_replication_when_indivisible(): + assert parameter_spec((15, 15), fsdp_size=2, min_shard_size=16) == P() + + +def test_build_mesh_rejects_bad_fsdp_size(): + with pytest.raises(ValueError): + build_mesh(fsdp_size=3) + + +def test_build_mesh_axes(): + mesh = build_mesh(fsdp_size=2) + assert mesh.shape['data'] == jax.device_count() // 2 + assert mesh.shape['fsdp'] == 2 + + +# -------------------------------------------------------------------------- +# Batch placement and prefetch +# -------------------------------------------------------------------------- + +def test_shard_batch_splits_across_all_devices(): + mesh = build_mesh(fsdp_size=2) + batch = {"image": np.zeros((jax.device_count(), 4), np.float32)} + sharded = shard_batch(batch_sharding(mesh), batch)["image"] + assert len(sharded.addressable_shards) == jax.device_count() + assert sharded.addressable_shards[0].data.shape == (1, 4) + + +def test_prefetch_iterator_preserves_order_and_terminates(): + mesh = build_mesh() + source = ({"x": np.full((jax.device_count(), 2), i, np.float32)} for i in range(5)) + it = DevicePrefetchIterator(source, batch_sharding(mesh), depth=2) + seen = [float(np.asarray(b["x"])[0, 0]) for b in it] + assert seen == [0.0, 1.0, 2.0, 3.0, 4.0] + with pytest.raises(StopIteration): + next(it) + + +def test_prefetch_iterator_surfaces_source_errors(): + mesh = build_mesh() + + def broken(): + yield {"x": np.zeros((jax.device_count(), 2), np.float32)} + raise ValueError("source exploded") + + it = DevicePrefetchIterator(broken(), batch_sharding(mesh), depth=2) + next(it) + with pytest.raises(ValueError, match="source exploded"): + next(it) + + +def test_prefetch_iterator_tracks_checkpointable_source_state(): + """The position handed out must be the consumed batch's, not the thread's.""" + import grain.python as pygrain + + loader = pygrain.DataLoader( + data_source=pygrain.RangeDataSource(0, 64, 1), + sampler=pygrain.IndexSampler(num_records=64, shuffle=False, seed=0, num_epochs=1, + shard_options=pygrain.NoSharding()), + operations=[pygrain.Batch(jax.device_count(), drop_remainder=True)], + worker_count=0, + ) + mesh = build_mesh() + it = DevicePrefetchIterator(iter(loader), batch_sharding(mesh), depth=2) + next(it) + next(it) + state = it.source_state + expected = np.asarray(next(it)) + + resumed = DevicePrefetchIterator(iter(loader), batch_sharding(mesh), depth=2, + source_state=state) + assert np.array_equal(np.asarray(next(resumed)), expected) + + +# -------------------------------------------------------------------------- +# Numerical parity +# -------------------------------------------------------------------------- + +def test_single_and_multi_device_losses_agree(tmp_path): + """The whole point of the migration: partitioning must not change the maths.""" + steps = 20 + single = run_losses(make_trainer(tmp_path / "one", "one", distributed_training=False), steps) + multi = run_losses(make_trainer(tmp_path / "many", "many", distributed_training=True), steps) + assert jax.device_count() > 1 + np.testing.assert_allclose(single, multi, rtol=2e-4, atol=2e-5) + + +def test_fsdp_losses_match_replicated(tmp_path): + """Sharding the parameters must not change the loss trajectory.""" + steps = 20 + replicated = run_losses( + make_trainer(tmp_path / "dp", "dp", distributed_training=True, fsdp_size=1), steps) + fsdp = make_trainer(tmp_path / "fsdp", "fsdp", distributed_training=True, + fsdp_size=2, fsdp_min_param_size=TINY) + assert any('fsdp' in str(x.sharding.spec) for x in jax.tree.leaves(fsdp.state.params)) + np.testing.assert_allclose(replicated, run_losses(fsdp, steps), rtol=2e-4, atol=2e-5) + + +# -------------------------------------------------------------------------- +# FSDP actually shards +# -------------------------------------------------------------------------- + +def test_fsdp_shards_parameters_and_optimizer_state(tmp_path): + trainer = make_trainer(tmp_path, "fsdp-shapes", distributed_training=True, + fsdp_size=2, fsdp_min_param_size=TINY) + leaves = jax.tree.leaves(trainer.state.params) + sharded = [x for x in leaves if 'fsdp' in str(x.sharding.spec)] + assert sharded, "no parameter was sharded over the fsdp axis" + + for param in sharded: + biggest = max(param.shape) + local = param.addressable_shards[0].data + assert local.size == param.size // 2, "shard is not half the global param" + assert biggest // 2 in local.shape + + # Adam moments and the EMA copy must follow the params they track, without + # the optimizer or the model ever describing a layout. + mu = trainer.state.opt_state[0].mu + param_specs = [x.sharding.spec for x in jax.tree.leaves(trainer.state.params)] + mu_specs = [x.sharding.spec for x in jax.tree.leaves(mu)] + assert param_specs == mu_specs + + ema_specs = [x.sharding.spec for x in jax.tree.leaves(trainer.state.ema_params)] + assert param_specs == ema_specs + + +def test_replicated_run_shards_nothing(tmp_path): + trainer = make_trainer(tmp_path, "dp-shapes", distributed_training=True, fsdp_size=1) + for leaf in jax.tree.leaves(trainer.state.params): + assert leaf.sharding.spec == P() + + +# -------------------------------------------------------------------------- +# Checkpointing under sharding +# -------------------------------------------------------------------------- + +def test_sharded_checkpoint_roundtrips(tmp_path): + trainer = make_trainer(tmp_path, "ckpt", distributed_training=True, + fsdp_size=2, fsdp_min_param_size=TINY) + grads = jax.tree.map(jnp.ones_like, trainer.state.params) + trainer.state = trainer.state.apply_gradients(grads=grads).apply_ema(0.99) + trainer.save(epoch=0, step=1) + trainer.wait_for_checkpoints() + + restored = make_trainer(tmp_path, "ckpt", distributed_training=True, fsdp_size=2, + fsdp_min_param_size=TINY, + load_from_checkpoint=trainer.checkpoint_path()) + assert int(restored.state.step) == 1 + for before, after in zip(jax.tree.leaves(trainer.state.params), + jax.tree.leaves(restored.state.params)): + assert before.sharding.spec == after.sharding.spec + np.testing.assert_allclose(np.asarray(before), np.asarray(after)) + + +def test_checkpoint_restores_onto_a_different_mesh(tmp_path): + """A run saved with FSDP must be resumable on a replicated mesh.""" + trainer = make_trainer(tmp_path, "mesh", distributed_training=True, + fsdp_size=2, fsdp_min_param_size=TINY) + grads = jax.tree.map(jnp.ones_like, trainer.state.params) + trainer.state = trainer.state.apply_gradients(grads=grads) + trainer.save(epoch=0, step=1) + trainer.wait_for_checkpoints() + + restored = make_trainer(tmp_path, "mesh", distributed_training=True, fsdp_size=1, + load_from_checkpoint=trainer.checkpoint_path()) + assert int(restored.state.step) == 1 + for leaf in jax.tree.leaves(restored.state.params): + assert leaf.sharding.spec == P() + for before, after in zip(jax.tree.leaves(trainer.state.params), + jax.tree.leaves(restored.state.params)): + np.testing.assert_allclose(np.asarray(before), np.asarray(after)) From 85dd63d616326677cb6cc6389ae80b5e3b252434 Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Mon, 27 Jul 2026 02:01:59 -0400 Subject: [PATCH 2/5] feat(models): rematerialize transformer blocks behind --remat Activation memory is what caps trainable model size, so make the DiT-family blocks rematerializable. ModulatedBlock and MMDiTBlock are wrapped through one helper, with dots_with_no_batch_dims_saveable so the big matmuls' outputs survive and the recompute stays cheap. Two things the tests forced out: - `train` picks a Python branch inside the blocks, so remat traced it and blew up on a bool conversion. It has to be static, which means the block call sites pass their arguments positionally now. - The S5 mixer carries complex64 residuals, and saving a residual goes through jax.lax.reduce_precision, which rejects complex dtypes. Those blocks get policy=None (recompute everything) rather than a silently degraded policy. Also adds --grad_accum_steps, which wraps the optimizer in optax.MultiSteps, and --fsdp_size / --fsdp_min_param_size / --profile_steps / --compilation_cache_dir / --log_every to reach the trainer's new knobs. Remat is asserted to leave the parameter tree, forward values and gradients unchanged, so turning it on cannot invalidate a checkpoint. Co-Authored-By: Claude Fable 5 --- flaxdiff/models/dit_common.py | 26 +++++++++++ flaxdiff/models/simple_dit.py | 7 +-- flaxdiff/models/simple_mmdit.py | 14 +++--- flaxdiff/models/simple_vit.py | 15 ++++--- flaxdiff/models/ssm_dit.py | 9 ++-- flaxdiff/models/video_dit.py | 9 ++-- tests/test_remat.py | 77 +++++++++++++++++++++++++++++++++ training.py | 22 +++++++++- 8 files changed, 153 insertions(+), 26 deletions(-) create mode 100644 tests/test_remat.py diff --git a/flaxdiff/models/dit_common.py b/flaxdiff/models/dit_common.py index 848d1b9..d18903d 100644 --- a/flaxdiff/models/dit_common.py +++ b/flaxdiff/models/dit_common.py @@ -9,6 +9,8 @@ module owns the sandwich; the model files just arrange blocks. """ +import inspect + import jax import jax.numpy as jnp from flax import linen as nn @@ -170,6 +172,30 @@ def __call__(self, tokens, inv_idx, H, W, conditioning=None): H_P=H // self.patch_size, W_P=W // self.patch_size) +def remat_block(block_cls, enabled: bool, policy='dots'): + """Optionally rematerialize a block class. + + Recomputing a block during the backward pass trades extra compute for a + large drop in activation memory, which is what caps trainable model size. + The default policy keeps the big matmul outputs so the recompute stays + cheap. Blocks carrying complex intermediates (the S5 mixer) must pass + policy=None: saving a residual goes through jax.lax.reduce_precision, + which only accepts floating dtypes. + + `train` selects a Python branch, so it has to stay static; that also means + callers must pass it positionally for jax to see it as such. + """ + if not enabled: + return block_cls + names = list(inspect.signature(block_cls.__call__).parameters) + return nn.remat( + block_cls, + static_argnums=tuple(i for i, name in enumerate(names) if name == 'train'), + policy=(jax.checkpoint_policies.dots_with_no_batch_dims_saveable + if policy == 'dots' else None), + ) + + class ModulatedBlock(nn.Module): """adaLN-Zero modulated residual block with a pluggable token mixer. diff --git a/flaxdiff/models/simple_dit.py b/flaxdiff/models/simple_dit.py index bc59f09..bc2a355 100644 --- a/flaxdiff/models/simple_dit.py +++ b/flaxdiff/models/simple_dit.py @@ -5,7 +5,7 @@ from .dit_common import ( PatchSequenceEmbed, ConditioningEmbed, PatchSequenceOutput, - ModulatedBlock, neutralized_rope_freqs, + ModulatedBlock, remat_block, neutralized_rope_freqs, ) from .vit_common import RotaryEmbedding @@ -26,6 +26,7 @@ class SimpleDiT(nn.Module): learn_sigma: bool = False qk_norm: bool = False attention_impl: Optional[str] = None + remat: bool = False use_hilbert: bool = False use_zigzag: bool = False @@ -52,7 +53,7 @@ def setup(self): self.rope = RotaryEmbedding( dim=self.emb_features // self.num_heads, max_seq_len=4096, dtype=self.dtype) self.blocks = [ - ModulatedBlock( + remat_block(ModulatedBlock, self.remat)( features=self.emb_features, num_heads=self.num_heads, rope_emb=self.rope, @@ -85,6 +86,6 @@ def __call__(self, x, temb, textcontext=None, train: bool = False): freqs_cis = neutralized_rope_freqs(self.rope, x_seq.shape[1], self.scan_order) for block in self.blocks: - x_seq = block(x_seq, conditioning=cond_emb, freqs_cis=freqs_cis, train=train) + x_seq = block(x_seq, cond_emb, freqs_cis, train) return self.output(x_seq, inv_idx, H, W) diff --git a/flaxdiff/models/simple_mmdit.py b/flaxdiff/models/simple_mmdit.py index 6ebc43e..0c320aa 100644 --- a/flaxdiff/models/simple_mmdit.py +++ b/flaxdiff/models/simple_mmdit.py @@ -18,7 +18,7 @@ from .dit_common import ( PatchSequenceEmbed, ConditioningEmbed, PatchSequenceOutput, - neutralized_rope_freqs, + neutralized_rope_freqs, remat_block, ) from .attention import scaled_dot_product_attention from .vit_common import RotaryEmbedding, AdaLNParams, apply_rotary_embedding @@ -150,6 +150,7 @@ class SimpleMMDiT(nn.Module): learn_sigma: bool = False qk_norm: bool = False attention_impl: Optional[str] = None + remat: bool = False use_hilbert: bool = False use_zigzag: bool = False @@ -180,7 +181,7 @@ def setup(self): self.rope = RotaryEmbedding( dim=self.emb_features // self.num_heads, max_seq_len=4096, dtype=self.dtype) self.blocks = [ - MMDiTBlock( + remat_block(MMDiTBlock, self.remat)( features=self.emb_features, num_heads=self.num_heads, mlp_ratio=self.mlp_ratio, @@ -215,7 +216,7 @@ def __call__(self, x, temb, textcontext, train: bool = False): # textcontext is freqs_cis = neutralized_rope_freqs(self.rope, img.shape[1], self.scan_order) for block in self.blocks: - img, txt = block(img, txt, conditioning=cond_emb, freqs_cis=freqs_cis, train=train) + img, txt = block(img, txt, cond_emb, freqs_cis, train) return self.output(img, inv_idx, H, W, conditioning=cond_emb) @@ -337,6 +338,7 @@ class HierarchicalMMDiT(nn.Module): learn_sigma: bool = False qk_norm: bool = False attention_impl: Optional[str] = None + remat: bool = False def setup(self): assert len(self.emb_features) == len(self.num_layers) == len(self.num_heads), \ @@ -377,7 +379,7 @@ def setup(self): def stage_blocks(stage, prefix): return [ - MMDiTBlock( + remat_block(MMDiTBlock, self.remat)( features=self.emb_features[stage], num_heads=self.num_heads[stage], mlp_ratio=self.mlp_ratio, @@ -455,7 +457,7 @@ def __call__(self, x, temb, textcontext, train: bool = False): freqs_cis = self.ropes[stage](seq_len=img.shape[1]) txt = txts[stage] for block in self.encoder_blocks[stage]: - img, txt = block(img, txt, conditioning=conds[stage], freqs_cis=freqs_cis, train=train) + img, txt = block(img, txt, conds[stage], freqs_cis, train) skips[stage] = img if stage < num_stages - 1: img, H_P, W_P = self.patch_mergers[stage](img, H_P, W_P) @@ -467,6 +469,6 @@ def __call__(self, x, temb, textcontext, train: bool = False): freqs_cis = self.ropes[stage](seq_len=img.shape[1]) txt = txts[stage] for block in self.decoder_blocks[i]: - img, txt = block(img, txt, conditioning=conds[stage], freqs_cis=freqs_cis, train=train) + img, txt = block(img, txt, conds[stage], freqs_cis, train) return self.output(img, None, H, W, conditioning=conds[0]) diff --git a/flaxdiff/models/simple_vit.py b/flaxdiff/models/simple_vit.py index 1554afb..e18f15c 100644 --- a/flaxdiff/models/simple_vit.py +++ b/flaxdiff/models/simple_vit.py @@ -12,7 +12,7 @@ from functools import partial from .hilbert import hilbert_indices, inverse_permutation, hilbert_patchify, hilbert_unpatchify from .vit_common import unpatchify, PatchEmbedding, RotaryEmbedding, RoPEAttention, AdaLNParams -from .dit_common import ModulatedBlock +from .dit_common import ModulatedBlock, remat_block class UViT(nn.Module): @@ -262,6 +262,7 @@ class SimpleUDiT(nn.Module): precision: PrecisionLike = None force_fp32_for_softmax: bool = True # Passed to DiTBlock -> RoPEAttention attention_impl: Optional[str] = None + remat: bool = False norm_epsilon: float = 1e-5 learn_sigma: bool = False use_hilbert: bool = False @@ -309,7 +310,7 @@ def setup(self): ) self.down_blocks = [ - ModulatedBlock( + remat_block(ModulatedBlock, self.remat)( features=self.emb_features, num_heads=self.num_heads, mlp_ratio=self.mlp_ratio, @@ -324,7 +325,7 @@ def setup(self): ) for i in range(half_layers) ] - self.mid_block = ModulatedBlock( + self.mid_block = remat_block(ModulatedBlock, self.remat)( features=self.emb_features, num_heads=self.num_heads, mlp_ratio=self.mlp_ratio, @@ -347,7 +348,7 @@ def setup(self): ) for i in range(half_layers) ] self.up_blocks = [ - ModulatedBlock( + remat_block(ModulatedBlock, self.remat)( features=self.emb_features, num_heads=self.num_heads, mlp_ratio=self.mlp_ratio, @@ -406,16 +407,16 @@ def __call__(self, x, temb, textcontext=None, train: bool = False): skips = [] for i in range(self.num_layers // 2): - x_seq = self.down_blocks[i](x_seq, conditioning=cond_emb, freqs_cis=None, train=train) + x_seq = self.down_blocks[i](x_seq, cond_emb, None, train) skips.append(x_seq) - x_seq = self.mid_block(x_seq, conditioning=cond_emb, freqs_cis=None, train=train) + x_seq = self.mid_block(x_seq, cond_emb, None, train) for i in range(self.num_layers // 2): skip_conn = skips.pop() x_seq = jnp.concatenate([x_seq, skip_conn], axis=-1) x_seq = self.up_dense[i](x_seq) - x_seq = self.up_blocks[i](x_seq, conditioning=cond_emb, freqs_cis=None, train=train) + x_seq = self.up_blocks[i](x_seq, cond_emb, None, train) x_out = self.final_norm(x_seq) x_out = self.final_proj(x_out) diff --git a/flaxdiff/models/ssm_dit.py b/flaxdiff/models/ssm_dit.py index 749aaf0..63f3033 100644 --- a/flaxdiff/models/ssm_dit.py +++ b/flaxdiff/models/ssm_dit.py @@ -11,7 +11,7 @@ from .dit_common import ( PatchSequenceEmbed, ConditioningEmbed, PatchSequenceOutput, - ModulatedBlock, neutralized_rope_freqs, + ModulatedBlock, remat_block, neutralized_rope_freqs, ) from .vit_common import RotaryEmbedding @@ -35,6 +35,7 @@ class HybridSSMAttentionDiT(nn.Module): learn_sigma: bool = False qk_norm: bool = False attention_impl: Optional[str] = None + remat: bool = False use_hilbert: bool = False use_zigzag: bool = False # ZigMa-style serpentine scan block_pattern: Optional[Sequence[str]] = None # e.g., ['ssm','ssm','ssm','attn'] @@ -85,7 +86,7 @@ def setup(self): blocks = [] for i, block_type in enumerate(pattern): if block_type == 'ssm': - blocks.append(ModulatedBlock( + blocks.append(remat_block(ModulatedBlock, self.remat, policy=None)( features=self.emb_features, num_heads=self.num_heads, rope_emb=self.rope, @@ -102,7 +103,7 @@ def setup(self): name=f"ssm_block_{i}" )) else: # 'attn' - blocks.append(ModulatedBlock( + blocks.append(remat_block(ModulatedBlock, self.remat)( features=self.emb_features, num_heads=self.num_heads, rope_emb=self.rope, @@ -136,6 +137,6 @@ def __call__(self, x, temb, textcontext=None, train: bool = False): freqs_cis = neutralized_rope_freqs(self.rope, x_seq.shape[1], self.scan_order) for block in self.blocks: - x_seq = block(x_seq, conditioning=cond_emb, freqs_cis=freqs_cis, train=train) + x_seq = block(x_seq, cond_emb, freqs_cis, train) return self.output(x_seq, inv_idx, H, W) diff --git a/flaxdiff/models/video_dit.py b/flaxdiff/models/video_dit.py index 8f2f7b6..069a5cc 100644 --- a/flaxdiff/models/video_dit.py +++ b/flaxdiff/models/video_dit.py @@ -16,7 +16,7 @@ from .dit_common import ( PatchSequenceEmbed, ConditioningEmbed, PatchSequenceOutput, - ModulatedBlock, neutralized_rope_freqs, + ModulatedBlock, remat_block, neutralized_rope_freqs, ) from .vit_common import RotaryEmbedding @@ -37,6 +37,7 @@ class VideoDiT(nn.Module): learn_sigma: bool = False qk_norm: bool = False attention_impl: Optional[str] = None + remat: bool = False use_hilbert: bool = False use_zigzag: bool = False @@ -67,7 +68,7 @@ def setup(self): dim=dim_head, max_seq_len=1024, dtype=self.dtype, name="temporal_rope") def block(name): - return ModulatedBlock( + return remat_block(ModulatedBlock, self.remat)( features=self.emb_features, num_heads=self.num_heads, mixer='attention', @@ -112,10 +113,10 @@ def __call__(self, x, temb, textcontext=None, train: bool = False): freqs_temporal = self.temporal_rope(seq_len=T) for spatial, temporal in zip(self.spatial_blocks, self.temporal_blocks): - tokens = spatial(tokens, conditioning=cond_spatial, freqs_cis=freqs_spatial, train=train) + tokens = spatial(tokens, cond_spatial, freqs_spatial, train) # [B*T, S, F] -> [B*S, T, F] tokens = tokens.reshape(B, T, S, -1).transpose(0, 2, 1, 3).reshape(B * S, T, -1) - tokens = temporal(tokens, conditioning=cond_temporal, freqs_cis=freqs_temporal, train=train) + tokens = temporal(tokens, cond_temporal, freqs_temporal, train) # back to [B*T, S, F] tokens = tokens.reshape(B, S, T, -1).transpose(0, 2, 1, 3).reshape(B * T, S, -1) diff --git a/tests/test_remat.py b/tests/test_remat.py new file mode 100644 index 0000000..c3a0cc6 --- /dev/null +++ b/tests/test_remat.py @@ -0,0 +1,77 @@ +"""Rematerialization must be invisible to everything except memory use. + +Same parameter tree, same forward values, same gradients - otherwise --remat +would silently invalidate checkpoints or change what a run converges to. +""" + +import jax +import jax.numpy as jnp +import pytest + +from flaxdiff.models.simple_dit import SimpleDiT +from flaxdiff.models.simple_mmdit import SimpleMMDiT, HierarchicalMMDiT +from flaxdiff.models.simple_vit import SimpleUDiT +from flaxdiff.models.ssm_dit import HybridSSMAttentionDiT +from flaxdiff.models.video_dit import VideoDiT + +RES = 32 + +BUILDERS = { + 'simple_dit': lambda remat: SimpleDiT( + patch_size=4, emb_features=64, num_layers=2, num_heads=2, mlp_ratio=2, remat=remat), + 'simple_udit': lambda remat: SimpleUDiT( + patch_size=4, emb_features=64, num_layers=2, num_heads=2, mlp_ratio=2, remat=remat), + 'simple_mmdit': lambda remat: SimpleMMDiT( + patch_size=4, emb_features=64, num_layers=2, num_heads=2, mlp_ratio=2, remat=remat), + 'hierarchical_mmdit': lambda remat: HierarchicalMMDiT( + base_patch_size=2, emb_features=(32, 64, 96), num_layers=(1, 1, 1), + num_heads=(2, 2, 2), mlp_ratio=2, remat=remat), + 'hybrid_dit': lambda remat: HybridSSMAttentionDiT( + patch_size=4, emb_features=64, num_layers=2, num_heads=2, mlp_ratio=2, remat=remat), +} + + +def image_inputs(rng): + return (jax.random.normal(rng, (2, RES, RES, 3)), jnp.ones((2,)), + jnp.ones((2, 77, 768), jnp.float32)) + + +@pytest.mark.parametrize('arch', sorted(BUILDERS)) +def test_remat_keeps_parameter_tree_identical(rng, arch): + x, temb, ctx = image_inputs(rng) + plain = BUILDERS[arch](False).init(rng, x, temb, ctx) + remat = BUILDERS[arch](True).init(rng, x, temb, ctx) + + def paths(tree): + return [jax.tree_util.keystr(p) for p, _ in jax.tree_util.tree_leaves_with_path(tree)] + + assert paths(plain) == paths(remat), "remat changed the checkpoint layout" + + +@pytest.mark.parametrize('arch', sorted(BUILDERS)) +def test_remat_preserves_outputs_and_gradients(rng, arch): + x, temb, ctx = image_inputs(rng) + plain, remat = BUILDERS[arch](False), BUILDERS[arch](True) + params = plain.init(rng, x, temb, ctx) + + def loss(model, p): + return jnp.sum(model.apply(p, x, temb, ctx) ** 2) + + assert jnp.allclose(loss(plain, params), loss(remat, params), rtol=1e-5, atol=1e-4) + + g_plain = jax.grad(lambda p: loss(plain, p))(params) + g_remat = jax.grad(lambda p: loss(remat, p))(params) + for a, b in zip(jax.tree.leaves(g_plain), jax.tree.leaves(g_remat)): + assert jnp.allclose(a, b, rtol=1e-3, atol=1e-4) + + +def test_video_dit_remat_matches(): + rng = jax.random.PRNGKey(0) + x = jax.random.normal(rng, (1, 3, 16, 16, 3)) + temb, ctx = jnp.ones((1,)), jnp.ones((1, 77, 768), jnp.float32) + plain = VideoDiT(patch_size=4, emb_features=32, num_layers=1, num_heads=2, mlp_ratio=1) + remat = VideoDiT(patch_size=4, emb_features=32, num_layers=1, num_heads=2, mlp_ratio=1, + remat=True) + params = plain.init(rng, x, temb, ctx) + assert jnp.allclose(plain.apply(params, x, temb, ctx), + remat.apply(params, x, temb, ctx), rtol=1e-5, atol=1e-4) diff --git a/training.py b/training.py index 9ede2e7..f34a72b 100644 --- a/training.py +++ b/training.py @@ -29,7 +29,7 @@ import warnings import traceback -from flaxdiff.utils import defaultTextEncodeModel +from flaxdiff.utils import DEFAULT_MIN_SHARD_SIZE, defaultTextEncodeModel from flaxdiff.inputs import DiffusionInputConfig, ConditionalInputConfig warnings.filterwarnings("ignore") @@ -137,6 +137,13 @@ def boolean_string(s): parser.add_argument('--precision', type=str, default='default', help='precision to use', choices=['high', 'default', 'highest', 'None', None]) parser.add_argument('--distributed_training', type=boolean_string, default=True, help='Should use distributed training or not') +parser.add_argument('--fsdp_size', type=int, default=1, help='Shard parameters over this many devices (FSDP). 1 replicates them (pure data parallelism)') +parser.add_argument('--fsdp_min_param_size', type=int, default=DEFAULT_MIN_SHARD_SIZE, help='Only shard parameters with at least this many elements') +parser.add_argument('--grad_accum_steps', type=int, default=1, help='Accumulate gradients over this many micro-batches before updating') +parser.add_argument('--remat', type=boolean_string, default=False, help='Rematerialize transformer blocks to trade compute for activation memory') +parser.add_argument('--profile_steps', type=int, default=0, help='Write a jax profiler trace covering this many steps of the first epoch') +parser.add_argument('--compilation_cache_dir', type=str, default=None, help='Directory for the persistent XLA compilation cache') +parser.add_argument('--log_every', type=int, default=100, help='Steps between throughput/loss logs') parser.add_argument('--experiment_name', type=str, default=None, help='Experiment name, would be generated if not provided') parser.add_argument('--load_from_checkpoint', type=str, default=None, help='Load from the best previously stored checkpoint. The checkpoint path should be provided') @@ -313,6 +320,7 @@ def main(args): "precision": PRECISION, "output_channels": INPUT_CHANNELS, "attention_impl": args.attention_impl, + "remat": args.remat, } @@ -535,13 +543,18 @@ def main(args): decay_steps=batches * args.learning_rate_decay_epochs, end_value=args.learning_rate_end, ) solver = optimizer(learning_rate, **optimizer_opts) - + if args.clip_grads > 0: solver = optax.chain( optax.clip_by_global_norm(args.clip_grads), solver, ) + if args.grad_accum_steps > 1: + # Accumulate gradients over several micro-batches so the effective batch + # can exceed what fits in device memory at once. + solver = optax.MultiSteps(solver, every_k_schedule=args.grad_accum_steps) + wandb_config = { "project": args.wandb_project, "entity": args.wandb_entity, @@ -572,6 +585,11 @@ def main(args): eval_metrics=eval_metrics, best_tracker_metric=args.best_tracker_metric, ema_decay=args.ema_decay, + fsdp_size=args.fsdp_size, + fsdp_min_param_size=args.fsdp_min_param_size, + compilation_cache_dir=args.compilation_cache_dir, + profile_steps=args.profile_steps, + log_every=args.log_every, ) if trainer.distributed_training: From a644bf5d9e98c656cd232d285ba684356d3d7b94 Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Mon, 27 Jul 2026 02:12:14 -0400 Subject: [PATCH 3/5] test(trainer): cover throughput accounting, divergence and the profiler The performance work is only evaluable if the instrumentation is trustworthy, so it gets tested like the training maths: FLOPs come back positive from the compiled step, samples/sec and step_time_ms are internally consistent, MFU scales the right way with time and device count and is skipped rather than guessed at on hardware with no published peak, the profiler actually writes a trace, and the compilation cache directory is really configured. Divergence detection is covered both ways: a run whose loss is always NaN stops with the consecutive-step count in the message, and a healthy run runs to completion without tripping it. That check now accumulates on device and is read at the logging cadence, so it costs no per-step host sync. Co-Authored-By: Claude Fable 5 --- tests/test_instrumentation.py | 140 ++++++++++++++++++++++++++++++++++ 1 file changed, 140 insertions(+) create mode 100644 tests/test_instrumentation.py diff --git a/tests/test_instrumentation.py b/tests/test_instrumentation.py new file mode 100644 index 0000000..8b4c81f --- /dev/null +++ b/tests/test_instrumentation.py @@ -0,0 +1,140 @@ +"""Throughput accounting, divergence detection and the profiler hook. + +None of the performance work is evaluable without these numbers, so they get +the same treatment as the training maths. +""" + +import os + +import jax +import jax.numpy as jnp +import numpy as np +import optax +import pytest + +from flaxdiff.inputs import DiffusionInputConfig +from flaxdiff.models.simple_dit import SimpleDiT +from flaxdiff.predictors import get_diffusion_preset +from flaxdiff.trainer import GeneralDiffusionTrainer +from flaxdiff.utils import ( + DevicePrefetchIterator, enable_compilation_cache, model_flops_utilization, step_flops, +) + +RES = 8 +BATCH = 8 + + +def make_trainer(tmp_path, **kwargs): + train_schedule, _, transform = get_diffusion_preset("edm") + return GeneralDiffusionTrainer( + model=SimpleDiT(patch_size=4, emb_features=16, num_layers=1, num_heads=2, mlp_ratio=1), + optimizer=optax.adam(1e-3), + noise_schedule=train_schedule, + model_output_transform=transform, + input_config=DiffusionInputConfig( + sample_data_key="image", sample_data_shape=(RES, RES, 3), conditions=[]), + rngs=jax.random.PRNGKey(0), + name="instr", + wandb_config=None, + distributed_training=False, + checkpoint_base_path=str(tmp_path), + **kwargs, + ) + + +def batches(): + images = np.tile(np.linspace(0, 255, RES, dtype=np.float32)[None, :, None, None], + (BATCH, 1, RES, 3)) + while True: + yield {"image": images} + + +def data_dict(): + return {"train": batches, "train_len": BATCH * 8, + "local_batch_size": BATCH, "global_batch_size": BATCH} + + +def test_step_flops_reports_a_positive_count(tmp_path): + trainer = make_trainer(tmp_path) + step = trainer._define_train_step(batch_size=BATCH) + source = DevicePrefetchIterator(batches(), trainer.batch_sharding) + flops = step_flops(step, trainer.state, trainer.rngstate, next(source)) + assert flops is not None and flops > 0 + + +def test_throughput_metrics_are_consistent(tmp_path): + trainer = make_trainer(tmp_path) + trainer.global_batch_size = 64 + metrics = trainer._throughput_metrics(elapsed=2.0, steps=10) + assert metrics["train/step_time_ms"] == pytest.approx(200.0) + assert metrics["train/samples_per_sec"] == pytest.approx(320.0) + + +def test_throughput_metrics_ignore_a_zero_interval(tmp_path): + assert make_trainer(tmp_path)._throughput_metrics(elapsed=0.0, steps=0) == {} + + +def test_mfu_is_skipped_on_unknown_hardware(): + # CPU is deliberately absent from the peak-FLOPs table + assert model_flops_utilization(1e12, 1.0, 8) is None + + +def test_mfu_scales_with_time_and_devices(monkeypatch): + import flaxdiff.utils as utils + monkeypatch.setitem(utils.PEAK_FLOPS_PER_DEVICE, jax.devices()[0].device_kind, 100.0) + assert utils.model_flops_utilization(50.0, 1.0, 1) == pytest.approx(0.5) + assert utils.model_flops_utilization(50.0, 1.0, 2) == pytest.approx(0.25) + assert utils.model_flops_utilization(50.0, 2.0, 1) == pytest.approx(0.25) + + +def test_fit_reports_throughput_to_wandb(tmp_path): + """The logging tick must actually carry the numbers, not just the loss.""" + logged = [] + + class FakeWandb: + def log(self, payload, step=None): + logged.append(payload) + + def define_metric(self, *args, **kwargs): + pass + + trainer = make_trainer(tmp_path, log_every=1) + trainer.wandb = FakeWandb() + trainer.fit(data_dict(), training_steps_per_epoch=3, epochs=1, val_steps_per_epoch=0) + + ticks = [p for p in logged if "train/samples_per_sec" in p] + assert ticks, "no throughput was logged" + assert all(p["train/step_time_ms"] > 0 for p in ticks) + assert all(p["train/samples_per_sec"] > 0 for p in ticks) + + +def test_compilation_cache_directory_is_configured(tmp_path): + path = str(tmp_path / "xla-cache") + enable_compilation_cache(path) + assert os.path.isdir(path) + assert jax.config.jax_compilation_cache_dir == path + + +def test_profiler_writes_a_trace(tmp_path): + trainer = make_trainer(tmp_path, profile_steps=2) + trainer.fit(data_dict(), training_steps_per_epoch=3, epochs=1, val_steps_per_epoch=0) + assert os.path.isdir(trainer.profile_path()) + assert any(files for _, _, files in os.walk(trainer.profile_path())) + + +# -------------------------------------------------------------------------- +# Divergence +# -------------------------------------------------------------------------- + +def test_sustained_non_finite_loss_stops_the_run(tmp_path): + trainer = make_trainer( + tmp_path, log_every=1, max_bad_loss_steps=3, + loss_fn=lambda pred, target: jnp.full_like(pred, jnp.nan)) + with pytest.raises(RuntimeError, match="non-finite"): + trainer.fit(data_dict(), training_steps_per_epoch=8, epochs=1, val_steps_per_epoch=0) + + +def test_healthy_run_does_not_trip_the_detector(tmp_path): + trainer = make_trainer(tmp_path, log_every=1, max_bad_loss_steps=3) + trainer.fit(data_dict(), training_steps_per_epoch=6, epochs=1, val_steps_per_epoch=0) + assert int(trainer.state.step) == 6 From aef3728c00a50f0f4f1667c71e5809b5fda84092 Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Mon, 27 Jul 2026 02:38:28 -0400 Subject: [PATCH 4/5] fix(trainer): restore the data iterator position, and drop code the migration orphaned Resuming rebuilt the template from self.dataset_state, which is always None on a fresh trainer, so the iterator position was written to the checkpoint and then never read back - every resume still restarted the epoch. The template is built from the checkpoint's own structure now, which also keeps checkpoints written before iterator tracking (and those from iterators that cannot report a position) restorable. Grain reports its position as JSON bytes and tensorstore has no dtype for that, so the bytes ride along as a uint8 array and are decoded on restore. Saving is async now, so tests that read a checkpoint back in the same process have to pass through wait_for_checkpoints(); one of them was racing. Also removes what the shard_map migration orphaned: shutil in simple_trainer (its only user was the deleted move_contents_to_subdir), the unused distributed_training local in the train step, the device-count locals in both validation loops, and jnp in dataloaders. Co-Authored-By: Claude Fable 5 --- flaxdiff/data/dataloaders.py | 1 - flaxdiff/trainer/general_diffusion_trainer.py | 3 - flaxdiff/trainer/simple_trainer.py | 37 ++++++----- flaxdiff/utils.py | 2 +- tests/test_parallelism.py | 61 +++++++++++++++++++ tests/test_trainer.py | 3 + 6 files changed, 88 insertions(+), 19 deletions(-) diff --git a/flaxdiff/data/dataloaders.py b/flaxdiff/data/dataloaders.py index b3ce58a..ccc1df1 100644 --- a/flaxdiff/data/dataloaders.py +++ b/flaxdiff/data/dataloaders.py @@ -1,4 +1,3 @@ -import jax.numpy as jnp import grain.python as pygrain from typing import Dict, Any, Optional, Union, List, Callable import numpy as np diff --git a/flaxdiff/trainer/general_diffusion_trainer.py b/flaxdiff/trainer/general_diffusion_trainer.py index a0a7f13..ece09d0 100644 --- a/flaxdiff/trainer/general_diffusion_trainer.py +++ b/flaxdiff/trainer/general_diffusion_trainer.py @@ -246,7 +246,6 @@ def _define_train_step(self, batch_size): model = self.model model_output_transform = self.model_output_transform loss_fn = self.loss_fn - distributed_training = self.distributed_training autoencoder = self.autoencoder unconditional_prob = self.unconditional_prob @@ -438,8 +437,6 @@ def validation_loop( """ Run validation and log samples for both image and video diffusion. """ - global_device_count = jax.device_count() - local_device_count = jax.local_device_count() process_index = jax.process_index() generate_samples = val_step_fn diff --git a/flaxdiff/trainer/simple_trainer.py b/flaxdiff/trainer/simple_trainer.py index 8b57104..a5ad219 100644 --- a/flaxdiff/trainer/simple_trainer.py +++ b/flaxdiff/trainer/simple_trainer.py @@ -20,7 +20,6 @@ ) from flax.training import dynamic_scale as dynamic_scale_lib from dataclasses import dataclass -import shutil PROCESS_COLOR_MAP = { 0: "green", @@ -229,24 +228,30 @@ def checkpoint_path(self): os.makedirs(path) return path - def _checkpoint_template(self): + def _checkpoint_template(self, stored_keys): """Restore template plus the per-leaf args that place arrays on the mesh. Shapes and types come from the freshly built state, so a checkpoint written on one mesh restores onto whatever mesh this run is using. Restoring untyped used to silently discard opt_state and reset the step counter (and with it the lr schedule) on every resume. + + The template must name exactly the keys the checkpoint holds: asking for + one it lacks is an error, and checkpoints predating iterator tracking - + or written from an iterator that cannot report a position - have no + dataset_state. """ - abstract_state = jax.eval_shape(lambda: self.state) template = { 'rngs': self.get_rngstate(), - 'state': abstract_state, + 'state': jax.eval_shape(lambda: self.state), 'best_state': self.best_state, 'best_loss': np.array(self.best_loss), 'epoch': 0, } - if self.dataset_state is not None: - template['dataset_state'] = self.dataset_state + if 'dataset_state' in stored_keys: + # Length varies with the iterator's position, so orbax takes the + # shape from the checkpoint rather than from this placeholder. + template['dataset_state'] = np.zeros((1,), np.uint8) restore_args = jax.tree.map(lambda _: ocp.RestoreArgs(), template) # Only the train state is placed onto the mesh; everything else is # bookkeeping that belongs on the host. @@ -255,25 +260,29 @@ def _checkpoint_template(self): return template, restore_args def load(self, checkpoint_path, checkpoint_step=None, load_directly_from_dir=False): + # The handler has to be registered for item_metadata to report the + # checkpoint's structure, which is what the template is built against. manager = ocp.CheckpointManager( - checkpoint_path, options=ocp.CheckpointManagerOptions(max_to_keep=4, create=False)) + checkpoint_path, options=ocp.CheckpointManagerOptions(max_to_keep=4, create=False), + item_handlers=ocp.PyTreeCheckpointHandler()) step = manager.latest_step() if checkpoint_step is None else checkpoint_step print("Loading model from checkpoint at step ", step) self.loaded_checkpoint_path = os.path.join( checkpoint_path if checkpoint_path else self.checkpoint_path(), f"{step}") - template, restore_args = self._checkpoint_template() - # A checkpoint written before iterator state was tracked simply has no - # such key, so ask for it only when the run can already produce one. target = checkpoint_path if load_directly_from_dir else step + template, restore_args = self._checkpoint_template(manager.item_metadata(target).keys()) ckpt = manager.restore( target, args=ocp.args.PyTreeRestore(item=template, restore_args=restore_args)) self.state = ckpt['state'] self.best_state = ckpt['best_state'] self.rngstate = ckpt['rngs'] - self.dataset_state = ckpt.get('dataset_state') + stored_position = ckpt.get('dataset_state') + self.dataset_state = ( + None if stored_position is None + else np.asarray(stored_position, np.uint8).tobytes()) self.best_loss = float(ckpt['best_loss']) if self.best_loss == 0: # It cant be zero as that must have been some problem @@ -294,7 +303,9 @@ def save(self, epoch=0, step=0, state=None, rngstate=None): 'epoch': epoch, } if self.dataset_state is not None: - ckpt['dataset_state'] = self.dataset_state + # Grain reports its position as JSON bytes, which tensorstore has no + # dtype for; the raw bytes ride along as a uint8 array instead. + ckpt['dataset_state'] = np.frombuffer(self.dataset_state, np.uint8) try: self.checkpointer.save(step, args=ocp.args.PyTreeSave(ckpt), force=True) except Exception as e: @@ -327,8 +338,6 @@ def validation_loop( val_steps_per_epoch, current_step, ): - global_device_count = jax.device_count() - local_device_count = jax.local_device_count() process_index = jax.process_index() val_ds = iter(val_ds()) if val_ds else None diff --git a/flaxdiff/utils.py b/flaxdiff/utils.py index c8b54e3..1d63026 100644 --- a/flaxdiff/utils.py +++ b/flaxdiff/utils.py @@ -2,7 +2,7 @@ import jax.numpy as jnp import flax.struct as struct import flax.linen as nn -from typing import Any, Iterator, Optional +from typing import Iterator, Optional import numpy as np import os import queue diff --git a/tests/test_parallelism.py b/tests/test_parallelism.py index b007140..d0edae3 100644 --- a/tests/test_parallelism.py +++ b/tests/test_parallelism.py @@ -236,6 +236,67 @@ def test_sharded_checkpoint_roundtrips(tmp_path): np.testing.assert_allclose(np.asarray(before), np.asarray(after)) +def grain_image_loader(num_records=256): + """A checkpointable source that yields distinguishable image batches.""" + import grain.python as pygrain + + class ToImage(pygrain.MapTransform): + def map(self, index): + return {"image": np.full((RES, RES, 3), index, np.float32)} + + return pygrain.DataLoader( + data_source=pygrain.RangeDataSource(0, num_records, 1), + sampler=pygrain.IndexSampler(num_records=num_records, shuffle=False, seed=0, + num_epochs=4, shard_options=pygrain.NoSharding()), + operations=[ToImage(), pygrain.Batch(BATCH, drop_remainder=True)], + worker_count=0, + ) + + +def test_resume_continues_mid_epoch(tmp_path): + """A resumed run must not replay the batches it already trained on.""" + steps = 3 + data = {"train": grain_image_loader, "train_len": BATCH * 64, + "local_batch_size": BATCH, "global_batch_size": BATCH} + + trainer = make_trainer(tmp_path, "resume", distributed_training=True) + trainer.fit(data, training_steps_per_epoch=steps, epochs=1, val_steps_per_epoch=0) + assert trainer.dataset_state is not None, "iterator position was never captured" + trainer.wait_for_checkpoints() + + resumed = make_trainer(tmp_path, "resume", distributed_training=True, + load_from_checkpoint=trainer.checkpoint_path()) + assert resumed.dataset_state == trainer.dataset_state + + # The next batch after resuming is the one the first run would have seen next + reference = DevicePrefetchIterator( + iter(grain_image_loader()), trainer.batch_sharding, + source_state=trainer.dataset_state) + resumed_iter = DevicePrefetchIterator( + iter(grain_image_loader()), resumed.batch_sharding, + source_state=resumed.dataset_state) + np.testing.assert_array_equal(np.asarray(next(resumed_iter)["image"]), + np.asarray(next(reference)["image"])) + + # and it is past the batches already consumed + fresh = DevicePrefetchIterator(iter(grain_image_loader()), resumed.batch_sharding) + first = np.asarray(next(fresh)["image"]) + assert not np.array_equal(np.asarray(next(resumed_iter)["image"]), first) + + +def test_load_tolerates_checkpoints_without_iterator_state(tmp_path): + """Checkpoints written before iterator tracking must still restore.""" + trainer = make_trainer(tmp_path, "legacy", distributed_training=True) + assert trainer.dataset_state is None + trainer.save(epoch=0, step=1) + trainer.wait_for_checkpoints() + + restored = make_trainer(tmp_path, "legacy", distributed_training=True, + load_from_checkpoint=trainer.checkpoint_path()) + assert int(restored.state.step) == 0 + assert restored.dataset_state is None + + def test_checkpoint_restores_onto_a_different_mesh(tmp_path): """A run saved with FSDP must be resumable on a replicated mesh.""" trainer = make_trainer(tmp_path, "mesh", distributed_training=True, diff --git a/tests/test_trainer.py b/tests/test_trainer.py index 0f21004..21540c3 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -59,6 +59,8 @@ def test_save_writes_checkpoint(tmp_path): """save() swallows exceptions, so assert the checkpoint actually landed.""" trainer = make_trainer(tmp_path) trainer.save(epoch=0, step=1) + # Saving is async to keep it off the training loop; this is the barrier + trainer.wait_for_checkpoints() assert trainer.checkpointer.latest_step() == 1 @@ -69,6 +71,7 @@ def test_restore_preserves_optimizer_state(tmp_path): grads = jax.tree.map(jnp.ones_like, trainer.state.params) trainer.state = trainer.state.apply_gradients(grads=grads).apply_ema(0.99) trainer.save(epoch=0, step=1) + trainer.wait_for_checkpoints() restored = make_trainer(tmp_path, load_from_checkpoint=trainer.checkpoint_path()) assert int(restored.state.step) == 1, "optimizer step counter was reset" From 1587b09798858d570943ee2abd023ccac32d7596 Mon Sep 17 00:00:00 2001 From: Ashish Kumar Singh Date: Mon, 27 Jul 2026 10:03:19 -0400 Subject: [PATCH 5/5] test(trainer): cover gradient accumulation across the sharding heuristic Co-Authored-By: Claude Fable 5 --- tests/test_parallelism.py | 33 +++++++++++++++++++++++++++++++-- 1 file changed, 31 insertions(+), 2 deletions(-) diff --git a/tests/test_parallelism.py b/tests/test_parallelism.py index d0edae3..931d639 100644 --- a/tests/test_parallelism.py +++ b/tests/test_parallelism.py @@ -27,11 +27,12 @@ TINY = 256 -def make_trainer(tmp_path, name, distributed_training, fsdp_size=1, **kwargs): +def make_trainer(tmp_path, name, distributed_training, fsdp_size=1, + optimizer=None, **kwargs): train_schedule, _, transform = get_diffusion_preset("edm") return GeneralDiffusionTrainer( model=SimpleDiT(patch_size=4, emb_features=32, num_layers=1, num_heads=2, mlp_ratio=1), - optimizer=optax.adam(1e-3), + optimizer=optax.adam(1e-3) if optimizer is None else optimizer, noise_schedule=train_schedule, model_output_transform=transform, input_config=DiffusionInputConfig( @@ -236,6 +237,34 @@ def test_sharded_checkpoint_roundtrips(tmp_path): np.testing.assert_allclose(np.asarray(before), np.asarray(after)) +def test_gradient_accumulation_updates_only_on_the_boundary(tmp_path): + """MultiSteps must hold the params still until k micro-batches have run. + + Also covers the accumulator surviving the sharding heuristic: its buffers + are param-shaped, so they pick up the param specs. + """ + accum = 3 + trainer = make_trainer(tmp_path, "accum", distributed_training=True, fsdp_size=2, + fsdp_min_param_size=TINY, + optimizer=optax.MultiSteps(optax.sgd(0.5), every_k_schedule=accum)) + + train_step = trainer._define_train_step(batch_size=BATCH) + source = DevicePrefetchIterator(batches(), trainer.batch_sharding) + state, rng = trainer.state, trainer.rngstate + + def snapshot(s): + return [np.asarray(x).copy() for x in jax.tree.leaves(s.params)] + + reference = snapshot(state) + for micro in range(1, accum * 2 + 1): + state, _, rng, _ = train_step(state, rng, next(source)) + moved = any(not np.array_equal(a, b) for a, b in zip(reference, snapshot(state))) + at_boundary = micro % accum == 0 + assert moved == at_boundary, f"micro-step {micro}: moved={moved}" + if at_boundary: + reference = snapshot(state) + + def grain_image_loader(num_records=256): """A checkpointable source that yields distinguishable image batches.""" import grain.python as pygrain