support bs > 1. Issue where all gens except for 1st coming out bad.

Author: jfacevedo-google

src/maxdiffusion/configs/base_wan_14b.yml

@@ -53,14 +53,14 @@ split_head_dim: True
 attention: 'flash' # Supported attention: dot_product, flash, cudnn_flash_te
 
 flash_block_sizes: {
-  "block_q" : 512,
-  "block_kv_compute" : 512,
-  "block_kv" : 512,
-  "block_q_dkv" : 512,
-  "block_kv_dkv" : 512,
-  "block_kv_dkv_compute" : 512,
-  "block_q_dq" : 512,
-  "block_kv_dq" : 512
+  "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
 }
 # GroupNorm groups
 norm_num_groups: 32

src/maxdiffusion/generate_wan.py

@@ -23,7 +23,7 @@
 def run(config):
   pipeline = WanPipeline.from_pretrained(config)
   s0 = time.perf_counter()
-  video = pipeline(
+  videos = pipeline(
     prompt=config.prompt,
     negative_prompt=config.negative_prompt,
     height=config.height,
@@ -34,10 +34,11 @@ def run(config):
   )
 
   print("compile time: ", (time.perf_counter() - s0))
-  export_to_video(video[0], "jax_output.mp4", fps=16)
+  for i in range(len(videos)):
+    export_to_video(videos[i], f"wan_output_{i}.mp4", fps=16)
   s0 = time.perf_counter()
   with jax.profiler.trace("/tmp/trace/"):
-    video = pipeline(
+    videos = pipeline(
       prompt=config.prompt,
       negative_prompt=config.negative_prompt,
       height=config.height,
@@ -47,7 +48,8 @@ def run(config):
       guidance_scale=config.guidance_scale,
     )
   print("generation time: ", (time.perf_counter() - s0))
-  export_to_video(video[0], "jax_output.mp4", fps=16)
+  for i in range(len(videos)):
+    export_to_video(videos[i], f"wan_output_{i}.mp4", fps=16)
 
 
 def main(argv: Sequence[str]) -> None:

src/maxdiffusion/models/wan/autoencoder_kl_wan.py

@@ -1131,12 +1131,17 @@ def _decode(
         # Ideally shouldn't need to do this however, can't find where the frame is going out of sync.
         # Most likely due to an incorrect reshaping in the decoder.
         fm1, fm2, fm3, fm4 = out_[:, 0, :, :, :], out_[:, 1, :, :, :], out_[:, 2, :, :, :], out_[:, 3, :, :, :]
-        if len(fm1.shape) == 4:
-          fm1 = jnp.expand_dims(fm1, axis=0)
-          fm2 = jnp.expand_dims(fm2, axis=0)
-          fm3 = jnp.expand_dims(fm3, axis=0)
-          fm4 = jnp.expand_dims(fm4, axis=0)
+        # When batch_size is 0, expand batch dim for contatenation
+        # else, expand frame dim for concatenation so that batch dim stays intact.
+        axis=0
+        if fm1.shape[0] > 1:
+          axis=1
 
+        if len(fm1.shape) == 4:
+          fm1 = jnp.expand_dims(fm1, axis=axis)
+          fm2 = jnp.expand_dims(fm2, axis=axis)
+          fm3 = jnp.expand_dims(fm3, axis=axis)
+          fm4 = jnp.expand_dims(fm4, axis=axis)
         out = jnp.concatenate([out, fm1, fm3, fm2, fm4], axis=1)
     out = jnp.clip(out, min=-1.0, max=1.0)
     feat_cache.clear_cache()

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

@@ -398,7 +398,7 @@ def __init__(
         dtype=dtype,
         param_dtype=weights_dtype,
         precision=precision,
-        kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), ("batch",)),
+        kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), (None, None, None, None, "conv_out",)),
     )
 
     # 2. Condition embeddings

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -17,13 +17,14 @@
 import numpy as np
 import jax
 import jax.numpy as jnp
-from jax.sharding import Mesh, PositionalSharding
+from jax.sharding import Mesh, PositionalSharding, PartitionSpec as P
 import flax
 import flax.linen as nn
 from flax import nnx
 from ...pyconfig import HyperParameters
 from ... import max_logging
 from ... import max_utils
+from ...max_utils import get_flash_block_sizes, get_precision
 from ...models.wan.wan_utils import load_wan_transformer, load_wan_vae
 from ...models.wan.transformers.transformer_wan import WanModel
 from ...models.wan.autoencoder_kl_wan import AutoencoderKLWan, AutoencoderKLWanCache
