add custom vjp over repeat scan

Author: NuojCheng

src/maxtext/configs/base.yml

@@ -299,7 +299,7 @@ pipeline_fsdp_ag_per_repeat: False
 # It may be useful to do the reverse when the layers_per_stage is very large.
 # The below settings only have effect when using pipeline parallelism.
 scan_pipeline_iterations: True
-scan_pipeline_repeats: True
+scan_pipeline_repeats: False
 scan_layers_per_stage: False
 set_remat_policy_on_pipeline_iterations: True
 set_remat_policy_on_layers_per_stage: False

src/maxtext/layers/pipeline.py

@@ -1389,8 +1389,6 @@ def run_iteration_scannable(model, loop_state, bsw):
     # base scannable function used twice for real and bubble runs
     base_scannable = functools.partial(
         pipeline_utils.create_rematerialized_pipeline_stage,
-        model=self,
-        run_iteration_scannable=run_iteration_scannable,
         deterministic=deterministic,
         model_mode=model_mode,
         logical_partition_spec=logical_partition_spec,
@@ -1404,9 +1402,15 @@ def run_iteration_scannable(model, loop_state, bsw):
     run_bubbles_scannable = base_scannable(length=bubble_iterations)
 
     run_repeats_scanned = pipeline_utils.create_flax_pipeline_scan(
-        pipeline_stage_fn=run_one_repeat_scannable, length=self.config.num_pipeline_repeats
+        pipeline_stage_fn=run_one_repeat_scannable,
+        length=self.config.num_pipeline_repeats,
+        use_scan=self.config.scan_pipeline_repeats,
+    )
+    run_bubbles_scanned = pipeline_utils.create_flax_pipeline_scan(
+        pipeline_stage_fn=run_bubbles_scannable,
+        length=1,
+        use_scan=self.config.scan_pipeline_repeats,
     )
-    run_bubbles_scanned = pipeline_utils.create_flax_pipeline_scan(pipeline_stage_fn=run_bubbles_scannable, length=1)
     (loop_state, w_curr), _ = run_repeats_scanned(self, (loop_state, bsw[0]))
     (loop_state, _), _ = run_bubbles_scanned(self, (loop_state, w_curr))
 

src/maxtext/utils/pipeline_utils.py

@@ -19,6 +19,7 @@
 from jax.sharding import PartitionSpec as P
 from flax import linen as nn
 from flax.linen.spmd import LogicallyPartitioned
+import jax.numpy as jnp
 
 
 def get_mesh_axis_dim_indices(physical_partition_spec, axis_name="fsdp"):
@@ -248,8 +249,6 @@ def run_pipeline_microbatches_custom_bwd(residuals, g_final_state):
 
 
 def create_rematerialized_pipeline_stage(
-    model,
-    run_iteration_scannable,
     length,
     deterministic,
     model_mode,
@@ -285,42 +284,62 @@ def create_rematerialized_pipeline_stage(
     the updated `loop_state`.
   """
 
-  def execute_pipeline_stage(model, loop_state_and_bsw):
-    loop_state, w_curr = loop_state_and_bsw
-    # # Retrieve the specific weights needed for this pipeline chunk
-    # bsw = model.both_weight_prefetching(pipeline_weights, physical_partition_spec, loop_state["loop_iteration"])
-    w_next = jax.remat(
-        model.one_weight_prefetching,
-        static_argnums=(1,),
-        policy=jax.checkpoint_policies.nothing_saveable,
-    )(
-        pipeline_weights,
-        physical_partition_spec,
-        loop_state["loop_iteration"],
-    )
-    bsw = (w_curr, w_next)
+  def execute_pipeline_stage_outer(model, loop_state_and_bsw):
+
     scan_microbatches_fn = create_gradient_accumulation_scan(
         model=model,
         length=length,
         deterministic=deterministic,
         model_mode=model_mode,
         logical_partition_spec=logical_partition_spec,
     )
-    loop_state, bsw = scan_microbatches_fn(loop_state, bsw, positions, segment_ids)
-    w_curr, w_next = bsw
-    del w_curr
-    return (loop_state, w_next), None
 
-  return execute_pipeline_stage
+    remat_weight_prefetching = model.one_weight_prefetching
+
+    @jax.custom_vjp
+    def execute_pipeline_stage(loop_state_and_bsw, pipeline_weights):
+      return execute_pipeline_stage_custom_fwd(loop_state_and_bsw, pipeline_weights)[0]
+
+    def execute_pipeline_stage_custom_fwd(loop_state_and_bsw, pipeline_weights):
+      loop_state, w_curr = loop_state_and_bsw
+      # # Retrieve the specific weights needed for this pipeline chunk
+      w_next = remat_weight_prefetching(
+          pipeline_weights,
+          physical_partition_spec,
+          loop_state["loop_iteration"],
+      )
+      bsw = (w_curr, w_next)
+      p_remat_weight_prefetching = functools.partial(
+          remat_weight_prefetching,
+          physical_partition_spec=physical_partition_spec,
+          loop_iteration=loop_state["loop_iteration"],
+      )
+      remat_weight_prefetching_t = jax.linear_transpose(
+          p_remat_weight_prefetching,
+          pipeline_weights,
+      )
+      (loop_state, bsw), scan_fn_vjp = jax.vjp(scan_microbatches_fn, loop_state, bsw, positions, segment_ids)
+      w_curr, w_next = bsw
+      return (loop_state, w_next), (scan_fn_vjp, remat_weight_prefetching_t)
+
+    def execute_pipeline_stage_custom_bwd(residuals, g_outputs):
+      g_loop_state, g_w_next = g_outputs
+      scan_fn_vjp, remat_weight_prefetching_t = residuals
+      g_w_curr = jax.tree.map(jnp.zeros_like, g_w_next)
+      g_bsw = (g_w_curr, g_w_next)
+      g_loop_state, g_bsw, _, _ = scan_fn_vjp((g_loop_state, g_bsw))
+      g_w_curr, g_w_next = g_bsw
+      (g_pipeline_weights,) = remat_weight_prefetching_t(g_w_next)
+      return (g_loop_state, g_w_curr), g_pipeline_weights
+
+    execute_pipeline_stage.defvjp(execute_pipeline_stage_custom_fwd, execute_pipeline_stage_custom_bwd)
+
+    return execute_pipeline_stage(loop_state_and_bsw, pipeline_weights), 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_outer
 
 
-def create_flax_pipeline_scan(pipeline_stage_fn, length):
+def create_flax_pipeline_scan(pipeline_stage_fn, length, use_scan=True):
   """Wraps the pipeline stage execution in a `flax.linen.scan`.
 
   This lifts the pipeline stage function so it can be repeated sequentially over
@@ -332,10 +351,12 @@ def create_flax_pipeline_scan(pipeline_stage_fn, length):
     pipeline_stage_fn: The function representing a single pipeline stage
                        (usually created by `create_rematerialized_pipeline_stage`).
     length: The total number of pipeline stages/repeats to scan over.
+    use_scan: Either scan over repeats or unroll the scan.
 
   Returns:
     A Flax scanned function that executes the full pipeline schedule.
   """
+  unroll_length = 1 if use_scan else length
   return nn.scan(
       pipeline_stage_fn,
       variable_axes={
@@ -346,5 +367,5 @@ def create_flax_pipeline_scan(pipeline_stage_fn, length):
       },
       split_rngs={"random": True},
       length=length,
-      unroll=length,
+      unroll=unroll_length,
   )