add jaxpr dump support

Author: NuojCheng

src/MaxText/configs/base.yml

@@ -696,7 +696,7 @@ profile_periodically_period: -1 # If set to a positive integer, profile every pr
 managed_mldiagnostics: False # Whether to enable the managed diagnostics
 managed_mldiagnostics_run_group: ""  # Optional. Used to group multiple runs.
 
-# Dump HLO options
+# Dump HLO and jaxpr options
 dump_hlo: False
 dump_step: -1 # Dump modules at the given step if set to a positive integer.
 dump_hlo_local_dir: "/tmp/xla_dump/"
@@ -708,6 +708,10 @@ dump_hlo_xla_flags: "" # Defaults to "--xla_dump_to={dump_hlo_local_dir} --xla_d
 dump_hlo_upload_all: False # If true all hosts dump HLO, false only jax.process_index()==0
 # All hosts should have identical HLO for SPMD programs, however we have encountered some bugs
 # where this is not the case and it is helpful to compare HLO across hosts.
+dump_jaxpr: False
+dump_jaxpr_local_dir: "/tmp/jaxpr_dump/"
+dump_jaxpr_delete_local_after: True
+dump_jaxpr_gcs_dir: "" # Defaults to {base_output_directory}/{run_name}/jaxpr_dump
 
 # When dropout is false the model is a deterministic function of the
 # data_shuffle_seed and init_weights_seed (i.e. reproducible losses)

src/MaxText/configs/types.py

@@ -1289,6 +1289,13 @@ class HloDump(BaseModel):
   dump_hlo_local_module_name: str = Field("jit_train_step", description="Filter modules to save locally by this name.")
   dump_hlo_xla_flags: str = Field("", description="Pass custom XLA flags for HLO dumping.")
   dump_hlo_upload_all: bool = Field(False, description="Upload HLO from all hosts.")
+  dump_jaxpr: bool = Field(False, description="Enable jaxpr dumping.")
+  dump_jaxpr_local_dir: PathStr = Field(
+      os.path.join(gettempdir(), "jaxpr_dump", ""),
+      description="Local directory to dump jaxpr.",
+  )
+  dump_jaxpr_delete_local_after: bool = Field(True, description="Delete local jaxpr dump after uploading to GCS.")
+  dump_jaxpr_gcs_dir: PathStr = Field("", description="GCS directory to upload jaxpr dumps.")
 
 
 class StackTrace(BaseModel):
@@ -1837,6 +1844,10 @@ def validate_and_set_hlo_dump_defaults():
         self.dump_hlo_gcs_dir = os.path.join(self.base_output_directory, self.run_name, "xla_dump")
       else:
         self.dump_hlo_gcs_dir = gcs_utils.add_trailing_slash(self.dump_hlo_gcs_dir)
+      if not self.dump_jaxpr_gcs_dir:
+        self.dump_jaxpr_gcs_dir = os.path.join(self.base_output_directory, self.run_name, "jaxpr_dump")
+      else:
+        self.dump_jaxpr_gcs_dir = gcs_utils.add_trailing_slash(self.dump_jaxpr_gcs_dir)
       if not os.environ.get("XLA_FLAGS"):
         os.environ["XLA_FLAGS"] = self.dump_hlo_xla_flags
 

src/MaxText/maxtext_utils.py

@@ -17,6 +17,7 @@
 
 import functools
 import pickle
+import os
 
 from flax import linen as nn
 from flax.linen import partitioning as nn_partitioning
@@ -41,6 +42,7 @@
 from MaxText import multimodal_utils
 from MaxText import sharding
 from MaxText.configs import types
+from MaxText.utils import gcs_utils
 from MaxText.common_types import DecoderBlockType, MODEL_MODE_PREFILL, MODEL_MODE_AUTOREGRESSIVE
 from MaxText.inference.page_manager import PageState
 
@@ -1213,3 +1215,33 @@ def print_shardings_params(params, params_sharding, mesh):
     shape = jax.typeof(leaf_val)
     pspec = sharding.remove_size_one_mesh_axis(leaf_sharding.spec, mesh)
     max_logging.log(f"{path_str:.<80} {shape} {tuple(pspec)}")
