fix(flux): make MLPerf validation measure what it claims to

examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8_mlperf.yaml

@@ -73,9 +73,29 @@ modules:
       max_samples_per_sequence: null
 
       # Training iterations
-      train_iters: 5000
+      #
+      # This is a safety cap, not a target: the run stops at the quality
+      # threshold, and reaching it is the result. The reference convergence
+      # points for flux1 at GBS 512 (mlperf_logging rcp_checker
+      # training_6.0.0/rcps_flux1.json, 20 NVIDIA BF16 runs) span 7077888 to
+      # 7602176 samples, i.e. 13824 to 14848 steps. The previous 5000 steps
+      # was 2560000 samples -- roughly a third of the way -- so every run
+      # exhausted its budget below target and no run could ever have counted.
+      # 16000 steps (8192000 samples) clears the slowest reference run with
+      # margin for a run that converges later than the reference spread.
+      train_iters: 16000
       eval_interval: 512
-      eval_iters: 10
+      # Cover the whole MLPerf validation set. eval_iters is derived from this
+      # (29696 / 512 = 58); setting both is rejected. The previous eval_iters of
+      # 10 read 5120 samples, and even at 58 the training num_workers of 16
+      # would have left 1920 of them unread, so val_num_workers is pinned to 0.
+      eval_samples: 29696
+      val_num_workers: 0
+      # MLPerf assigns each val image one fixed timestep, carried in the Arrow
+      # 'timestep' column. Requires a val split ingested with that column; shards
+      # ingested earlier carry {"key": ...} only and will now fail loudly rather
+      # than silently fall back to injecting t = index % 8.
+      eval_timestep_source: dataset
       log_interval: 10
       save_interval: 10000
 
@@ -104,7 +124,9 @@ modules:
 
       # Learning rate scheduler (warmup-hold, no decay)
       lr_warmup_iters: 1600
-      lr_decay_iters: 4000
+      # Inert while lr_decay_style is constant; kept equal to the budget so
+      # it cannot be misread as a decay that ends a quarter of the way in.
+      lr_decay_iters: 16000
       lr_decay_style: constant
 
       # ==========================================
@@ -185,7 +207,9 @@ modules:
       use_turbo_attention: true
       use_dual_fp8_output_projection: false
 
-      seed: 42
+      # Each run in a submission campaign needs its own seed, supplied from
+      # outside so the value in the log is the value the launcher chose.
+      seed: ${PRIMUS_SEED:42}
       # MLPerf-aligned per-step CUDA RNG reseed (defaults off elsewhere; MLPerf
       # reproduction must opt in for run-to-run determinism).
       per_step_rng_reseed: true

examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8_mlperf.yaml

@@ -86,9 +86,28 @@ modules:
       num_workers: 8
       max_samples_per_sequence: null
 
-      train_iters: 5000
+      # This is a safety cap, not a target: the run stops at the quality
+      # threshold, and reaching it is the result. The reference convergence
+      # points for flux1 at GBS 512 (mlperf_logging rcp_checker
+      # training_6.0.0/rcps_flux1.json, 20 NVIDIA BF16 runs) span 7077888 to
+      # 7602176 samples, i.e. 13824 to 14848 steps. The previous 5000 steps
+      # was 2560000 samples -- roughly a third of the way -- so every run
+      # exhausted its budget below target and no run could ever have counted.
+      # 16000 steps (8192000 samples) clears the slowest reference run with
+      # margin for a run that converges later than the reference spread.
+      train_iters: 16000
       eval_interval: 512
-      eval_iters: 10
+      # Cover the whole MLPerf validation set. eval_iters is derived from this
+      # (29696 / 512 = 58); setting both is rejected. The previous eval_iters of
+      # 10 read 5120 samples, and even at 58 the training num_workers of 8 would
+      # have left 768 of them unread, so val_num_workers is pinned to 0.
+      eval_samples: 29696
+      val_num_workers: 0
+      # MLPerf assigns each val image one fixed timestep, carried in the Arrow
+      # 'timestep' column. Requires a val split ingested with that column; shards
+      # ingested earlier carry {"key": ...} only and will now fail loudly rather
+      # than silently fall back to injecting t = index % 8.
+      eval_timestep_source: dataset
       log_interval: 10
       save_interval: 10000
 
@@ -120,7 +139,9 @@ modules:
       clip_grad: 1.0
 
       lr_warmup_iters: 1600
-      lr_decay_iters: 4000
+      # Inert while lr_decay_style is constant; kept equal to the budget so
+      # it cannot be misread as a decay that ends a quarter of the way in.
+      lr_decay_iters: 16000
       lr_decay_style: constant
 
       # ==========================================
