Merge branch 'jax_0.8.0_py3.12' into jax_py3.12_nightly.

Author: andersensam

Dockerfile

@@ -139,7 +139,7 @@ FROM base AS gpu
 # Enable the CUDA repository and install the required libraries (libnvrtc.so)
 RUN curl -o cuda-keyring_1.1-1_all.deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb && \
     dpkg -i cuda-keyring_1.1-1_all.deb && \
-    apt-get update && apt-get install -y cuda-libraries-dev-12-8 ibverbs-utils && \
+    apt-get update && apt-get install -y cuda-libraries-dev-12-9 ibverbs-utils && \
     apt clean -y
 COPY pyproject.toml README.md /root/
 RUN uv pip install -qq .[core,gpu] && uv cache clean

axlearn/cloud/gcp/jobset_utils.py

@@ -685,7 +685,7 @@ def _build_pod(self) -> Nested[Any]:
         # TODO(samos123) support using reservation when using local launch
         # the local launch command automatically sets tier=disabled.
         logging.info("Found tier=%s in env. Using reservation=%s", tier, cfg.reservation)
-        if tier == "0" and cfg.reservation is not None:
+        if tier in ["0", "disabled"] and cfg.reservation is not None:
             logging.info("Found tier=%s in env. Using reservation=%s", tier, cfg.reservation)
             selector.update({"cloud.google.com/reservation-name": cfg.reservation})
             if cfg.reservation_project is not None:

axlearn/common/array_serialization.py

@@ -473,7 +473,6 @@ async def _async_deserialize(
     context = serialization.TS_CONTEXT
     if os.getenv("ENABLE_GCS_GRPC", "false") == "true":
         tensorstore_spec, context = use_gcs_grpc(tensorstore_spec)
-
     t = await ts.open(
         tensorstore_spec,
         open=True,

axlearn/common/array_serialization_test.py

@@ -483,7 +483,6 @@ def test_shard_info_partially_replicated(
         devices = mesh_utils.create_device_mesh((8,))
         mesh = jax.sharding.Mesh(devices.reshape((4, 2)), ("x", "y"))
         sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(None, "y"))
-
         arr = jax.device_put(single_device_arr, sharding)
 
         replica_count = _num_replicas_per_shard(arr)
@@ -504,7 +503,6 @@ def test_shard_info_fully_sharded(self, max_data_shard_degree: int, shard_thresh
         devices = mesh_utils.create_device_mesh((8,))
         mesh = jax.sharding.Mesh(devices.reshape((4, 2)), ("x", "y"))
         sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec("x", "y"))
-
         arr = jax.device_put(single_device_arr, sharding)
 
         replica_count = _num_replicas_per_shard(arr)
@@ -564,6 +562,5 @@ def test_shard_coordinate(self, index, expected):
             expected,
         )
 
-
 if __name__ == "__main__":
     absltest.main()

axlearn/common/compiler_options.py

@@ -151,6 +151,37 @@ def default_xla_options(
             xla_tpu_enable_ici_ag_pipelining="true",
         )
 
+        # 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",
+            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",
+            xla_tpu_enable_sparse_core_collective_offload_3d_all_gather="true",
+            xla_tpu_use_single_sparse_core_for_all_gather_offload="true",
+            xla_sc_disable_megacore_partitioning="true",
+            xla_tpu_use_tc_device_shape_on_sc="true",
+        )
+    if version == "v7x":
+        options.update(
+            # Many of the v7x flags are similar to v6e
+            xla_tpu_scoped_vmem_limit_kib=65536,
+            xla_tpu_enable_all_gather_offload_tracing="true",
+            xla_tpu_enable_async_collective_fusion_fuse_all_gather="false",
+            xla_enable_async_all_gather="true",
+            xla_tpu_prefer_async_allgather_to_allreduce="true",
+            xla_tpu_bf16_emission_mode="NATIVE_EMISSION",
+        )
+
+        # Ensure pipelining is properly configured
+        options.update(
+            xla_should_allow_loop_variant_parameter_in_chain="true",
+            xla_should_add_loop_invariant_op_in_chain="true",
+            xla_tpu_enable_ici_ag_pipelining="true",
+        )
+
         # v7x flags from MaxText, inclusing SparseCore configs
         # TODO(samuel-andersen): Move v7x SparseCore config below with v6e
         options.update(

axlearn/common/flash_attention/layer.py

@@ -414,7 +414,6 @@ def enter_fn(_, value, default_kv):
         cfg.visit(visit_fn=visit_fn, enter_fn=enter_fn)
         return cfg
 
-
 class BackendOverrideModifier(ConfigModifier):
     """Modifies the backend_overrides config of Flash Attention."""
 

axlearn/common/layers.py

@@ -882,7 +882,9 @@ class Config(BaseLayer.Config):
     @classmethod
     def default_config(cls):
         cfg = super().default_config()
-        cfg.param_partition_spec = (None, "model")
+        # Prevent OOM on 405b
+        # cfg.param_partition_spec = (None, "model")
+        cfg.param_partition_spec = ("fsdp", "model")
         # By default, initialize to Gaussian with std=1/sqrt(dim), e.g., 0.036 when dim=768.
         #
         # This is the same as:

axlearn/common/optimizers.py

@@ -38,7 +38,6 @@
 from absl import logging
 from jax import numpy as jnp
 
-# from jax._src.sharding_impls import TransferToMemoryKind
 from optax._src import numerics
 
 from axlearn.common import flax_struct, schedule

axlearn/common/state_builder_test.py

@@ -293,6 +293,7 @@ def _mock_bert_trainer_config(
                 embedding_cfg=bert.bert_embedding_config(max_position_embeddings=max_len),
                 stack_cfg=bert.bert_transformer_config(num_layers=3, num_heads=2),
             ),
+            mesh_axis_names=("fsdp", "model")
         )
 
     def _mock_bert_trainer_config_and_state(

axlearn/common/trainer.py

@@ -630,7 +630,7 @@ def run(
                         )
                         self.vlog(3, "Done step %s", self.step)
                         num_steps += 1
-                        if num_steps % 100 == 0:
+                        if num_steps % 1 == 0:
                             now = time.perf_counter()
                             average_step_time = (now - start_time) / num_steps
                             self._step_log("Average step time: %s seconds", average_step_time)
@@ -1348,6 +1348,7 @@ def _maybe_stop_or_start_tracing(
                 cfg.start_trace_process_indices == "all"
                 or jax.process_index() in cfg.start_trace_process_indices
             )
+            self._step_log(f"{should_start_tracing=}")
         if should_start_tracing:
             self._step_log("Start profiler tracing")
             jax.profiler.start_trace(self.summary_writer.config.dir)