Helper function to initialize MLA with different model names.
(self, config_arguments, rope_type)
| 64 | self.mesh = Mesh(devices_array, self.cfg.mesh_axes) |
| 65 | |
| 66 | def init_mla(self, config_arguments, rope_type): |
| 67 | """Helper function to initialize MLA with different model names.""" |
| 68 | cfg = pyconfig.initialize( |
| 69 | [sys.argv[0], os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml")], |
| 70 | **config_arguments, |
| 71 | rope_type=rope_type, |
| 72 | ) |
| 73 | |
| 74 | devices_array = maxtext_utils.create_device_mesh(cfg) |
| 75 | mesh = Mesh(devices_array, cfg.mesh_axes) |
| 76 | |
| 77 | dummy_inputs_q = jnp.ones(( |
| 78 | cfg.global_batch_size_to_train_on, |
| 79 | cfg.max_target_length, |
| 80 | cfg.base_emb_dim, |
| 81 | )) |
| 82 | dummy_inputs_kv = jnp.ones(( |
| 83 | cfg.global_batch_size_to_train_on, |
| 84 | cfg.max_target_length, |
| 85 | cfg.base_emb_dim, |
| 86 | )) |
| 87 | |
| 88 | mla = MLA( |
| 89 | config=cfg, |
| 90 | num_query_heads=cfg.num_query_heads, |
| 91 | num_kv_heads=cfg.num_kv_heads, |
| 92 | head_dim=cfg.head_dim, |
| 93 | inputs_q_shape=dummy_inputs_q.shape, |
| 94 | inputs_kv_shape=dummy_inputs_kv.shape, |
| 95 | max_target_length=cfg.max_target_length, |
| 96 | max_prefill_predict_length=cfg.max_prefill_predict_length, |
| 97 | mesh=mesh, |
| 98 | attention_kernel="dot_product", |
| 99 | dtype=cfg.dtype, |
| 100 | dropout_rate=cfg.dropout_rate, |
| 101 | attention_type=cfg.attention_type, |
| 102 | q_lora_rank=cfg.q_lora_rank, |
| 103 | kv_lora_rank=cfg.kv_lora_rank, |
| 104 | qk_nope_head_dim=cfg.qk_nope_head_dim, |
| 105 | qk_rope_head_dim=cfg.qk_rope_head_dim, |
| 106 | v_head_dim=cfg.v_head_dim, |
| 107 | model_mode=MODEL_MODE_PREFILL, |
| 108 | rngs=self.nnx_rng, |
| 109 | ) |
| 110 | |
| 111 | return cfg, mla |
| 112 | |
| 113 | def get_data(self, cfg, dtype): |
| 114 | """get data""" |
no test coverage detected