[CI] Fix head dim check (#1091)

Author: kyuyeunk

tests/models/jax/test_llama3.py

@@ -107,15 +107,14 @@ def test_llama32_1b(self, mock_vllm_config, rng, mesh, mock_model_inputs):
         num_heads = hf_config.num_attention_heads
         num_kv_heads = hf_config.num_key_value_heads
         rope_theta = hf_config.rope_theta
-        original_head_dim = hf_config.head_dim
-        head_dim = 128
+        head_dim = hf_config.head_dim
         intermediate_size = hf_config.intermediate_size
 
         assert attn.hidden_size == hidden_size
         assert attn.num_heads == num_heads
         assert attn.num_kv_heads == num_kv_heads
         assert attn.rope_theta == rope_theta
-        assert attn.head_dim_original == original_head_dim
+        assert attn.head_dim_original == head_dim
         assert attn.head_dim == head_dim
         assert attn.q_proj.kernel.shape == (hidden_size, num_heads, head_dim)
         assert attn.k_proj.kernel.shape == (hidden_size, num_kv_heads,

tests/test_utils.py

@@ -168,7 +168,7 @@ def test_hbm_usage_bytes_pathways_no_arrays(mock_devices, mock_live_arrays):
     "head_dim, expected_padded_head_dim",
     [
         (1, 128),
-        (64, 128),
+        (64, 64),
         (127, 128),
         (128, 128),
         (129, 256),