@@ -59,11 +60,12 @@ def _add_sharding_rule(vs: nnx.VariableState, logical_axis_rules) -> nnx.Variabl
 
 partial(nnx.jit, static_argnums=(3,))
 def create_sharded_logical_transformer(devices_array: np.array, mesh: Mesh, rngs: nnx.Rngs, config: HyperParameters):
-  # breakpoint()
+
   def create_model(rngs: nnx.Rngs, wan_config: dict):
     wan_transformer = WanModel(**wan_config, rngs=rngs)
     return wan_transformer
 
+  # 1. Load config.
   wan_config = WanModel.load_config(
     config.pretrained_model_name_or_path,
     subfolder="transformer"
@@ -72,32 +74,39 @@ def create_model(rngs: nnx.Rngs, wan_config: dict):
   wan_config["dtype"] = config.activations_dtype
   wan_config["weights_dtype"] = config.weights_dtype
   wan_config["attention"] = config.attention
+  wan_config["precision"] = get_precision(config)
+  wan_config["flash_block_sizes"] = get_flash_block_sizes(config)
+
+  # 2. eval_shape - will not use flops or create weights on device
+  # thus not using HBM memory.
   p_model_factory = partial(create_model, wan_config=wan_config)
   wan_transformer = nnx.eval_shape(p_model_factory, rngs=rngs)
   graphdef, state, rest_of_state = nnx.split(wan_transformer, nnx.Param, ...)
-  #breakpoint()
+
+  # 3. retrieve the state shardings, mapping logical names to mesh axis names.
   logical_state_spec = nnx.get_partition_spec(state)
   logical_state_sharding = nn.logical_to_mesh_sharding(logical_state_spec, mesh, config.logical_axis_rules)
   logical_state_sharding = dict(nnx.to_flat_state(logical_state_sharding))
   params = state.to_pure_dict()
   state = dict(nnx.to_flat_state(state))
-  # del state
+
+  # 4. Load pretrained weights and move them to device using the state shardings from (3) above.
+  # This helps with loading sharded weights directly into the accelerators without fist copying them
+  # all to one device and then distributing them, thus using low HBM memory.
   params = load_wan_transformer(config.pretrained_model_name_or_path, params, "cpu")
   params = jax.tree_util.tree_map(lambda x: x.astype(config.weights_dtype), params)
   for path, val in flax.traverse_util.flatten_dict(params).items():
     sharding = logical_state_sharding[path].value
-    state[path].value = jax.device_put(val, sharding)
+    try:
+      state[path].value = jax.device_put(val, sharding)
+    except:
+      breakpoint()
   state = nnx.from_flat_state(state)
-  p_add_sharding_rule = partial(_add_sharding_rule, logical_axis_rules=config.logical_axis_rules)
-  state = jax.tree.map(p_add_sharding_rule, state, is_leaf=lambda x: isinstance(x, nnx.VariableState))
-  pspecs = nnx.get_partition_spec(state)
-  #breakpoint()
-  sharded_state = jax.lax.with_sharding_constraint(state, pspecs)
-  #breakpoint()
-  #wan_transformer = jax.jit(nnx.merge(graphdef, sharded_state, rest_of_state), in_shardings=None, out_shardings=sharded_state)
-  wan_transformer = nnx.merge(graphdef, sharded_state, rest_of_state)
+
+  wan_transformer = nnx.merge(graphdef, state, rest_of_state)
   return wan_transformer
 
+
 partial(nnx.jit, static_argnums=(1,))
 def create_sharded_logical_model(model, logical_axis_rules):
   graphdef, state, rest_of_state = nnx.split(model, nnx.Param, ...)
@@ -108,6 +117,7 @@ def create_sharded_logical_model(model, logical_axis_rules):
   wan_transformer = nnx.merge(graphdef, sharded_state, rest_of_state)
   return wan_transformer
 
+
 class WanPipeline:
   r"""
   Pipeline for text-to-video generation using Wan.
@@ -155,6 +165,7 @@ def __init__(
 
     self.p_run_inference = None
 
+
   @classmethod
   def load_text_encoder(cls, config: HyperParameters):
     text_encoder = UMT5EncoderModel.from_pretrained(
@@ -163,6 +174,7 @@ def load_text_encoder(cls, config: HyperParameters):
     )
     return text_encoder
 
+
   @classmethod
   def load_tokenizer(cls, config: HyperParameters):
     tokenizer = AutoTokenizer.from_pretrained(
@@ -171,6 +183,7 @@ def load_tokenizer(cls, config: HyperParameters):
     )
     return tokenizer
 
+
   @classmethod
   def load_vae(cls, devices_array: np.array, mesh: Mesh, rngs: nnx.Rngs, config: HyperParameters):
     wan_vae = AutoencoderKLWan.from_config(
@@ -196,33 +209,14 @@ def load_vae(cls, devices_array: np.array, mesh: Mesh, rngs: nnx.Rngs, config: H
       wan_vae = p_create_sharded_logical_model(model=wan_vae)
     return wan_vae, vae_cache
 
+
   @classmethod
   def load_transformer(cls, devices_array: np.array, mesh: Mesh, rngs: nnx.Rngs, config: HyperParameters):
     with mesh:
       wan_transformer = create_sharded_logical_transformer(devices_array=devices_array, mesh=mesh, rngs=rngs, config=config)
-    # wan_transformer = WanModel.from_config(
-    #   config.pretrained_model_name_or_path,
-    #   subfolder="transformer",
-    #   rngs=rngs,
-    #   attention=config.attention,
-    #   mesh=mesh,
-    #   dtype=config.activations_dtype,
-    #   weights_dtype=config.weights_dtype
-    # )
-    # graphdef, state, rest_of_state = nnx.split(wan_transformer, nnx.Param, ...)
-    # breakpoint()
-    # params = state.to_pure_dict()
-    # del state
-    # #params = load_wan_transformer(config.pretrained_model_name_or_path, params, "cpu")
-    # params = jax.tree_util.tree_map(lambda x: x.astype(config.weights_dtype), params)
-    # #params = jax.device_put(params, PositionalSharding(devices_array).replicate())
-    # wan_transformer = nnx.merge(graphdef, params, rest_of_state)
-    # # Shard
-    # p_create_sharded_logical_model = partial(create_sharded_logical_model, logical_axis_rules=config.logical_axis_rules)
-    # with mesh:
-    #   wan_transformer = p_create_sharded_logical_model(model=wan_transformer)
     return wan_transformer
 
+
   @classmethod
   def load_scheduler(cls, config):
     scheduler, scheduler_state = FlaxUniPCMultistepScheduler.from_pretrained(
@@ -232,6 +226,7 @@ def load_scheduler(cls, config):
     )
     return scheduler, scheduler_state
 
+
   @classmethod  
   def from_pretrained(cls, config: HyperParameters, vae_only=False):
     devices_array = max_utils.create_device_mesh(config)
@@ -268,6 +263,7 @@ def from_pretrained(cls, config: HyperParameters, vae_only=False):
       config=config
     )
 
+
   def _get_t5_prompt_embeds(
     self,
     prompt: Union[str, List[str]] = None,
@@ -302,6 +298,7 @@ def _get_t5_prompt_embeds(
 
     return prompt_embeds
 
+
   def encode_prompt(
     self,
     prompt: Union[str, List[str]],
@@ -333,6 +330,7 @@ def encode_prompt(
 
     return prompt_embeds, negative_prompt_embeds
 
+
   def prepare_latents(
     self,
     batch_size: int,
@@ -356,6 +354,7 @@ def prepare_latents(
 
     return latents
 
+
   def __call__(
     self,
     prompt: Union[str, List[str]] = None,
@@ -382,9 +381,9 @@ def __call__(
 
       # 2. Define call parameters
       if prompt is not None and isinstance(prompt, str):
-          batch_size = 1
-      elif prompt is not None and isinstance(prompt, list):
-          batch_size = len(prompt)
+        prompt = [prompt]
+      
+      batch_size = len(prompt)
 
       prompt_embeds, negative_prompt_embeds = self.encode_prompt(
         prompt=prompt,
@@ -406,12 +405,13 @@ def __call__(
           num_channels_latents=num_channel_latents
         )
 
-      prompt_embeds = jnp.concatenate([prompt_embeds] * latents.shape[0], dtype=self.config.weights_dtype)
-      negative_prompt_embeds = jnp.concatenate([negative_prompt_embeds] * latents.shape[0], dtype=self.config.weights_dtype)
-      
-      latents = jax.device_put(latents, PositionalSharding(self.devices_array).replicate())
-      prompt_embeds = jax.device_put(prompt_embeds, PositionalSharding(self.devices_array).replicate())
-      negative_prompt_embeds = jax.device_put(negative_prompt_embeds, PositionalSharding(self.devices_array).replicate())
+      data_sharding = PositionalSharding(self.devices_array).replicate()
+      if len(prompt) % jax.device_count() == 0:
+        data_sharding = jax.sharding.NamedSharding(self.mesh, P(*self.config.data_sharding))
+        
+      latents = jax.device_put(latents, data_sharding)
+      prompt_embeds = jax.device_put(prompt_embeds, data_sharding)
+      negative_prompt_embeds = jax.device_put(negative_prompt_embeds, data_sharding)
 
       scheduler_state = self.scheduler.set_timesteps(
         self.scheduler_state, num_inference_steps=num_inference_steps, shape=latents.shape