Wire MaxText RL profiling and metrics onto tunix's targeted RL profiler and ML Diagnostics.

src/maxtext/common/managed_mldiagnostics.py

@@ -12,60 +12,71 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Create the managed mldiagnostics run."""
+"""Managed ML Diagnostics run lifecycle management for MaxText."""
 
 import json
 from typing import Any
 
 from maxtext.common.gcloud_stub import mldiagnostics_modules
+from maxtext.utils import max_logging
 
-mldiag, _ = mldiagnostics_modules()
+mldiag, _mldiag_is_stub = mldiagnostics_modules()
 
 from maxtext.configs.pyconfig import KEYS_NO_LOGGING
 
 
-class ManagedMLDiagnostics:
-  """ML Diagnostics Run, implemented with the Singleton pattern.
+def mldiagnostics_available() -> bool:
+  """Returns whether the real google_cloud_mldiagnostics SDK is loaded (not the stub)."""
+  return not _mldiag_is_stub
 
-  Ensures that only one instance of the class can exist.
-  """
 
-  _instance = None  # Class attribute to hold the single instance
+def _config_items(config: Any) -> dict[str, Any]:
+  """Extracts configuration key-value pairs as a flat dictionary."""
+  if hasattr(config, "get_keys"):
+    return config.get_keys()
+  return config.model_dump(mode="json")
 
-  def __new__(cls, *args: Any, **kwargs: Any):
-    """Overrides the instance creation method.
 
