Set offload checkpoint policy

google · Mar 12, 2024 · ad18fb0 · ad18fb0
1 parent 378c6ad
commit ad18fb0
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/paxml/tasks/lm/model_params.py b/paxml/tasks/lm/model_params.py
@@ -663,6 +663,12 @@ def task(self) -> pax_fiddle.Config[tasks_lib.SingleTask]:
           self.CHECKPOINT_POLICY)
     else:
       model_p.lm_tpl.stacked_transformer_tpl = stacked_transformer_tpl
+      if (self.CHECKPOINT_POLICY ==
+          layers.AutodiffCheckpointType.OFFLOAD_DOT_WITH_NO_BATCH_DIM):
+        model_p.lm_tpl.stacked_transformer_tpl.checkpoint_policy = (
+            self.CHECKPOINT_POLICY)
+        model_p.lm_tpl.stacked_transformer_tpl.remat = True
+
 
     # Enable bf16.
     model_p.fprop_dtype = self.FPROP_DTYPE