Update default XLA flags for TPU 7x to mitigate loss lifting and add allow_split_physical_axes to split physical axes

GitOrigin-RevId: df14d6e

Author: Liang (AIML-Compute) He

axlearn/common/compiler_options.py

@@ -157,8 +157,11 @@ def default_xla_options(
         # v7x flags from MaxText, inclusing SparseCore configs
         # TODO(samuel-andersen): Move v7x SparseCore config below with v6e
         options.update(
-            xla_tpu_enable_sparse_core_collective_offload_all_reduce="true",
-            xla_tpu_enable_sparse_core_collective_offload_reduce_scatter="true",
+            # TODO(lianghe): Re-enable after JAX 0.8.3 issues are resolved.
+            # - all_reduce: causes abnormal loss in pretraining.
+            # - reduce_scatter: causes abnormal loss in post-training (SFT).
+            # xla_tpu_enable_sparse_core_collective_offload_all_reduce="true",
+            # xla_tpu_enable_sparse_core_collective_offload_reduce_scatter="true",
             xla_tpu_enable_sparse_core_collective_offload_all_gather="true",
             xla_tpu_enable_sparse_core_collective_offload_2d_all_gather="true",
             xla_tpu_enable_sparse_core_reduce_scatter_v2="true",

axlearn/common/utils.py

@@ -1693,7 +1693,13 @@ def build_standard_mesh(mesh_shape: MeshShape, *, devices: np.ndarray) -> np.nda
             mesh = mesh.reshape(mesh_shape)
             logging.log_first_n(logging.INFO, "Using custom mesh: %s", 1, str(mesh))
             return mesh
-        return mesh_utils.create_device_mesh(mesh_shape, devices=devices)
+        return mesh_utils.create_device_mesh(
+            mesh_shape,
+            devices=devices,
+            # Set allow_split_physical_axes to True to split physical axes.
+            # Reference: jax.experimental.mesh_utils.create_device_mesh docs.
+            allow_split_physical_axes=True,
+        )
     except NotImplementedError as e:
         logging.warning(
             "mesh_utils.create_device_mesh cannot handle shape %s: %s. "