Cleanup for v7x

Author: andersensam

Dockerfile

@@ -98,14 +98,13 @@ ARG INSTALL_PATHWAYS_JAXLIB=false
 
 # Ensure we install the TPU version, even if building locally.
 # Jax will fallback to CPU when run on a machine without TPU.
-RUN uv pip install --prerelease=allow .[core,tpu] && uv cache clean
+RUN uv pip install -qq --prerelease=allow .[core,tpu] && uv cache clean
 RUN if [ -n "$EXTRAS" ]; then uv pip install -qq .[$EXTRAS] && uv cache clean; fi
 RUN if [ "$INSTALL_PATHWAYS_JAXLIB" = "true" ]; then \
       uv pip install --prerelease=allow "jaxlib==0.5.3.dev20250918" \
         --find-links https://storage.googleapis.com/axlearn-wheels/wheels.html; \
     fi
-COPY --from=libtpu-target:latest /wheels /wheels
-RUN uv pip install --no-deps /wheels/*.whl && uv cache clean
+RUN uv pip install -qq --no-deps libtpu==0.0.26.dev20251022+nightly && uv cache clean
 COPY . .
 
 ################################################################################

axlearn/common/compiler_options.py

@@ -136,13 +136,12 @@ def default_xla_options(
     if version == "v7x":
         options.update(
             # 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_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_impure_enable_packed_bf16_math_ops="true",
+            xla_tpu_bf16_emission_mode="NATIVE_EMISSION",
         )
 
         # Ensure pipelining is properly configured
@@ -153,6 +152,7 @@ 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",
@@ -209,7 +209,9 @@ def default_xla_options(
             continue
         elif isinstance(v, str):
             # Allow numeric strings, time-based strings (e.g., "10m", "30s", "60m"), and bool str.
-            if v.isdigit() or re.match(r"^\d+[ms]$", v.strip()) or v.strip() in ["true", "false"]:
+            if v.isdigit() or re.match(r"^\d+[ms]$", v.strip()) or v.strip() in [
+                "true", "false", "NATIVE_EMISSION"
+            ]:
                 continue
             # Allow paths.
             if v.startswith("/"):

axlearn/experiments/text/gpt/fuji.py

@@ -727,11 +727,52 @@ def get_trainer_kwargs(
                     ),
                 ),
                 (
-                    "tpu-v7x-.*",
+                    "tpu-v7x-128",
                     ChainConfigModifier.default_config().set(
                         config_modifiers=[
                             MeshShapeModifier.default_config().set(
-                                mesh_shape=mesh_shape_from_axes(data=-1, fsdp=-1)
+                                mesh_shape=mesh_shape_from_axes(data=-1, fsdp=128)
+                            ),
+                            # Ensure we set the default tpu_block_size=2048 on v7x
+                            # With the 70B model, MaxText also modifies block_q
+                            # and block_kv_compute to 4096 and 1024 respectively
+                            V7xFlashConfigModifier.default_config(),
+                            SplashAttentionConfigModifier.default_config().set(
+                                splash_block_q=4096,
+                                splash_block_kv_compute=1024,
+                            ),
+                            RematSpecModifier.default_config().set(
+                                remat_policies={
+                                    "model.decoder.transformer.layer": RematSpec(
+                                        prevent_cse=False,
+                                        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_offloaded="|".join(
+                                                [
+                                                    RematRegexSavePatterns.INPUT.value,
+                                                ]
+                                            ),
+                                            offload_src="device",
+                                            offload_dst="pinned_host",
+                                        ),
+                                    ),
+                                }
+                            ),
+                        ],
+                    ),
+                ),
+                (
+                    "tpu-v7x-256",
+                    ChainConfigModifier.default_config().set(
+                        config_modifiers=[
+                            MeshShapeModifier.default_config().set(
+                                mesh_shape=mesh_shape_from_axes(data=-1, fsdp=128)
                             ),
                             # Ensure we set the default tpu_block_size=2048 on v7x
                             # With the 70B model, MaxText also modifies block_q
@@ -748,11 +789,6 @@ 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=None,
                                             names_which_can_be_offloaded="|".join(
                                                 [