+
+
+def maybe_dump_jaxpr(config, p_train_step, train_step_inputs):
+  """Dump jaxpr to local then upload to GCS."""
+  if not config.dump_jaxpr:
+    return
+  max_logging.log("Tracing train_step to jaxpr...")
+
+  # We use the p_train_step (the JIT-decorated function)
+  p_train_jaxpr = jax.make_jaxpr(p_train_step)(*train_step_inputs)
+
+  local_filename = "train_step.jaxpr"
+  local_path = os.path.join(config.dump_jaxpr_local_dir, local_filename)
+
+  os.makedirs(config.dump_jaxpr_local_dir, exist_ok=True)
+
+  # pylint: disable=unspecified-encoding
+  with open(local_path, "w") as f:
+    f.write(str(p_train_jaxpr))
+
+  max_logging.log(f"Jaxpr dumped locally to {local_path}")
+
+  if config.dump_jaxpr_gcs_dir:
+    gcs_utils.upload_dump(
+        config.dump_jaxpr_local_dir,
+        config.dump_jaxpr_gcs_dir,
+        module_name=local_filename,
+        delete_local_after=config.dump_jaxpr_delete_local_after,  # Keeping local for debugging
+        all_host_upload=False,  # Only upload from lead host (Host 0)
+    )

src/MaxText/train.py

@@ -427,6 +427,7 @@ def train_loop(config, recorder, state=None):
     shaped_batch = maxtext_utils.get_shaped_batch(config)
     if config.shard_optimizer_over_data:
       state = sharding.maybe_shard_with_name(state, state_mesh_shardings, config.shard_mode)
+    maxtext_utils.maybe_dump_jaxpr(config, p_train_step, (state, shaped_batch, init_rng))
     if config.compiled_trainstep_file == "":  # compile only when there is no pre-compiled file loaded
       compiled = p_train_step.lower(state, shaped_batch, init_rng).compile()
       compiled_stats = compiled.memory_analysis()

src/MaxText/train_compile.py

@@ -121,6 +121,7 @@ def jit_and_compile(
     out_shardings,
     static_argnums,
     donate_argnums,
+    config,
     logical_axis_rules,
 ):
   """Jit, lower, and compile func."""
@@ -132,6 +133,7 @@ def jit_and_compile(
         static_argnums=static_argnums,
         donate_argnums=donate_argnums,
     )
+    maxtext_utils.maybe_dump_jaxpr(config, jitted, func_input_args)
     lowered = jitted.lower(*func_input_args, **func_input_kwargs)
   compiled = lowered.compile()
   return compiled
@@ -180,6 +182,7 @@ def is_oom(argv: Sequence[str]) -> bool:
         out_shard,
         static_argnums,
         donate_argnums,
+        config,
         nn_partitioning.axis_rules(config.logical_axis_rules),
     )
     return False
@@ -241,6 +244,7 @@ def main(argv: Sequence[str]) -> None:
       out_shard,
       static_argnums,
       donate_argnums,
+      config,
       nn_partitioning.axis_rules(config.logical_axis_rules),
   )
   print("Jitting and compilation complete!", flush=True)

src/MaxText/utils/gcs_utils.py

@@ -79,7 +79,7 @@ def upload_dump(local_dir, target_dir, module_name=None, delete_local_after=True
     hostname = socket.gethostname()  # Alternatively can use jax.process_id()
     prefix_name = os.path.join(prefix_name, hostname)
     target_dir = os.path.join(target_dir, hostname)
-  max_logging.log(f"Uploading HLO Dump to {target_dir}...")
+  max_logging.log(f"Uploading Dump to {target_dir}...")
   for root, _, files in os.walk(local_dir):
     for file in files:
       if module_name and module_name not in file:
@@ -91,7 +91,7 @@ def upload_dump(local_dir, target_dir, module_name=None, delete_local_after=True
       blob_name = os.path.join(prefix_name, relative_path)
       blob = bucket.blob(blob_name)
       blob.upload_from_filename(local_path)
-  max_logging.log(f"HLO Dump Uploaded to {target_dir}!")
+  max_logging.log(f"Dump Uploaded to {target_dir}!")
   if delete_local_after:
     shutil.rmtree(local_dir)