More logging for XLA compile options

GitOrigin-RevId: 3fad67f

Author: changlan

axlearn/common/compiler_options.py

@@ -375,6 +375,14 @@ def infer_xla_performance_flags(
     *, mesh_shape: "MeshShape", mesh_axis_names: Sequence[str], device_kind: str
 ) -> dict[str, str]:
     """Performs automatic XLA flag tuning based on mesh shape and device kind."""
+
+    logging.info(
+        "Inferring XLA flags for mesh %s (%s) and device %s",
+        mesh_shape,
+        mesh_axis_names,
+        device_kind,
+    )
+
     if device_kind == "TPU v6 lite":
         device_kind = "TPU v6e"
 

axlearn/common/trainer.py

@@ -1057,6 +1057,7 @@ def _get_compiled_train_step_fn(
         options = infer_xla_performance_flags(
             mesh_shape=cfg.mesh_shape, mesh_axis_names=cfg.mesh_axis_names, device_kind=device_kind
         )
+        logging.log_first_n(logging.INFO, "Compiler options: %s", 1, options)
         if not with_xsc:
             self._maybe_record_event(
                 measurement.Event.START_CUSTOM_BADPUT_EVENT,