Refactor: Transformer blocks

src/mflux/models/transformer/ada_layer_norm_zero.py

@@ -8,7 +8,7 @@ class AdaLayerNormZero(nn.Module):
         self.linear = nn.Linear(3072, 18432)
         self.norm = nn.LayerNorm(dims=3072, eps=1e-6, affine=False)
 
-    def __call__(self, x: mx.array, text_embeddings: mx.array):
+    def __call__(self, hidden_states: mx.array, text_embeddings: mx.array):
         text_embeddings = self.linear(nn.silu(text_embeddings))
         chunk_size = 18432 // 6
         shift_msa = text_embeddings[:, 0 * chunk_size : 1 * chunk_size]
@@ -17,5 +17,5 @@ class AdaLayerNormZero(nn.Module):
         shift_mlp = text_embeddings[:, 3 * chunk_size : 4 * chunk_size]
         scale_mlp = text_embeddings[:, 4 * chunk_size : 5 * chunk_size]
         gate_mlp = text_embeddings[:, 5 * chunk_size : 6 * chunk_size]
-        x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
+        hidden_states = self.norm(hidden_states) * (1 + scale_msa[:, None]) + shift_msa[:, None]
-        return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
+        return hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp

src/mflux/models/transformer/ada_layer_norm_zero_single.py

@@ -8,11 +8,11 @@ class AdaLayerNormZeroSingle(nn.Module):
         self.linear = nn.Linear(3072, 3 * 3072)
         self.norm = nn.LayerNorm(dims=3072, eps=1e-6, affine=False)
 
-    def __call__(self, x: mx.array, text_embeddings: mx.array):
+    def __call__(self, hidden_states: mx.array, text_embeddings: mx.array) -> mx.array:
         text_embeddings = self.linear(nn.silu(text_embeddings))
         chunk_size = 9216 // 3
         shift_msa = text_embeddings[:, 0 * chunk_size : 1 * chunk_size]
         scale_msa = text_embeddings[:, 1 * chunk_size : 2 * chunk_size]
         gate_msa = text_embeddings[:, 2 * chunk_size : 3 * chunk_size]
-        x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
+        hidden_states = self.norm(hidden_states) * (1 + scale_msa[:, None]) + shift_msa[:, None]
-        return x, gate_msa
+        return hidden_states, gate_msa

src/mflux/models/transformer/joint_transformer_block.py

@@ -11,12 +11,12 @@ class JointTransformerBlock(nn.Module):
         super().__init__()
         self.layer = layer
         self.norm1 = AdaLayerNormZero()
-        self.norm2 = nn.LayerNorm(dims=3072, eps=1e-6, affine=False)
-        self.ff = FeedForward(activation_function=nn.gelu)
-        self.attn = JointAttention()
         self.norm1_context = AdaLayerNormZero()
-        self.ff_context = FeedForward(activation_function=nn.gelu_approx)
+        self.attn = JointAttention()
+        self.norm2 = nn.LayerNorm(dims=3072, eps=1e-6, affine=False)
         self.norm2_context = nn.LayerNorm(dims=1536, eps=1e-6, affine=False)
+        self.ff = FeedForward(activation_function=nn.gelu)
+        self.ff_context = FeedForward(activation_function=nn.gelu_approx)
 
     def __call__(
         self,
@@ -25,30 +25,67 @@ class JointTransformerBlock(nn.Module):
         text_embeddings: mx.array,
         rotary_embeddings: mx.array,
     ) -> (mx.array, mx.array):
-        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, text_embeddings)
+        # 1a. Compute norm for hidden_states
+        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
+            hidden_states=hidden_states,
+            text_embeddings=text_embeddings
+        )  # fmt: off
 
+        # 1b. Compute norm for encoder_hidden_states
         norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
-            x=encoder_hidden_states, text_embeddings=text_embeddings
+            hidden_states=encoder_hidden_states,
-        )
+            text_embeddings=text_embeddings
+        )  # fmt: off
 
+        # 2. Compute attention
         attn_output, context_attn_output = self.attn(
             hidden_states=norm_hidden_states,
             encoder_hidden_states=norm_encoder_hidden_states,
             image_rotary_emb=rotary_embeddings,
         )
 
+        # 3a. Apply norm and feed forward for hidden states
+        hidden_states = JointTransformerBlock._apply_norm_and_feed_forward(
+            hidden_states=hidden_states,
+            attn_output=attn_output,
+            gate_mlp=gate_mlp,
+            gate_msa=gate_msa,
+            scale_mlp=scale_mlp,
+            shift_mlp=shift_mlp,
+            norm_layer=self.norm2,
+            ff_layer=self.ff,
+        )
+
+        # 3b. Apply norm and feed forward for encoder hidden states
+        encoder_hidden_states = JointTransformerBlock._apply_norm_and_feed_forward(
+            hidden_states=encoder_hidden_states,
+            attn_output=context_attn_output,
+            gate_mlp=c_gate_mlp,
+            gate_msa=c_gate_msa,
+            scale_mlp=c_scale_mlp,
+            shift_mlp=c_shift_mlp,
+            norm_layer=self.norm2_context,
+            ff_layer=self.ff_context,
+        )
+
+        return encoder_hidden_states, hidden_states
+
+    @staticmethod
+    def _apply_norm_and_feed_forward(
+        hidden_states: mx.array,
+        attn_output: mx.array,
+        gate_mlp: mx.array,
+        gate_msa: mx.array,
+        scale_mlp: mx.array,
+        shift_mlp: mx.array,
+        norm_layer: nn.Module,
+        ff_layer: nn.Module,
+    ) -> mx.array:
         attn_output = mx.expand_dims(gate_msa, axis=1) * attn_output
         hidden_states = hidden_states + attn_output
-        norm_hidden_states = self.norm2(hidden_states)
+        norm_hidden_states = norm_layer(hidden_states)
         norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
-        ff_output = self.ff(norm_hidden_states)
+        ff_output = ff_layer(norm_hidden_states)
         ff_output = mx.expand_dims(gate_mlp, axis=1) * ff_output
         hidden_states = hidden_states + ff_output
-
+        return hidden_states
-        context_attn_output = mx.expand_dims(c_gate_msa, axis=1) * context_attn_output
-        encoder_hidden_states = encoder_hidden_states + context_attn_output
-        norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
-        norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
-        context_ff_output = self.ff_context(norm_encoder_hidden_states)
-        encoder_hidden_states = encoder_hidden_states + mx.expand_dims(c_gate_mlp, axis=1) * context_ff_output
-        return encoder_hidden_states, hidden_states

src/mflux/models/transformer/single_transformer_block.py

@@ -12,8 +12,8 @@ class SingleTransformerBlock(nn.Module):
         super().__init__()
         self.layer = layer
         self.norm = AdaLayerNormZeroSingle()
-        self.proj_mlp = nn.Linear(3072, 4 * 3072)
         self.attn = SingleBlockAttention()
+        self.proj_mlp = nn.Linear(3072, 4 * 3072)
         self.proj_out = nn.Linear(3072 + 4 * 3072, 3072)
 
     def __call__(
@@ -22,15 +22,39 @@ class SingleTransformerBlock(nn.Module):
         text_embeddings: mx.array,
         rotary_embeddings: mx.array,
     ) -> (mx.array, mx.array):
+        # 0. Establish residual connection
         residual = hidden_states
-        norm_hidden_states, gate = self.norm(x=hidden_states, text_embeddings=text_embeddings)
+
-        mlp_hidden_states = nn.gelu_approx(self.proj_mlp(norm_hidden_states))
+        # 1. Compute norm for hidden_states
+        norm_hidden_states, gate = self.norm(
+            hidden_states=hidden_states,
+            text_embeddings=text_embeddings
+        )  # fmt: off
+
+        # 2. Compute attention
         attn_output = self.attn(
             hidden_states=norm_hidden_states,
             image_rotary_emb=rotary_embeddings,
         )
+
+        # 3. Apply norm and feed forward for hidden states
+        hidden_states = self._apply_feed_forward_and_projection(
+            norm_hidden_states=norm_hidden_states,
+            attn_output=attn_output,
+            gate=gate,
+        )
+
+        return residual + hidden_states
+
+    def _apply_feed_forward_and_projection(
+        self,
+        norm_hidden_states: mx.array,
+        attn_output: mx.array,
+        gate: mx.array,
+    ) -> mx.array:
+        feed_forward = self.proj_mlp(norm_hidden_states)
+        mlp_hidden_states = nn.gelu_approx(feed_forward)
         hidden_states = mx.concatenate([attn_output, mlp_hidden_states], axis=2)
         gate = mx.expand_dims(gate, axis=1)
         hidden_states = gate * self.proj_out(hidden_states)
-        hidden_states = residual + hidden_states
         return hidden_states