Support MLPerf training logging compliance for MaxText

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

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

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)

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.

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,

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