add fp32 layer norm

Author: jfacevedo-google

src/maxdiffusion/models/normalization_flax.py

@@ -17,6 +17,7 @@
 import jax
 import jax.numpy as jnp
 import flax.linen as nn
+from flax import nnx
 
 
 class AdaLayerNormContinuous(nn.Module):
@@ -147,3 +148,20 @@ def __call__(self, x, emb):
     else:
       raise ValueError(f"Unsupported `norm_type` ({self.norm_type}) provided. Supported ones are: 'layer_norm'.")
     return x, gate_msa
+
+class FP32LayerNorm(nnx.Module):
+  def __init__(self, rngs: nnx.Rngs, dim: int, eps : float, elementwise_affine: bool):
+    self.layer_norm = nnx.LayerNorm(
+      rngs=rngs,
+      num_features=dim,
+      epsilon=eps,
+      use_bias=elementwise_affine,
+      use_scale=elementwise_affine,
+      param_dtype=jnp.float32,
+      dtype=jnp.float32
+    )
+  def __call__(self, inputs: jax.Array) -> jax.Array:
+    origin_dtype = inputs.dtype
+    return self.layer_norm(
+      inputs.astype(dtype=jnp.float32)
+    ).astype(dtype=origin_dtype)

src/maxdiffusion/tests/wan_transformer_test.py

@@ -22,6 +22,7 @@
 
 from ..models.wan.transformers.transformer_wan import WanRotaryPosEmbed, WanTimeTextImageEmbedding
 from ..models.embeddings_flax import NNXTimestepEmbedding, NNXPixArtAlphaTextProjection
+from ..models.normalization_flax import FP32LayerNorm
 
 class WanTransformerTest(unittest.TestCase):
   def setUp(self):
@@ -68,6 +69,21 @@ def test_nnx_timestep_embedding(self):
     dummy_output = layer(dummy_sample)
     assert dummy_output.shape == (1, 5120)
 
+  def test_fp32_layer_norm(self):
+    key = jax.random.key(0)
+    rngs = nnx.Rngs(key)
+    batch_size = 1
+    dummy_hidden_states = jnp.ones((batch_size, 75600, 5120))
+    # expected same output shape with same dtype
+    layer = FP32LayerNorm(
+      rngs=rngs,
+      dim=5120,
+      eps=1e-6,
+      elementwise_affine=False
+    )
+    dummy_output = layer(dummy_hidden_states)
+    assert dummy_output.shape == dummy_hidden_states.shape    
+
   def test_wan_time_text_embedding(self):
     key = jax.random.key(0)
     rngs = nnx.Rngs(key)