From 4b48f3a845f0d959bf29479259b45a79bc4c8c68 Mon Sep 17 00:00:00 2001 From: maxtext authors Date: Wed, 9 Sep 2026 16:09:17 -0700 Subject: [PATCH] Wire MaxText RL profiling and metrics onto tunix's targeted RL profiler and ML Diagnostics. The Tunix-based RL trainer used tunix's SFT-era step profiler, which only traces the actor on process 0 and never the vLLM rollout, and it had no ML Diagnostics integration. Tunix now provides a per-role, invocation-counted `RLProfileConfig` and an `MLDiagScalarBackend`; this change maps the MaxText RL config onto them. * `profiler=xplane` builds a tunix `RLProfileConfig` from the new `rl.profiler_start_invocation`, `rl.profiler_num_invocations` and `rl.profiler_timeout_secs` keys, producing one `_invocation_` trace per role. Non-managed traces land in `tensorboard_dir`; with `managed_mldiagnostics=true` the tunix ML Diagnostics backend writes under the MLRun's `gcs_path` (`managed_mldiagnostics_dir`), so `RLProfileConfig.mldiagnostics_dir` is not forwarded (tunix's JAX backend would otherwise prefer it over `tensorboard_dir`). * `managed_mldiagnostics=true` creates the ML Diagnostics run before the `RLCluster` (tunix's `RLProfiler` preflights the active run), uploads the trainer config (plus `sampler.` for any sampler-side differences on multi-slice topologies), and installs `MLDiagScalarBackend` next to `TensorboardBackend` as the tunix metrics backends. `RLConfig` now derives `managed_mldiagnostics_dir` like `MaxTextConfig` does and rejects a non-`gs://` directory or on-demand profiling combined with `profiler=xplane` at config time. * `common/managed_mldiagnostics.py` changes accordingly. `ManagedMLDiagnostics` stays a process-wide singleton (the RL trainer constructs it once at startup; pretrain's `common/profiler.py` is unchanged). It gains an optional `sampler_config`: when it is the same object as the trainer config (single-slice, colocated) the merge is skipped and the run config is a flat dict with no `sampler.` keys; when it is a distinct object (multi-slice, disaggregated) only keys whose values differ are added as `sampler.`. The JSON-serializability filter now also catches `ValueError` so NaN/Inf values are dropped instead of failing the upload. Because `gcloud_stub.mldiagnostics_modules()` returns a stub object rather than `None` when the SDK is absent, availability is detected via the returned stub flag (`mldiagnostics_available()`); `ManagedMLDiagnostics` logs a warning when the SDK is in stub mode, while pretrain keeps its previous warn-and-continue behavior. * With `managed_mldiagnostics=false` the metrics backends are left to tunix's defaults, so non-managed runs log exactly as before. With `managed_mldiagnostics=true` the backends are listed explicitly (`MLDiagScalarBackend` plus `TensorboardBackend` when `enable_tensorboard` is set); a non-empty explicit list replaces tunix's defaults, so managed runs do not get tunix's W&B backend in OSS (or the CLU backend internally). * The step-window profiling flags (`profiler_steps`, `skip_first_n_steps_for_profiler`, `profile_periodically_period`, ...) are removed from `RLConfig` via a `ProfilingCore` mixin that only keeps `profiler`; they have no consumer on the Tunix RL path and are rejected with the list of valid fields. Pretrain/SFT keep the full `Profiling` mixin. The in-tree RL scripts are updated to the new keys. * On Pathways, `pathwaysutils` traces worker 0 only (`max_num_hosts=1`), so a disaggregated sampler never appears in the rollout trace. A new `rl.profiler_max_num_hosts` key (default 1) parameterizes a wrapper around `jax.profiler.start_trace` to trace worker hosts across all slices. * Format `tfds_data_processing_c4_mlperf.py` to satisfy repo-wide `pyink` pre-commit checks. PiperOrigin-RevId: 978806139 --- src/maxtext/common/managed_mldiagnostics.py | 55 +++-- src/maxtext/configs/post_train/rl.yml | 14 ++ src/maxtext/configs/types.py | 132 ++++++++++- .../tfds_data_processing_c4_mlperf.py | 7 +- .../trainers/post_train/rl/train_rl.py | 114 +++++++-- tests/post_training/unit/train_rl_test.py | 3 +- tests/unit/managed_mldiagnostics_test.py | 167 ++++++++++++- tests/unit/rl_profiling_config_test.py | 224 ++++++++++++++++++ 8 files changed, 657 insertions(+), 59 deletions(-) create mode 100644 tests/unit/rl_profiling_config_test.py diff --git a/src/maxtext/common/managed_mldiagnostics.py b/src/maxtext/common/managed_mldiagnostics.py index 5f9b29691d..8a9998733b 100644 --- a/src/maxtext/common/managed_mldiagnostics.py +++ b/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}", diff --git a/src/maxtext/configs/post_train/rl.yml b/src/maxtext/configs/post_train/rl.yml index 4485736612..5b3f2b197c 100644 --- a/src/maxtext/configs/post_train/rl.yml +++ b/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 + # `_invocation_`. 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/ 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 diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index c64f49f9f8..c11e8b10fb 100644 --- a/src/maxtext/configs/types.py +++ b/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 `_invocation_`." + ), + ) + 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`, + 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) diff --git a/src/maxtext/input_pipeline/tfds_data_processing_c4_mlperf.py b/src/maxtext/input_pipeline/tfds_data_processing_c4_mlperf.py index e90f524d6b..8d82b4865e 100644 --- a/src/maxtext/input_pipeline/tfds_data_processing_c4_mlperf.py +++ b/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 ) diff --git a/src/maxtext/trainers/post_train/rl/train_rl.py b/src/maxtext/trainers/post_train/rl/train_rl.py index 5306d587b7..e8e3068658 100644 --- a/src/maxtext/trainers/post_train/rl/train_rl.py +++ b/src/maxtext/trainers/post_train/rl/train_rl.py @@ -70,12 +70,13 @@ from tunix.rl import rl_cluster as rl_cluster_lib from tunix.rl.rollout import base_rollout from tunix.rl.grpo.grpo_learner import GrpoConfig, GrpoLearner -from tunix.sft import metrics_logger, profiler +from tunix.sft import metrics_logger os.environ["TOKENIZERS_PARALLELISM"] = "0" from maxtext.common.common_types import DecoderBlockType +from maxtext.common.managed_mldiagnostics import ManagedMLDiagnostics from maxtext.configs import pyconfig, types from maxtext.common.goodput import ( GoodputEvent, @@ -94,7 +95,6 @@ from maxtext.utils import max_logging, max_utils, model_creation_utils from maxtext.utils.model_creation_utils import get_rollout_kwargs_for_parallelism - _RECURRENT_ROLLOUT_DECODER_BLOCKS = frozenset((DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5)) @@ -358,6 +358,45 @@ def _kwargs_supported_by(cls: Any, **kwargs: Any) -> dict[str, Any]: return {k: v for k, v in kwargs.items() if k in params} +def build_metrics_logging_options( + trainer_config: Any, +) -> metrics_logger.MetricsLoggerOptions: + """Returns Tunix MetricsLoggerOptions configured for ML Diagnostics or defaults.""" + backends = [] + if trainer_config.managed_mldiagnostics: + from tunix.rl.metrics import mldiag_backend # pylint: disable=import-outside-toplevel + + backends.append(mldiag_backend.MLDiagScalarBackend) + if trainer_config.enable_tensorboard: + backends.append( + functools.partial( + metrics_logger.TensorboardBackend, + log_dir=trainer_config.tensorboard_dir, + flush_every_n_steps=trainer_config.log_period, + ) + ) + return metrics_logger.MetricsLoggerOptions( + log_dir=trainer_config.tensorboard_dir, + flush_every_n_steps=trainer_config.log_period, + backend_kwargs={"custom_backend": backends}, + ) + + +def build_rl_profiler_config( + trainer_config: Any, +) -> rl_cluster_lib.RLProfileConfig | None: + """Returns an RLProfileConfig if xplane profiling is enabled, else None.""" + if trainer_config.profiler != "xplane": + return None + return rl_cluster_lib.RLProfileConfig( + start_invocation=trainer_config.rl.profiler_start_invocation, + num_invocations=trainer_config.rl.profiler_num_invocations, + output_dir=trainer_config.tensorboard_dir, + managed_mldiagnostics=trainer_config.managed_mldiagnostics, + timeout_secs=trainer_config.rl.profiler_timeout_secs, + ) + + def create_rl_components( # pylint: disable=too-many-positional-arguments trainer_config, sampler_config, @@ -389,24 +428,11 @@ def create_rl_components( # pylint: disable=too-many-positional-arguments None if trainer_config.rollout_micro_batch_size == -1 else trainer_config.rollout_micro_batch_size ) - # Setup metrics logging - metrics_logging_options = metrics_logger.MetricsLoggerOptions( - log_dir=trainer_config.tensorboard_dir, flush_every_n_steps=trainer_config.log_period - ) - - profiler_options = None - if trainer_config.profiler == "xplane": - profiler_options = profiler.ProfilerOptions( - log_dir=trainer_config.tensorboard_dir, - skip_first_n_steps=trainer_config.skip_first_n_steps_for_profiler, - profiler_steps=trainer_config.profiler_steps, - # Skip setting tracer levels. - set_profile_options=False, - **_kwargs_supported_by( - profiler.ProfilerOptions, - enable_continuous_profiling=trainer_config.enable_continuous_profiling, - ), - ) + # Setup metrics logging and targeted RL profiling. `profiler_options` stays + # None: tunix `RLTrainingConfig` forbids combining the SFT step profiler with + # `rl_profiler_config`, and only the latter traces the rollout role. + metrics_logging_options = build_metrics_logging_options(trainer_config) + rl_profiler_config = build_rl_profiler_config(trainer_config) # Parse vllm_additional_config rollout_additional_config = None @@ -464,7 +490,8 @@ def create_rl_components( # pylint: disable=too-many-positional-arguments rollout_micro_batch_size=rollout_micro_batch_size, max_seq_token_per_tpu=trainer_config.max_seq_token_per_tpu or None, metrics_logging_options=metrics_logging_options, - profiler_options=profiler_options, + profiler_options=None, + rl_profiler_config=rl_profiler_config, checkpoint_root_directory=checkpoint_dir, checkpointing_options=checkpointing_options, ), @@ -662,6 +689,15 @@ def rl_train(argv: Sequence[str], kwargs: dict): kwargs, config_class=types.RLConfig, ) + install_pathways_multihost_profiler_patch(trainer_config.rl.profiler_max_num_hosts) + + # Create the ML Diagnostics run before any tunix object: `RLCluster` builds + # `RLProfiler`, whose managed backend preflights `Xprof._ensure_initialized()` + # and raises if no MLRun is active. Both configs are uploaded (sampler keys + # that differ from the trainer's are recorded as `sampler.`). + if trainer_config.managed_mldiagnostics: + ManagedMLDiagnostics(trainer_config, sampler_config=sampler_config) + goodput_recorder = create_goodput_recorder(trainer_config) record_goodput(goodput_recorder, RECORD_JOB_START_TIME) _job_completed_gracefully = False @@ -812,6 +848,42 @@ def _rl_train_impl( ) +def install_pathways_multihost_profiler_patch(max_hosts: int) -> None: + """Injects max_num_hosts into start_trace when unset by caller on Pathways.""" + if max_hosts <= 1: + return + + orig_start_trace = jax.profiler.start_trace + if getattr(orig_start_trace, "_is_pathways_multihost_patched", False): + orig_start_trace._max_hosts = max_hosts # pylint: disable=protected-access + return + + try: + sig = inspect.signature(orig_start_trace) + if "max_num_hosts" not in sig.parameters: + max_logging.warning( + "rl.profiler_max_num_hosts > 1 requested, but jax.profiler.start_trace " + "does not accept max_num_hosts. Skipping monkeypatch." + ) + return + except (ValueError, TypeError): + return + + @functools.wraps(orig_start_trace) + def start_trace_wrapper(*args, **kwargs): + if "max_num_hosts" not in kwargs and len(args) <= 4: + kwargs["max_num_hosts"] = getattr(start_trace_wrapper, "_max_hosts", max_hosts) + return orig_start_trace(*args, **kwargs) + + start_trace_wrapper._is_pathways_multihost_patched = True # pylint: disable=protected-access + start_trace_wrapper._max_hosts = max_hosts # pylint: disable=protected-access + + jax.profiler.start_trace = start_trace_wrapper + jax_src_profiler = getattr(getattr(jax, "_src", None), "profiler", None) + if jax_src_profiler is not None: + jax_src_profiler.start_trace = start_trace_wrapper + + def main(argv: Sequence[str], kwargs: dict = None) -> None: """Main function to run RL training. diff --git a/tests/post_training/unit/train_rl_test.py b/tests/post_training/unit/train_rl_test.py index 4899c44fed..735e4a14f4 100644 --- a/tests/post_training/unit/train_rl_test.py +++ b/tests/post_training/unit/train_rl_test.py @@ -1070,6 +1070,7 @@ def _run_rl_train_with_goodput(self, enable_goodput_recording: bool): checkpoint_dir="/tmp/rl_ckpt", tensorboard_dir="/tmp/rl_tb", profiler="", + managed_mldiagnostics=False, steps=1, batch_size=1, num_batches=1, @@ -1084,7 +1085,7 @@ def _run_rl_train_with_goodput(self, enable_goodput_recording: bool): enable_single_controller=False, reward_functions_path="", reward_functions="", - rl=SimpleNamespace(use_agentic_rollout=False), + rl=SimpleNamespace(use_agentic_rollout=False, profiler_max_num_hosts=1), ) sampler_config = SimpleNamespace(debug=SimpleNamespace(rl=False)) diff --git a/tests/unit/managed_mldiagnostics_test.py b/tests/unit/managed_mldiagnostics_test.py index 6fd113a011..d75e7e6ec2 100644 --- a/tests/unit/managed_mldiagnostics_test.py +++ b/tests/unit/managed_mldiagnostics_test.py @@ -14,13 +14,16 @@ """Unit tests for ManagedMLDiagnostics.""" -import unittest +import json from unittest import mock +from absl.testing import absltest +from maxtext.common import managed_mldiagnostics from maxtext.common.managed_mldiagnostics import ManagedMLDiagnostics, mldiag +from maxtext.configs import types -class ManagedMLDiagnosticsTest(unittest.TestCase): +class ManagedMLDiagnosticsTest(absltest.TestCase): # pylint: disable=protected-access def setUp(self): @@ -78,6 +81,164 @@ def test_enabled_populated_region_passes_region(self): region="us-east1", ) + def test_sampler_config_merging_differing_keys(self): + mock_config = mock.MagicMock() + mock_config.managed_mldiagnostics = True + mock_config.managed_mldiagnostics_region = "" + mock_config.run_name = "test_run" + mock_config.managed_mldiagnostics_run_group = "test_group" + mock_config.managed_mldiagnostics_dir = "gs://test_dir" + mock_config.managed_mldiagnostics_on_demand_profiling = False + mock_config.get_keys.return_value = {"key1": "val1", "num_slices": 1} + + mock_sampler_config = mock.MagicMock() + mock_sampler_config.get_keys.return_value = { + "key1": "val1", + "num_slices": 2, + } + + with mock.patch.object(mldiag, "machinelearning_run") as mock_run: + ManagedMLDiagnostics(mock_config, sampler_config=mock_sampler_config) + mock_run.assert_called_once_with( + name="test_run", + run_group="test_group", + configs={"key1": "val1", "num_slices": 1, "sampler.num_slices": 2}, + gcs_path="gs://test_dir", + on_demand_xprof=False, + region=None, + ) + + def test_sampler_config_same_object_no_prefix(self): + mock_config = mock.MagicMock() + mock_config.managed_mldiagnostics = True + mock_config.managed_mldiagnostics_region = "" + mock_config.run_name = "test_run" + mock_config.managed_mldiagnostics_run_group = "test_group" + mock_config.managed_mldiagnostics_dir = "gs://test_dir" + mock_config.managed_mldiagnostics_on_demand_profiling = False + mock_config.get_keys.return_value = {"key1": "val1"} + + with mock.patch.object(mldiag, "machinelearning_run") as mock_run: + ManagedMLDiagnostics(mock_config, sampler_config=mock_config) + mock_run.assert_called_once_with( + name="test_run", + run_group="test_group", + configs={"key1": "val1"}, + gcs_path="gs://test_dir", + on_demand_xprof=False, + region=None, + ) + + def test_stub_sdk_warns_and_still_no_ops(self): + mock_config = mock.MagicMock() + mock_config.managed_mldiagnostics = True + mock_config.managed_mldiagnostics_region = "" + mock_config.managed_mldiagnostics_dir = "gs://test_dir" + mock_config.get_keys.return_value = {} + + with ( + mock.patch.object(managed_mldiagnostics, "_mldiag_is_stub", True), + mock.patch.object(managed_mldiagnostics.max_logging, "warning") as mock_warning, + mock.patch.object(mldiag, "machinelearning_run") as mock_run, + ): + self.assertFalse(managed_mldiagnostics.mldiagnostics_available()) + ManagedMLDiagnostics(mock_config) + mock_warning.assert_called_once() + self.assertIn("not available", mock_warning.call_args.args[0]) + # The stub's machinelearning_run is a no-op; pretrain behavior is unchanged. + mock_run.assert_called_once() + + def test_mldiagnostics_available_when_real_sdk_loaded(self): + with mock.patch.object(managed_mldiagnostics, "_mldiag_is_stub", False): + self.assertTrue(managed_mldiagnostics.mldiagnostics_available()) + + def test_gcs_path_is_passed_through_unchanged(self): + mock_config = mock.MagicMock() + mock_config.managed_mldiagnostics = True + mock_config.managed_mldiagnostics_region = "" + mock_config.run_name = "test_run" + mock_config.managed_mldiagnostics_run_group = "" + mock_config.managed_mldiagnostics_dir = "gs://bucket/run/managed-mldiagnostics" + mock_config.managed_mldiagnostics_on_demand_profiling = False + mock_config.get_keys.return_value = {} + + with mock.patch.object(mldiag, "machinelearning_run") as mock_run: + ManagedMLDiagnostics(mock_config) + self.assertEqual( + mock_run.call_args.kwargs["gcs_path"], + "gs://bucket/run/managed-mldiagnostics", + ) + + def test_sampler_keys_skipped_when_not_json_serializable(self): + mock_config = mock.MagicMock() + mock_config.managed_mldiagnostics = True + mock_config.managed_mldiagnostics_region = "" + mock_config.run_name = "test_run" + mock_config.managed_mldiagnostics_run_group = "" + mock_config.managed_mldiagnostics_dir = "gs://test_dir" + mock_config.managed_mldiagnostics_on_demand_profiling = False + mock_config.get_keys.return_value = {"a": 1} + + mock_sampler_config = mock.MagicMock() + mock_sampler_config.get_keys.return_value = { + "a": 1, + "b": object(), + "c": float("nan"), + "d": 2, + } + + with mock.patch.object(mldiag, "machinelearning_run") as mock_run: + ManagedMLDiagnostics(mock_config, sampler_config=mock_sampler_config) + self.assertEqual(mock_run.call_args.kwargs["configs"], {"a": 1, "sampler.d": 2}) + + def test_config_with_get_keys(self): + class PydanticConfigMock: + managed_mldiagnostics = True + managed_mldiagnostics_region = "" + run_name = "pydantic_run" + managed_mldiagnostics_run_group = "pydantic_group" + managed_mldiagnostics_dir = "gs://pydantic_dir" + managed_mldiagnostics_on_demand_profiling = False + + def get_keys(self): + return {"pydantic_key": "pydantic_val", "epochs": 10} + + mock_config = PydanticConfigMock() + with mock.patch.object(mldiag, "machinelearning_run") as mock_run: + ManagedMLDiagnostics(mock_config) + mock_run.assert_called_once_with( + name="pydantic_run", + run_group="pydantic_group", + configs={"pydantic_key": "pydantic_val", "epochs": 10}, + gcs_path="gs://pydantic_dir", + on_demand_xprof=False, + region=None, + ) + + def test_pydantic_rl_config_without_get_keys(self): + """The RL trainer passes pydantic `RLConfig` objects (initialize_pydantic), which have no get_keys().""" + trainer = types.RLConfig( + model_name="llama2-7b", + tokenizer_path="meta-llama/Llama-2-7b", + run_name="rl_run", + base_output_directory="gs://bucket/out", + managed_mldiagnostics=True, + managed_mldiagnostics_on_demand_profiling=False, + managed_mldiagnostics_region="us-central1", + ) + sampler = trainer.model_copy(update={"debug": not trainer.debug}) + with mock.patch.object(mldiag, "machinelearning_run") as mock_run: + ManagedMLDiagnostics(trainer, sampler_config=sampler) + kwargs = mock_run.call_args.kwargs + self.assertEqual(kwargs["name"], "rl_run") + self.assertEqual(kwargs["gcs_path"], "gs://bucket/out/rl_run/managed-mldiagnostics") + self.assertEqual(kwargs["region"], "us-central1") + self.assertFalse(kwargs["on_demand_xprof"]) + self.assertEqual(kwargs["configs"]["model_name"], "llama2-7b") + self.assertEqual(kwargs["configs"]["sampler.debug"], sampler.debug) + self.assertNotIn("sampler.model_name", kwargs["configs"]) + json.dumps(kwargs["configs"], allow_nan=False) # everything uploaded is JSON-serializable + if __name__ == "__main__": - unittest.main() + absltest.main() diff --git a/tests/unit/rl_profiling_config_test.py b/tests/unit/rl_profiling_config_test.py new file mode 100644 index 0000000000..e61cc1e94e --- /dev/null +++ b/tests/unit/rl_profiling_config_test.py @@ -0,0 +1,224 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Config-level tests for RL profiling and managed ML Diagnostics on `RLConfig`. + +These build configs through the real `pyconfig` path (so the pydantic validators +and the unknown-key check run) and need no accelerator or network. +""" + +import sys +from unittest import mock + +from absl.testing import absltest +from absl.testing import parameterized +from maxtext.configs import pyconfig +from maxtext.configs import types +from tests.utils.test_helpers import get_test_config_path +import pydantic + +# Tiny-model overrides that let RLConfig validate quickly on CPU. +_MODEL_OVERRIDES = { + "enable_checkpointing": False, + "base_num_decoder_layers": 1, + "attention": "dot_product", + "max_target_length": 8, + "base_emb_dim": 128, + "base_num_query_heads": 2, + "base_num_kv_heads": 2, + "base_mlp_dim": 256, + "max_prefill_predict_length": 4, + "model_name": "llama2-7b", + "override_model_config": True, + "weight_dtype": "bfloat16", + "tokenizer_path": "meta-llama/Llama-2-7b", +} + +# Step-window profiling flags that exist on the pretrain `Profiling` mixin and +# must be rejected by `RLConfig`. +_LEGACY_PROFILING_FLAGS = ( + "skip_first_n_steps_for_profiler=1", + "profiler_steps=2", + "profile_periodically_period=10", + "profile_cleanly=false", + "upload_all_profiler_results=true", + "enable_continuous_profiling=true", +) + + +def _rl_config(*cli_args: str, **overrides): + return pyconfig.initialize_pydantic( + [sys.argv[0], get_test_config_path("post_train/rl.yml"), *cli_args], + config_class=types.RLConfig, + **{**_MODEL_OVERRIDES, **overrides}, + ) + + +class RLProfilingConfigTest(parameterized.TestCase): + + def test_rl_yml_declares_invocation_window_defaults(self): + cfg = _rl_config(run_name="rl_profiling_test") + self.assertEqual(cfg.rl.profiler_start_invocation, 2) + self.assertEqual(cfg.rl.profiler_num_invocations, 1) + self.assertEqual(cfg.rl.profiler_timeout_secs, 60.0) + self.assertEqual(cfg.rl.profiler_max_num_hosts, 1) + self.assertEqual(cfg.profiler, types.ProfilerType.NONE) + + def test_cli_overrides_reach_nested_rl_block(self): + with mock.patch("pathwaysutils.is_pathways_backend_used", return_value=True): + cfg = _rl_config( + "rl.profiler_start_invocation=5", + "rl.profiler_num_invocations=3", + "rl.profiler_timeout_secs=120", + "rl.profiler_max_num_hosts=4", + "profiler=xplane", + run_name="rl_profiling_test", + ) + self.assertEqual(cfg.rl.profiler_start_invocation, 5) + self.assertEqual(cfg.rl.profiler_num_invocations, 3) + self.assertEqual(cfg.rl.profiler_timeout_secs, 120.0) + self.assertEqual(cfg.rl.profiler_max_num_hosts, 4) + self.assertEqual(cfg.profiler, types.ProfilerType.XPLANE) + + def test_profiler_max_num_hosts_rejects_non_xplane_profiler(self): + with mock.patch("pathwaysutils.is_pathways_backend_used", return_value=True): + with self.assertRaisesRegex(ValueError, "requires profiler=xplane"): + _rl_config("rl.profiler_max_num_hosts=4", run_name="rl_profiling_test") + + def test_profiler_max_num_hosts_rejects_non_pathways_backend(self): + with mock.patch("pathwaysutils.is_pathways_backend_used", return_value=False): + with self.assertRaisesRegex(ValueError, "Pathways backend"): + _rl_config( + "profiler=xplane", + "rl.profiler_max_num_hosts=4", + run_name="rl_profiling_test", + ) + + def test_profiler_max_num_hosts_allows_managed_mldiagnostics(self): + with mock.patch("pathwaysutils.is_pathways_backend_used", return_value=True): + cfg = _rl_config( + "profiler=xplane", + "rl.profiler_max_num_hosts=4", + run_name="run1", + base_output_directory="gs://bucket/out", + managed_mldiagnostics=True, + managed_mldiagnostics_on_demand_profiling=False, + ) + self.assertEqual(cfg.rl.profiler_max_num_hosts, 4) + self.assertTrue(cfg.managed_mldiagnostics) + + @parameterized.parameters( + {"rl": {"profiler_start_invocation": -1}}, + {"rl": {"profiler_num_invocations": 0}}, + {"rl": {"profiler_timeout_secs": 0}}, + ) + def test_invocation_window_bounds(self, **bad): + with self.assertRaises(pydantic.ValidationError): + types.RLConfig(model_name="llama2-7b", tokenizer_path="meta-llama/Llama-2-7b", **bad) + + @parameterized.parameters(*_LEGACY_PROFILING_FLAGS) + def test_legacy_step_profiling_flags_are_rejected_by_rl_config(self, flag): + key = flag.split("=", 1)[0] + with self.assertRaisesRegex(ValueError, key): + _rl_config(flag, run_name="rl_profiling_test") + + def test_legacy_step_profiling_flags_still_exist_on_pretrain_config(self): + for flag in _LEGACY_PROFILING_FLAGS: + self.assertIn(flag.split("=", 1)[0], types.MaxTextConfig.model_fields) + self.assertIn("profiler", types.RLConfig.model_fields) + self.assertNotIn("rl", types.MaxTextConfig.model_fields) + + def test_managed_mldiagnostics_dir_is_derived_without_trailing_slash(self): + cfg = _rl_config(run_name="run1", base_output_directory="gs://bucket/out") + self.assertEqual( + cfg.managed_mldiagnostics_dir, + "gs://bucket/out/run1/managed-mldiagnostics", + ) + cfg = _rl_config( + run_name="run1", + base_output_directory="gs://bucket/out", + managed_mldiagnostics_storage_path="gs://telemetry", + ) + self.assertEqual( + cfg.managed_mldiagnostics_dir, + "gs://telemetry/run1/managed-mldiagnostics", + ) + self.assertEqual( + cfg.managed_mldiagnostics_dir, + types.derive_managed_mldiagnostics_dir("gs://bucket/out", "run1", "gs://telemetry"), + ) + + def test_managed_mldiagnostics_dir_matches_pretrain_derivation(self): + pretrain = pyconfig.initialize_pydantic( + [sys.argv[0], get_test_config_path("base.yml")], + run_name="run1", + base_output_directory="gs://bucket/out", + **_MODEL_OVERRIDES, + ) + rl = _rl_config(run_name="run1", base_output_directory="gs://bucket/out") + self.assertEqual(pretrain.managed_mldiagnostics_dir, rl.managed_mldiagnostics_dir) + + def test_managed_requires_gcs_run_directory(self): + with self.assertRaisesRegex(ValueError, "gs://"): + _rl_config(run_name="run1", managed_mldiagnostics=True) # no base_output_directory + with self.assertRaisesRegex(ValueError, "gs://"): + _rl_config( + run_name="run1", + base_output_directory="/tmp/local", + managed_mldiagnostics=True, + ) + with self.assertRaisesRegex(ValueError, "gs://"): + _rl_config( + run_name="", + base_output_directory="gs://bucket/out", + managed_mldiagnostics=True, + ) + + def test_managed_xplane_rejects_on_demand_profiling(self): + with self.assertRaisesRegex(ValueError, "on_demand_profiling"): + _rl_config( + "profiler=xplane", + run_name="run1", + base_output_directory="gs://bucket/out", + managed_mldiagnostics=True, + managed_mldiagnostics_on_demand_profiling=True, + ) + cfg = _rl_config( + "profiler=xplane", + run_name="run1", + base_output_directory="gs://bucket/out", + managed_mldiagnostics=True, + managed_mldiagnostics_on_demand_profiling=False, + ) + self.assertTrue(cfg.managed_mldiagnostics) + # Without xplane, on-demand profiling remains allowed. + cfg = _rl_config( + run_name="run1", + base_output_directory="gs://bucket/out", + managed_mldiagnostics=True, + managed_mldiagnostics_on_demand_profiling=True, + ) + self.assertTrue(cfg.managed_mldiagnostics_on_demand_profiling) + + def test_unmanaged_config_is_not_validated_for_gcs(self): + cfg = _rl_config( + run_name="run1", + base_output_directory="/tmp/local", + managed_mldiagnostics=False, + ) + self.assertEqual(cfg.managed_mldiagnostics_dir, "/tmp/local/run1/managed-mldiagnostics") + + +if __name__ == "__main__": + absltest.main()