Update

Author: hsuan-lun-chiang

src/maxtext/configs/base.yml

@@ -315,7 +315,7 @@ attention_out: 'remat'
 
 optimizer_memory_host_offload: False
 parameter_memory_host_offload: False
-scan_layers: False # We recommend setting this to false when using pipeline parallelism, instead scanning the PP iterations.
+scan_layers: True # We recommend setting this to false when using pipeline parallelism, instead scanning the PP iterations.
 param_scan_axis: 1
 
 # The attention parameter dictates the specific algorithm/methodology used to compute the attention scores