reduces memory significantly when loading transformer. Needs clean up.

Author: jfacevedo-google

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

@@ -34,6 +34,38 @@
 
 BlockSizes = common_types.BlockSizes
 
+def get_frequencies(max_seq_len: int, theta: int, attention_head_dim: int):
+  h_dim = w_dim = 2 * (attention_head_dim // 6)
+  t_dim = attention_head_dim - h_dim - w_dim
+  freqs = []
+  for dim in [t_dim, h_dim, w_dim]:
+    freq = get_1d_rotary_pos_embed(
+      dim,
+      max_seq_len,
+      theta,
+      freqs_dtype=jnp.float64,
+      use_real=False
+    )
+    freqs.append(freq)
+  freqs = jnp.concatenate(freqs, axis=1)
+  # sizes = jnp.array([
+  #     attention_head_dim // 2 - 2 * (attention_head_dim // 6),
+  #     attention_head_dim // 6,
+  #     attention_head_dim // 6,
+  # ])
+  # cumulative_sizes = jnp.cumsum(jnp.array(sizes))
+  # split_indices = cumulative_sizes[:-1]
+  t_size = attention_head_dim // 2 - 2 * (attention_head_dim // 6)
+  hw_size = attention_head_dim // 6
+  
+  dims = [t_size, hw_size, hw_size]
+  
+  # Calculate split indices as a static list of integers
+  cumulative_sizes = np.cumsum(dims)
+  split_indices = cumulative_sizes[:-1].tolist()
+  freqs_split = jnp.split(freqs, split_indices, axis=1)
+  return freqs_split
+
 class WanRotaryPosEmbed(nnx.Module):
   def __init__(
     self,
@@ -45,44 +77,23 @@ def __init__(
     self.attention_head_dim = attention_head_dim
     self.patch_size = patch_size
     self.max_seq_len = max_seq_len
-
-    h_dim = w_dim = 2 * (attention_head_dim // 6)
-    t_dim = attention_head_dim - h_dim - w_dim
-
-    freqs = []
-    for dim in [t_dim, h_dim, w_dim]:
-      freq = get_1d_rotary_pos_embed(
-        dim,
-        self.max_seq_len,
-        theta,
-        freqs_dtype=jnp.float64,
-        use_real=False
-      )
-      freqs.append(freq)
-    freqs = jnp.concatenate(freqs, axis=1)
-
-    sizes = [
-        self.attention_head_dim // 2 - 2 * (self.attention_head_dim // 6),
-        self.attention_head_dim // 6,
-        self.attention_head_dim // 6,
-    ]
-    cumulative_sizes = jnp.cumsum(jnp.array(sizes))
-    split_indices = cumulative_sizes[:-1]
-    self.freqs_split = jnp.split(freqs, split_indices, axis=1)
+    self.theta = theta
 
   def __call__(self, hidden_states: jax.Array) -> jax.Array:
     _, num_frames, height, width, _ = hidden_states.shape
     p_t, p_h, p_w = self.patch_size
     ppf, pph, ppw = num_frames // p_t, height // p_h, width // p_w
 
-    freqs_f = jnp.expand_dims(jnp.expand_dims(self.freqs_split[0][:ppf], axis=1), axis=1)
-    freqs_f = jnp.broadcast_to(freqs_f, (ppf, pph, ppw, self.freqs_split[0].shape[-1]))
+    freqs_split = get_frequencies(self.max_seq_len, self.theta, self.attention_head_dim)
+
+    freqs_f = jnp.expand_dims(jnp.expand_dims(freqs_split[0][:ppf], axis=1), axis=1)
+    freqs_f = jnp.broadcast_to(freqs_f, (ppf, pph, ppw, freqs_split[0].shape[-1]))
 
-    freqs_h = jnp.expand_dims(jnp.expand_dims(self.freqs_split[1][:pph], axis=0), axis=2)
-    freqs_h = jnp.broadcast_to(freqs_h, (ppf, pph, ppw, self.freqs_split[1].shape[-1]))
+    freqs_h = jnp.expand_dims(jnp.expand_dims(freqs_split[1][:pph], axis=0), axis=2)
+    freqs_h = jnp.broadcast_to(freqs_h, (ppf, pph, ppw, freqs_split[1].shape[-1]))
 
-    freqs_w = jnp.expand_dims(jnp.expand_dims(self.freqs_split[2][:ppw], axis=0), axis=1)
-    freqs_w = jnp.broadcast_to(freqs_w, (ppf, pph, ppw, self.freqs_split[2].shape[-1]))
+    freqs_w = jnp.expand_dims(jnp.expand_dims(freqs_split[2][:ppw], axis=0), axis=1)
+    freqs_w = jnp.broadcast_to(freqs_w, (ppf, pph, ppw, freqs_split[2].shape[-1]))
 
     freqs_concat = jnp.concatenate([freqs_f, freqs_h, freqs_w], axis=-1)
     freqs_final = jnp.reshape(freqs_concat, (1, 1, ppf * pph * ppw, -1))
@@ -362,7 +373,7 @@ def __init__(
       qk_norm: Optional[str] = "rms_norm_across_heads",
       eps: float = 1e-6,
       image_dim: Optional[int] = None,
-      added_kn_proj_dim: Optional[int] = None,
+      added_kv_proj_dim: Optional[int] = None,
       rope_max_seq_len: int = 1024,
       pos_embed_seq_len: Optional[int] = None,
       flash_min_seq_length: int = 4096,

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -18,6 +18,8 @@
 import jax
 import jax.numpy as jnp
 from jax.sharding import Mesh, PositionalSharding
+import flax
+import flax.linen as nn
 from flax import nnx
 from ...pyconfig import HyperParameters
 from ... import max_logging
@@ -54,6 +56,48 @@ def _add_sharding_rule(vs: nnx.VariableState, logical_axis_rules) -> nnx.Variabl
   vs.sharding_rules = logical_axis_rules
   return vs
 
+
+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
+
+  wan_config = WanModel.load_config(
+    config.pretrained_model_name_or_path,
+    subfolder="transformer"
+  )
+  wan_config["mesh"] = mesh
+  wan_config["dtype"] = config.activations_dtype
+  wan_config["weights_dtype"] = config.weights_dtype
+  wan_config["attention"] = config.attention
+  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()
+  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
+  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)
+  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)
+  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, ...)
@@ -154,26 +198,29 @@ def load_vae(cls, devices_array: np.array, mesh: Mesh, rngs: nnx.Rngs, config: H
 
   @classmethod
   def load_transformer(cls, devices_array: np.array, mesh: Mesh, rngs: nnx.Rngs, config: HyperParameters):
-    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, ...)
-    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)
+      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