jitting and sharding vae, refactored for loops in jitted VAE, 132 sec on 16 TPUs

Author: eltsai

src/maxdiffusion/models/attention_flax.py

@@ -334,15 +334,15 @@ def wrap_flash_attention(query, key, value):
           mask=mask,
           q_seq_shards=1,  # the sizes of the axis is sharding over seq_len
           config=convert_to_tokamax_splash_config(block_sizes, residual_checkpoint_name=residual_checkpoint_name),
-          save_residuals=True if "ring" in attention_kernel else False,
+          save_residuals=False,
       )
     elif attention_kernel == "tokamax_ring":
       mask = tokamax_splash_attention_mask.FullMask(_shape=(query.shape[2], key.shape[2]),)
       splash_kernel = tokamax_ring_attention_kernel.make_ring_attention(
           mask=mask,
           is_mqa=False,
           config=convert_to_tokamax_splash_config(block_sizes, residual_checkpoint_name=residual_checkpoint_name),
-          save_residuals=True,
+          save_residuals=False,
           ring_axis="fsdp",
       )
     else:

src/maxdiffusion/models/vae_flax.py

@@ -22,11 +22,13 @@
 import flax.linen as nn
 import jax
 import jax.numpy as jnp
+from jax import tree_util
 from flax.core.frozen_dict import FrozenDict
 
 from ..configuration_utils import ConfigMixin, flax_register_to_config
 from ..utils import BaseOutput
 from .modeling_flax_utils import FlaxModelMixin
+ 
 
 
 @flax.struct.dataclass
@@ -931,3 +933,29 @@ def __call__(self, sample, sample_posterior=False, deterministic: bool = True, r
       return (sample,)
 
     return FlaxDecoderOutput(sample=sample)
+
+class WanDiagonalGaussianDistribution(FlaxDiagonalGaussianDistribution):
+  pass
+
+
+def _wan_diag_gauss_dist_flatten(dist):
+  return (dist.mean, dist.logvar, dist.std, dist.var), (dist.deterministic,)
+
+
+def _wan_diag_gauss_dist_unflatten(aux, children):
+  mean, logvar, std, var = children
+  deterministic = aux[0]
+  obj = WanDiagonalGaussianDistribution.__new__(WanDiagonalGaussianDistribution)
+  obj.mean = mean
+  obj.logvar = logvar
+  obj.std = std
+  obj.var = var
+  obj.deterministic = deterministic
+  return obj
+
+
+tree_util.register_pytree_node(
+    WanDiagonalGaussianDistribution,
+    _wan_diag_gauss_dist_flatten,
+    _wan_diag_gauss_dist_unflatten,
+)