Merge pull request #2811 from AI-Hypercomputer:qinwen/add_ckpt_mla

PiperOrigin-RevId: 843274677

Author: Google-ML-Automation

src/MaxText/configs/base.yml

@@ -298,6 +298,8 @@ key_proj: 'remat'
 value_proj: 'remat'
 qkv_proj: 'remat'
 out_proj: 'remat'
+mla_q: 'remat'
+mla_kv: 'remat'
 
 optimizer_memory_host_offload: False
 parameter_memory_host_offload: False

src/MaxText/configs/types.py

@@ -805,6 +805,14 @@ class RematAndOffload(BaseModel):
       RematLocation.REMAT,
       description="Remat policy for the attention output projection.",
   )
+  mla_q: RematLocation = Field(
+      RematLocation.REMAT,
+      description="Remat policy for the mla's query projectiont.",
+  )
+  mla_kv: RematLocation = Field(
+      RematLocation.REMAT,
+      description="Remat policy for the mla's key and value projection.",
+  )
   optimizer_memory_host_offload: bool = Field(False, description="Offload optimizer state to host memory.")
   parameter_memory_host_offload: bool = Field(False, description="Offload parameters to host memory.")
 
@@ -1855,6 +1863,8 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
           "query_proj",
           "key_proj",
           "value_proj",
+          "mla_kv",
+          "mla_q",
           "qkv_proj",
           "out_proj",
       ]

src/MaxText/layers/attention_mla.py

@@ -530,6 +530,7 @@ def mla_query_projection(self, inputs_q: Array, inputs_positions: Array, model_m
       # LoRA path
       low_rank_q = self.wq_a(inputs_q, out_sharding=wqa_out_sharding)  # [B, L, q_lora_rank]
       low_rank_q = self.q_norm(low_rank_q)  # RMSNorm on low rank
+      low_rank_q = checkpoint_name(low_rank_q, "mla_q")
       q = self.wq_b(low_rank_q, out_sharding=query_sharding)  # [B, L, n_heads * qk_head_dim]
 
     # Split into non-positional and rotary parts.
@@ -668,7 +669,7 @@ def mla_kv_projection(self, inputs: Array, inputs_positions: Array, decoder_segm
     low_rank = self.wkv_a(inputs, out_sharding=wkva_out_sharding)
     low_rank_main, low_rank_rope = jnp.split(low_rank, [self.kv_lora_rank], axis=-1)
     low_rank_main = self.kv_norm(low_rank_main)
-
+    low_rank_main = checkpoint_name(low_rank_main, "mla_kv")
     # Apply rotary embedding to key_rope.
     key_rope = jnp.expand_dims(low_rank_rope, axis=2)
     key_rope = self.apply_rotary_embedding(key_rope, inputs_positions=inputs_positions)