Allow user-specified num_slices

Author: xuefgu

src/MaxText/checkpointing.py

@@ -427,7 +427,9 @@ def _restore_grain_iterator(
   elif expansion_factor_real_data > 1 and process_count_stored == process_count_jax // expansion_factor_real_data:
     # Scaling up to a larger number of hosts.(e.g., 32 files -> 64 processes)
     # In this case, a subset of hosts restore the data iterator.
-    assert not isinstance(data_iterator, list), "when expansion_factor_real_data > 1, the data iterator should not be a list."
+    assert not isinstance(
+        data_iterator, list
+    ), "when expansion_factor_real_data > 1, the data iterator should not be a list."
     grain_restore_args = GrainCheckpointRestore(
         data_iterator.local_iterator, process_index=jax.process_index(), process_count=process_count_stored
     )

src/MaxText/configs/base.yml

@@ -471,7 +471,8 @@ ici_expert_parallelism: 1
 # Enable ZeRO-1 optimizer sharding over data axis
 shard_optimizer_over_data: False
 
-# The number of TPU slices is automatically determined, you should not set this explicitly. For ahead of time compilation,
+# Unless explicitly specified, the number of TPU slices is automatically determined. It should only be set for
+# disaggregated reinforcement learning workloads using multiple slices. For ahead of time compilation,
 # you should set compile_toplogy_num_slices, which will in turn set this value. For non-TPU environments this is set to 1.
 num_slices: -1
 

src/MaxText/max_utils.py

@@ -289,6 +289,9 @@ def _retrieve_jax_init_info(raw_keys):
 
 def get_num_slices(raw_keys):
   """Calculate num_slices based on number of devices."""
+  if raw_keys["num_slices"] != -1:
+    max_logging.log(f"Using num_slices={raw_keys['num_slices']} per user request.")
+    return raw_keys["num_slices"]
   if raw_keys["hardware"] == "cpu":
     max_logging.log(" Setting num_slices=1 for CPU hardware type")
     return 1