add lax.cond and dtype control

Author: NuojCheng

src/maxtext/configs/custom_mesh_and_rule/pipeline-large-moe.yml

@@ -41,6 +41,8 @@ logical_axis_rules: [
                       ['activation_attn_length', ['context', 'expert']],
                       ['activation_q_length', ['context', 'expert']],
                       ['activation_attn_embed', ['tensor']],
+                      ['activation_norm_length', ['context']],
+                      ['activation_norm_length_moe', ['context']],
                       ['activation_embed', ['tensor']],
                       ['activation_embed_moe', ['tensor']],
                       ['activation_mlp', ['tensor']],

src/maxtext/layers/pipeline.py

@@ -970,8 +970,8 @@ def init_states(self, inputs):
     def _init_empty_bsw_buffers(variables):
       # BSW requires two buffers (current and next) for the sliding window
       return (
-          jax.tree.map(lambda x: jnp.zeros_like(x[0]), variables),
-          jax.tree.map(lambda x: jnp.zeros_like(x[0]), variables),
+          jax.tree.map(lambda x: jnp.zeros_like(x[0], dtype=jnp.bfloat16), variables),
+          jax.tree.map(lambda x: jnp.zeros_like(x[0], dtype=jnp.bfloat16), variables),
       )
 
     if self.is_initializing():

src/maxtext/utils/pipeline_utils.py

@@ -22,6 +22,20 @@
 import jax.numpy as jnp
 
 
+def to_bf16(tree):
+  """Casts all floating-point leaves in a PyTree to bfloat16."""
+  return jax.tree.map(
+      lambda t: t.astype(jnp.bfloat16) if hasattr(t, "dtype") and jnp.issubdtype(t.dtype, jnp.floating) else t, tree
+  )
+
+
+def to_f32(tree):
+  """Casts all floating-point leaves in a PyTree to float32."""
+  return jax.tree.map(
+      lambda t: t.astype(jnp.float32) if hasattr(t, "dtype") and jnp.issubdtype(t.dtype, jnp.floating) else t, tree
+  )
+
+
 def get_mesh_axis_dim_indices(physical_partition_spec, axis_name="fsdp"):
   """Finds the tensor dimension index sharded across a specific physical mesh axis.
 
@@ -310,7 +324,7 @@ def execute_pipeline_stage_pure(loop_state, w_curr, pipeline_weights):
     def execute_pipeline_stage_pure_fwd(loop_state, w_curr, pipeline_weights):
       # Prefetch FSDP-sharded weights for the upcoming pipeline repeat
       w_next = model.weight_prefetching(
-          pipeline_weights,
+          to_bf16(pipeline_weights),
           physical_partition_spec,
           loop_state["loop_iteration"],
       )
@@ -332,9 +346,8 @@ def execute_pipeline_stage_pure_fwd(loop_state, w_curr, pipeline_weights):
       )
       # Execute the forward pass of the microbatches and generate its VJP.
       # The VJP captures necessary checkpoints to evaluate gradients later.
-      (loop_state, bsw), scan_microbatches_vjp = jax.vjp(scan_microbatches_fn, loop_state, bsw, positions, segment_ids)
+      (loop_state, _), scan_microbatches_vjp = jax.vjp(scan_microbatches_fn, loop_state, bsw, positions, segment_ids)
       # Discard the old weights (w_curr) and advance w_next to act as the current weights in the next iteration
-      _, w_next = bsw
       return (loop_state, w_next), (scan_microbatches_vjp, weight_prefetching_t)
 
     def execute_pipeline_stage_pure_bwd(residuals, g_outputs):
@@ -349,7 +362,7 @@ def execute_pipeline_stage_pure_bwd(residuals, g_outputs):
       # Apply the linear transpose of the weight prefetch to execute the reduce-scatter
       # This maps the gradients of the gathered weights back to the FSDP-sharded parameter space
       g_w_curr, g_w_next = g_bsw
-      (g_pipeline_weights,) = weight_prefetching_t(g_w_next)
+      (g_pipeline_weights,) = weight_prefetching_t(to_f32(g_w_next))
       # Return gradients corresponding to the three original inputs of execute_pipeline_stage_pure
       return g_loop_state, g_w_curr, g_pipeline_weights