add unroll repeat scan

Author: NuojCheng

src/maxtext/utils/pipeline_utils.py

@@ -311,11 +311,13 @@ def execute_pipeline_stage(model, loop_state_and_bsw):
     del w_curr
     return (loop_state, w_next), None
 
-  return nn.remat(
-      execute_pipeline_stage,
-      prevent_cse=not model.config.scan_pipeline_iterations,
-      policy=model.get_pipeline_remat_policy(),
-  )
+  return execute_pipeline_stage
+
+  # return nn.remat(
+  #     execute_pipeline_stage,
+  #     prevent_cse=not model.config.scan_pipeline_iterations,
+  #     policy=model.get_pipeline_remat_policy(),
+  # )
 
 
 def create_flax_pipeline_scan(pipeline_stage_fn, length):
@@ -344,4 +346,5 @@ def create_flax_pipeline_scan(pipeline_stage_fn, length):
       },
       split_rngs={"random": True},
       length=length,
+      unroll=length,
   )