add bsw checkpoint

Author: NuojCheng

src/MaxText/layers/pipeline.py

@@ -590,8 +590,9 @@ def _all_gather_invariant(x, i):
         return x
 
       new_variables = jax.tree.map(_all_gather_invariant, sharded_weights, fsdp_idx)
+      new_variables = jax.ad_checkpoint.checkpoint_name(new_variables, "bsw_gathered_weights")
 
-      return (cur_bsw[1], new_variables)
+      return jax.ad_checkpoint.checkpoint_name((cur_bsw[1], new_variables), "bsw")
 
     return _all_gather_inner(repeat_weights, bsw, fsdp_idx)
 
@@ -984,7 +985,7 @@ def run_iteration_scannable(model, loop_state):
     if self.config.set_remat_policy_on_pipeline_iterations:
       run_iteration_scannable = nn.remat(
           run_iteration_scannable,
-          prevent_cse=not self.config.scan_pipeline_iterations,
+          prevent_cse=True,  # not self.config.scan_pipeline_iterations,
           policy=self.get_pipeline_remat_policy(),
       )
 
@@ -1017,7 +1018,7 @@ def run_one_repeat_scannable(model, loop_state):
 
     run_one_repeat_scannable = nn.remat(
         run_one_repeat_scannable,
-        prevent_cse=not self.config.scan_pipeline_iterations,
+        prevent_cse=True,
         policy=self.get_pipeline_remat_policy(),
     )