Merge pull request #77 from hx89:checkpoint_fix

PiperOrigin-RevId: 647034147
google · Jun 26, 2024 · f0149a3 · f0149a3
2 parents 7293a44 + 4878086
commit f0149a3
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/praxis/layers/attentions.py b/praxis/layers/attentions.py
@@ -1698,6 +1698,9 @@ def __call__(
       query_proj = self.query(query_vec)
       key_proj = self.key(key_vec)
       value_proj = self.value(value_vec)
+      query_proj = checkpoint_name(query_proj, 'query_proj')
+      key_proj = checkpoint_name(key_proj, 'key_proj')
+      value_proj = checkpoint_name(value_proj, 'value_proj')
 
     if not self.consolidate_rope_key_state:
       self._fprop_update_decode_state('key_state', key_proj)

diff --git a/praxis/layers/grok.py b/praxis/layers/grok.py
@@ -350,5 +350,6 @@ def GrokUniTransformerLmHParams(
         num_pipeline_stages=num_pipeline_stages,
         num_pipeline_microbatches=num_pipeline_microbatches,
         stream_io=True,
+        checkpoint_policy=checkpoint_policy,
     )
   return p