diff --git a/flaxdiff/data/dataloaders.py b/flaxdiff/data/dataloaders.py
index a388bec..ccc1df1 100644
--- a/flaxdiff/data/dataloaders.py
+++ b/flaxdiff/data/dataloaders.py
@@ -1,87 +1,15 @@
-import jax.numpy as jnp
 import grain.python as pygrain
 from typing import Dict, Any, Optional, Union, List, Callable
 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 +335,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 +365,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 +384,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 +430,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 +467,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 +482,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 +525,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/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/flaxdiff/trainer/general_diffusion_trainer.py b/flaxdiff/trainer/general_diffusion_trainer.py
index b651d9d..ece09d0 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']
@@ -247,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
         
@@ -263,13 +261,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 +322,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):
         """
@@ -451,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
         
@@ -465,9 +449,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 +705,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..a5ad219 100644
--- a/flaxdiff/trainer/simple_trainer.py
+++ b/flaxdiff/trainer/simple_trainer.py
@@ -1,32 +1,25 @@
-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
 
 PROCESS_COLOR_MAP = {
     0: "green",
@@ -49,80 +42,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 +66,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 +144,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 +176,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 +228,61 @@ 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, 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.
+        """
         template = {
             'rngs': self.get_rngstate(),
-            'state': self.get_state(),
-            'best_state': self.get_best_state(),
+            'state': jax.eval_shape(lambda: self.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 '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.
+        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):
+        # 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),
+            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}")
+
+        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']
+        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
@@ -328,25 +292,24 @@ 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:
+            # 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:
-            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")
@@ -375,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
@@ -386,9 +347,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 +374,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 +547,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 +562,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..1d63026 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 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_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
diff --git a/tests/test_parallelism.py b/tests/test_parallelism.py
new file mode 100644
index 0000000..931d639
--- /dev/null
+++ b/tests/test_parallelism.py
@@ -0,0 +1,345 @@
+"""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,
+                 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) if optimizer is None else optimizer,
+        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_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
+
+    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,
+                           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))
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/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"
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: