Reduce numerical difference in llama4 vision encoder

Author: aireenmei

MaxText/layers/llama4.py

@@ -61,7 +61,11 @@ def setup(self):
     cfg = self.config
     # Linear projection layer using DenseGeneral
     self.linear = linears.DenseGeneral(
-        features=cfg.hidden_size_for_vit, dtype=cfg.dtype_mm, name="vit_unfold_linear", use_bias=False
+        features=cfg.hidden_size_for_vit,
+        dtype=cfg.dtype_mm,
+        name="vit_unfold_linear",
+        use_bias=False,
+        matmul_precision=cfg.matmul_precision,
     )
 
   def __call__(self, inputs: Array) -> Array:
@@ -137,10 +141,18 @@ class Llama4VisionMLP(nn.Module):
   def setup(self):
     cfg = self.config
     self.fc1 = linears.DenseGeneral(
-        features=cfg.intermediate_size_for_vit, dtype=cfg.dtype_mm, name="vit_encoder_layer_mlp_fc1", use_bias=True
+        features=cfg.intermediate_size_for_vit,
+        dtype=cfg.dtype_mm,
+        name="vit_encoder_layer_mlp_fc1",
+        use_bias=True,
+        matmul_precision=cfg.matmul_precision,
     )
     self.fc2 = linears.DenseGeneral(
-        features=cfg.hidden_size_for_vit, dtype=cfg.dtype_mm, name="vit_encoder_layer_mlp_fc2", use_bias=True
+        features=cfg.hidden_size_for_vit,
+        dtype=cfg.dtype_mm,
+        name="vit_encoder_layer_mlp_fc2",
+        use_bias=True,
+        matmul_precision=cfg.matmul_precision,
     )
 
   def __call__(self, hidden_states: Array) -> Array:
@@ -170,10 +182,18 @@ class Llama4VisionMLP2(nn.Module):
   def setup(self):
     cfg = self.config
     self.fc1 = linears.DenseGeneral(
-        features=cfg.projector_input_dim_for_vit, dtype=cfg.dtype_mm, name="vit_pixel_shuffle_mlp_fc1", use_bias=False
+        features=cfg.projector_input_dim_for_vit,
+        dtype=cfg.dtype_mm,
+        name="vit_pixel_shuffle_mlp_fc1",
+        use_bias=False,
+        matmul_precision=cfg.matmul_precision,
     )
     self.fc2 = linears.DenseGeneral(
-        features=cfg.projector_output_dim_for_vit, dtype=cfg.dtype_mm, name="vit_pixel_shuffle_mlp_fc2", use_bias=False
+        features=cfg.projector_output_dim_for_vit,
+        dtype=cfg.dtype_mm,
+        name="vit_pixel_shuffle_mlp_fc2",
+        use_bias=False,
+        matmul_precision=cfg.matmul_precision,
     )
     self.dropout = nn.Dropout(rate=cfg.projector_dropout_for_vit)
 
@@ -252,6 +272,7 @@ def setup(self):
         dtype=cfg.dtype_mm,
         name="vit_multi_modal_projector",
         use_bias=False,
+        matmul_precision=cfg.matmul_precision,
     )
 
   def __call__(self, image_features: Array) -> Array:
@@ -579,6 +600,8 @@ def __call__(
         head_dim=self.config.hidden_size_for_vit // self.config.num_attention_heads_for_vit,
         max_target_length=(self.config.image_size_for_vit // self.config.patch_size_for_vit) ** 2 + 1,
         attention_kernel="dot_product",
+        float32_qk_product=self.config.float32_qk_product,
+        float32_logits=self.config.float32_logits,
         mesh=self.mesh,
         dropout_rate=0,
         name="self_attention_vision",

MaxText/tests/check_llama4_layers.py

@@ -842,13 +842,15 @@ class Config(NamedTuple):
     attention_dropout: int = 0
 
   config_arguments = {
-      "per_device_batch_size": 4.0,
       "run_name": "test",
       "enable_checkpointing": False,
       "model_name": "llama4-17b-16e",
       "scan_layers": False,
-      "num_hidden_layers_for_vit": 6,
+      "num_hidden_layers_for_vit": 34,
       "dtype": "float32",
+      "matmul_precision": "float32",
+      "float32_qk_product": True,
+      "float32_logits": True,
   }
 
   def setUp(self):
@@ -885,6 +887,7 @@ def test_vision_encoder(self):
     # Create test input using config dimensions
     batch_size = 4
     inputs = jnp.ones((batch_size, self.seq_len_for_vit, self.cfg.hidden_size_for_vit), dtype=jnp.float32)
+    inputs /= 10
 
     # Initialize JAX parameters
     params = jax_model.init(self.rng, inputs, deterministic=True)
@@ -909,7 +912,7 @@ def test_vision_encoder(self):
     jax_outputs = jax_model.apply(params, inputs, deterministic=True)
 
     # Compare outputs
-    np.testing.assert_allclose(jax_outputs, to_jax(pt_outputs), rtol=1e-3, atol=0.05)
+    np.testing.assert_allclose(jax_outputs, to_jax(pt_outputs), rtol=0.01, atol=0.05)
 
 
 if __name__ == "__main__":