[DRAFT] feat(mlperf): add MLPerf RCP logging support for distributed GSM8K GRPO

Dockerfile

@@ -108,8 +108,11 @@ COPY . .
 RUN uv pip install grpcio-tools
 RUN cd /app && find tunix/experimental/distributed -name "*.proto" -exec python -m grpc_tools.protoc -I/app --python_out=/app --grpc_python_out=/app {} +
 
-# Install Tunix in editable mode
-RUN uv pip install --no-deps -e .
+# Install Tunix in editable mode and register MaxText vLLM adapter + Raiden
+RUN uv pip install --no-deps -e . && \
+    RAIDEN_WHEEL_DIR="/app/raiden_wheels" bash /app/scripts/install_raiden.sh && \
+    uv pip install --no-deps /opt/venv/lib/python3.12/site-packages/maxtext/integration/vllm && \
+    uv pip install git+https://github.com/mlcommons/logging.git
 
 
 

tunix/experimental/examples/math_gsm8k_dist/k8s_launcher.sh

@@ -107,6 +107,9 @@ export WANDB_RUN_NAME=${WANDB_RUN_NAME:-}
 export WANDB_API_KEY=${WANDB_API_KEY:-}
 export LOG_DIR=${LOG_DIR:-}
 export TRAJECTORY_LOG_DIR=${TRAJECTORY_LOG_DIR:-}
+export RCP_LOGGING=${RCP_LOGGING:-false}
+export METRIC_LOGGER_DIR=${METRIC_LOGGER_DIR:-}
+export TARGET_ACCURACY=${TARGET_ACCURACY:-0.69}
 export TFDS_DATA_DIR=${TFDS_DATA_DIR:-"artifacts/data"}
 export TFDS_SPLIT=${TFDS_SPLIT:-train}
 export FLUSH_METRICS_EVERY_N_STEPS=${FLUSH_METRICS_EVERY_N_STEPS:-1}
@@ -180,6 +183,10 @@ start_orchestrator() {
   if [[ "${DEBUG}" == "1" || "${DEBUG}" == "true" || "${DEBUG}" == "True" ]]; then
     debug_flag="--debug"
   fi
+  local rcp_flag=""
+  if [[ "${RCP_LOGGING}" == "1" || "${RCP_LOGGING}" == "true" || "${RCP_LOGGING}" == "True" ]]; then
+    rcp_flag="--rcp_logging"
+  fi
 
   "$PYTHON" "$YAML_GEN" \
     "$YAML_DIR/jobset.cpu.yaml" \
@@ -221,6 +228,20 @@ start_orchestrator() {
         ${MAX_SEQ_TOKEN_PER_TPU:+--max_seq_token_per_tpu=${MAX_SEQ_TOKEN_PER_TPU}} \
         ${MAX_SEGMENTS_PER_PACKED_ROW:+--max_segments_per_packed_row=${MAX_SEGMENTS_PER_PACKED_ROW}} \
         ${TRAINER_MESH_FSDP:+--trainer_fsdp=${TRAINER_MESH_FSDP}} \
+        --eval_every_n_steps=${EVAL_EVERY_N_STEPS} \
+        --learning_rate=${LEARNING_RATE} \
+        --b1=${ADAM_B1} \
+        --b2=${ADAM_B2} \
+        --weight_decay=${WEIGHT_DECAY} \
+        --max_grad_norm=${MAX_GRAD_NORM} \
+        --train_mesh_tp=${TRAINER_MESH_TP} \
+        --train_mesh_expert=${TRAINER_MESH_EXPERT} \
+        --rollout_mesh_tp=${ROLLOUT_MESH_TP} \
+        --rollout_engine=${SAMPLER} \
+        --tpu_topology="${TRAINER_TPU_SLICE}+${ROLLOUT_TPU_SLICE}" \
+        --target_accuracy=${TARGET_ACCURACY} \
+        ${METRIC_LOGGER_DIR:+--metric_logger_dir="${METRIC_LOGGER_DIR}"} \
+        ${rcp_flag} \
         ${debug_flag} \
     " \
     | apply_manifest

tunix/experimental/examples/math_gsm8k_dist/run_gsm8k_dist_grpo.py

@@ -57,6 +57,7 @@
 from tunix.experimental.worker import remote_execution  # pylint: disable=g-import-not-at-top
 from tunix.rl import algorithm_config  # pylint: disable=g-import-not-at-top
 from tunix.sft import metrics_logger as metrics_logger_lib  # pylint: disable=g-import-not-at-top
+from tunix.utils import mllog_utils  # pylint: disable=g-import-not-at-top
 
 ProcessContext = runtime_context.ProcessContext
 
@@ -288,6 +289,79 @@ def _parse_args(argv: list[str]) -> argparse.Namespace:
       action="store_true",
       help="Enable debug logging and print full sampler responses.",
   )
+  parser.add_argument(
+      "--rcp_logging",
+      action="store_true",
+      default=False,
+      help="Enable MLPerf RCP (mllog) compliance logging.",
+  )
+  parser.add_argument(
+      "--metric_logger_dir",
+      type=str,
+      default=os.getenv("METRIC_LOGGER_DIR", None),
+      help="Directory or GCS URI for MLPerf RCP output (seed_<seed>.out).",
+  )
+  parser.add_argument(
+      "--target_accuracy",
+      type=float,
+      default=float(os.getenv("TARGET_ACCURACY", "0.69")),
+      help="Target evaluation accuracy for MLPerf RCP compliance logging.",
+  )
+  parser.add_argument(
+      "--eval_every_n_steps",
+      type=int,
+      default=int(os.getenv("EVAL_EVERY_N_STEPS", "1000000")),
+  )
+  parser.add_argument(
+      "--learning_rate",
+      type=float,
+      default=float(os.getenv("LEARNING_RATE", "2.0e-7")),
+  )
+  parser.add_argument(
+      "--b1",
+      type=float,
+      default=float(os.getenv("ADAM_B1", "0.9")),
+  )
+  parser.add_argument(
+      "--b2",
+      type=float,
+      default=float(os.getenv("ADAM_B2", "0.999")),
+  )
+  parser.add_argument(
+      "--weight_decay",
+      type=float,
+      default=float(os.getenv("WEIGHT_DECAY", "0.01")),
+  )
+  parser.add_argument(
+      "--max_grad_norm",
+      type=float,
+      default=float(os.getenv("MAX_GRAD_NORM", "1.0")),
+  )
+  parser.add_argument(
+      "--train_mesh_tp",
+      type=int,
+      default=int(os.getenv("TRAINER_MESH_TP", "1")),
+  )
+  parser.add_argument(
+      "--train_mesh_expert",
+      type=int,
+      default=int(os.getenv("TRAINER_MESH_EXPERT", "1")),
+  )
+  parser.add_argument(
+      "--rollout_mesh_tp",
+      type=int,
+      default=int(os.getenv("ROLLOUT_MESH_TP", "1")),
+  )
+  parser.add_argument(
+      "--rollout_engine",
+      type=str,
+      default=os.getenv("SAMPLER", "vllm"),
+  )
+  parser.add_argument(
+      "--tpu_topology",
+      type=str,
+      default=os.getenv("TPU_TOPOLOGY", None),
+  )
   return parser.parse_args(argv)
 
 
@@ -389,6 +463,8 @@ def main(argv: list[str], context: ProcessContext | None = None) -> None:
   )
 
   args = _parse_args(argv)
