From 8afc957acd95d05f56f3c6038f74a094fe4b18c7 Mon Sep 17 00:00:00 2001 From: Zhengxian He Date: Thu, 3 Sep 2026 01:43:51 +0000 Subject: [PATCH] Support MLPerf training logging compliance for MaxText --- .../extra_deps/pre_train_github_deps.txt | 2 +- .../requirements/requirements.txt | 2 +- src/maxtext/common/metric_logger.py | 7 +- src/maxtext/configs/base.yml | 7 + src/maxtext/configs/types.py | 20 ++ src/maxtext/trainers/pre_train/train.py | 25 +- src/maxtext/utils/mllog_utils.py | 306 ++++++++++++++++++ 7 files changed, 363 insertions(+), 6 deletions(-) create mode 100644 src/maxtext/utils/mllog_utils.py diff --git a/src/dependencies/extra_deps/pre_train_github_deps.txt b/src/dependencies/extra_deps/pre_train_github_deps.txt index 676f2e58e7..87e569cbe1 100644 --- a/src/dependencies/extra_deps/pre_train_github_deps.txt +++ b/src/dependencies/extra_deps/pre_train_github_deps.txt @@ -1,2 +1,2 @@ google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/29329e8e73820993f77cfc8efe34eb2a73f5de98.zip -mlperf-logging @ https://github.com/mlcommons/logging/archive/38ab22670527888c8eb7825a4ece176fcc36a95d.zip +mlperf-logging @ https://github.com/mlcommons/logging/archive/refs/tags/6.0.0-rc6.zip diff --git a/src/dependencies/requirements/requirements.txt b/src/dependencies/requirements/requirements.txt index 81168e0ae0..f4dc649133 100644 --- a/src/dependencies/requirements/requirements.txt +++ b/src/dependencies/requirements/requirements.txt @@ -43,4 +43,4 @@ tiktoken tokamax>=0.0.4 transformers google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/29329e8e73820993f77cfc8efe34eb2a73f5de98.zip -mlperf-logging @ https://github.com/mlcommons/logging/archive/38ab22670527888c8eb7825a4ece176fcc36a95d.zip +mlperf-logging @ https://github.com/mlcommons/logging/archive/refs/tags/6.0.0-rc6.zip diff --git a/src/maxtext/common/metric_logger.py b/src/maxtext/common/metric_logger.py index 687c35963c..9c7ca64ced 100644 --- a/src/maxtext/common/metric_logger.py +++ b/src/maxtext/common/metric_logger.py @@ -25,12 +25,12 @@ import numpy as np import jax - from maxtext.utils.globals import EPS from maxtext.common.gcloud_stub import mldiagnostics_modules from maxtext.common.gcloud_stub import workload_monitor from maxtext.common.managed_mldiagnostics import ManagedMLDiagnostics from maxtext.utils import exceptions +from maxtext.utils import mllog_utils from maxtext.utils import gcs_utils from maxtext.utils import max_logging from maxtext.utils import max_utils @@ -93,13 +93,14 @@ class MetricLogger: Logger for saving metrics to a local file, GCS and TensorBoard. """ - def __init__(self, config, learning_rate_schedule): + def __init__(self, config, learning_rate_schedule, start_step=0): self.writer = max_utils.initialize_summary_writer(config.tensorboard_dir, config.run_name, config.enable_tensorboard) self.config = config self.metadata = {} self.running_gcs_metrics = [] if config.gcs_metrics else None self.performance_metric_queue = self.get_performance_metric_queue(config) self.learning_rate_schedule = learning_rate_schedule + self.start_step = start_step self.cumulative_eval_metrics = {"scalar": defaultdict(float)} # self.buffered_metrics is a polymorphic deferred-write queue. Entries are one of: # ("train", train_step, metrics, step_time_delta) @@ -493,6 +494,7 @@ def _finalize_eval_metrics(self, train_step): self.write_metrics(self.cumulative_eval_metrics, train_step, metric_type="eval") self._pending_eval_step_count = 0 + mllog_utils.check_eval(self.config, train_step + 1, eval_loss, self.start_step) if self.config.target_eval_loss and eval_loss <= self.config.target_eval_loss: raise exceptions.StopTraining(f"Target loss {self.config.target_eval_loss=} is achieved.") @@ -507,4 +509,5 @@ def flush_metrics_and_cleanup(self): self._flush_one_buffered_entry(entry) self.buffered_metrics = [] + mllog_utils.flush_and_sync(force=True) max_utils.close_summary_writer(self.writer) diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 5669e012ae..5f9ffe4c08 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -490,6 +490,13 @@ record_internal_nn_metrics: 0 # Create a GCS bucket, e.g. my-maxtext-outputs and set this to "gs://my-maxtext-outputs/" base_output_directory: "" +# Whether to enable MLPerf logging (mllog). Default is false. +enable_mllog: false + +# File name for exporting MLPerf mllog entries in base_output_directory (e.g. "mllog.log" or "seed_1.out"). +# If base_output_directory is set and mllog_file is empty, defaults to "mllog.log" inside base_output_directory/run_name. +mllog_file: "" + # Multi-tier checkpointing is an experimental Orbax feature that: periodically saves to persistent storage(GCS bucket) dictated by `multi_tier_checkpointing_backup_interval_minutes` and, # saves to a local directory for smaller checkpoint intervals(local_checkpoint_period). # The local checkpoint directory must be specified when enabling multi-tier checkpointing. diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 353656c665..26a636586d 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -324,6 +324,11 @@ class RunInfo(BaseModel): ) debug_sharding: bool = Field(False, description="If True, print model weight sharding details.") base_output_directory: PathStr = Field("", description="Base directory for all outputs, typically a GCS path.") + enable_mllog: bool = Field(False, description="If True, enables MLPerf logging (mllog).") + mllog_file: None | PathStr = Field( + "", + description="Optional filename or path for mllog export in base_output_directory (defaults to 'mllog.log').", + ) sharding_strategy: None | Literal["experimental"] = Field( None, description="Experimental sharding strategy used for some inference configs.", @@ -3546,6 +3551,21 @@ def set_derived_and_validate_values(self) -> "MaxTextConfig": # 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") + if self.enable_mllog: + if not self.mllog_file: + self.mllog_file = os.path.join(output_dir, "mllog.log") + elif not self.mllog_file.startswith("gs://") and not os.path.isabs(self.mllog_file): + self.mllog_file = os.path.join(output_dir, self.mllog_file) + else: + self.mllog_file = "" + elif self.base_output_directory: + if self.enable_mllog: + if not self.mllog_file: + self.mllog_file = os.path.join(self.base_output_directory, "mllog.log") + elif not self.mllog_file.startswith("gs://") and not os.path.isabs(self.mllog_file): + self.mllog_file = os.path.join(self.base_output_directory, self.mllog_file) + else: + self.mllog_file = "" else: self.checkpoint_dir, self.metrics_dir, self.tensorboard_dir = ( None, diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py index 692690db76..b283fe3d66 100644 --- a/src/maxtext/trainers/pre_train/train.py +++ b/src/maxtext/trainers/pre_train/train.py @@ -76,6 +76,7 @@ from maxtext.utils import sharding from maxtext.utils import maxtext_utils_nnx from maxtext.utils import train_utils +from maxtext.utils import mllog_utils from maxtext.utils.gradient_accumulation import gradient_accumulation_loss_and_grad from maxtext.utils.vocabulary_tiling import vocab_tiling_linen_loss, vocab_tiling_nnx_loss @@ -884,6 +885,15 @@ def training_loop_iteration( step_time_delta = datetime.datetime.now() - last_step_completion last_step_completion = datetime.datetime.now() + completed_step = step + 1 + mllog_utils.tracked_stats( + config, + completed_step, + step_time_delta.total_seconds(), + metrics["scalar"]["learning/loss"], + start_step=start_step, + ) + checkpointing.maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step) if dump_hlo and step == (dump_step if dump_step >= 0 else start_step): @@ -906,6 +916,7 @@ def training_loop_iteration( # Explicitly reset the eval iterator and counters before starting the eval loop eval_data_iterator.reset() metric_logger_instance.reset_eval_metrics() + mllog_utils.eval_start(config, completed_step, start_step=start_step) max_logging.log(f"Starting eval after train step {step}") eval_step_count = 0 @@ -1017,7 +1028,9 @@ def train_loop(config, recorder, state=None): compiled_stats = compiled.memory_analysis() max_utils.print_compiled_memory_stats(compiled_stats) prof = profiler.Profiler(config, offset_step=start_step) - metric_logger_instance = metric_logger.MetricLogger(config=config, learning_rate_schedule=learning_rate_schedule) + metric_logger_instance = metric_logger.MetricLogger( + config=config, learning_rate_schedule=learning_rate_schedule, start_step=start_step + ) # Write train config params, num model params, and XLA flags to tensorboard if isinstance(model, nn.Module): @@ -1079,6 +1092,11 @@ def train_loop(config, recorder, state=None): try: python_vars["last_step_completion"] = datetime.datetime.now() + mllog_utils.init_print(config, start_step) + mllog_utils.init_stop() + mllog_utils.run_start() + mllog_utils.block_start(config, start_step) + # Using while loop to allow for potential dynamic 'steps' adjustment in future while python_vars["step"] < immutable_data["steps"]: training_loop_iteration(jax_device_state, python_vars, immutable_data) @@ -1089,7 +1107,6 @@ def train_loop(config, recorder, state=None): if immutable_data["save_checkpoint_on_completion"]: checkpointing.maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator) - if checkpoint_manager is not None: # in case the last checkpoint_period checkpoint is still in progress checkpointing.wait_until_finished(checkpoint_manager) @@ -1101,6 +1118,9 @@ def train_loop(config, recorder, state=None): finally: if _job_completed_gracefully: record_goodput(recorder, RECORD_JOB_END_TIME) + samples_count = (python_vars["step"] - immutable_data["start_step"]) * config.global_batch_size_to_train_on + mllog_utils.run_stop(status="success", current_epoch_num=samples_count) + mllog_utils.flush_and_sync() metric_logger_instance.flush_metrics_and_cleanup() train_utils.maybe_cleanup_dcn_throttling(config) @@ -1135,6 +1155,7 @@ def initialize(argv: Sequence[str]) -> tuple[pyconfig.HyperParameters, Any]: max_utils.bootstrap_transformer_engine_cgemm(config) # Create the Goodput recorder + mllog_utils.init_start(config) recorder = create_goodput_recorder(config) return config, recorder diff --git a/src/maxtext/utils/mllog_utils.py b/src/maxtext/utils/mllog_utils.py new file mode 100644 index 0000000000..b6017fb054 --- /dev/null +++ b/src/maxtext/utils/mllog_utils.py @@ -0,0 +1,306 @@ +""" +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. +""" + +"""Utils for MLPerf submission compliance.""" + +import os +import jax +from maxtext.utils import max_logging + +try: + from mlperf_logging import mllog + + mllogger = mllog.get_mllogger() +except ImportError: + mllog = None + mllogger = None + +_destination_path = None +_local_staging_file = None +_is_configured = False +_run_stopped = False +_enabled = False + + +def is_mllog_enabled(): + return _enabled + + +def setup_mllog(config): + """Configures mllogger to output to a file in base_output_directory (or custom mllog_file).""" + global _destination_path, _local_staging_file, _is_configured, _enabled + if _is_configured or mllog is None: + return + + _enabled = bool(getattr(config, "enable_mllog", False)) + if not _enabled: + _is_configured = True + return + + if jax.process_index() != 0: + _is_configured = True + return + + target_path = getattr(config, "mllog_file", "") or "" + if not target_path and getattr(config, "base_output_directory", ""): + run_name = getattr(config, "run_name", "") + target_path = ( + os.path.join(config.base_output_directory, run_name, "mllog.log") + if run_name + else os.path.join(config.base_output_directory, "mllog.log") + ) + + if not target_path: + _is_configured = True + return + + _destination_path = target_path + if target_path.startswith("gs://"): + run_name = getattr(config, "run_name", "") or "maxtext" + _local_staging_file = f"/tmp/mllog_{run_name}.log" + try: + with open(_local_staging_file, "w", encoding="utf8"): + pass + except Exception: # pylint: disable=broad-exception-caught + pass + mllog.config(filename=_local_staging_file) + max_logging.log(f"Configured mllog to staging file {_local_staging_file} (destination: {_destination_path})") + else: + os.makedirs(os.path.dirname(os.path.abspath(target_path)), exist_ok=True) + mllog.config(filename=target_path) + _local_staging_file = None + max_logging.log(f"Configured mllog to file {_destination_path}") + + _is_configured = True + + +_last_sync_time = 0.0 +_min_sync_interval = 5.0 + + +def flush_and_sync(force=False): + """Flushes mllog handlers and uploads the local staging file to GCS if needed.""" + global _last_sync_time + if not _enabled or mllogger is None or jax.process_index() != 0: + return + + if hasattr(mllogger, "logger") and mllogger.logger: + for handler in mllogger.logger.handlers: + try: + handler.flush() + except Exception: # pylint: disable=broad-exception-caught + pass + + if _destination_path and _local_staging_file and _destination_path.startswith("gs://"): + import time # pylint: disable=import-outside-toplevel + + now = time.time() + if force or (now - _last_sync_time >= _min_sync_interval): + _upload_file_to_gcs(_destination_path, _local_staging_file) + _last_sync_time = now + + +def _upload_file_to_gcs(dest_gcs: str, src_local: str): + """Uploads a local file to GCS using google-cloud-storage client or etils.epath fallback.""" + if not os.path.exists(src_local): + return + try: + from google.cloud import storage # pylint: disable=import-outside-toplevel + + path_parts = dest_gcs.replace("gs://", "").split("/") + bucket_name = path_parts.pop(0) + blob_name = "/".join(path_parts) + client = storage.Client() + bucket = client.bucket(bucket_name) + blob = bucket.blob(blob_name) + blob.upload_from_filename(src_local) + except Exception: # pylint: disable=broad-exception-caught + try: + from etils import epath # pylint: disable=import-outside-toplevel + + epath.Path(dest_gcs).write_bytes(epath.Path(src_local).read_bytes()) + except Exception as e: # pylint: disable=broad-exception-caught + max_logging.log(f"Warning: Failed to sync mllog to GCS destination {dest_gcs}: {e}") + + +def init_start(config=None): + if config is not None: + setup_mllog(config) + if not _enabled or mllogger is None or jax.process_index() != 0: + return + mllogger.event(mllog.constants.CACHE_CLEAR) + mllogger.start(mllog.constants.INIT_START) + flush_and_sync() + + +def init_stop(): + if not _enabled or mllogger is None or jax.process_index() != 0: + return + mllogger.end(mllog.constants.INIT_STOP) + flush_and_sync() + + +def run_start(): + if not _enabled or mllogger is None or jax.process_index() != 0: + return + mllogger.start(mllog.constants.RUN_START) + flush_and_sync() + + +def block_start(config, step=0): + """Logs BLOCK_START for an MLPerf evaluation block.""" + if not _enabled or mllogger is None or jax.process_index() != 0: + return + eval_frequency_samples = config.eval_interval * config.global_batch_size_to_train_on + mllogger.start( + mllog.constants.BLOCK_START, + metadata={ + "samples_count": eval_frequency_samples, + "step": step, + }, + ) + flush_and_sync() + + +def init_print(config, start_step): + """The initial mllog for mlperf submission compliance check.""" + setup_mllog(config) + if not _enabled or mllogger is None or jax.process_index() != 0: + return + # General + mllogger.event(mllog.constants.SUBMISSION_ORG, "Google") + mllogger.event(mllog.constants.SUBMISSION_PLATFORM, "TPU-Ironwood") + mllogger.event(mllog.constants.SUBMISSION_STATUS, mllog.constants.CLOUD) + mllogger.event(mllog.constants.SUBMISSION_DIVISION, mllog.constants.CLOSED) + + # Model specific + mllogger.event(mllog.constants.SUBMISSION_BENCHMARK, mllog.constants.DEEPSEEKV3_671B) + mllogger.event(mllog.constants.SEED, config.data_shuffle_seed) + mllogger.event(mllog.constants.MAX_STEPS, config.steps) + mllogger.event(mllog.constants.GLOBAL_BATCH_SIZE, config.global_batch_size_to_train_on) + mllogger.event(mllog.constants.GRADIENT_ACCUMULATION_STEPS, config.gradient_accumulation_steps) + mllogger.event(mllog.constants.MAX_SEQUENCE_LENGTH, config.max_target_length) + mllogger.event(mllog.constants.EVAL_SAMPLES, config.global_batch_size_to_eval_on * config.eval_steps) + mllogger.event(mllog.constants.TRAIN_SAMPLES, 1574207408) + mllogger.event(mllog.constants.INIT_CHECKPOINT_STEP, 0) + mllogger.event(mllog.constants.OPT_NAME, mllog.constants.ADAMW) + mllogger.event(mllog.constants.OPT_BASE_LR, config.learning_rate) + mllogger.event(mllog.constants.OPT_ADAMW_BETA_1, config.adam_b1) + mllogger.event(mllog.constants.OPT_ADAMW_BETA_2, config.adam_b2) + mllogger.event(mllog.constants.OPT_ADAMW_EPSILON, config.adam_eps) + mllogger.event(mllog.constants.OPT_ADAMW_WEIGHT_DECAY, config.adam_weight_decay) + mllogger.event(mllog.constants.OPT_GRADIENT_CLIP_NORM, config.gradient_clipping_threshold) + mllogger.event(mllog.constants.MOE_AUX_LOSS_COEFF, config.load_balance_loss_weight) + mllogger.event(mllog.constants.OPT_END_LR, config.learning_rate * config.learning_rate_final_fraction) + mllogger.event( + mllog.constants.OPT_LR_WARMUP_STEPS, int(config.learning_rate_schedule_steps * config.warmup_steps_fraction) + ) + mllogger.event( + mllog.constants.OPT_LR_DECAY_STEPS, + int(config.learning_rate_schedule_steps * (1 - config.warmup_steps_fraction) + 1), + ) + mllogger.event(mllog.constants.OPT_LR_DECAY_SCHEDULE, "cosine with linear warmup") + mllogger.event("target_accuracy", config.target_eval_loss) + flush_and_sync() + + +def run_stop(status="success", current_epoch_num=None): + """Logs RUN_STOP for MLPerf completion.""" + global _run_stopped + if not _enabled or _run_stopped: + return + if mllogger is not None and jax.process_index() == 0: + metadata = {"status": status} + if current_epoch_num is not None: + metadata["samples_count"] = current_epoch_num + mllogger.end(mllog.constants.RUN_STOP, metadata=metadata) + flush_and_sync(force=True) + _run_stopped = True + + +def eval_start(config, step, start_step=0): + """Logs BLOCK_STOP and EVAL_START before the evaluation loop.""" + if not _enabled or mllogger is None or jax.process_index() != 0: + return + eval_frequency_samples = config.eval_interval * config.global_batch_size_to_train_on + samples_count = (step - start_step) * config.global_batch_size_to_train_on + mllogger.end( + mllog.constants.BLOCK_STOP, + metadata={ + "samples_count": eval_frequency_samples, + "step": step, + }, + ) + mllogger.start( + mllog.constants.EVAL_START, + metadata={ + "samples_count": samples_count, + "step": step, + }, + ) + flush_and_sync() + + +def eval_stop(config, step, eval_loss, start_step=0): + """Logs EVAL_ACCURACY, EVAL_STOP, and starts a new BLOCK_START or triggers RUN_STOP.""" + if not _enabled or mllogger is None or jax.process_index() != 0: + return + samples_count = (step - start_step) * config.global_batch_size_to_train_on + eval_frequency_samples = config.eval_interval * config.global_batch_size_to_train_on + is_early_stop = bool(config.target_eval_loss and eval_loss <= config.target_eval_loss) + + mllogger.event( + mllog.constants.EVAL_ACCURACY, + float(eval_loss), + metadata={"samples_count": samples_count}, + ) + mllogger.end( + mllog.constants.EVAL_STOP, + metadata={ + "samples_count": samples_count, + "step": step, + }, + ) + if is_early_stop: + run_stop(status="success", current_epoch_num=samples_count) + else: + mllogger.start( + mllog.constants.BLOCK_START, + metadata={ + "samples_count": eval_frequency_samples, + "step": step, + }, + ) + flush_and_sync() + + +def check_eval(config, step, eval_loss, start_step=0): + """Logs an MLPerf evaluation block completion, checks for early stopping, and starts a new block if continuing.""" + eval_stop(config, step, eval_loss, start_step) + + +def tracked_stats(config, step, step_time, loss, start_step=0): + """Logs tracked_stats for MLPerf training compliance.""" + if not _enabled or mllogger is None or jax.process_index() != 0: + return + loss_val = loss.item() if hasattr(loss, "item") else float(loss) + samples_count = (step - start_step) * config.global_batch_size_to_train_on + value = {"reduced_train_loss": loss_val} + if step_time is not None: + value["train_step_time"] = step_time + mllogger.event( + key="tracked_stats", + metadata={mllog.constants.SAMPLES_COUNT: samples_count}, + value=value, + ) + flush_and_sync()