Update

Author: andersensam

axlearn/cloud/gcp/jobset_utils.py

@@ -635,7 +635,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/compiler_options.py

@@ -133,15 +133,11 @@ def default_xla_options(
             # Many of the v7x flags are similar to v6e
             xla_tpu_dvfs_p_state=7,
             xla_tpu_scoped_vmem_limit_kib=65536,
-            xla_tpu_enable_sparse_core_reduce_scatter_v2="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_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_enable_sparse_core_collective_offload_all_reduce="true",
-            xla_tpu_enable_sparse_core_collective_offload_reduce_scatter="true",
+            xla_tpu_impure_enable_packed_bf16_math_ops="true",
         )
 
         # Ensure pipelining is properly configured
@@ -153,11 +149,15 @@ def default_xla_options(
 
         # v7x flags from MaxText, inclusing SparseCore configs
         options.update(
-            xla_tpu_use_single_sparse_core_for_all_gather_offload="true",
+            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",
-            xla_tpu_impure_enable_packed_bf16_math_ops="true",
         )
     if num_slices > 1:
         # Support multiple TPU slices connected over a data center network.

axlearn/common/trainer.py

@@ -626,7 +626,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)
@@ -1338,6 +1338,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)

axlearn/experiments/text/gpt/fuji.py

@@ -748,11 +748,12 @@ def get_trainer_kwargs(
                                         policy=config_for_function(
                                             save_and_offload_only_these_names_regex
                                         ).set(
-                                            names_which_can_be_saved="|".join(
-                                                [
-                                                    RematRegexSavePatterns.QKV_PROJ.value
-                                                ]
-                                            ),
+                                            # names_which_can_be_saved="|".join(
+                                            #     [
+                                            #         RematRegexSavePatterns.QKV_PROJ.value
+                                            #     ]
+                                            # ),
+                                            names_which_can_be_saved=None,
                                             names_which_can_be_offloaded="|".join(
                                                 [
                                                     RematRegexSavePatterns.INPUT.value,