+  if args.rcp_logging:
+    mllog_utils.init_start(args)
   if args.debug:
     # Enable canonical debug logging and print full sampler responses
     logging.getLogger().setLevel(logging.DEBUG)
@@ -520,16 +596,39 @@ def main(argv: list[str], context: ProcessContext | None = None) -> None:
           step,
           step,
       ),
-      on_step_end=lambda step, result: logging.info(
-          "<<< Step %d finished | Advanced to Policy Version: %d",
-          step,
-          step + 1,
+      on_step_end=lambda step, result: (
+          logging.info(
+              "<<< Step %d finished | Advanced to Policy Version: %d",
+              step,
+              step + 1,
+          ),
+          mllog_utils.log_rcp_step_stats(
+              program.metrics_logger,
+              args=args,
+              step=step + 1,
+          )
+          if args.rcp_logging
+          else None,
       ),
   )
 
+  if args.rcp_logging:
+    train_ds = gsm8k.load_gsm8k_dataset(
+        split=args.tfds_split,
+        data_dir=args.tfds_data_dir,
+        shuffle=args.shuffle,
+        seed=args.seed,
+    )
+    mllog_utils.init_print(
+        args,
+        train_dataset=train_ds,
+    )
+
   try:
     logging.info("Bringing up remote workers through ClusterOrchestrator...")
     cluster.bring_up_workers(dummy_data=None)
+    if args.rcp_logging:
+      mllog_utils.train_start(args, step=0)
     logging.info(
         "Cluster workers ready: %s. Starting StandardRLProgram execution...",
         [w.worker_id for w in cluster.worker_infos()],
@@ -539,7 +638,21 @@ def main(argv: list[str], context: ProcessContext | None = None) -> None:
         num_steps=args.max_steps,
         bring_up=False,
     )
+    if args.rcp_logging:
+      completed_steps = (
+          program.last_step_result.step + 1
+          if program.last_step_result is not None
+          else args.max_steps
+      )
+      mllog_utils.train_stop(args, step=completed_steps, status="success")
   except BaseException as exc:
