pre-shard with pjit

fattorib · Jun 6, 2023 · fc7c8b7 · fc7c8b7
1 parent 496adf4
commit fc7c8b7
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/main.py b/main.py
@@ -308,6 +308,10 @@ def compute_tokens_seen(absolute_step, max_context):
             import gc
             gc.collect()
 
+            params = pjit(lambda x:x, out_axis_resources=param_spec)(params)
+            opt_state = pjit(lambda x:x, out_axis_resources=opt_state_spec)(opt_state)
+
+
         else:
             raise NotImplementedError(
                 "Checkpointing not currently implemented for GPU."