Merge pull request #4312 from AI-Hypercomputer:fix_engram_test

PiperOrigin-RevId: 941189378

Author: Google-ML-Automation

tests/integration/deepseek_scan_engram_test.py

@@ -127,9 +127,6 @@ def batch_decode(self, token_ids, *args, **kwargs):
         ]
     )
 
-    if base_num_decoder_layers >= 4 and jax.device_count() <= 8 and any("TPU7x" in d.device_kind for d in jax.devices()):
-      pytest.skip("Compiling 4+ DeepSeek MoE/Engram scanned layers requires >= 8" " physical TPU chips on TPU7x")
-
     devices_array = maxtext_utils.create_device_mesh(config)
     mesh = Mesh(devices_array, config.mesh_axes)
 
@@ -176,16 +173,14 @@ def test_decoder_init_engram_2_8(self, mock_from_pretrained):
     """Test engram layers at indices 2 and 8."""
     self._test_engram_pattern(
         mock_from_pretrained,
-        "1,4",
+        "1",
         [
             "dense_layers_0_0",
             "dense_layers_engram_1",
             "dense_layers_2_2",
-            "moe_layers_3_3",
-            "moe_layers_engram_4",
         ],
         first_num_dense_layers=3,
-        base_num_decoder_layers=5,
+        base_num_decoder_layers=3,
     )
 
   @pytest.mark.tpu_only
@@ -206,8 +201,8 @@ def test_decoder_init_engram_4_9(self, mock_from_pretrained):
     """Test engram layers at indices 4 and 9 - last engram layer of block."""
     self._test_engram_pattern(
         mock_from_pretrained,
-        "1,3",
-        ["dense_layers_0_0", "dense_layers_engram_1", "moe_layers_2_2", "moe_layers_engram_3"],
+        "1,2",
+        ["dense_layers_0_0", "dense_layers_engram_1", "moe_layers_engram_2"],
         first_num_dense_layers=2,
-        base_num_decoder_layers=4,
+        base_num_decoder_layers=3,
     )