+    if args.rcp_logging:
+      completed_steps = (
+          program.last_step_result.step + 1
+          if program.last_step_result is not None
+          else 0
+      )
+      mllog_utils.train_stop(args, step=completed_steps, status="aborted")
     logging.exception("FATAL: StandardRLProgram execution failed: %s", exc)
     raise
   finally:

tunix/experimental/orchestrator/rl_program.py

@@ -625,6 +625,9 @@ def _collect_and_log_step_metrics(
       consumed_policy_version: int,
       log_step: int,
       sampler_agreement: dict[str, tuple[Any, list[float]]] | None = None,
+      policy_training_time: float = 0.0,
+      exposed_generation_time: float = 0.0,
+      weight_sync_time: float = 0.0,
   ) -> dict[str, Any]:
     """Logs rollout, reward, trainer, and orchestrator metrics.
 
@@ -744,6 +747,29 @@ def _collect_and_log_step_metrics(
           self.mode,
           log_step,
       )
+      self.metrics_logger.log(
+          self.metrics_prefix,
+          "rollout/global_valid_toks",
+          float(np.sum(total_lengths)),
+          self.mode,
+          log_step,
+      )
+    elif completion_lengths:
+      self.metrics_logger.log(
+          self.metrics_prefix,
+          "rollout/global_valid_toks",
+          float(np.sum(completion_lengths)),
+          self.mode,
+          log_step,
+      )
+    if all_step_items:
+      self.metrics_logger.log(
+          self.metrics_prefix,
+          "rollout/global_valid_seqs",
+          float(len(all_step_items)),
+          self.mode,
+          log_step,
+      )
     if turns_list:
       self.metrics_logger.log(
           self.metrics_prefix,
@@ -839,6 +865,9 @@ def _collect_and_log_step_metrics(
         "num_rollouts": float(num_rollouts),
         "num_microbatches": float(num_microbatches),
         "step_time_sec": float(step_time_sec),
+        "policy_training_time": float(policy_training_time),
+        "exposed_generation_time": float(exposed_generation_time),
+        "weight_sync_time": float(weight_sync_time),
     }
     for tag, val in orchestrator_stats.items():
       self.metrics_logger.log(
@@ -1101,6 +1130,9 @@ async def train_stage(self) -> None:
       groups_consumed = 0
       checkpoint_saved = False
       final_minibatch_completed = False
+      policy_training_time = 0.0
+      exposed_generation_time = 0.0
+      weight_sync_time = 0.0
 
       async def _maybe_save_checkpoint() -> None:
         nonlocal checkpoint_saved
@@ -1123,7 +1155,9 @@ async def _maybe_save_checkpoint() -> None:
         checkpoint_saved = True
 
       while groups_consumed < self.full_batch_size:
+        _t_gen = time.monotonic()
         scored_items = await self.scored_q.get_batch(num_groups=1)
+        exposed_generation_time += time.monotonic() - _t_gen
         if not scored_items:
           assembled_batches = self.assembler.flush()
         else:
@@ -1184,16 +1218,20 @@ async def _maybe_save_checkpoint() -> None:
               len(mb.trajectory_ids),
               logging_utils.summarize_list(list(mb.trajectory_ids)),
           )
+          _t_train = time.monotonic()
           step_result = await self.engine.train_step(
               batch,
               role=datatypes.Role.ACTOR,
               accumulate_gradients=True,
               apply_optimizer=mb.is_final_batch,
           )
+          policy_training_time += time.monotonic() - _t_train
           if mb.is_final_batch:
+            _t_metrics = time.monotonic()
             trainer_metrics = await self.engine.get_metrics(
                 role=datatypes.Role.ACTOR
             )
+            policy_training_time += time.monotonic() - _t_metrics
             final_minibatch_completed = True
             # TODO(tunix-dev): Configurable checkpointing frequency. Today we
             # checkpoint at the same frequency as the weight update.
@@ -1221,7 +1259,9 @@ async def _maybe_save_checkpoint() -> None:
         break
 
       if self.sync_weights:
+        _t_sync = time.monotonic()
         new_version = await self.engine.sync_weights(role=datatypes.Role.ACTOR)
+        weight_sync_time = time.monotonic() - _t_sync
         self.policy_version = (
             new_version if new_version is not None else self.policy_version + 1
         )
@@ -1251,6 +1291,9 @@ async def _maybe_save_checkpoint() -> None:
           consumed_policy_version=consumed_policy_version,
           log_step=current_step,
           sampler_agreement=step_sampler_agreement,
+          policy_training_time=policy_training_time,
+          exposed_generation_time=exposed_generation_time,
+          weight_sync_time=weight_sync_time,
       )
       self._log_consumed_trajectories(
           all_step_items,