Support MLPerf training logging compliance for MaxText - #5118
Open
zxhe-sean wants to merge 1 commit into
Open
Conversation
zxhe-sean
requested review from
SurbhiJainUSC,
bvandermoon,
parambole,
richjames0 and
shralex
as code owners
September 3, 2026 01:58
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
zxhe-sean
requested review from
A9isha,
NuojCheng,
RissyRan,
abhinavclemson,
aireenmei,
darisoy,
dipannita08,
gagika,
gobbleturk,
hengtaoguo,
huytransformer,
igorts-git,
jiangjy1982,
khatwanimohit,
shuningjin,
vipannalla and
xibinliu
as code owners
September 3, 2026 01:58
zxhe-sean
force-pushed
the
mlperf_logging
branch
from
September 3, 2026 02:03
32149fc to
8afc957
Compare
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
Contributor
|
🤖 Hi @RissyRan, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
This Pull Request introduces end-to-end MLPerf Training logging (mllog) compliance support in MaxText, aligning with MLCommons v6.0 schemas. It correctly structures rank-0 isolated logging and implements continuous synchronization of local compliance logs to GCS. The integration with the core training and evaluation loops is robust, well-sequenced, and introduces zero overhead when disabled.
🔍 General Feedback
- Architecture and Integration: The integration with the core
train.pylifecycle andMetricLoggeris highly cohesive and does not disrupt standard execution paths. - Dependency Management: The update of the
mlperf-loggingdependency to the latest6.0.0-rc6tag is correctly aligned across package requirements and GitHub deps. - Robustness: The fallback synchronization using
etils.epathifgoogle-cloud-storageis absent ensures excellent runtime robustness.
Comment on lines
+87
to
+133
| _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}") |
Contributor
There was a problem hiding this comment.
🟠 High - Synchronous blocking network I/O in distributed training loop.
Currently, flush_and_sync is called on every step (via tracked_stats), and every 5 seconds it synchronously calls _upload_file_to_gcs to upload the local staging file to GCS. Because this is a blocking network call on Rank 0, it will stall/halt the entire distributed training cluster at the next collective sync.
Additionally, _upload_file_to_gcs creates a new GCS storage.Client instance on every sync, which incurs substantial connection and auth overhead.
Recommendation: Run the GCS upload in a background daemon thread to avoid blocking the training loop, and cache/reuse the GCS client.
Suggested change
| _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}") | |
| _last_sync_time = 0.0 | |
| _min_sync_interval = 5.0 | |
| _gcs_client = None | |
| _upload_thread = None | |
| def flush_and_sync(force=False): | |
| """Flushes mllog handlers and uploads the local staging file to GCS if needed.""" | |
| global _last_sync_time, _upload_thread | |
| 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): | |
| if _upload_thread is not None and _upload_thread.is_alive(): | |
| if force: | |
| _upload_thread.join() # If forced (e.g. final shutdown), block to complete | |
| else: | |
| return # Skip this interval to prevent overlap with a running upload | |
| import threading # pylint: disable=import-outside-toplevel | |
| _upload_thread = threading.Thread( | |
| target=_upload_file_to_gcs, | |
| args=(_destination_path, _local_staging_file), | |
| daemon=True, | |
| ) | |
| _upload_thread.start() | |
| _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.""" | |
| global _gcs_client | |
| 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) | |
| if _gcs_client is None: | |
| _gcs_client = storage.Client() | |
| bucket = _gcs_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}") |
| limitations under the License. | ||
| """ | ||
|
|
||
| """Utils for MLPerf submission compliance.""" |
Contributor
There was a problem hiding this comment.
🟢 Low - Missing unit tests for the newly added compliance logging utility.
The PR introduces a comprehensive logging utility in mllog_utils.py without any associated test coverage.
Recommendation: Add a unit test file (e.g. tests/unit/mllog_utils_test.py) to verify configurations, GCS staging, synchronization intervals, and compliance lifecycle logging methods using mock handlers.
Comment on lines
+71
to
+76
| try: | ||
| with open(_local_staging_file, "w", encoding="utf8"): | ||
| pass | ||
| except Exception: # pylint: disable=broad-exception-caught | ||
| pass | ||
| mllog.config(filename=_local_staging_file) |
Contributor
There was a problem hiding this comment.
🟡 Medium - Silent failure and subsequent crash if `/tmp` or staging log file is not writeable.
In setup_mllog, the try-except block checks if the file can be created but silently catches the exception and proceeds to call mllog.config(filename=_local_staging_file). If the file is indeed not writeable, mllog.config or subsequent logging calls will crash.
Recommendation: Safely disable logging and return early if staging file creation fails.
Suggested change
| try: | |
| with open(_local_staging_file, "w", encoding="utf8"): | |
| pass | |
| except Exception: # pylint: disable=broad-exception-caught | |
| pass | |
| mllog.config(filename=_local_staging_file) | |
| try: | |
| with open(_local_staging_file, "w", encoding="utf8"): | |
| pass | |
| mllog.config(filename=_local_staging_file) | |
| except Exception as e: # pylint: disable=broad-exception-caught | |
| max_logging.log(f"Warning: Staging file {_local_staging_file} is not writeable: {e}. MLPerf logging is disabled.") | |
| _enabled = False | |
| _is_configured = True | |
| return |
Comment on lines
+181
to
+194
| 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) |
Contributor
There was a problem hiding this comment.
🟡 Medium - Hardcoded compliance metadata limits benchmark portability.
The submission organization ("Google"), platform ("TPU-Ironwood"), benchmark ("DEEPSEEKV3_671B"), and train samples (1574207408) are hardcoded. This restricts the compliance utility to a single hardware and model configuration, which goes against MaxText's support for multiple platforms and models (e.g. Llama, Gemma, etc.).
Recommendation: Make these configurable via config with sensible default values.
Suggested change
| 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) | |
| submission_org = getattr(config, "submission_org", "Google") | |
| submission_platform = getattr(config, "submission_platform", "TPU-Ironwood") | |
| submission_benchmark = getattr(config, "submission_benchmark", mllog.constants.DEEPSEEKV3_671B) | |
| train_samples = getattr(config, "mllog_train_samples", 1574207408) | |
| mllogger.event(mllog.constants.SUBMISSION_ORG, submission_org) | |
| mllogger.event(mllog.constants.SUBMISSION_PLATFORM, submission_platform) | |
| 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, submission_benchmark) | |
| 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, train_samples) |
susanbao
reviewed
Sep 4, 2026
|
|
||
| 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) |
Collaborator
There was a problem hiding this comment.
Check whether train_step has +1 before _finalize_eval_metrics.
| mllog_utils.init_print(config, start_step) | ||
| mllog_utils.init_stop() | ||
| mllog_utils.run_start() | ||
| mllog_utils.block_start(config, start_step) |
Collaborator
There was a problem hiding this comment.
We have the block_start at the beginning. The block_stop is only in eval_start. If we do not have the eval. The block_stop will not be added.
| @@ -0,0 +1,306 @@ | |||
| """ | |||
Collaborator
There was a problem hiding this comment.
Could you please add unit tests for rcp logging?
| ) | ||
| mllogger.event( | ||
| mllog.constants.OPT_LR_DECAY_STEPS, | ||
| int(config.learning_rate_schedule_steps * (1 - config.warmup_steps_fraction) + 1), |
Collaborator
There was a problem hiding this comment.
Why have +1? opt_learning_rate_decay_steps requires " v['value'] == 12000 - s['opt_learning_rate_warmup_steps'] ".
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This PR introduces end-to-end MLPerf training logging (
mllog) compliance support into MaxText for MLPerf Training submissions.When enabled (
enable_mllog=true), MaxText logs all required initialization metadata, training step metrics (tracked_stats), evaluation intervals (block_start/block_stop,eval_start/eval_stop), evaluation accuracy milestones, and early stopping / run completion markers strictly following MLCommons Training compliance rules.Key Changes
src/maxtext/utils/mllog_utils.py):mllogcalls are gated strictly tojax.process_index() == 0.mllogrecords locally throughmllog.config(filename=...)and periodically/synchronously syncs the staging log tobase_output_directoryon GCS usinggoogle-cloud-storage(client.bucket(...)avoiding bucket metadata permission requirements) withetils.epathfallback.submission_org,submission_platform,submission_division,submission_benchmark(deepseekv3_671b), andsubmission_status.opt_adamw_beta_1,opt_adamw_beta_2,opt_adamw_epsilon, andopt_adamw_weight_decay.moe_aux_loss_coeff(mapped fromconfig.load_balance_loss_weight).opt_base_learning_rate,opt_learning_rate_warmup_steps,opt_learning_rate_decay_steps, andopt_learning_rate_decay_schedule.eval_samples,train_samples,init_checkpoint_step, andmax_sequence_length.init_start,cache_clear,init_stop,run_start, and initialblock_start.tracked_statson each completed training step withreduced_train_lossandtrain_step_time.block_stop,eval_start,eval_accuracy,eval_stop, and conditionalrun_stopupon reachingtarget_eval_loss.train.py,metric_logger.py):train.py:mllog_utils.init_start(config)during initialization.init_print,init_stop,run_start, andblock_startbefore entering the step loop.tracked_statsandeval_startat evaluation checkpoints.run_stop(status='success')upon job completion.metric_logger.py:mllog_utils.check_evalfromfinalize_eval_metricswith eval loss and step count.mllog_utils.flush_and_sync(force=True)on writer close.base.yml,types.py):enable_mllog: bool = False: Flag to enable MLPerf logging (defaults toFalsefor zero overhead during standard runs).mllog_file: None | PathStr = "": Path or filename for exporting mllog entries (defaults toos.path.join(base_output_directory, run_name, "mllog.log")). SupportsNone | PathStrto handle empty CLI overrides (mllog_file="").requirements.txt,pre_train_github_deps.txt):mlperf-loggingdependency to release6.0.0-rc6.BUGS: b/556312967
Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.
Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.
Tests
Had a successful e2e run, the generated mllog file, also did a local RCP check with compliance_checker. It didn't pass the check as a few flag naming issue, and the mismatch of the mlperf_logging version.
I made the necessary change and verified locally for the missing flags.
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.