improves performance by 14% on v5p.

Author: jfacevedo-google

src/maxdiffusion/configs/base_wan_14b.yml

@@ -52,16 +52,7 @@ from_pt: True
 split_head_dim: True
 attention: 'flash' # Supported attention: dot_product, flash, cudnn_flash_te
 
-flash_block_sizes: {
-  "block_q" : 1024,
-  "block_kv_compute" : 1024,
-  "block_kv" : 1024,
-  "block_q_dkv" : 1024,
-  "block_kv_dkv" : 1024,
-  "block_kv_dkv_compute" : 1024,
-  "block_q_dq" : 1024,
-  "block_kv_dq" : 1024
-}
+flash_block_sizes: {}
 # GroupNorm groups
 norm_num_groups: 32
 
@@ -127,6 +118,7 @@ logical_axis_rules: [
                       ['mlp','tensor'],
                       ['embed','fsdp'],
                       ['heads', 'tensor'],
+                      ['norm', 'fsdp'],
                       ['conv_batch', ['data','fsdp']],
                       ['out_channels', 'tensor'],
                       ['conv_out', 'fsdp'],

src/maxdiffusion/generate_wan.py

@@ -35,7 +35,7 @@ def run(config):
 
   print("compile time: ", (time.perf_counter() - s0))
   for i in range(len(videos)):
-    export_to_video(videos[i], f"wan_output_{i}.mp4", fps=16)
+    export_to_video(videos[i], f"wan_output_{config.seed}_{i}.mp4", fps=16)
   s0 = time.perf_counter()
   with jax.profiler.trace("/tmp/trace/"):
     videos = pipeline(
@@ -49,7 +49,7 @@ def run(config):
     )
   print("generation time: ", (time.perf_counter() - s0))
   for i in range(len(videos)):
-    export_to_video(videos[i], f"wan_output_{i}.mp4", fps=16)
+    export_to_video(videos[i], f"wan_output_{config.seed}_{i}.mp4", fps=16)
 
 
 def main(argv: Sequence[str]) -> None:

src/maxdiffusion/models/attention_flax.py

@@ -19,6 +19,7 @@
 import flax.linen as nn
 from flax import nnx
 import jax
+from jax.sharding import PartitionSpec
 import jax.numpy as jnp
 from jax.experimental import shard_map
 from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask
@@ -139,21 +140,23 @@ def _tpu_flash_attention(
   heads: int,
   mesh: Mesh,
   flash_axis_names: AxisNames,
-  flash_block_sizes: BlockSizes) -> jax.Array:
+  flash_block_sizes: BlockSizes,
+  dtype: jnp.dtype = jnp.float32) -> jax.Array:
   """TPU Flash Attention"""
 
+  max_block_size = 1024 if dtype == jnp.bfloat16 else 512
   if flash_block_sizes:
     block_sizes = flash_block_sizes
   else:
     block_sizes = splash_attention_kernel.BlockSizes(
-        block_q=min(512, query.shape[2]),
-        block_kv_compute=min(512, key.shape[2]),
-        block_kv=min(512, key.shape[2]),
-        block_q_dkv=min(512, query.shape[2]),
-        block_kv_dkv=min(512, key.shape[2]),
-        block_kv_dkv_compute=min(512, query.shape[2]),
-        block_q_dq=min(512, query.shape[2]),
-        block_kv_dq=min(512, query.shape[2]),
+        block_q=min(max_block_size, query.shape[2]),
+        block_kv_compute=min(max_block_size, key.shape[2]),
+        block_kv=min(max_block_size, key.shape[2]),
+        block_q_dkv=min(max_block_size, query.shape[2]),
+        block_kv_dkv=min(max_block_size, key.shape[2]),
+        block_kv_dkv_compute=min(max_block_size, query.shape[2]),
+        block_q_dq=min(max_block_size, query.shape[2]),
+        block_kv_dq=min(max_block_size, query.shape[2]),
     )
 
   query, kv_size, query_seq_len = _reshape_data_for_flash(query, heads, block_sizes.block_q)