-    If an instance already exists, it is returned instead of creating a new one.
-    """
+class ManagedMLDiagnostics:
+  """Centralized ML Diagnostics Run manager implemented as a Singleton."""
+
+  _instance = None
+
+  def __new__(cls, *args: Any, **kwargs: Any):
     if cls._instance is None:
       cls._instance = super(ManagedMLDiagnostics, cls).__new__(cls)
-
     return cls._instance
 
-  def __init__(self, config):
-    """Initializes the ManagedMLDiagnostics, ensuring this method runs only once."""
-    # We need a flag to ensure __init__ only runs once,
-    # as the object is returned multiple times by __new__.
+  def __init__(self, config: Any, sampler_config: Any | None = None):
     if hasattr(self, "_initialized"):
       return
     self._initialized = True
     if not config.managed_mldiagnostics:
       return
 
-    # Set up the managed mldiagnostics for profiling and metrics uploading.
-    def should_log_key(key, value):
+    if not mldiagnostics_available():
+      max_logging.warning(
+          "managed_mldiagnostics=True, but the google_cloud_mldiagnostics SDK"
+          " is not available (not installed, or MaxText is running in decoupled"
+          " mode); no ML Diagnostics run will be created."
+      )
+
+    def should_log_key(key: str, value: Any) -> bool:
       if key in KEYS_NO_LOGGING:
         return False
       try:
-        # Verify the value can be serialized to json. If not, we'll skip it.
         json.dumps(value, allow_nan=False)
-      except TypeError:
+      except (TypeError, ValueError):
         return False
       return True
 
-    config_dict = {key: value for key, value in config.get_keys().items() if should_log_key(key, value)}
+    config_dict = {key: value for key, value in _config_items(config).items() if should_log_key(key, value)}
+
+    if sampler_config is not None and sampler_config is not config:
+      for key, value in _config_items(sampler_config).items():
+        if should_log_key(key, value) and config_dict.get(key) != value:
+          config_dict[f"sampler.{key}"] = value
 
-    # Create a run for the managed mldiagnostics, and upload the configuration.
     region = config.managed_mldiagnostics_region if config.managed_mldiagnostics_region else None
     mldiag.machinelearning_run(
         name=f"{config.run_name}",

src/maxtext/configs/post_train/rl.yml

@@ -70,6 +70,20 @@ rl:
   # If null, the entire model is resharded at once.
   reshard_chunk_size: null
 
+  # ====== Targeted profiling (profiler=xplane) ======
+  # One trace per RL role (`actor` = update_actor, `rollout` = generate), written as
+  # `<role>_invocation_<n>`. The window counts role invocations, not training steps:
+  # update_actor runs once per micro-batch and eval actor passes share the actor counter.
+  # Traces go to tensorboard_dir, or under managed_mldiagnostics_dir/<run id> when
+  # managed_mldiagnostics=true (set managed_mldiagnostics_on_demand_profiling=false then).
+  # On Pathways only worker 0 is traced unless profiler_max_num_hosts > 1
+  # (e.g. 1000 to include the sampler worker on a disaggregated topology).
+  profiler_max_num_hosts: 1
+  profiler_start_invocation: 2
+  profiler_num_invocations: 1
+  # Seconds to wait for the background trace stop/upload before the next trace or shutdown.
+  profiler_timeout_secs: 60.0
+
 # ====== Models ======
 # for MaxText
 # Model and Tokenizer Configuration

src/maxtext/configs/types.py

@@ -3154,10 +3154,22 @@ def _clean_empty_string_for_list(cls, v: Any) -> Any:
   _validate_config = field_validator("constant_bound_config", mode="before")(_clean_empty_string_for_list)
 
 
-class Profiling(BaseModel):
-  """Configuration for performance profiling."""
+class ProfilingCore(BaseModel):
+  """Profiler backend selection shared by every trainer.
+
+  The Tunix-based RL trainer (`trainers/post_train/rl/train_rl.py`) only needs
+  this field: its profiling window is expressed in role invocations via
+  `RL.profiler_start_invocation` / `RL.profiler_num_invocations` and handled by
+  tunix `RLProfileConfig`, so the step-based knobs in `Profiling` are not
+  accepted there.
+  """
 
   profiler: ProfilerType = Field(ProfilerType.NONE, description="Profiler to use ('xplane', 'nsys').")
+
+
+class Profiling(ProfilingCore):
+  """Configuration for step-window performance profiling (`common/profiler.py`)."""
+
   upload_all_profiler_results: bool = Field(False, description="Upload profiler results from all hosts.")
   skip_first_n_steps_for_profiler: int = Field(1, description="Number of initial steps to skip for profiling.")
   profiler_steps: int = Field(5, description="Number of steps to profile.")
@@ -3588,6 +3600,50 @@ class VLLM(BaseModel):
 class RL(BaseModel):
   """Configuration for RL algorithms like Group Relative Policy Optimization (GRPO) among others."""
 
+  # Targeted profiling of the Tunix RL roles (tunix `RLProfileConfig`). Active
+  # only with `profiler=xplane`.
+  profiler_start_invocation: int = Field(
+      2,
+      ge=0,
+      description=(
+          "Zero-based invocation index at which tracing starts for each RL"
+          " role: `update_actor` for the actor and `generate` for the rollout"
+          " engine (tunix `RLProfileConfig.start_invocation`). This counts role"
+          " function calls, not global training steps: `update_actor` runs once"
+          " per micro-batch (mini-batch x gradient accumulation x"
+          " num_iterations per step) and eval actor passes share the actor"
+          " counter. Default 2 skips JIT compilation and warm-up."
+      ),
+  )
+  profiler_num_invocations: int = Field(
+      1,
+      ge=1,
+      description=(
+          "Number of consecutive invocations traced per role once"
+          " `profiler_start_invocation` is reached (tunix"
+          " `RLProfileConfig.num_invocations`). One session is written per"
+          " invocation as `<role>_invocation_<n>`."
+      ),
+  )
+  profiler_timeout_secs: float = Field(
+      60.0,
+      gt=0,
+      description=(
+          "Seconds to wait for tunix's background trace stop and upload thread"
+          " before starting the next trace or closing the cluster (tunix"
+          " `RLProfileConfig.timeout_secs`)."
+      ),
+  )
+  profiler_max_num_hosts: PositiveInt = Field(
+      1,
+      description=(
+          "Pathways only: maximum number of worker hosts to trace with the"
+          " xplane profiler across all RL roles (trainer and sampler). Defaults"
+          " to 1 (worker 0 only). Values > 1 trace worker hosts across all"
+          " slices on the Pathways backend."
+      ),
+  )
+
   num_generations: int = Field(2, description="Number of responses to generate per prompt (G in GRPO paper).")
   num_iterations: int = Field(1, ge=1, description="Number of iterations per batch (μ in GRPO paper).")
   grpo_beta: float = Field(0.08, description="Coefficient for the KL divergence penalty (β).")
@@ -3901,6 +3957,32 @@ class DerivedValues(BaseModel):
 # ----------------------------------------------------------------------------
 
 
+def derive_managed_mldiagnostics_dir(
+    base_output_directory: str | None,
+    run_name: str | None,
+    managed_mldiagnostics_storage_path: str | None,
+) -> str | None:
+  """Returns the managed ML Diagnostics root for a run, or None if it cannot be derived.
+
+  Layout: `<managed_mldiagnostics_storage_path or base_output_directory>/<run_name>/managed-mldiagnostics`,
+  without a trailing slash (the SDK mis-joins paths that end in one, b/454725283). The SDK appends its
+  own run identifier below this directory for profiles.
+
+  Args:
+    base_output_directory: The run's base output directory; used when no dedicated storage path is set.
+    run_name: The run name; the directory cannot be derived without it.
+    managed_mldiagnostics_storage_path: Optional dedicated storage root that takes priority over
+      `base_output_directory`.
+
+  Returns:
+    The derived directory, or None when `run_name` or both base paths are empty.
+  """
+  telemetry_base = managed_mldiagnostics_storage_path or base_output_directory
+  if telemetry_base and run_name:
+    return os.path.join(telemetry_base, run_name, "managed-mldiagnostics")
+  return None
+
+
 # Decoder blocks that support router replay (`forced_routed_experts`). All are
 # homogeneous-MoE, so a 4D layer axis is just the decoder layer index;
 # interleaved architectures (Llama4/Envy) would need a MoE-only counter.
@@ -4713,9 +4795,11 @@ def set_derived_and_validate_values(self) -> "MaxTextConfig":
       self.checkpoint_dir = os.path.join(output_dir, "checkpoints", "")
       self.metrics_dir = os.path.join(output_dir, "metrics", "")
       self.tensorboard_dir = os.path.join(output_dir, "tensorboard", "")
-      # To work around SDK bug b/454725283, remove the trailing back slash from the managed_mldiagnostics_dir.
-      telemetry_base = getattr(self, "managed_mldiagnostics_storage_path", "") or self.base_output_directory
-      self.managed_mldiagnostics_dir = os.path.join(telemetry_base, self.run_name, "managed-mldiagnostics")
+      self.managed_mldiagnostics_dir = derive_managed_mldiagnostics_dir(
+          self.base_output_directory,
+          self.run_name,
+          self.managed_mldiagnostics_storage_path,
+      )
       if self.enable_mllog:
         if not self.mllog_file:
           self.mllog_file = os.path.join(output_dir, f"mllog_{self.data_shuffle_seed}.log")
@@ -6427,7 +6511,7 @@ class RLConfig(
     VisionProjector,
     AudioEncoder,
     DevelopmentAndDebugging,
-    Profiling,
+    ProfilingCore,
     AOT,
     Metrics,
     Tensorboard,
@@ -6604,6 +6688,42 @@ def set_derived_values_and_validate(self) -> "RLConfig":
       tensorboard_dir = os.path.join(os.path.abspath("maxtext_output"), "tensorboard", "")
     object.__setattr__(self, "tensorboard_dir", tensorboard_dir)
 
+    # Managed ML Diagnostics: derive the run directory and fail at config time
+    # on combinations that would only surface mid-run.
+    object.__setattr__(
+        self,
+        "managed_mldiagnostics_dir",
+        derive_managed_mldiagnostics_dir(
+            self.base_output_directory,
+            self.run_name,
+            self.managed_mldiagnostics_storage_path,
+        ),
+    )
+    if self.managed_mldiagnostics:
+      if not self.managed_mldiagnostics_dir or not self.managed_mldiagnostics_dir.startswith("gs://"):
+        raise ValueError(
+            "managed_mldiagnostics=True requires a gs:// run directory: set"
+            " run_name and either managed_mldiagnostics_storage_path or"
+            " base_output_directory to a gs:// path (derived"
+            f" managed_mldiagnostics_dir={self.managed_mldiagnostics_dir!r})."
+        )
+      if self.profiler == ProfilerType.XPLANE and self.managed_mldiagnostics_on_demand_profiling:
+        raise ValueError(
+            "managed_mldiagnostics_on_demand_profiling=True cannot be combined"
+            " with profiler=xplane on the RL path: an on-demand capture during"
+            " a programmatic trace session fails with 'profile already"
+            " started'. Set managed_mldiagnostics_on_demand_profiling=False (or"
+            " profiler='')."
+        )
+
+    if self.rl.profiler_max_num_hosts > 1:
+      if self.profiler != ProfilerType.XPLANE:
+        raise ValueError(f"rl.profiler_max_num_hosts > 1 requires profiler=xplane, got profiler={self.profiler!r}.")
+      import pathwaysutils  # pylint: disable=import-outside-toplevel
+
+      if not pathwaysutils.is_pathways_backend_used():
+        raise ValueError("rl.profiler_max_num_hosts > 1 is only supported on the Pathways backend (JAX_PLATFORMS=proxy).")
+
     # Slice configuration
     if not (
         (self.cluster.num_trainer_slices == -1 and self.cluster.num_samplers_slices == -1)

src/maxtext/input_pipeline/tfds_data_processing_c4_mlperf.py

@@ -756,12 +756,7 @@ def make_c4_mlperf_eval_iterator(
   # sequences), several hosts must share a file and read disjoint sequences of it instead
   # of each taking the first `local_num_examples` sequences of file `host_index`.
   skip_examples = 0
-  if (
-      eval_shard_in_read
-      and effective_stream_chunking
-      and num_examples > 0
-      and num_examples % len(process_indices) == 0
-  ):
+  if eval_shard_in_read and effective_stream_chunking and num_examples > 0 and num_examples % len(process_indices) == 0:
     sequences_per_file = _sequences_per_eval_file(
         config.eval_dataset_name, eval_split, eval_data_dir, eval_col, config.max_target_length
     )