@@ -183,7 +204,9 @@ modules:
       enable_primus_turbo: false
       use_turbo_attention: false
 
-      seed: 2025
+      # Each run in a submission campaign needs its own seed, supplied from
+      # outside so the value in the log is the value the launcher chose.
+      seed: ${PRIMUS_SEED:2025}
       te_rng_tracker: true
       # MLPerf-aligned per-step CUDA RNG reseed (defaults off elsewhere; MLPerf
       # reproduction must opt in for run-to-run determinism).

examples/megatron/prepare.py

@@ -176,6 +176,18 @@ def prepare_dataset_if_needed(
         )
         return
 
+    # An external dataloader means the trainer supplies its own data pipeline
+    # -- Energon, for the diffusion recipes -- reading from data_path rather
+    # than from a tokenised corpus. Tokenising bookcorpus for one of those
+    # builds a dataset it will never open, and demands HF_TOKEN for a
+    # tokenizer it never loads.
+    if getattr(pre_trainer_cfg, "dataloader_type", None) == "external":
+        log_info(
+            "dataloader_type=external detected, skipping bookcorpus tokenisation "
+            "(the trainer supplies its own dataloader)."
+        )
+        return
+
     tokenizer_type = pre_trainer_cfg.tokenizer_type
     if (
         pre_trainer_cfg.full_validation or pre_trainer_cfg.eval_iters > 0

examples/mlperf/flux1/megatron/README.md

@@ -0,0 +1,69 @@
+# FLUX.1-Schnell MLPerf Training — Megatron backend
+
+The sibling directory (`examples/mlperf/flux1/`) runs FLUX.1-Schnell through the
+`diffusion` backend. This one runs it through Megatron, which is the path the
+MXFP6 work sits on, and produces runtime logs meant to be read by
+`mlperf_logging.compliance_checker` rather than by a person.
+
+## What the launcher supplies, and why it has to
+
+The logging patch refuses to invent any value that ends up in the submitted
+log. Startup fails, loudly, if any of these is missing:
+
+| Variable | What it decides |
+| --- | --- |
+| `MLLOG_OUTPUT_FILE` | Where the result file is written. Rank zero writes it directly, so the artifact the checker reads is the artifact the run produced — not a filtered copy of stdout. |
+| `MLLOG_SUBMISSION_ORG` / `_DIVISION` / `_PLATFORM` | Which division the log is judged in. A wrong default here is a wrong submission. |
+| `MLLOG_LOWEST_NUMERICAL_PRECISION_IN_LINEAR` / `_ATTN` / `_COMM` | The numerics disclosure. |
+| `MLPERF_CLEAR_CACHES` | Whether the machine actually started cold. |
+| `EXP` | The recipe, which a reviewer has to be able to find in the submission. |
+
+`run_and_time.sh` sets all of them, drops the page cache, and reports
+`cache_clear=false` if it could not — a run without the privileges to drop
+caches is still a valid run, just not a cold one.
+
+## One run
+
+```bash
+RUN_INDEX=0 RESULTS_DIR=/results \
+bash examples/mlperf/flux1/megatron/run_and_time.sh
+```
+
+This writes `/results/result_0.txt` and immediately runs the compliance checker
+against it. Checking at the end of each run is deliberate: the alternative is
+discovering after ten runs that none of the logs parse.
+
+## A submission campaign
+
+```bash
+RESULTS_DIR=/results bash examples/mlperf/flux1/megatron/run_campaign.sh
+```
+
+Ten runs with seeds `42..51`, then an RCP comparison over the collected
+results. Ten is what `mlperf_logging/rcp_checker/rcp_checker.py` requires for
+`flux1`. Runs that finish without reaching the target still produce a result
+file and still count — they are part of the distribution being compared, and
+discarding them would bias it.
+
+## Where the numbers come from
+
+The reference convergence points for `flux1` at global batch size 512
+(`rcp_checker/training_6.0.0/rcps_flux1.json`, 20 NVIDIA BF16 runs) span
+7,077,888 to 7,602,176 samples, which is 13,824 to 14,848 steps. The recipe's
+`train_iters` is a safety cap set above that range, not a target: a run ends
+when `eval_accuracy` reaches 0.586, and how many samples that took is the
+result.
+
+`closed_flux1.yaml` also pins several hyperparameters exactly — AdamW betas
+0.9/0.95, epsilon 1e-8, weight decay 0.1, gradient clip 1.0, and
+`evaluation_frequency` at exactly 262,144 samples (`eval_interval: 512` at
+GBS 512). Changing any of them in the recipe makes the log fail the checker.
+
+## MXFP6 and the disclosure vocabulary
+
+`training_6.0.0/common.yaml` accepts a fixed set of values for
+`lowest_numerical_precision_in_*`, and `mxfp6` is not in it. An MXFP6 run
+therefore produces a structurally valid log that the checker rejects until the
+format is accepted upstream. The logger emits the configured value anyway and
+warns — describing an MXFP6 run as `fp8` would be a false disclosure, which is
+worse than a log that has to wait for approval.

examples/mlperf/flux1/megatron/run_and_time.sh

@@ -0,0 +1,120 @@
+#!/bin/bash
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+#
+# One timed MLPerf Training run of FLUX.1-Schnell on the Megatron backend.
+#
+# The logging patch fails closed on every value it would otherwise have to
+# guess -- submission identity, precision disclosures, whether caches were
+# dropped, where the log goes, which recipe produced it. This script is where
+# those answers come from, which is why it is part of the submission rather
+# than a convenience wrapper.
+#
+# One invocation produces one result file. A submission needs ten of them,
+# each with its own seed; see run_campaign.sh.
+
+set -euo pipefail
+
+: "${PRIMUS_PATH:=$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../.." && pwd)}"
+export PRIMUS_PATH
+
+# --- What is being submitted ------------------------------------------------
+: "${EXP:=${PRIMUS_PATH}/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8_mlperf.yaml}"
+: "${MLLOG_SUBMISSION_ORG:=AMD}"
+: "${MLLOG_SUBMISSION_DIVISION:=closed}"
+: "${MLLOG_SUBMISSION_PLATFORM:=MI355X}"
+: "${MLLOG_SUBMISSION_STATUS:=onprem}"
+export EXP MLLOG_SUBMISSION_ORG MLLOG_SUBMISSION_DIVISION
+export MLLOG_SUBMISSION_PLATFORM MLLOG_SUBMISSION_STATUS
+
+# --- Numerics disclosure ----------------------------------------------------
+# The compliance checker accepts a fixed vocabulary here (see
+# mlperf_logging/compliance_checker/training_6.0.0/common.yaml). mxfp6 is not
+# in it yet, so an MXFP6 run needs the format approved upstream before its log
+# can pass; describing the run as anything else would be a false disclosure.
+: "${MLLOG_LOWEST_NUMERICAL_PRECISION_IN_LINEAR:=fp8}"
+: "${MLLOG_LOWEST_NUMERICAL_PRECISION_IN_ATTN:=bfloat16}"
+: "${MLLOG_LOWEST_NUMERICAL_PRECISION_IN_COMM:=bfloat16}"
+export MLLOG_LOWEST_NUMERICAL_PRECISION_IN_LINEAR
+export MLLOG_LOWEST_NUMERICAL_PRECISION_IN_ATTN
+export MLLOG_LOWEST_NUMERICAL_PRECISION_IN_COMM
+
+# --- This run ---------------------------------------------------------------
+: "${RESULTS_DIR:=/results}"
+: "${RUN_INDEX:=0}"
+: "${PRIMUS_SEED:=$((42 + RUN_INDEX))}"
+: "${MLLOG_OUTPUT_FILE:=${RESULTS_DIR}/result_${RUN_INDEX}.txt}"
+export RESULTS_DIR RUN_INDEX PRIMUS_SEED MLLOG_OUTPUT_FILE
+
+mkdir -p "${RESULTS_DIR}"
+
+# --- Cold start -------------------------------------------------------------
+# cache_clear is a claim about the machine, so drop the caches here and report
+# what actually happened rather than what was requested. Dropping caches needs
+# privileges the container may not have; a run that could not do it says so.
+: "${MLPERF_CLEAR_CACHES:=true}"
+if [[ "${MLPERF_CLEAR_CACHES}" == "true" ]]; then
+    if sync && echo 3 > /proc/sys/vm/drop_caches 2>/dev/null; then
+        echo "[MLPerf] Dropped page cache"
+    else
+        echo "[MLPerf] WARNING: could not drop page cache; reporting cache_clear=false"
+        MLPERF_CLEAR_CACHES=false
+    fi
+fi
+export MLPERF_CLEAR_CACHES
+
+echo "============================================"
+echo "MLPerf FLUX.1-Schnell Training (Megatron)"
+echo "============================================"
+echo "Recipe:    ${EXP}"
+echo "Run index: ${RUN_INDEX}"
+echo "Seed:      ${PRIMUS_SEED}"
+echo "Result:    ${MLLOG_OUTPUT_FILE}"
+echo "Division:  ${MLLOG_SUBMISSION_DIVISION} (${MLLOG_SUBMISSION_ORG} / ${MLLOG_SUBMISSION_PLATFORM})"
+echo "Precision: linear=${MLLOG_LOWEST_NUMERICAL_PRECISION_IN_LINEAR}" \
+     "attn=${MLLOG_LOWEST_NUMERICAL_PRECISION_IN_ATTN}" \
+     "comm=${MLLOG_LOWEST_NUMERICAL_PRECISION_IN_COMM}"
+echo "============================================"
+
+start=$(date +%s)
+start_fmt=$(date +%Y-%m-%d\ %r)
+echo "STARTING TIMING RUN AT ${start_fmt}"
+
+set +e
+"${PRIMUS_PATH}/primus-cli" direct -- \
+    train pretrain \
+    --config "${EXP}" \
+    2>&1 | tee "${RESULTS_DIR}/train_flux1_${RUN_INDEX}.log"
+ret_code=${PIPESTATUS[0]}
+set -e
+
+end=$(date +%s)
+end_fmt=$(date +%Y-%m-%d\ %r)
+echo "ENDING TIMING RUN AT ${end_fmt}"
+
+result=$(( end - start ))
+echo "RESULT,FLUX1,${PRIMUS_SEED},${result},${MLLOG_SUBMISSION_ORG},${start_fmt}"
+
+if [[ ${ret_code} != 0 ]]; then
+    echo "Training failed with exit code: ${ret_code}"
+    exit "${ret_code}"
+fi
+
+# --- Check the artifact before it is treated as a result --------------------
+# Finding out at submission time that a week of runs produced unparseable logs
+# is the failure this guards against.
+: "${MLPERF_RULESET:=6.0.0}"
+: "${CHECK_COMPLIANCE:=1}"
+if [[ "${CHECK_COMPLIANCE}" == "1" ]]; then
+    python3 -m mlperf_logging.compliance_checker \
+        --usage training \
+        --ruleset "${MLPERF_RULESET}" \
+        --log_output "${RESULTS_DIR}/compliance_${RUN_INDEX}.out" \
+        "${MLLOG_OUTPUT_FILE}" \
+        || echo "[MLPerf] WARNING: compliance check failed for run ${RUN_INDEX}"
+fi
+
+exit 0

