remove additional AG

Author: NuojCheng

src/maxtext/layers/pipeline.py

@@ -1239,7 +1239,7 @@ def _apply_sharding_hint(weight, pspec):
       return _from_repeat_weights_to_bsw_shardmap(repeat_weights, physical_partition_spec, axes_to_gather=axes_to_gather)
     return _from_repeat_weights_to_bsw_hint(repeat_weights)
 
-  def weight_prefetching(self, weights, physical_partition_spec, loop_iteration):
+  def both_weight_prefetching(self, weights, physical_partition_spec, loop_iteration):
     """Triggers asynchronous FSDP-like all-gathers for the current and next pipeline steps.
 
     By gathering weights for `loop_iteration + 1` right now, the network communication
@@ -1250,7 +1250,16 @@ def weight_prefetching(self, weights, physical_partition_spec, loop_iteration):
     nxt_repeat_weights = self.from_all_variables_to_repeat_weights(weights, loop_iteration + 1)
     bsw_0 = self.from_repeat_weights_to_bsw(cur_repeat_weights, physical_partition_spec)
     bsw_1 = self.from_repeat_weights_to_bsw(nxt_repeat_weights, physical_partition_spec)
-    return jax.ad_checkpoint.checkpoint_name((bsw_0, bsw_1), "bsw")
+    return bsw_0, bsw_1
+
+  def one_weight_prefetching(self, weights, physical_partition_spec, loop_iteration):
+    """Triggers asynchronous FSDP-like all-gathers for the next pipeline steps.
+
+    By gathering weights for `loop_iteration + 1` right now, the network communication
+    can overlap with the compute happening in `loop_iteration`.
+    """
+    repeat_weights = self.from_all_variables_to_repeat_weights(weights, loop_iteration + 1)
+    return self.from_repeat_weights_to_bsw(repeat_weights, physical_partition_spec)
 
   def run_one_iteration(self, loop_state, bsw, positions, segment_ids, deterministic, model_mode, logical_partition_spec):
     """Executes the forward/backward logic for a single microbatch inside the pipeline.
@@ -1389,18 +1398,12 @@ def run_iteration_scannable(model, loop_state, bsw):
     run_one_repeat_scannable = base_scannable(length=self.config.num_pipeline_microbatches)
     run_bubbles_scannable = base_scannable(length=bubble_iterations)
 
-    if self.config.scan_pipeline_repeats:
-      run_repeats_scanned = pipeline_utils.create_flax_pipeline_scan(
-          pipeline_stage_fn=run_one_repeat_scannable, length=self.config.num_pipeline_repeats
-      )
-      run_bubbles_scanned = pipeline_utils.create_flax_pipeline_scan(pipeline_stage_fn=run_bubbles_scannable, length=1)
-      (loop_state, bsw), _ = run_repeats_scanned(self, (loop_state, bsw))
-      (loop_state, bsw), _ = run_bubbles_scanned(self, (loop_state, bsw))
-    else:
-      for _ in range(self.config.num_pipeline_repeats):
-        (loop_state, bsw), _ = run_one_repeat_scannable(self, (loop_state, bsw))
-      for _ in range(bubble_iterations):
-        (loop_state, bsw), _ = run_iteration_scannable(self, loop_state, bsw)
+    run_repeats_scanned = pipeline_utils.create_flax_pipeline_scan(
+        pipeline_stage_fn=run_one_repeat_scannable, length=self.config.num_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))
 
     final_output = self.realign_output_microbatches(loop_state["state_io"])
     final_output = jnp.reshape(

src/maxtext/utils/pipeline_utils.py

@@ -286,23 +286,30 @@ def create_rematerialized_pipeline_stage(
   """
 
   def execute_pipeline_stage(model, loop_state_and_bsw):
-    loop_state, bsw = loop_state_and_bsw
-    # Retrieve the specific weights needed for this pipeline chunk
-    bsw = model.weight_prefetching(pipeline_weights, physical_partition_spec, loop_state["loop_iteration"])
-
-    if model.config.scan_pipeline_iterations:
-      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)
-    else:
-      for _ in range(length):
-        (loop_state, bsw), _ = run_iteration_scannable(model, loop_state, bsw)
-    return (loop_state, bsw), None
+    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)
+    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 nn.remat(
       execute_pipeline_stage,