Phase 3: real FSDP, jit+NamedSharding, measurement, mid-epoch resume

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 {

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.
 

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)

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])

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)