examples/mlperf/flux1/megatron/run_campaign.sh

@@ -0,0 +1,57 @@
+#!/bin/bash
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+#
+# Run the full flux1 submission campaign: ten timed runs, each with its own
+# seed, then check the collected logs as a set.
+#
+# Ten is not a convention -- mlperf_logging/rcp_checker/rcp_checker.py maps
+# flux1 to 10 runs, and the RCP comparison is made over that many results. A
+# run that ends without converging still produces a result file and still
+# counts as one of the ten; dropping it would bias the set.
+
+set -euo pipefail
+
+: "${PRIMUS_PATH:=$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../.." && pwd)}"
+: "${RESULTS_DIR:=/results}"
+: "${NUM_RUNS:=10}"
+: "${SEED_BASE:=42}"
+: "${MLPERF_RULESET:=6.0.0}"
+export PRIMUS_PATH RESULTS_DIR
+
+mkdir -p "${RESULTS_DIR}"
+
+failed_runs=()
+for (( index = 0; index < NUM_RUNS; index++ )); do
+    echo
+    echo "########## flux1 run ${index} of ${NUM_RUNS} ##########"
+    # Each run is checked on its own inside run_and_time.sh; a failure here is
+    # recorded and the campaign continues, because nine good runs plus a
+    # diagnosis beats stopping on the first bad one after hours of compute.
+    if ! RUN_INDEX="${index}" \
+         PRIMUS_SEED="$(( SEED_BASE + index ))" \
+         bash "${PRIMUS_PATH}/examples/mlperf/flux1/megatron/run_and_time.sh"; then
+        failed_runs+=("${index}")
+    fi
+done
+
+echo
+echo "########## campaign summary ##########"
+if (( ${#failed_runs[@]} > 0 )); then
+    echo "Runs that exited non-zero: ${failed_runs[*]}"
+else
+    echo "All ${NUM_RUNS} runs completed."
+fi
+
+echo
+echo "Comparing the collected results against the reference convergence points:"
+python3 -m mlperf_logging.rcp_checker \
+    --rcp_usage training \
+    --rcp_version "${MLPERF_RULESET}" \
+    --log_output "${RESULTS_DIR}/rcp_checker.out" \
+    --verbose \
+    "${RESULTS_DIR}" \
+    || echo "[MLPerf] WARNING: RCP comparison failed; see the output above"

examples/mlperf/flux1/requirements.txt

@@ -1,4 +1,4 @@
 datasets>=3.6.0
 webdataset==1.0.2
-git+https://github.com/mlcommons/logging.git@6.0.0-rc5
+git+https://github.com/mlcommons/logging.git@6.0.0-rc6
 git+https://github.com/NVIDIA/mlperf-common.git@b86d175a05849d650a8ff69c1e2c37b9f4e61d51