Tests that MLA and Attention layers initialize the correct projection weights.
(self)
| 1236 | # self.assertTrue(jax.numpy.allclose(mla_full_this_idx, mla_idx, rtol=1e-02, atol=1e-02, equal_nan=False)) |
| 1237 | |
| 1238 | def test_projection_initialization(self): |
| 1239 | """Tests that MLA and Attention layers initialize the correct projection weights.""" |
| 1240 | # 1. Initialize a standard Attention layer for comparison |
| 1241 | # Create a copy of the arguments and override the attention_type for the base model |
| 1242 | attention_config_args = self.config_arguments.copy() |
| 1243 | attention_config_args["attention_type"] = AttentionType.GLOBAL.value |
| 1244 | attention_cfg = pyconfig.initialize( |
| 1245 | [sys.argv[0], os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml")], |
| 1246 | **attention_config_args, |
| 1247 | ) |
| 1248 | dummy_inputs_q = jnp.ones( |
| 1249 | (attention_cfg.global_batch_size_to_train_on, attention_cfg.max_target_length, attention_cfg.base_emb_dim) |
| 1250 | ) |
| 1251 | dummy_inputs_kv = jnp.ones( |
| 1252 | (attention_cfg.global_batch_size_to_train_on, attention_cfg.max_target_length, attention_cfg.base_emb_dim) |
| 1253 | ) |
| 1254 | |
| 1255 | base_attention = Attention( |
| 1256 | config=attention_cfg, |
| 1257 | num_query_heads=attention_cfg.num_query_heads, |
| 1258 | num_kv_heads=attention_cfg.num_kv_heads, |
| 1259 | head_dim=attention_cfg.head_dim, |
| 1260 | max_target_length=attention_cfg.max_target_length, |
| 1261 | max_prefill_predict_length=attention_cfg.max_prefill_predict_length, |
| 1262 | inputs_q_shape=dummy_inputs_q.shape, |
| 1263 | inputs_kv_shape=dummy_inputs_kv.shape, |
| 1264 | mesh=self.mesh, |
| 1265 | attention_kernel="dot_product", |
| 1266 | dtype=attention_cfg.dtype, |
| 1267 | rngs=self.nnx_rng, |
| 1268 | ) |
| 1269 | |
| 1270 | # 2. Assert that the base Attention layer HAS all its standard projections |
| 1271 | self.assertTrue(hasattr(base_attention, "query"), "Base Attention should have 'query' projection.") |
| 1272 | self.assertTrue(hasattr(base_attention, "key"), "Base Attention should have 'key' projection.") |
| 1273 | self.assertTrue(hasattr(base_attention, "value"), "Base Attention should have 'value' projection.") |
| 1274 | self.assertTrue(hasattr(base_attention, "out"), "Base Attention should have 'out' projection.") |
| 1275 | |
| 1276 | # 3. Initialize the MLA layer |
| 1277 | _, mla_layer = self.init_mla(self.config_arguments, rope_type="default") |
| 1278 | |
| 1279 | # 4. Assert that the MLA layer DOES NOT HAVE the base projections |
| 1280 | self.assertFalse(hasattr(mla_layer, "query"), "MLA should not have 'query' projection.") |
| 1281 | self.assertFalse(hasattr(mla_layer, "key"), "MLA should not have 'key' projection.") |
| 1282 | self.assertFalse(hasattr(mla_layer, "value"), "MLA should not have 'value' projection.") |
| 1283 | |
| 1284 | # 5. Assert that the MLA layer HAS all of its own specific projections AND the common 'out' projection |
| 1285 | self.assertTrue(hasattr(mla_layer, "wq_a"), "MLA should have 'wq_a' projection.") |
| 1286 | self.assertTrue(hasattr(mla_layer, "wq_b"), "MLA should have 'wq_b' projection.") |
| 1287 | self.assertTrue(hasattr(mla_layer, "wkv_a"), "MLA should have 'wkv_a' projection.") |
| 1288 | self.assertTrue(hasattr(mla_layer, "wkv_b"), "MLA should have 'wkv_b' projection.") |
| 1289 | self.assertTrue(hasattr(mla_layer, "q_norm"), "MLA should have 'q_norm' projection.") |
| 1290 | self.assertTrue(hasattr(mla_layer, "kv_norm"), "MLA should have 'kv_norm' projection.") |
| 1291 | self.assertTrue(hasattr(mla_layer, "out"), "MLA should have 'out' projection.") |
| 1292 | |
| 1293 | @parameterized.named_parameters( |
| 1294 | { |
nothing calls this directly
no test coverage detected