Refactor maxtext_rl DAG and organize tasks in a group (GoogleCloudPlatform#1125)

Author: SurbhiJainUSC

dags/post_training/maxtext_rl.py

@@ -11,6 +11,7 @@
 
 from airflow import models
 from airflow.models.baseoperator import chain
+from airflow.utils.task_group import TaskGroup
 
 from dags import composer_env
 from dags.common import test_owner
@@ -69,7 +70,6 @@
       accelerator="v5p-128",
       slices=[1],  # Single slice for RL training
       model_name="llama3.1-70b",
-      short_id="mrl",
       base_dir=(
           f"{test_config_util.DEFAULT_BUCKET}/llama3.1-70b-Instruct/outputs"
       ),
@@ -94,79 +94,78 @@
 
     for loss_algo in training_config.loss_algos:
       for slice_num in training_config.slices:
-        task_suffix = f"{mode.value}_{loss_algo.value}_{slice_num}"
-        run_name = validation_util.generate_posttraining_run_name.override(
-            task_id=f"run_name_{task_suffix}"
-        )(
-            short_id=training_config.short_id,
-            checkpointing_type=loss_algo.value,
-            slice_number=slice_num,
-            mode=mode.value,
-        )
+        current_datetime = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
+        run_name = f"{loss_algo.value}-{mode.value}-{current_datetime}"
 
         rl_training_command = training_config.generate_rl_training_command(
             loss_algo=loss_algo,
             run_name=run_name,
             hf_token=HF_TOKEN_LLAMA3_1,
         )
 
-        start_time = validation_util.generate_timestamp.override(
-            task_id=f"start_time_{task_suffix}"
-        )()
-
-        test_name = f"{training_config.short_id[:3]}{loss_algo.value[:3]}"
-
-        training_task = gke_config.get_gke_config(
-            num_slices=slice_num,
-            cluster=training_config.cluster,
-            time_out_in_min=30,
-            test_name=test_name,
-            run_model_cmds=rl_training_command,
-            docker_image=image.value,
-            test_owner=test_owner.JACKY_F,
-        ).run(
-            use_pathways=True,
-            xpk_branch=MAIN_BRANCH,
-            skip_post_process=True,
-        )
-
-        end_time = validation_util.generate_timestamp.override(
-            task_id=f"end_time_{task_suffix}"
-        )()
-
-        validate_loss_algo = validation_util.validate_log_exist.override(
-            task_id=f"validate_loss_{task_suffix}"
-        )(
-            project_id=training_config.cluster.project,
-            location=zone_to_region(training_config.cluster.zone),
-            cluster_name=training_config.cluster.name,
-            text_filter=f'"Config param loss_algo: {loss_algo.value}"',
-            namespace="default",
-            container_name="jax-tpu",
-            pod_pattern=f"{test_name}.*",
-            start_time=start_time,
-            end_time=end_time,
-        )
-
-        validate_training = validation_util.validate_log_exist.override(
-            task_id=f"validate_training_{task_suffix}"
-        )(
-            project_id=training_config.cluster.project,
-            location=zone_to_region(training_config.cluster.zone),
-            cluster_name=training_config.cluster.name,
-            text_filter='"Post RL Training"',
-            namespace="default",
-            container_name="jax-tpu",
-            pod_pattern=f"{test_name}.*",
-            start_time=start_time,
-            end_time=end_time,
-        )
-
-        chain(
-            run_name,
-            start_time,
-            training_task,
-            end_time,
-            validate_loss_algo,
-            validate_training,
-        )
+        with TaskGroup(
+            group_id=f"{loss_algo.value}-{mode.value}-{slice_num}x{training_config.accelerator}"
+        ) as group:
+          with TaskGroup(group_id="run_training") as training_group:
+            start_time = validation_util.generate_timestamp.override(
+                task_id="generate_start_time"
+            )()
+
+            training_task = gke_config.get_gke_config(
+                num_slices=slice_num,
+                cluster=training_config.cluster,
+                time_out_in_min=30,
+                test_name=loss_algo.value,
+                run_model_cmds=rl_training_command,
+                docker_image=image.value,
+                test_owner=test_owner.JACKY_F,
+            ).run_model(
+                use_pathways=True,
+                xpk_branch=MAIN_BRANCH,
+            )
+
+            end_time = validation_util.generate_timestamp.override(
+                task_id="generate_end_time"
+            )()
+
+            chain(
+                start_time,
+                training_task,
+                end_time,
+            )
+
+          with TaskGroup(group_id="validate_training") as validation_group:
+            validate_loss_algo = validation_util.validate_log_exist.override(
+                task_id="validate_loss_algo"
+            )(
+                project_id=training_config.cluster.project,
+                location=zone_to_region(training_config.cluster.zone),
+                cluster_name=training_config.cluster.name,
+                text_filter=f'"Config param loss_algo: {loss_algo.loss_name}"',
+                namespace="default",
+                container_name="jax-tpu",
+                pod_pattern=f"{loss_algo.value}.*",
+                start_time=start_time,
+                end_time=end_time,
+            )
+
+            validate_training_logs = (
+                validation_util.validate_log_exist.override(
+                    task_id="validate_training_logs"
+                )(
+                    project_id=training_config.cluster.project,
+                    location=zone_to_region(training_config.cluster.zone),
+                    cluster_name=training_config.cluster.name,
+                    text_filter="Post RL Training",
+                    namespace="default",
+                    container_name="jax-tpu",
+                    pod_pattern=f"{loss_algo.value}.*",
+                    start_time=start_time,
+                    end_time=end_time,
+                )
+            )
+
+          chain(
+              training_group,
+              validation_group,
+          )

dags/post_training/util/test_config_util.py

@@ -23,6 +23,14 @@ class LossAlgo(Enum):
   GRPO = "grpo"
   GSPO = "gspo"
 
+  @property
+  def loss_name(self) -> str:
+    """Returns the specific loss algorithm string used for computation."""
+    return {
+        LossAlgo.GRPO: "grpo",
+        LossAlgo.GSPO: "gspo-token",
+    }[self]
+
 
 @dataclass
 class RLTestConfig:
@@ -38,7 +46,6 @@ class RLTestConfig:
     slices: List of slice numbers to test with.
     model_name: The name of the model being trained
         (e.g., llama3.1-70b).
-    short_id: A short identifier for the test run.
     base_dir: Base GCS directory for outputs.
     tokenizer_path: Path to the tokenizer (HuggingFace model
         path or local path).
@@ -53,7 +60,6 @@ class RLTestConfig:
   accelerator: str
   slices: list[int]
   model_name: str
-  short_id: str
   base_dir: str
   tokenizer_path: str
   load_parameters_path: str
@@ -66,7 +72,6 @@ def __init__(
       accelerator: str,
       slices: list[int],
       model_name: str,
-      short_id: str,
       base_dir: str,
       tokenizer_path: str,
       load_parameters_path: str,
@@ -82,7 +87,6 @@ def __init__(
       slices: The number of slices to be used.
       model_name: The name of the base model being tested
           (e.g., llama3.1-70b).
-      short_id: A short identifier for the test run.
       base_dir: The base GCS directory for storing outputs.
       tokenizer_path: Path to the tokenizer (HuggingFace
           model path).
@@ -96,7 +100,6 @@ def __init__(
     self.accelerator = accelerator
     self.slices = slices
     self.model_name = model_name
-    self.short_id = short_id
     self.base_dir = base_dir
     self.tokenizer_path = tokenizer_path
     self.load_parameters_path = load_parameters_path
@@ -129,11 +132,9 @@ def generate_rl_training_command(
         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"base_output_directory={self.base_dir} "
+        f"loss_algo={loss_algo.loss_name}"
     )
 
-    if loss_algo == LossAlgo.GSPO:
-      command += f" loss_algo={loss_algo.value}-token"
-
     # Return as tuple for k8s yaml compatibility.
     return (command,)