fix: Support multi-slice training and add experiment name (GoogleCloudPlatform#1161)

This change introduces enhanced multi-slice support for RL and SFT training workflows, improves experiment tracking with Vertex AI Tensorboard integration, and updates core utility functions and configuration defaults. The changes primarily focus on enabling more flexible distributed training, improving experiment reproducibility, and ensuring compatibility with the latest `xpk` tooling which fixed the version comparison issue.

Author: RexBearIU

dags/post_training/maxtext_rl.py

@@ -66,7 +66,7 @@
   training_config = test_config_util.RLTestConfig(
       cluster=XpkClusters.TPU_V5P_128_CLUSTER,
       accelerator="v5p-128",
-      slices=[1],  # Single slice for RL training
+      slices=[1, 2],  # Multi-slice support
       model_name="llama3.1-70b",
       base_dir=(
           f"{test_config_util.DEFAULT_BUCKET}/llama3.1-70b-Instruct/outputs"
@@ -99,6 +99,7 @@
             loss_algo=loss_algo,
             run_name=run_name,
             hf_token=HF_TOKEN_LLAMA3_1,
+            num_slices=slice_num,
         )
 
         with TaskGroup(
@@ -122,6 +123,8 @@
                 test_owner=test_owner.JACKY_F,
             ).run_model(
                 use_pathways=True,
+                use_vertex_tensorboard=True,
+                experiment_name=loss_algo.value,
                 xpk_branch=MAIN_BRANCH,
             )
 

dags/post_training/maxtext_sft.py

@@ -49,6 +49,8 @@ def run_training(
         test_owner=test_owner.JACKY_F,
     ).run_model(
         use_pathways=True,
+        use_vertex_tensorboard=True,
+        experiment_name="sft",
         xpk_branch=MAIN_BRANCH,
     )
 
@@ -128,7 +130,7 @@ def validate_training(
   training_config = test_config_util.SFTTestConfig(
       cluster=XpkClusters.TPU_V5P_128_CLUSTER,
       accelerator="v5p-128",
-      slices=[1],  # Single slice for SFT training
+      slices=[2],
       model_name="llama3.1-70b",
       steps=training_steps,
       short_id="msft",

dags/post_training/util/test_config_util.py

@@ -107,40 +107,56 @@ def __init__(
     self.rl_config_path = rl_config_path
 
   def generate_rl_training_command(
-      self, loss_algo: LossAlgo, run_name: str, hf_token: str
+      self,
+      loss_algo: LossAlgo,
+      run_name: str,
+      hf_token: str,
+      num_slices: int = 1,
   ) -> tuple[str]:
-    """Generates the RL training command as a tuple for GKE compatibility.
+    """Generates the RL training command as a tuple for GKE compatibility."""
+    command_params = [
+        f"run_name={run_name}",
+        f"model_name={self.model_name}",
+        f"tokenizer_path={self.tokenizer_path}",
+        f"load_parameters_path={self.load_parameters_path}",
+        f"base_output_directory={self.base_dir}",
+        f"rl.loss_algo={loss_algo.loss_name}",
+    ]
 
-    Args:
-      loss_algo: The loss algorithm to use (e.g., LossAlgo.GRPO).
-      run_name: The run name for the training job.
-      hf_token: The HuggingFace token for authentication.
+    num_trainer = 1
+    num_samplers = max(1, num_slices - num_trainer)
+
+    if num_slices > 1:
+      command_params.extend([
+          f"num_trainer_slices={num_trainer}",
+          f"num_samplers_slices={num_samplers}",
+          f"rollout_data_parallelism={num_samplers * 2}",
+      ])
 
-    Returns:
-      A tuple containing the RL training command string.
-    """
     rl_command = (
         "python -m src.MaxText.rl.train_rl "
-        f"{self.rl_config_path} run_name={run_name} "
-        f"model_name={self.model_name} "
-        f"tokenizer_path={self.tokenizer_path} "
-        f"load_parameters_path={self.load_parameters_path} "
-        f"base_output_directory={self.base_dir} "
-        f"rl.loss_algo={loss_algo.loss_name}"
+        f"{self.rl_config_path} " + " ".join(command_params)
     )
-    command = " && ".join([
+
+    environment_variables = [
         f"export HF_TOKEN={hf_token}",
         "export TPU_MIN_LOG_LEVEL=0",
         "export TF_CPP_MIN_LOG_LEVEL=0",
         "export TPU_STDERR_LOG_LEVEL=0",
         "export JAX_PLATFORMS=proxy,cpu",
         "export JAX_BACKEND_TARGET=grpc://127.0.0.1:29000",
         "export ENABLE_PATHWAYS_PERSISTENCE='1'",
-        rl_command,
-    ])
+    ]
 
-    # Return as tuple for k8s yaml compatibility.
-    return (command,)
+    if num_slices > 1:
+      environment_variables.extend([
+          f"export NUM_SLICES={num_samplers}",
+          "export JAX_RANDOM_WEIGHTS=true",
+          "export VLLM_ENABLE_V1_MULTIPROCESSING=0",
+      ])
+
+    full_command = " && ".join(environment_variables + [rl_command])
+    return (full_command,)
 
 
 @dataclass

xlml/apis/task.py

@@ -699,6 +699,7 @@ def run_model(
       self,
       gcs_location: Optional[airflow.XComArg] = None,
       use_vertex_tensorboard: bool = False,
+      experiment_name: str = "",
       use_pathways: bool = False,
       ramdisk_directory: str = "",
       mtc_enabled: bool = False,
@@ -711,6 +712,7 @@ def run_model(
       gcs_location: GCS path for all artifacts of the test.
       use_vertex_tensorboard: Set to True to view workload data on
         Vertex AI Tensorboard.
+      experiment_name: The name of the experiment for Vertex AI Tensorboard.
 
     Returns:
       A DAG node that executes the model test.
@@ -728,6 +730,7 @@ def run_model(
           workload_id,
           gcs_path,
           use_vertex_tensorboard,
+          experiment_name,
           use_pathways,
           ramdisk_directory,
           mtc_enabled,
@@ -764,6 +767,7 @@ def launch_workload(
       workload_id: str,
       gcs_path: str,
       use_vertex_tensorboard: bool,
+      experiment_name: str = "",
       use_pathways: bool = False,
       ramdisk_directory: str = "",
       mtc_enabled: bool = False,
@@ -787,6 +791,7 @@ def launch_workload(
           run_cmds=self.task_test_config.test_script,
           num_slices=self.task_test_config.num_slices,
           use_vertex_tensorboard=use_vertex_tensorboard,
+          experiment_name=experiment_name,
           use_pathways=use_pathways,
           ramdisk_directory=ramdisk_directory,
           mtc_enabled=mtc_enabled,

xlml/utils/xpk.py

@@ -33,7 +33,7 @@
 
 # NOTE: This version needs to be pinned to ensure compatibility when using
 # xpk.py for workload creation.
-MAIN_BRANCH = "v0.17.3"
+MAIN_BRANCH = "v1.1.2"
 
 # Duration = past 7 days
 LOGGING_URL_FORMAT = (
@@ -105,6 +105,7 @@ def run_workload(
     run_cmds: str,
     num_slices: int = 1,
     use_vertex_tensorboard: bool = False,
+    experiment_name: str = "",
     use_pathways: bool = False,
     # Directory for enabling emergency checkpointing
     ramdisk_directory: str = "",
@@ -179,6 +180,8 @@ def run_workload(
       workload_create_cmd += " --scheduler=gke.io/topology-aware-auto"
     if use_vertex_tensorboard:
       workload_create_cmd += " --use-vertex-tensorboard"
+      if experiment_name:
+        workload_create_cmd += f" --experiment-name={experiment_name}"
       vertex_ai_dependency = (
           "pip install -U google-cloud-aiplatform cloud-accelerator-diagnostics"
       )