@@ -340,7 +343,7 @@ def _apply_attention(
   if attention_kernel == "dot_product" or use_memory_efficient_attention or not can_use_flash_attention:
     return _apply_attention_dot(query, key, value, dtype, heads, dim_head, scale, split_head_dim, float32_qk_product, use_memory_efficient_attention)
   elif attention_kernel == "flash":
-    return _tpu_flash_attention(query, key * scale, value, heads, mesh, flash_axis_names, flash_block_sizes)
+    return _tpu_flash_attention(query, key * scale, value, heads, mesh, flash_axis_names, flash_block_sizes, dtype)
   elif attention_kernel == "cudnn_flash_te":
     return _cudnn_flash_attention(query, key, value, heads, mesh, dpa_layer)
   else:
@@ -668,15 +671,15 @@ def __init__(
         rngs=rngs,
         epsilon=eps,
         dtype=dtype,
-        scale_init=nnx.with_partitioning(nnx.initializers.ones, ("heads", )),
+        scale_init=nnx.with_partitioning(nnx.initializers.ones, ("norm", )),
         param_dtype=weights_dtype
       )
 
       self.norm_k = nnx.RMSNorm(
         num_features=self.inner_dim,
         rngs=rngs,
         dtype=dtype,
-        scale_init=nnx.with_partitioning(nnx.initializers.ones, ("heads", )),
+        scale_init=nnx.with_partitioning(nnx.initializers.ones, ("norm", )),
         param_dtype=weights_dtype
       )
 
@@ -702,9 +705,12 @@ def __call__(
     encoder_hidden_states: jax.Array = None,
     rotary_emb: Optional[jax.Array] = None
   ) -> jax.Array:
+    hidden_states = jax.lax.with_sharding_constraint(hidden_states, PartitionSpec('data', 'fsdp','tensor'))
+    encoder_hidden_states = jax.lax.with_sharding_constraint(encoder_hidden_states, PartitionSpec('data', 'fsdp','tensor'))
     dtype = hidden_states.dtype
     if encoder_hidden_states is None:
       encoder_hidden_states = hidden_states
+
     query_proj = self.query(hidden_states)
     key_proj = self.key(encoder_hidden_states)
     value_proj = self.value(encoder_hidden_states)
@@ -717,8 +723,13 @@ def __call__(
       key_proj = _unflatten_heads(key_proj, self.heads)
       value_proj = _unflatten_heads(value_proj, self.heads)
       query_proj, key_proj = self._apply_rope(query_proj, key_proj, rotary_emb)
-    
+      query_proj = jax.lax.with_sharding_constraint(query_proj, PartitionSpec('data', 'tensor', None, None))
+      key_proj = jax.lax.with_sharding_constraint(key_proj, PartitionSpec('data', 'tensor', None, None))
+      value_proj = jax.lax.with_sharding_constraint(value_proj, PartitionSpec('data', 'tensor', None, None))
+
     attn_output = self.attention_op.apply_attention(query_proj, key_proj, value_proj)
+    attn_output = jax.lax.with_sharding_constraint(attn_output, PartitionSpec('data', None, None))
+
     attn_output = attn_output.astype(dtype=dtype)
 
     hidden_states = self.proj_attn(attn_output)

src/maxdiffusion/models/wan/transformers/transformer_wan.py

@@ -179,8 +179,6 @@ def __init__(
       dtype=dtype,
       param_dtype=weights_dtype,
       precision=precision,
-      kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), ("embed", "mlp",)),
-      bias_init=nnx.with_partitioning(nnx.initializers.zeros, ("mlp",)),
     )
 
   def __call__(self, x: jax.Array) -> jax.Array:
@@ -231,7 +229,6 @@ def __init__(
       param_dtype=weights_dtype,
       precision=precision,
       kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), ("mlp", "embed",)),
-      bias_init=nnx.with_partitioning(nnx.initializers.zeros, ("embed",)),
     )
 
   def __call__(self, hidden_states: jax.Array) -> jax.Array:
@@ -338,7 +335,7 @@ def __call__(
 
     # 1. Self-attention
     norm_hidden_states = (self.norm1(hidden_states.astype(jnp.float32)) * (1 + scale_msa) + shift_msa).astype(hidden_states.dtype)
-    attn_output = self.attn1(hidden_states=norm_hidden_states, rotary_emb=rotary_emb)
+    attn_output = self.attn1(hidden_states=norm_hidden_states, encoder_hidden_states=norm_hidden_states, rotary_emb=rotary_emb)
     hidden_states = (hidden_states.astype(jnp.float32) + attn_output * gate_msa).astype(hidden_states.dtype)
 
     # 2. Cross-attention
@@ -443,11 +440,13 @@ def __init__(
       dtype=dtype,
       param_dtype=weights_dtype,
       precision=precision,
-      kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), ("embed", "mlp",)),
-      bias_init=nnx.with_partitioning(nnx.initializers.zeros, ("mlp",)),
+      kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), ("embed", None)),
     )
     key = rngs.params()
-    self.scale_shift_table = nnx.Param(jax.random.normal(key, (1, 2, inner_dim)) / inner_dim**0.5)
+    self.scale_shift_table = nnx.Param(
+      jax.random.normal(key, (1, 2, inner_dim)) / inner_dim**0.5,
+      kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), (None, None, "embed"))
+    )
 
   def __call__(
     self,