restore norm

configs_3b/main.yaml

@@ -11,7 +11,7 @@ dit:
     vid_in_channels: 33
     vid_out_channels: 16
     vid_dim: 2560
-    vid_out_norm: layer
+    vid_out_norm: fusedrms
     txt_in_dim: 5120
     txt_in_norm: fusedln
     txt_dim: ${.vid_dim}
@@ -19,11 +19,11 @@ dit:
     heads: 20
     head_dim: 128 # llm-like
     expand_ratio: 4
-    norm: layer
+    norm: fusedrms
     norm_eps: 1.0e-05
     ada: single
     qk_bias: False
-    qk_norm: layer
+    qk_norm: fusedrms
     patch_size: [ 1,2,2 ]
     num_layers: 32 # llm-like
     mm_layers: 10

configs_7b/main.yaml

@@ -17,12 +17,12 @@ dit:
     heads: 24
     head_dim: 128 # llm-like
     expand_ratio: 4
-    norm: layer
+    norm: fusedrms
     norm_eps: 1e-5
     ada: single
     qk_bias: False
     qk_rope: True
-    qk_norm: layer
+    qk_norm: fusedrms
     patch_size: [ 1,2,2 ]
     num_layers: 36 # llm-like
     shared_mlp: False

models/dit_v2/normalization.py

@@ -30,7 +30,7 @@ def get_norm_layer(norm_type: Optional[str]) -> norm_layer_type:
             return nn.LayerNorm(
                 normalized_shape=dim,
                 eps=eps,
-                elementwise_affine=False,
+                elementwise_affine=elementwise_affine,
             )
 
         if norm_type == "rms":