fix ag from vmap/scan.

Author: entrpn

src/maxdiffusion/data_preprocessing/wan_pusav1_to_tfrecords.py

@@ -20,18 +20,11 @@
 """
 
 import os
-import functools
 from absl import app
-from typing import Sequence, Union, List
-from datasets import load_dataset
+from typing import Sequence
 import csv
-import numpy as np
-import jax
 import jax.numpy as jnp
-from jax.sharding import Mesh
-from maxdiffusion import pyconfig, max_utils
-from maxdiffusion.pipelines.wan.wan_pipeline import WanPipeline
-from maxdiffusion.video_processor import VideoProcessor
+from maxdiffusion import pyconfig
 
 import torch
 import tensorflow as tf
@@ -72,6 +65,7 @@ def create_example(latent, hidden_states):
   example = tf.train.Example(features=tf.train.Features(feature=feature))
   return example.SerializeToString()
 
+
 def generate_dataset(config):
 
   tfrecords_dir = config.tfrecords_dir
@@ -88,7 +82,7 @@ def generate_dataset(config):
 
   # Load dataset
   metadata_path = os.path.join(config.train_data_dir, "metadata.csv")
-  with open(metadata_path, 'r', newline='') as file:
+  with open(metadata_path, "r", newline="") as file:
     # Create a csv.reader object
     csv_reader = csv.reader(file)
     next(csv_reader)
@@ -99,11 +93,11 @@ def generate_dataset(config):
     # Iterate over each row in the CSV file
     for row in csv_reader:
       video_name = row[0]
-      pth_path = os.path.join(config.train_data_dir,"train", f"{video_name}.tensors.pth")
-      loaded_state_dict = torch.load(pth_path, map_location=torch.device('cpu'))
+      pth_path = os.path.join(config.train_data_dir, "train", f"{video_name}.tensors.pth")
+      loaded_state_dict = torch.load(pth_path, map_location=torch.device("cpu"))
       prompt_embeds = loaded_state_dict["prompt_emb"]["context"].squeeze()
       latent = loaded_state_dict["latents"]
-      
+
       # Format we want(Batch, channels, Frames, Height, Width)
       # Save them as float32 because numpy cannot read bfloat16.
       latent = jnp.array(latent.float().numpy(), dtype=jnp.float32)
@@ -120,6 +114,7 @@ def generate_dataset(config):
         )
         shard_record_count = 0
 
+
 def run(config):
   generate_dataset(config)
 

src/maxdiffusion/models/attention_flax.py

@@ -674,8 +674,10 @@ def __init__(
         dtype=dtype,
         quant=quant,
     )
-
-    kernel_axes = ("embed", "heads")
+    # None axes corresponds to the stacked weights across all blocks
+    # because of the use of nnx.vmap and nnx.scan.
+    # Dims are [num_blocks, embed, heads]
+    kernel_axes = (None, "embed", "heads")
     qkv_init_kernel = nnx.with_partitioning(nnx.initializers.lecun_normal(), kernel_axes)
 
     self.query = nnx.Linear(
@@ -686,7 +688,7 @@ def __init__(
         dtype=dtype,
         param_dtype=weights_dtype,
         precision=precision,
-        bias_init=nnx.with_partitioning(nnx.initializers.zeros, ("embed",)),
+        bias_init=nnx.with_partitioning(nnx.initializers.zeros, (None, "embed",)),
     )
 
     self.key = nnx.Linear(
@@ -697,7 +699,7 @@ def __init__(
         dtype=dtype,
         param_dtype=weights_dtype,
         precision=precision,
-        bias_init=nnx.with_partitioning(nnx.initializers.zeros, ("embed",)),
+        bias_init=nnx.with_partitioning(nnx.initializers.zeros, (None, "embed",)),
     )
 
     self.value = nnx.Linear(
@@ -708,14 +710,14 @@ def __init__(
         dtype=dtype,
         param_dtype=weights_dtype,
         precision=precision,
-        bias_init=nnx.with_partitioning(nnx.initializers.zeros, ("embed",)),
+        bias_init=nnx.with_partitioning(nnx.initializers.zeros, (None, "embed",)),
     )
 
     self.proj_attn = nnx.Linear(
         rngs=rngs,
         in_features=self.inner_dim,
         out_features=self.inner_dim,
-        kernel_init=nnx.with_partitioning(nnx.initializers.lecun_normal(), ("heads", "embed")),
+        kernel_init=nnx.with_partitioning(nnx.initializers.lecun_normal(), (None, "heads", "embed")),
         dtype=dtype,
         param_dtype=weights_dtype,
         precision=precision,
@@ -729,15 +731,15 @@ def __init__(
           rngs=rngs,
           epsilon=eps,
           dtype=dtype,
-          scale_init=nnx.with_partitioning(nnx.initializers.ones, ("norm",)),
+          scale_init=nnx.with_partitioning(nnx.initializers.ones, (None, "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, ("norm",)),
+          scale_init=nnx.with_partitioning(nnx.initializers.ones, (None, "norm",)),
           param_dtype=weights_dtype,
       )
 

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

@@ -398,7 +398,7 @@ def __init__(
 
     # 3. Transformer blocks
     @nnx.split_rngs(splits=num_layers)
-    @nnx.vmap
+    @nnx.vmap(in_axes=0, out_axes=0)
     def init_block(rngs):
       return WanTransformerBlock(
           rngs=rngs,
@@ -416,6 +416,7 @@ def init_block(rngs):
           precision=precision,
           attention=attention,
       )
+
     self.blocks = init_block(rngs)
 
     self.norm_out = FP32LayerNorm(rngs=rngs, dim=inner_dim, eps=eps, elementwise_affine=False)
@@ -471,10 +472,10 @@ def scan_fn(carry, block):
 
     initial_carry = (hidden_states, encoder_hidden_states, timestep_proj, rotary_emb)
     final_carry = nnx.scan(
-      scan_fn,
-      length=self.num_layers,
-      in_axes=(nnx.Carry, 0),
-      out_axes=nnx.Carry,
+        scan_fn,
+        length=self.num_layers,
+        in_axes=(nnx.Carry, 0),
+        out_axes=nnx.Carry,
     )(initial_carry, self.blocks)
 
     hidden_states = final_carry[0]

src/maxdiffusion/models/wan/wan_utils.py

@@ -87,7 +87,9 @@ def load_fusionx_transformer(pretrained_model_name_or_path: str, eval_shapes: di
           new_key = ("blocks",) + pt_tuple_key[2:]
           block_index = int(pt_tuple_key[1])
           pt_tuple_key = new_key
-        flax_key, flax_tensor = rename_key_and_reshape_tensor(pt_tuple_key, tensor, random_flax_state_dict, model_type=WAN_MODEL)
+        flax_key, flax_tensor = rename_key_and_reshape_tensor(
+            pt_tuple_key, tensor, random_flax_state_dict, model_type=WAN_MODEL
+        )
         flax_key = rename_for_nnx(flax_key)
         flax_key = _tuple_str_to_int(flax_key)
 
@@ -133,11 +135,12 @@ def load_causvid_transformer(pretrained_model_name_or_path: str, eval_shapes: di
           new_key = ("blocks",) + pt_tuple_key[2:]
           block_index = int(pt_tuple_key[1])
           pt_tuple_key = new_key
-        flax_key, flax_tensor = rename_key_and_reshape_tensor(pt_tuple_key, tensor, random_flax_state_dict, model_type=WAN_MODEL)
+        flax_key, flax_tensor = rename_key_and_reshape_tensor(
+            pt_tuple_key, tensor, random_flax_state_dict, model_type=WAN_MODEL
+        )
         flax_key = rename_for_nnx(flax_key)
         flax_key = _tuple_str_to_int(flax_key)
 
-        
         if "blocks" in flax_key:
           if flax_key in flax_state_dict:
             new_tensor = flax_state_dict[flax_key]
@@ -224,7 +227,9 @@ def load_base_wan_transformer(pretrained_model_name_or_path: str, eval_shapes: d
         new_key = ("blocks",) + pt_tuple_key[2:]
         block_index = int(pt_tuple_key[1])
         pt_tuple_key = new_key
-      flax_key, flax_tensor = rename_key_and_reshape_tensor(pt_tuple_key, tensor, random_flax_state_dict, model_type=WAN_MODEL)
+      flax_key, flax_tensor = rename_key_and_reshape_tensor(
+          pt_tuple_key, tensor, random_flax_state_dict, model_type=WAN_MODEL
+      )
       flax_key = rename_for_nnx(flax_key)
       flax_key = _tuple_str_to_int(flax_key)