Support subslice shapes in maxtext when using Pathways

MaxText/configs/base.yml

@@ -745,4 +745,8 @@ projector_output_dim_for_vit: 4096
 rope_theta_for_vit: 10000
 vision_output_dim_for_vit: 4096
 pixel_shuffle_ratio_for_vit: 0.5
-projector_dropout_for_vit: 0.0
+projector_dropout_for_vit: 0.0
+
+# Subslice shape in the form of "x,y,z" when using pathways (single controller). 
+# Example: "8,8" to use a 8x8 subgrid (64 chips) of a full pod (16x16) of trillium.
+subslice_shape: ""

MaxText/maxtext_utils.py

@@ -921,6 +921,31 @@ def create_device_mesh(config, devices=None):
   """Creates a device mesh with each slice in its own data parallel group. If there is only one slice, uses two replicas"""
   if devices is None:
     devices = jax.devices()
+  if (
+      config.subslice_shape is not None
+      and config.enable_single_controller
+      and config.num_slices == 1
+  ):
+    max_logging.log(
+        f"Trying to create a subslice with shape: {config.subslice_shape}"
+    )
+    subslice_shape = tuple(int(x) for x in config.subslice_shape.split(","))
+    device_coords = [device.coords for device in devices]
+    device_coords_np = np.array(device_coords)
+
+    # Find the minimum coordinates to start the subslice
+    min_coords = device_coords_np.min(axis=0)
+
+    subslice_devices = []
+    for device in devices:
+      coords = device.coords
+      if all(
+          min_coords[i] <= coords[i] < min_coords[i] + subslice_shape[i]
+          for i in range(len(subslice_shape))
+      ):
+        subslice_devices.append(device)
+    devices = subslice_devices
+
   num_devices = len(devices)
   num_slices = 1 if config.inference_benchmark_test else config.num_slices
   num_devices_per_slice = num_devices // num_slices

MaxText/pyconfig.py

@@ -18,6 +18,7 @@
 # pylint: disable=missing-module-docstring, bare-except, consider-using-generator, missing-function-docstring
 from collections import OrderedDict
 from typing import Any, Union
+from math import prod
 import math
 import os
 import sys
@@ -1014,6 +1015,9 @@ def get_num_target_devices(raw_keys):
     compile_topology = accelerator_to_spec_map.get_system_characteristics(raw_keys["compile_topology"])
     devices_per_slice = compile_topology.devices_per_slice
     return int(devices_per_slice * raw_keys["compile_topology_num_slices"])
+  elif raw_keys.get("subslice_shape") and raw_keys.get("enable_single_controller"):
+    subslice_shape = tuple(int(x) for x in raw_keys["subslice_shape"].split(","))
+    return prod(subslice_shape)
   else:
     return len(jax.devices())