remove weights from most inputs

Author: NuojCheng

src/maxtext/configs/base.yml

@@ -923,7 +923,7 @@ xprof_e2e_enable_fw_power_level_event: False
 xprof_e2e_enable_fw_thermal_event: False
 profile_power_events: False # Set to True to enable TPU-specific power/thermal profiling events. Defaults to False to avoid breaking GPU xplane tracing.
 
-log_config: True # Prints the config (after defaults have been set by pyconfig logic)
+log_config: False # Prints the config (after defaults have been set by pyconfig logic)
 debug_sharding: False # Prints model weights sharding info
 
 # Checkpoint Structured logging

src/maxtext/layers/pipeline.py

@@ -1252,9 +1252,7 @@ def weight_prefetching(self, weights, physical_partition_spec, loop_iteration):
     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")
 
-  def run_one_iteration(
-      self, loop_state, bsw, weights, positions, segment_ids, deterministic, model_mode, logical_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.
 
     This acts as the core step function that our `jax.lax.scan` wrappers call. It routes
@@ -1353,12 +1351,11 @@ def __call__(
 
     logical_partition_spec = pipeline_utils.strip_pipeline_repeat_logical_axis(logical_partition_spec)
 
-    def run_iteration_scannable(model, loop_state, bsw, weights):
+    def run_iteration_scannable(model, loop_state, bsw):
       return (
           model.run_one_iteration(
               loop_state,
               bsw,
-              weights,
               positions,
               segment_ids,
               deterministic,
@@ -1386,6 +1383,7 @@ def run_iteration_scannable(model, loop_state, bsw, weights):
         physical_partition_spec=physical_partition_spec,
         positions=positions,
         segment_ids=segment_ids,
+        pipeline_weights=weights,
     )
 
     run_one_repeat_scannable = base_scannable(length=self.config.num_pipeline_microbatches)
@@ -1396,13 +1394,13 @@ def run_iteration_scannable(model, loop_state, bsw, weights):
           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, weights), _ = run_repeats_scanned(self, (loop_state, bsw, weights))
-      (loop_state, bsw, weights), _ = run_bubbles_scanned(self, (loop_state, bsw, weights))
+      (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, weights), _ = run_one_repeat_scannable(self, (loop_state, bsw, weights))
+        (loop_state, bsw), _ = run_one_repeat_scannable(self, (loop_state, bsw))
       for _ in range(bubble_iterations):
-        (loop_state, bsw, weights), _ = run_iteration_scannable(self, loop_state, bsw, weights)
+        (loop_state, bsw), _ = run_iteration_scannable(self, loop_state, bsw)
 
     final_output = self.realign_output_microbatches(loop_state["state_io"])
     final_output = jnp.reshape(

src/maxtext/utils/pipeline_utils.py

@@ -194,56 +194,54 @@ def create_gradient_accumulation_scan(
   """
 
   @functools.partial(jax.custom_vjp)
-  def run_single_microbatch_custom(lightweight_state, bsw, weights, pos_arg, seg_arg):
-    return run_single_microbatch_custom_fwd(lightweight_state, bsw, weights, pos_arg, seg_arg)[0]
+  def run_single_microbatch_custom(lightweight_state, bsw, pos_arg, seg_arg):
+    return run_single_microbatch_custom_fwd(lightweight_state, bsw, pos_arg, seg_arg)[0]
 
-  def run_single_microbatch_custom_fwd(lightweight_state, bsw, weights, pos_arg, seg_arg):
-    def _run(l, b, w):
+  def run_single_microbatch_custom_fwd(lightweight_state, bsw, pos_arg, seg_arg):
+    def _run(l, b):
       out = model.run_one_iteration(
-          l, b, w, pos_arg, seg_arg, deterministic, model_mode, logical_partition_spec=logical_partition_spec
+          l, b, pos_arg, seg_arg, deterministic, model_mode, logical_partition_spec=logical_partition_spec
       )
-      return out, b, w
+      return out, b
 
     # Rematerialize the inner step to save activation memory
     _run_remat = jax.remat(_run, prevent_cse=False, policy=model.get_pipeline_remat_policy())
-    out, vjp_fun = jax.vjp(_run_remat, lightweight_state, bsw, weights)
+    out, vjp_fun = jax.vjp(_run_remat, lightweight_state, bsw)
     return out, vjp_fun
 
   def run_single_microbatch_custom_bwd(res, g_out):
     vjp_fun = res
-    d_l, d_b, d_w = vjp_fun(g_out)
-    return d_l, d_b, d_w, None, None
+    d_l, d_b = vjp_fun(g_out)
+    return d_l, d_b, None, None
 
   run_single_microbatch_custom.defvjp(run_single_microbatch_custom_fwd, run_single_microbatch_custom_bwd)
 
   @functools.partial(jax.custom_vjp)
-  def run_pipeline_microbatches_custom(loop_state, bsw, weights, positions, segment_ids):
-    return run_pipeline_microbatches_custom_fwd(loop_state, bsw, weights, positions, segment_ids)[0]
+  def run_pipeline_microbatches_custom(loop_state, bsw, positions, segment_ids):
+    return run_pipeline_microbatches_custom_fwd(loop_state, bsw, positions, segment_ids)[0]
 
-  def run_pipeline_microbatches_custom_fwd(loop_state, bsw, weights, positions, segment_ids):
+  def run_pipeline_microbatches_custom_fwd(loop_state, bsw, positions, segment_ids):
     final_lightweight, scan_vjp_fun = jax.vjp(
-        lambda l, b, w: jax.lax.scan(
-            lambda carry, _: (run_single_microbatch_custom(carry, b, w, positions, segment_ids)[0], None),
+        lambda l, b: jax.lax.scan(
+            lambda carry, _: (run_single_microbatch_custom(carry, b, positions, segment_ids)[0], None),
             l,
             None,
             length=length,
         )[0],
         loop_state,
         bsw,
-        weights,
     )
 
-    return (final_lightweight, bsw, weights), scan_vjp_fun
+    return (final_lightweight, bsw), scan_vjp_fun
 
   def run_pipeline_microbatches_custom_bwd(residuals, g_final_state):
     scan_vjp_fun = residuals
-    g_lightweight, g_bsw, g_weights = g_final_state
-    d_init_lightweight, d_init_bsw, d_init_weights = scan_vjp_fun(g_lightweight)
+    g_lightweight, g_bsw = g_final_state
+    d_init_lightweight, d_init_bsw = scan_vjp_fun(g_lightweight)
 
     d_init_bsw = jax.tree.map(lambda d, g: d + g if hasattr(d, "shape") else d, d_init_bsw, g_bsw)
-    d_init_weights = jax.tree.map(lambda d, g: d + g if hasattr(d, "shape") else d, d_init_weights, g_weights)
 
-    return (d_init_lightweight, d_init_bsw, d_init_weights, None, None)
+    return (d_init_lightweight, d_init_bsw, None, None)
 
   run_pipeline_microbatches_custom.defvjp(run_pipeline_microbatches_custom_fwd, run_pipeline_microbatches_custom_bwd)
   return run_pipeline_microbatches_custom
@@ -259,6 +257,7 @@ def create_rematerialized_pipeline_stage(
     physical_partition_spec,
     positions,
     segment_ids,
+    pipeline_weights,
 ):
   """Builds a memory-checkpointed execution block for a single pipeline stage.
 
@@ -279,16 +278,17 @@ def create_rematerialized_pipeline_stage(
     physical_partition_spec: Rules for physical device mesh mappings (used in prefetching).
     positions: Position IDs for the sequence.
     segment_ids: Segment/Attention routing IDs for the sequence.
+    pipeline_weights: The fully gathered pipeline weights explicitly passed via closure.
 
   Returns:
     A function decorated with `nn.remat` that takes `(model, loop_state)` and returns
     the updated `loop_state`.
   """
 
-  def execute_pipeline_stage(model, loop_state_and_bsw_and_weights):
-    loop_state, bsw, weights = loop_state_and_bsw_and_weights
+  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(weights, physical_partition_spec, loop_state["loop_iteration"])
+    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(
@@ -298,11 +298,11 @@ def execute_pipeline_stage(model, loop_state_and_bsw_and_weights):
           model_mode=model_mode,
           logical_partition_spec=logical_partition_spec,
       )
-      loop_state, bsw, weights = scan_microbatches_fn(loop_state, bsw, weights, positions, segment_ids)
+      loop_state, bsw = scan_microbatches_fn(loop_state, bsw, positions, segment_ids)
     else:
       for _ in range(length):
-        (loop_state, bsw, weights), _ = run_iteration_scannable(model, loop_state, bsw, weights)
-    return (loop_state, bsw, weights), None
+        (loop_state, bsw), _ = run_iteration_scannable(model, loop_state, bsw)
+    return (loop_state, bsw), None
 
   return nn.remat(
       execute_pipeline_stage,