diff --git a/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8_mlperf.yaml b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8_mlperf.yaml index d67662f76..d933c49cf 100644 --- a/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8_mlperf.yaml +++ b/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 diff --git a/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8_mlperf.yaml b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8_mlperf.yaml index cd01468c8..a1f07de06 100644 --- a/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8_mlperf.yaml +++ b/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). diff --git a/examples/megatron/prepare.py b/examples/megatron/prepare.py index 86725f311..ef5158c41 100644 --- a/examples/megatron/prepare.py +++ b/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 diff --git a/examples/mlperf/flux1/megatron/README.md b/examples/mlperf/flux1/megatron/README.md new file mode 100644 index 000000000..48a668a90 --- /dev/null +++ b/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. diff --git a/examples/mlperf/flux1/megatron/run_and_time.sh b/examples/mlperf/flux1/megatron/run_and_time.sh new file mode 100755 index 000000000..ad0e1c0b0 --- /dev/null +++ b/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 diff --git a/examples/mlperf/flux1/megatron/run_campaign.sh b/examples/mlperf/flux1/megatron/run_campaign.sh new file mode 100755 index 000000000..7ef2dd372 --- /dev/null +++ b/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" diff --git a/examples/mlperf/flux1/requirements.txt b/examples/mlperf/flux1/requirements.txt index 4477174d6..f58617fab 100644 --- a/examples/mlperf/flux1/requirements.txt +++ b/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 diff --git a/primus/backends/megatron/data/diffusion/preprocessing/download.py b/primus/backends/megatron/data/diffusion/preprocessing/download.py index c58c2814e..a1ec9ff5b 100644 --- a/primus/backends/megatron/data/diffusion/preprocessing/download.py +++ b/primus/backends/megatron/data/diffusion/preprocessing/download.py @@ -139,18 +139,28 @@ def parse_md5_manifest( return entries -def fetch_manifest(manifest_url: str) -> Tuple[str, List[Tuple[str, str]]]: +def fetch_manifest( + manifest_url: str, + suffix_filter: Optional[str] = None, +) -> Tuple[str, List[Tuple[str, str]]]: """Fetch .uri and .md5 manifests, return (base_url, [(md5, filename)]). The manifest_url should end with '.uri' or '.md5'. The function derives the complementary URL by replacing the suffix. + + Args: + manifest_url: URL of either manifest. + suffix_filter: Keep only files with this suffix (e.g. ".arrow"). + MLCommons manifests list dataset metadata alongside the data, and + a caller that consumes one kind of file wants the other kind gone + before anything counts, indexes or numbers the entries. """ uri_url = manifest_url.replace(".md5", ".uri") md5_url = manifest_url.replace(".uri", ".md5") base_url = fetch_url_text(uri_url).strip() md5_text = fetch_url_text(md5_url) - entries = parse_md5_manifest(md5_text) + entries = parse_md5_manifest(md5_text, suffix_filter=suffix_filter) logger.info(f"Manifest: {len(entries)} files, base URL: {base_url}") return base_url, entries diff --git a/primus/backends/megatron/data/diffusion/preprocessing/pipelines/ingest.py b/primus/backends/megatron/data/diffusion/preprocessing/pipelines/ingest.py index 5c68bc572..60000f5c8 100644 --- a/primus/backends/megatron/data/diffusion/preprocessing/pipelines/ingest.py +++ b/primus/backends/megatron/data/diffusion/preprocessing/pipelines/ingest.py @@ -48,6 +48,11 @@ def _arrow_to_tar( ``__key__`` column from the Arrow file when available, otherwise generates sequential keys. + The val split additionally carries a ``timestep`` column, which MLPerf + validation is defined in terms of (sigma = t / 8). It is copied into the + JSON sidecar rather than ``ARROW_COLUMNS`` because it is scalar metadata, + not a tensor entry. The train split has no such column and is unaffected. + Returns the number of samples written. """ reader = pyarrow.ipc.open_stream(str(arrow_path)) @@ -55,6 +60,7 @@ def _arrow_to_tar( num_rows = table.num_rows has_key_col = "__key__" in table.schema.names + timestep_col = table.column("timestep") if "timestep" in table.schema.names else None tar_path.parent.mkdir(parents=True, exist_ok=True) with tarfile.open(str(tar_path), "w") as tar: @@ -76,7 +82,13 @@ def _arrow_to_tar( info.size = len(data) tar.addfile(info, io.BytesIO(data)) - meta = json.dumps({"key": base_name}).encode("utf-8") + metadata = {"key": base_name} + if timestep_col is not None: + timestep = timestep_col[row_idx].as_py() + if timestep is not None: + metadata["timestep"] = timestep + + meta = json.dumps(metadata).encode("utf-8") meta_info = tarfile.TarInfo(name=f"{base_name}.json") meta_info.size = len(meta) tar.addfile(meta_info, io.BytesIO(meta)) @@ -128,7 +140,12 @@ def run(self, **kwargs) -> Dict[str, int]: Dict with 'files_processed', 'samples_written', 'shards_created', 'shards_skipped', 'files_failed'. """ - base_url, entries = fetch_manifest(self.manifest_url) + # Drop the manifest's metadata files (dataset_info.json, state.json) + # before anything counts or indexes the entries. They are not Arrow, so + # conversion fails on them and reports the run as having failed files; + # they would also consume max_files slots and, if one ever sorted ahead + # of a data file, shift every shard index after it. + base_url, entries = fetch_manifest(self.manifest_url, suffix_filter=".arrow") if self.max_files is not None: entries = entries[: self.max_files] diff --git a/primus/backends/megatron/data/diffusion/task_encoders/image.py b/primus/backends/megatron/data/diffusion/task_encoders/image.py index 837bbcfdb..55778a0ef 100644 --- a/primus/backends/megatron/data/diffusion/task_encoders/image.py +++ b/primus/backends/megatron/data/diffusion/task_encoders/image.py @@ -32,6 +32,10 @@ logger = logging.getLogger(__name__) +# MLPerf Flux validation evaluates each sample at a fixed timestep drawn from +# {0/8, ..., 7/8}; forward_step turns the stored integer into sigma = t / 8. +NUM_VALIDATION_TIMESTEPS = 8 + # ============================================================================ # Sample Definition (with proper Sample inheritance) @@ -68,6 +72,48 @@ class DiffusionSample(Sample): timestep: Optional[torch.Tensor] = None +def _collate_timesteps(samples) -> Optional[torch.Tensor]: + """Stack per-sample validation timesteps, validating every sample. + + Checking only ``samples[0]`` would let a partially-ingested shard drop the + field for a whole batch, silently downgrading evaluation to the positional + fallback in forward_step, or crash inside ``torch.stack`` with no + indication of which sample was at fault. + + Returns None when no sample carries a timestep (the training split). + """ + present = [s.timestep is not None for s in samples] + if not any(present): + return None + + if not all(present): + missing = [s.__key__ for s, ok in zip(samples, present) if not ok] + raise ValueError( + f"{len(missing)} of {len(samples)} samples lack a 'timestep' field: " + f"{missing[:8]}{' ...' if len(missing) > 8 else ''}. The batch mixes " + f"samples from shards ingested with and without the timestep column; " + f"re-ingest the validation split so every sidecar carries it." + ) + + stacked = torch.stack([s.timestep for s in samples]) + + if stacked.dtype.is_floating_point or stacked.dtype.is_complex: + raise ValueError(f"'timestep' must be an integer type, got {stacked.dtype}.") + + out_of_range = (stacked < 0) | (stacked >= NUM_VALIDATION_TIMESTEPS) + if bool(out_of_range.any()): + offenders = [ + (s.__key__, int(t)) for s, bad, t in zip(samples, out_of_range.tolist(), stacked.tolist()) if bad + ] + raise ValueError( + f"'timestep' must be in [0, {NUM_VALIDATION_TIMESTEPS - 1}]; " + f"{len(offenders)} sample(s) out of range: {offenders[:8]}" + f"{' ...' if len(offenders) > 8 else ''}." + ) + + return stacked + + # ============================================================================ # Cooker Functions # ============================================================================ @@ -328,8 +374,9 @@ def batch(self, samples: List[DiffusionSample]) -> Dict[str, torch.Tensor]: batch["mean"] = torch.stack([s.mean for s in samples]) batch["logvar"] = torch.stack([s.logvar for s in samples]) - if samples[0].timestep is not None: - batch["timestep"] = torch.stack([s.timestep for s in samples]) + timesteps = _collate_timesteps(samples) + if timesteps is not None: + batch["timestep"] = timesteps return batch diff --git a/primus/backends/megatron/data/energon_dataset_provider.py b/primus/backends/megatron/data/energon_dataset_provider.py index 4ac1728dc..cf3bdf704 100644 --- a/primus/backends/megatron/data/energon_dataset_provider.py +++ b/primus/backends/megatron/data/energon_dataset_provider.py @@ -14,7 +14,6 @@ from typing import Any, Callable, List, Optional, Tuple from megatron.core import parallel_state -from megatron.core.num_microbatches_calculator import get_num_microbatches from megatron.core.parallel_state import ( get_pipeline_model_parallel_rank, get_pipeline_model_parallel_world_size, @@ -32,6 +31,13 @@ from primus.backends.megatron.data.dataloader import MegatronDataloaderWrapper from primus.backends.megatron.data.dataset_provider import DatasetProvider +from primus.backends.megatron.training.eval_budget import ( + EvalCoverageError, + assert_mlperf_timestep_source, + assert_val_worker_divisibility, + get_eval_num_microbatches, + get_val_num_workers, +) from primus.core.utils.module_utils import log_rank_0 @@ -93,53 +99,96 @@ def create_dataloaders( task_encoder = self.task_encoder_factory() log_rank_0(f"Created task encoder: {type(task_encoder).__name__}") - # Create worker config for distributed loading + # Create worker config for distributed loading. Validation gets its own, + # because the worker count decides how many samples an eval actually + # reads (see eval_budget) and the training value is rarely a safe one. worker_config = self._create_worker_config(args) + val_worker_config = self._create_worker_config(args, num_workers=get_val_num_workers(args)) # Get data path data_path = self._get_data_path(args) - # Create training dataset using Energon - log_rank_0(f"Creating training dataset from: {data_path}") - train_dataset = get_train_dataset( - data_path, - batch_size=args.micro_batch_size, - task_encoder=task_encoder, - worker_config=worker_config, - virtual_epoch_length=getattr(args, "virtual_epoch_length", 1_000_000_000), - max_samples_per_sequence=getattr(args, "max_samples_per_sequence", 100), - shuffle_buffer_size=getattr(args, "shuffle_buffer_size", None), - handler=lambda *args: None, # Error handler (print errors but continue) - ) - - # Wrap in savable loader for checkpointing support prefetch_factor = getattr(args, "prefetch_factor", 2) log_rank_0(f"Dataloader prefetch_factor: {prefetch_factor}") - train_dataloader = get_savable_loader( - train_dataset, worker_config=worker_config, prefetch_factor=prefetch_factor - ) - train_dataloader = MegatronDataloaderWrapper(train_dataloader) - log_rank_0("Created training dataloader") + + # Megatron drops the train iterator entirely under --skip-train, so + # building one is wasted work. It is also fatal for a dataset that + # holds only a validation split, which is a legitimate shape for an + # evaluation-only run. + if getattr(args, "skip_train", False): + train_dataloader = None + log_rank_0("skip_train is set: not creating a training dataset") + else: + log_rank_0(f"Creating training dataset from: {data_path}") + train_dataset = get_train_dataset( + data_path, + batch_size=args.micro_batch_size, + task_encoder=task_encoder, + worker_config=worker_config, + virtual_epoch_length=getattr(args, "virtual_epoch_length", 1_000_000_000), + max_samples_per_sequence=getattr(args, "max_samples_per_sequence", 100), + shuffle_buffer_size=getattr(args, "shuffle_buffer_size", None), + handler=lambda *args: None, # Error handler (print errors but continue) + ) + + # Wrap in savable loader for checkpointing support + train_dataloader = get_savable_loader( + train_dataset, worker_config=worker_config, prefetch_factor=prefetch_factor + ) + train_dataloader = MegatronDataloaderWrapper(train_dataloader) + log_rank_0("Created training dataloader") + + # The patch that turns eval_samples into eval_iters runs in build_args, + # where a failure is logged and swallowed rather than raised. If it did + # not take effect, eval_iters is still 0 and the job would run to + # completion, exit 0, and report no validation at all -- while the + # recipe plainly asked for a specific number of samples. Refuse that. + if getattr(args, "eval_samples", None) and not args.eval_iters: + raise EvalCoverageError( + f"eval_samples={args.eval_samples} is configured but eval_iters is 0, " + f"so no evaluation would run. The megatron.args.eval_samples patch " + f"that derives one from the other did not take effect; look for its " + f"failure in the build_args phase of the log." + ) # Create validation dataloaders if evaluation is enabled valid_dataloaders = None if args.eval_iters > 0: + # Assert before construction: a shape that cannot read every sample + # should fail here rather than silently report a short evaluation. + eval_num_microbatches = get_eval_num_microbatches(args) + eval_samples = ( + args.eval_iters + * eval_num_microbatches + * args.micro_batch_size + * (parallel_state.get_data_parallel_world_size()) + ) + assert_val_worker_divisibility(args, eval_samples) + # Reading every sample is not enough if they are read at timesteps + # the split never paired them with. + assert_mlperf_timestep_source(args) + log_rank_0( + f"Validation budget: {args.eval_iters} iterations x " + f"{eval_num_microbatches} microbatches x {args.micro_batch_size} " + f"= {eval_samples} samples, val_num_workers={get_val_num_workers(args)}" + ) + try: log_rank_0("Creating validation dataloaders...") val_datasets = get_val_datasets( data_path, batch_size=args.micro_batch_size, task_encoder=task_encoder, - worker_config=worker_config, + worker_config=val_worker_config, handler=lambda *args: None, ) # Limit validation datasets to eval_iters * num_microbatches val_datasets_limited = [ LimitDataset( - RepeatDataset(val_ds, worker_config=worker_config), - length=args.eval_iters * get_num_microbatches(), - worker_config=worker_config, + RepeatDataset(val_ds, worker_config=val_worker_config), + length=args.eval_iters * eval_num_microbatches, + worker_config=val_worker_config, reset_after_epoch=True, ) for val_ds, _src_ds in val_datasets @@ -147,7 +196,7 @@ def create_dataloaders( valid_dataloaders = [ MegatronDataloaderWrapper( - get_loader(valid_ds, worker_config=worker_config, prefetch_factor=prefetch_factor) + get_loader(valid_ds, worker_config=val_worker_config, prefetch_factor=prefetch_factor) ) for valid_ds in val_datasets_limited ] @@ -209,20 +258,28 @@ def _is_dataloader_rank(self) -> bool: return is_first_tp_rank and is_valid_pp_stage - def _create_worker_config(self, args) -> WorkerConfig: + def _create_worker_config(self, args, num_workers: Optional[int] = None) -> WorkerConfig: """ Create Energon WorkerConfig for distributed loading. WorkerConfig tells Energon how to shard data across workers. + + Args: + num_workers: Override the worker count. Validation passes its own so + it does not inherit the training value, which controls how many + samples an evaluation reads (see eval_budget). """ rank = parallel_state.get_data_parallel_rank() world_size = parallel_state.get_data_parallel_world_size() data_parallel_group = parallel_state.get_data_parallel_group() + if num_workers is None: + num_workers = getattr(args, "num_workers", 4) + return WorkerConfig( rank=rank, world_size=world_size, - num_workers=getattr(args, "num_workers", 4), + num_workers=num_workers, data_parallel_group=data_parallel_group, ) diff --git a/primus/backends/megatron/diffusion_trainer.py b/primus/backends/megatron/diffusion_trainer.py index dd91f3285..ffa71efc1 100644 --- a/primus/backends/megatron/diffusion_trainer.py +++ b/primus/backends/megatron/diffusion_trainer.py @@ -50,6 +50,10 @@ def __init__(self, *args, **kwargs): self._compiled_loss_fn = None self._forward_step_count = 0 self._forward_step_count_initialized = False + # None rather than 0: the session counter also starts at 0, and a run + # that never opens a session must still reset on its first evaluation. + self._eval_session = None + self._eval_microbatch_index = 0 # Composition pattern: avoids recreating the provider on each call use_mock_data = getattr(self.backend_args, "mock_data", False) @@ -225,6 +229,7 @@ def forward_step(self, data_iterator, model, return_schedule_plan=False): Tuple of (noise_pred, loss_func_callable) """ from primus.backends.megatron.training.diffusion.forward_step import ( + EQUIDISTANT_TIMESTEPS, flux_forward_step_func, ) @@ -232,10 +237,21 @@ def forward_step(self, data_iterator, model, return_schedule_plan=False): # validation steps would shift the next training step's per-step seed # by eval_iters * num_microbatches per --eval-interval window, # defeating the goal of isolating training RNG from unrelated forward - # passes. Eval forward passes reuse the most recent training counter - # value, so the per-step CUDA reseed is a no-op replay during eval. + # passes. + # + # Validation instead gets its own advancing index, because reusing the + # frozen training counter reseeds every eval microbatch identically and + # so repeats one draw of VAE epsilon and flow noise across the whole + # evaluation. + per_step_rng_reseed = getattr(self, "per_step_rng_reseed", False) + eval_step_index = None if model.training: self._forward_step_count += 1 + elif per_step_rng_reseed: + # Only derived when reseeding will consume it; without reseeding the + # ambient generator already advances per batch, which is the + # behaviour both references have. + eval_step_index = self._next_eval_step_index() # Megatron's pattern: forward_step returns model output, loss_func computes loss noise_pred, clean_latents, noise, loss_mask, metrics, is_validation = flux_forward_step_func( @@ -251,8 +267,10 @@ def forward_step(self, data_iterator, model, return_schedule_plan=False): vae_scale=getattr(self, "vae_scale", None), vae_shift=getattr(self, "vae_shift", None), vae_latent_mode=getattr(self, "vae_latent_mode", "presampled"), - per_step_rng_reseed=getattr(self, "per_step_rng_reseed", False), + per_step_rng_reseed=per_step_rng_reseed, step_count=self._forward_step_count, + eval_step_index=eval_step_index, + eval_timestep_source=getattr(self, "eval_timestep_source", EQUIDISTANT_TIMESTEPS), ) # Store values needed for loss computation (will be used by loss function) @@ -296,6 +314,60 @@ def diffusion_loss_func(output_tensor, non_loss_data=False): return noise_pred, diffusion_loss_func + def _next_eval_step_index(self) -> int: + """Index identifying this validation microbatch within the run. + + Built from ``(training step under evaluation, microbatch index within + this evaluation)``, so it is a pure function of where the run is and + reproduces after a checkpoint resume without checkpointing a counter. + + The step is ``args.curr_iteration``, not ``args.iteration``. + ``args.iteration`` is the resume point: Megatron assigns it at setup + and at checkpoint load and never inside the training loop, so keying on + it made the index depend on where the run *started*. An evaluation at + step k then drew one set of noise when reached continuously and a + different set when reached after a resume. + + ``curr_iteration`` is the last step the loop completed, so it is one + below the step the evaluation reports. The offset is constant, and + constant is all reproducibility needs, but it does mean this index + cannot be lined up directly against the step in an MLLOG record. + """ + from megatron.training import get_args + + from primus.backends.megatron.training.diffusion.forward_step import ( + EVAL_RNG_ITERATION_STRIDE, + ) + from primus.backends.megatron.training.eval_session import current_eval_session + + args = get_args() + # Unset until the loop runs its first step, so an evaluation that + # precedes training (--skip-train) keys on the checkpoint step, which + # is the step under evaluation in that case. + step = getattr(args, "curr_iteration", None) + if step is None: + step = getattr(args, "iteration", 0) + + # Reset per evaluation rather than per step: the in-loop evaluation + # that ends a run and pretrain's post-training one share a step, and + # the second has to reproduce the first rather than continue it. + session = current_eval_session() + if session != self._eval_session: + self._eval_session = session + self._eval_microbatch_index = 0 + + index = self._eval_microbatch_index + self._eval_microbatch_index += 1 + + if index >= EVAL_RNG_ITERATION_STRIDE: + raise RuntimeError( + f"Evaluation ran {index + 1} microbatches, at or beyond the " + f"stride {EVAL_RNG_ITERATION_STRIDE} that keeps consecutive " + f"training steps' evaluation RNG streams disjoint." + ) + + return step * EVAL_RNG_ITERATION_STRIDE + index + def get_forward_step(self): """ Return forward step function for diffusion models. diff --git a/primus/backends/megatron/flux_pretrain_trainer.py b/primus/backends/megatron/flux_pretrain_trainer.py index 602a21d85..1b41ba896 100644 --- a/primus/backends/megatron/flux_pretrain_trainer.py +++ b/primus/backends/megatron/flux_pretrain_trainer.py @@ -20,6 +20,10 @@ import torch.nn as nn from primus.backends.megatron.diffusion_trainer import DiffusionPretrainTrainer +from primus.backends.megatron.training.diffusion.forward_step import ( + EQUIDISTANT_TIMESTEPS, + EVAL_TIMESTEP_SOURCES, +) from primus.backends.megatron.training.diffusion.schedulers import ( FlowMatchEulerDiscreteScheduler, ) @@ -151,6 +155,16 @@ def __init__(self, *args, **kwargs): else: log_rank_0(f"VAE latent mode: presampled (stored latents used directly)") + # Validated here rather than at the first forward so a typo fails at + # startup instead of after the first eval_interval steps. + self.eval_timestep_source = getattr(params, "eval_timestep_source", EQUIDISTANT_TIMESTEPS) + if self.eval_timestep_source not in EVAL_TIMESTEP_SOURCES: + raise ValueError( + f"eval_timestep_source must be one of {list(EVAL_TIMESTEP_SOURCES)}, " + f"got '{self.eval_timestep_source}'" + ) + log_rank_0(f"Validation timestep source: {self.eval_timestep_source}") + log_rank_0(f"Guidance embedding: {self.use_guidance_embed}") log_rank_0(f"Scheduler shift: {self.scheduler_shift}") log_rank_0(f"Dynamic shifting: {self.use_dynamic_shifting}") diff --git a/primus/backends/megatron/patches/args/__init__.py b/primus/backends/megatron/patches/args/__init__.py index 647847380..d4d11dd03 100644 --- a/primus/backends/megatron/patches/args/__init__.py +++ b/primus/backends/megatron/patches/args/__init__.py @@ -19,6 +19,7 @@ from . import ( # noqa: F401 checkpoint_path_patches, data_path_split_patches, + eval_samples_patches, hsdp_args_patches, iterations_to_skip_default_patches, logging_level_patches, @@ -36,6 +37,7 @@ "wandb_config_patches", "logging_level_patches", "data_path_split_patches", + "eval_samples_patches", "hsdp_args_patches", "mock_data_patches", "sequence_parallel_tp1_patches", diff --git a/primus/backends/megatron/patches/args/eval_samples_patches.py b/primus/backends/megatron/patches/args/eval_samples_patches.py new file mode 100644 index 000000000..0763ae68e --- /dev/null +++ b/primus/backends/megatron/patches/args/eval_samples_patches.py @@ -0,0 +1,102 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Derive the evaluation budget from a sample count instead of an iteration count. + +``eval_iters`` only expresses "cover the validation set" at one particular +global batch size, so it silently becomes wrong when the batch size changes. +``eval_samples`` says what the evaluation is actually meant to measure. + +This runs in the ``build_args`` phase, before ``pretrain()`` builds the Energon +validation dataloader, so both the provider and the evaluation loop observe the +corrected ``eval_iters``. +""" + +from primus.backends.megatron.training.eval_budget import ( + assert_val_worker_divisibility, + get_val_num_workers, + read_energon_split_sample_count, + resolve_eval_iters, +) +from primus.core.patches import PatchContext, register_patch +from primus.core.utils.module_utils import log_kv_rank_0 + +# Keys Megatron's parser does not define. MegatronArgBuilder drops anything +# absent from that parser ("Non-Megatron parameters are silently ignored"), and +# the runtime merges the leftover Primus-only params into backend_args only +# *after* the build_args phase. So this patch cannot read them off args yet and +# must take them from the module config, or it would see every one as unset and +# quietly leave eval_iters alone. +PRIMUS_ONLY_EVAL_KEYS = ("eval_samples", "val_num_workers") + + +def _hydrate_primus_only_keys(args, module_config): + """Copy Primus-only eval keys from the module config onto args.""" + params = getattr(module_config, "params", None) + if params is None: + return + for key in PRIMUS_ONLY_EVAL_KEYS: + if not hasattr(args, key) and hasattr(params, key): + setattr(args, key, getattr(params, key)) + + +@register_patch( + "megatron.args.eval_samples", + backend="megatron", + phase="build_args", + description="Derive eval_iters from eval_samples and assert validation coverage is exact", +) +def patch_eval_samples(ctx: PatchContext): + args = ctx.extra.get("backend_args", {}) + if not args: + return + + _hydrate_primus_only_keys(args, ctx.extra.get("module_config")) + + # full_validation has never done anything on the Energon path: the + # validation LimitDataset is always sized from eval_iters. Give it a + # meaning by reading the split's true size out of the dataset index. + if getattr(args, "full_validation", False) and getattr(args, "eval_samples", None) is None: + dataset_samples = read_energon_split_sample_count(getattr(args, "data_path", None)) + if dataset_samples is None: + raise ValueError( + "full_validation is set but the validation split size could not be read " + "from the dataset index (.nv-meta/.info.json). Set eval_samples explicitly." + ) + args.eval_samples = dataset_samples + args.eval_iters = 0 + log_kv_rank_0( + "[Patch:megatron.args.eval_samples] -full_validation", + f"eval_samples={dataset_samples} (read from dataset index)", + ) + + derived = resolve_eval_iters(args) + if derived is not None: + previous = getattr(args, "eval_iters", None) + args.eval_iters = derived + # Name the value being replaced: eval_iters is set in trainer_base.yaml + # for every module, so this almost always overrides something, and a + # silent override is how the budget drifted from the intent before. + replaced = "" if previous in (None, derived) else f", was {previous}" + log_kv_rank_0( + "[Patch:megatron.args.eval_samples] -eval_iters", + f"{derived} (from eval_samples={args.eval_samples}, " + f"global_batch_size={args.global_batch_size}{replaced})", + ) + return + + # No eval_samples given: leave eval_iters alone, but still check that the + # budget it implies can actually be read. Without this, only configs that + # opt into eval_samples get the coverage guarantee. + eval_iters = getattr(args, "eval_iters", 0) or 0 + if eval_iters > 0: + assert_val_worker_divisibility(args, eval_iters * args.global_batch_size) + log_kv_rank_0( + "[Patch:megatron.args.eval_samples] -val_num_workers", + f"{get_val_num_workers(args)} (coverage verified for " + f"{eval_iters * args.global_batch_size} samples)", + ) diff --git a/primus/backends/megatron/patches/mlperf_boundary.py b/primus/backends/megatron/patches/mlperf_boundary.py new file mode 100644 index 000000000..6cb89704e --- /dev/null +++ b/primus/backends/megatron/patches/mlperf_boundary.py @@ -0,0 +1,160 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +The single transition between MLPerf initialization and measured training. + +MLPerf requires the clock to start before the implementation touches the +dataset. Megatron's ``pretrain()`` runs ``setup_model_and_optimizer``, then +builds the data iterators, then calls ``train()``, and Primus' patch phases +(``setup`` / ``build_args`` / ``before_train`` / ``after_train``) all fire +before ``pretrain()`` -- there is no phase between model build and data build. +Megatron-LM is a pinned submodule, so a new phase cannot be inserted at the +call site either. + +This module therefore creates the missing seam by wrapping three Megatron +entry points from the ``before_train`` phase: + + ``pretrain`` capture ``forward_step_func`` + ``setup_model_and_optimizer`` capture model / optimizer / scheduler + ``build_train_valid_test_data_iterators`` fire the transition on entry + +On the first call into the data-iterator builder -- before any shard, worker +or prefetcher exists -- registered pre-run hooks execute (compile warmup), +all ranks synchronize, rank 0 emits ``init_stop`` / ``run_start``, and a +second barrier holds every rank until those records exist. + +Hooks are ordered so warmup always precedes the transition; the transition +itself is a separate registration so the logging patch owns what it emits. +""" + +import logging + +logger = logging.getLogger(__name__) + +_ORIGINALS: dict = {} +_HOOKS: list = [] +_TRANSITION: list = [] +_CAPTURED: dict = {} +_FIRED = [False] + + +def register_pre_run_hook(name: str, fn, order: int = 50) -> None: + """Register work that must finish before the clock starts. + + Hooks run in ascending ``order`` on every rank, inside the initialization + window, with the captured model/optimizer available via :func:`captured`. + """ + _HOOKS.append((order, name, fn)) + _HOOKS.sort(key=lambda entry: entry[0]) + + +def set_transition(fn) -> None: + """Register the callable that emits ``init_stop`` / ``run_start``.""" + _TRANSITION.clear() + _TRANSITION.append(fn) + + +def captured() -> dict: + """Objects captured from Megatron on the way to the boundary. + + Keys are present only once the corresponding entry point has run: + ``forward_step_func``, ``model``, ``optimizer``, ``opt_param_scheduler``. + """ + return _CAPTURED + + +def has_fired() -> bool: + return _FIRED[0] + + +def reset_for_tests() -> None: + """Drop all registrations and captures. Tests only.""" + _ORIGINALS.clear() + _HOOKS.clear() + _TRANSITION.clear() + _CAPTURED.clear() + _FIRED[0] = False + + +def fire() -> None: + """Run the pre-run hooks, synchronize, emit the transition, synchronize.""" + if _FIRED[0]: + return + _FIRED[0] = True + + for _order, name, fn in _HOOKS: + logger.debug("MLPerf boundary: running pre-run hook %s", name) + fn() + + _synchronize() + + for fn in _TRANSITION: + fn() + + # Hold every rank until the records exist, so no rank can begin opening + # the dataset while rank 0 is still writing run_start. + _synchronize() + + +def _synchronize() -> None: + import torch + import torch.distributed + + if torch.cuda.is_available(): + torch.cuda.synchronize() + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.barrier() + + +def install() -> None: + """Wrap the three Megatron entry points. Idempotent.""" + if _ORIGINALS: + return + + import megatron.training as megatron_training_pkg + import megatron.training.training as mt + + _ORIGINALS["pretrain"] = getattr(megatron_training_pkg, "pretrain", None) + _ORIGINALS["setup_model_and_optimizer"] = mt.setup_model_and_optimizer + _ORIGINALS["build_train_valid_test_data_iterators"] = mt.build_train_valid_test_data_iterators + + # forward_step_func is a pretrain() argument and is reachable nowhere else + # before train(). The trainer imports pretrain inside its run method, after + # this phase, so rebinding the package attribute is picked up. + if _ORIGINALS["pretrain"] is not None: + original_pretrain = _ORIGINALS["pretrain"] + + def _capturing_pretrain(*args, **kwargs): + if len(args) > 3: + _CAPTURED["forward_step_func"] = args[3] + elif "forward_step_func" in kwargs: + _CAPTURED["forward_step_func"] = kwargs["forward_step_func"] + return original_pretrain(*args, **kwargs) + + megatron_training_pkg.pretrain = _capturing_pretrain + + original_setup = _ORIGINALS["setup_model_and_optimizer"] + + def _capturing_setup_model_and_optimizer(*args, **kwargs): + model, optimizer, opt_param_scheduler = original_setup(*args, **kwargs) + _CAPTURED["model"] = model + _CAPTURED["optimizer"] = optimizer + _CAPTURED["opt_param_scheduler"] = opt_param_scheduler + return model, optimizer, opt_param_scheduler + + mt.setup_model_and_optimizer = _capturing_setup_model_and_optimizer + + original_build = _ORIGINALS["build_train_valid_test_data_iterators"] + + def _boundary_build_data_iterators(*args, **kwargs): + # Virtual pipelining calls this once per stage; only the first call is + # the boundary, and fire() is idempotent regardless. + fire() + return original_build(*args, **kwargs) + + _boundary_build_data_iterators._primus_mlperf_boundary = True + mt.build_train_valid_test_data_iterators = _boundary_build_data_iterators diff --git a/primus/backends/megatron/patches/mlperf_logging_patches.py b/primus/backends/megatron/patches/mlperf_logging_patches.py index b53173d8a..f29415cad 100644 --- a/primus/backends/megatron/patches/mlperf_logging_patches.py +++ b/primus/backends/megatron/patches/mlperf_logging_patches.py @@ -25,6 +25,88 @@ logger = logging.getLogger(__name__) +_PRECISION_DISCLOSURE_ENV = { + "lowest_numerical_precision_in_linear": "MLLOG_LOWEST_NUMERICAL_PRECISION_IN_LINEAR", + "lowest_numerical_precision_in_attn": "MLLOG_LOWEST_NUMERICAL_PRECISION_IN_ATTN", + "lowest_numerical_precision_in_comm": "MLLOG_LOWEST_NUMERICAL_PRECISION_IN_COMM", +} + +# The compliance checker rejects any lowest_numerical_precision_* value outside +# this set (training_6.0.0/common.yaml). mxfp6 is deliberately absent: adding it +# is the Training WG request tracked separately, and until it lands a run that +# discloses mxfp6 produces a structurally valid log that the checker refuses. +# Emitting anything else would misdescribe the run, so the value is passed +# through and the mismatch is surfaced loudly rather than silently corrected. +_CHECKER_PRECISION_VALUES = frozenset( + { + "fp64", + "fp32", + "tf32", + "fp16", + "fp8", + "nvfp4", + "mxfp4", + "bfloat16", + "Graphcore FLOAT 16.16", + "int8", + "uint8", + "int4", + "uint4", + } +) + +# Identity records that decide which division the log is judged in. None may +# fall back to a built-in guess. +_SUBMISSION_IDENTITY_ENV = { + "submission_org": "MLLOG_SUBMISSION_ORG", + "submission_division": "MLLOG_SUBMISSION_DIVISION", + "submission_platform": "MLLOG_SUBMISSION_PLATFORM", +} + + +def _is_rank_zero() -> bool: + return int(os.environ.get("RANK", "0")) == 0 + + +def _require_env(name: str, purpose: str) -> str: + value = os.environ.get(name, "").strip() + if not value: + raise RuntimeError(f"MLPerf mode requires {name} to be set explicitly ({purpose}).") + return value + + +def _precision_disclosures_from_env() -> dict[str, str]: + """Return the mandatory precision disclosures without guessing policy names.""" + values = { + key: os.environ.get(environment_name, "").strip() + for key, environment_name in _PRECISION_DISCLOSURE_ENV.items() + } + missing = [_PRECISION_DISCLOSURE_ENV[key] for key, value in values.items() if not value] + if missing: + raise RuntimeError( + "MLPerf mode requires explicit precision disclosures; missing " + ", ".join(missing) + ) + unaccepted = sorted({value for value in values.values() if value not in _CHECKER_PRECISION_VALUES}) + if unaccepted: + logger.warning( + "Precision disclosure(s) %s are not in the compliance checker's accepted set; " + "the resulting log will be rejected until the format is approved upstream.", + ", ".join(unaccepted), + ) + return values + + +def _submission_identity_from_env() -> dict[str, str]: + """Return org/division/platform, refusing to default any of them.""" + values = { + key: _require_env(environment_name, f"MLLOG {key}") + for key, environment_name in _SUBMISSION_IDENTITY_ENV.items() + } + division = values["submission_division"] + if division not in ("closed", "open"): + raise RuntimeError(f"MLLOG_SUBMISSION_DIVISION must be 'closed' or 'open', got {division!r}.") + return values + def _mlperf_logging_enabled(ctx: PatchContext) -> bool: args = get_args(ctx) @@ -94,6 +176,21 @@ def __init__( self.log_every_n_steps = log_every_n_steps self.timer = ThroughputTimer(global_batch_size) self._converged = False + self._run_started = False + self._run_stopped = False + + # A submission result is a file, not console scrollback: stdout is + # interleaved with every other rank's output and with framework noise + # that the parser then has to be trusted to ignore. Rank zero writes the + # log itself so the artifact the checker reads is the artifact produced. + if _is_rank_zero(): + mllog.config( + filename=_require_env("MLLOG_OUTPUT_FILE", "the path this run's result_*.txt is written to"), + # Frames from mllogger.event() back to the FluxMLPerfLogger + # method that called it, so every record reports a stable + # origin line. The seed checker compares those across runs. + default_stack_offset=int(os.environ.get("MLLOG_STACK_OFFSET", "3")), + ) self.profiler = os.getenv("PROFILER", "") self.profiler_warmup_steps = int(os.getenv("PROF_WARMUP_STEPS", "0")) @@ -121,38 +218,71 @@ def _end(self, key, value=None, metadata=None): self._mllogger.end(key=key, value=value, metadata=metadata) def log_init(self, seed: int): - if int(os.environ.get("RANK", "0")) == 0: - self._start(key=self._constants.INIT_START) - self._event(key=self._constants.SUBMISSION_BENCHMARK, value="flux1") - self._event( - key=self._constants.SUBMISSION_ORG, - value=os.environ.get("MLLOG_SUBMISSION_ORG", "AMD"), - ) - self._event( - key=self._constants.SUBMISSION_DIVISION, - value=os.environ.get("MLLOG_SUBMISSION_DIVISION", "closed"), - ) - self._event( - key=self._constants.SUBMISSION_PLATFORM, - value=os.environ.get("MLLOG_SUBMISSION_PLATFORM", "MI355X"), - ) - self._event(key=self._constants.SUBMISSION_STATUS, value="onprem") - self._event(key="target_accuracy", value=self.target_val_loss) - self._event(key=self._constants.SEED, value=seed) + if not _is_rank_zero(): + return + identity = _submission_identity_from_env() + # The launcher clears the page cache before it starts the container and + # reports what it did; defaulting this to false would let a run that + # never cleared claim a cold start. + clear_caches = ( + _require_env("MLPERF_CLEAR_CACHES", "whether the launcher dropped caches before this run").lower() + == "true" + ) + self._event(key="cache_clear", value=clear_caches) + self._start(key=self._constants.INIT_START) + self._event(key=self._constants.SUBMISSION_BENCHMARK, value="flux1") + self._event(key=self._constants.SUBMISSION_ORG, value=identity["submission_org"]) + self._event(key=self._constants.SUBMISSION_DIVISION, value=identity["submission_division"]) + self._event(key=self._constants.SUBMISSION_PLATFORM, value=identity["submission_platform"]) + self._event( + key=self._constants.SUBMISSION_STATUS, + value=os.environ.get("MLLOG_SUBMISSION_STATUS", "onprem"), + ) + self._event(key="target_accuracy", value=self.target_val_loss) + self._event(key=self._constants.SEED, value=seed) def log_hyperparams(self, args): - if int(os.environ.get("RANK", "0")) != 0: + if not _is_rank_zero(): return self._event(key=self._constants.GLOBAL_BATCH_SIZE, value=self.gbs) + for key, value in _precision_disclosures_from_env().items(): + self._event(key=key, value=value) + for key, value in ( + ("tensor_parallelism", getattr(args, "tensor_model_parallel_size", 1)), + ("pipeline_parallelism", getattr(args, "pipeline_model_parallel_size", 1)), + ("context_parallelism", getattr(args, "context_parallel_size", 1)), + ("expert_parallelism", getattr(args, "expert_model_parallel_size", 1)), + ("micro_batch_size", self.mbs), + # Names the recipe a reviewer has to be able to find in the + # submission's code/ directory, so "unknown" is not an answer. + ("config_filename", _require_env("EXP", "the recipe this run was launched from")), + ): + self._event(key=key, value=value) self._event( key=self._constants.TRAIN_SAMPLES, - value=getattr(args, "train_samples", 1099776), + value=getattr(args, "train_samples", None) or 1099776, ) + # EVAL_SAMPLES is emitted before any evaluation has run, so it can only + # ever state the configured budget. The check that the budget was + # actually read lives in the evaluation loop + # (evaluator._record_consumed_valid_samples), which raises rather than + # letting this number stand in for an unverified one. + eval_iters = getattr(args, "eval_iters", 0) or 0 self._event( key=self._constants.EVAL_SAMPLES, - value=getattr(args, "eval_samples", 29696), + value=getattr(args, "eval_samples", None) or eval_iters * self.gbs, + ) + # How often evaluation runs, in samples. EXACTLY_ONE in + # closed_flux1.yaml, and spelled as a literal because no + # mlperf_logging release defines a constant for it: the name appears + # only in the checker's own rulesets. A getattr against the constants + # module would therefore always take its fallback. + self._event( + key="evaluation_frequency", + value=getattr(args, "eval_interval", 0) * self.gbs, ) - gas = max(self.gbs // self.mbs, 1) + data_parallel_size = getattr(args, "data_parallel_size", 1) or 1 + gas = max(self.gbs // (self.mbs * data_parallel_size), 1) self._event(key=self._constants.GRADIENT_ACCUMULATION_STEPS, value=gas) self._event(key=self._constants.OPT_NAME, value="adamw") self._event( @@ -175,13 +305,31 @@ def log_hyperparams(self, args): key="opt_adamw_weight_decay", value=getattr(args, "weight_decay", 0.1), ) + self._event( + key="opt_learning_rate_warmup_steps", + value=getattr(args, "lr_warmup_iters", 0), + ) + self._event( + key="opt_gradient_clip_norm", + value=getattr(args, "clip_grad", 1.0), + ) def log_init_stop_run_start(self): - if int(os.environ.get("RANK", "0")) == 0: + if self._run_started: + return + self._run_started = True + if _is_rank_zero(): self._end(key=self._constants.INIT_STOP) self._start(key=self._constants.RUN_START) self._start(key=self._constants.EPOCH_START, metadata={"epoch_num": 0}) - self._start(key=self._constants.BLOCK_START, metadata={"first_epoch_num": 0}) + self.log_block_start(0) + + def log_block_start(self, global_step: int): + if _is_rank_zero(): + self._start( + key=self._constants.BLOCK_START, + metadata={"samples_count": global_step * self.gbs}, + ) def on_train_batch_end(self, global_step: int, loss: float, lr: float): self.timer.mark_training_start() @@ -189,7 +337,7 @@ def on_train_batch_end(self, global_step: int, loss: float, lr: float): self._handle_profiler(global_step) - if int(os.environ.get("RANK", "0")) != 0: + if not _is_rank_zero(): return if global_step % self.log_every_n_steps == 0: self._event( @@ -203,10 +351,11 @@ def on_train_batch_end(self, global_step: int, loss: float, lr: float): ) def on_validation_start(self, global_step: int): + self.log_init_stop_run_start() self.timer.update_samples(global_step) self.timer.pause_for_eval() - if int(os.environ.get("RANK", "0")) == 0: + if _is_rank_zero(): if global_step > 0: throughput = self.timer.compute_throughput() self._event( @@ -219,14 +368,17 @@ def on_validation_start(self, global_step: int): ) self._end( key=self._constants.BLOCK_STOP, - metadata={"first_epoch_num": 0}, + metadata={"samples_count": global_step * self.gbs}, + ) + self._start( + key=self._constants.EVAL_START, + metadata={"samples_count": global_step * self.gbs}, ) - self._start(key=self._constants.EVAL_START, metadata={"epoch_num": 0}) def on_validation_end(self, global_step: int, val_loss: float): self.timer.resume_after_eval() - if int(os.environ.get("RANK", "0")) == 0: + if _is_rank_zero(): self._event( key=self._constants.EVAL_ACCURACY, value=val_loss, @@ -265,10 +417,24 @@ def _handle_profiler(self, global_step: int): def converged(self): return self._converged + @property + def run_stopped(self): + """Whether run_stop has been emitted, however the run ended. + + Distinct from ``converged``, which is only the success path. Anything + guarding against records past the end of the run wants this one. + """ + return self._run_stopped + def log_run_stop(self, success: bool, global_step: int): + # run_stop is EXACTLY_ONE in the ruleset: a converged run that also hits + # the end-of-training path must not emit a second, contradictory record. + if self._run_stopped: + return if success: self._converged = True - if int(os.environ.get("RANK", "0")) == 0: + self._run_stopped = True + if _is_rank_zero(): status = "success" if success else "aborted" self._end( key=self._constants.RUN_STOP, @@ -323,6 +489,7 @@ def patch_mlperf_logging(ctx: PatchContext): mbs = getattr(args, "micro_batch_size", 64) target_val_loss = getattr(args, "target_val_loss", 0.586) log_interval = getattr(args, "log_interval", 10) + eval_purge_memory = getattr(args, "eval_purge_memory", False) mlperf_logger = FluxMLPerfLogger( global_batch_size=gbs, @@ -334,6 +501,20 @@ def patch_mlperf_logging(ctx: PatchContext): mlperf_logger.log_init(seed=seed) mlperf_logger.log_hyperparams(args) + # The clock has to start before Megatron opens the dataset, which happens + # inside pretrain() after this phase has already run. mlperf_boundary + # creates that seam; the call below is what it fires there. The + # first-training_log path further down stays as a backstop and turns into a + # no-op once this has run. + from primus.backends.megatron.patches import mlperf_boundary + + mlperf_boundary.set_transition(mlperf_logger.log_init_stop_run_start) + mlperf_boundary.install() + + # Reachable from the after_train phase, which closes out runs that finish + # their step budget without converging. + megatron_training._primus_mlperf_logger = mlperf_logger + # --- Suppress Megatron's built-in logging --- megatron_training.print_rank_last = lambda *a, **k: None @@ -386,6 +567,14 @@ def _mlperf_evaluate_and_print_results(*eval_args, **eval_kwargs): # evaluate_and_print_results(prefix, fwd, data, model, iteration[4], ...) iteration = eval_kwargs.get("iteration", eval_args[4] if len(eval_args) > 4 else 0) + # The run is over and run_stop is EXACTLY_ONE, so an evaluation reaching + # here is outside the measured region and contributes nothing to the + # submission. Clearing do_valid on convergence means nothing should get + # this far; this makes "the log ends at run_stop" hold structurally + # rather than by call ordering, for whatever call site comes next. + if mlperf_logger.run_stopped: + return _orig_eval(*eval_args, **eval_kwargs) + mlperf_logger.on_validation_start(iteration) # Temporarily wrap whatever `evaluate` is at call time (e.g. @@ -404,54 +593,40 @@ def _capture_wrapper(*a, **kw): megatron_training.evaluate = _capture_wrapper try: - import gc - result = _orig_eval(*eval_args, **eval_kwargs) - gc.collect() finally: megatron_training.evaluate = _current_eval - try: + # Reclaiming memory after every eval costs a full GC pause plus an + # allocator flush inside the measured window, and the allocator has to + # re-grow its pools on the next training step. Off by default; set + # eval_purge_memory to re-enable if a run proves it needs the headroom. + if eval_purge_memory: + import gc + import torch + gc.collect() torch.cuda.empty_cache() - except Exception: - pass val_loss = _extract_val_loss(_loss_capture) if val_loss is not None: - # Megatron's `evaluate()` (training.py:3178-3180) divides the - # per-rank accumulated loss locally and does NOT all-reduce - # across the data-parallel group — the result is intended for - # `print_rank_last` / TensorBoard which only read on a single - # rank. We must reduce here so every rank evaluates the same - # global validation loss against `target_val_loss`; otherwise - # ranks can disagree on the early-stop branch and the - # divergent `args.train_iters` mutation below desyncs - # collective ordering at the next training step, producing a - # NCCL watchdog deadlock (observed on FLUX 12B MLPerf at - # step 2560 when val_loss landed near target). - import torch - import torch.distributed as dist - - if dist.is_initialized(): - try: - from megatron.core import parallel_state as mpu - - dp_group = mpu.get_data_parallel_group() - except Exception: - dp_group = None - _vl = torch.tensor(val_loss, dtype=torch.float64, device="cuda") - dist.all_reduce(_vl, op=dist.ReduceOp.AVG, group=dp_group) - val_loss = _vl.item() - + # primus_evaluate already reduces over the data-parallel group and + # returns a value identical on every rank, which is what the + # early-stop comparison below requires: if ranks disagree near the + # target, one can exit train() alone while the others keep training, + # desyncing collectives into an NCCL watchdog deadlock (observed on + # FLUX 12B MLPerf at step 2560). A further ReduceOp.AVG here would + # average identical values -- a no-op costing one collective and one + # host sync -- so it is deliberately absent. mlperf_logger.on_validation_end(iteration, val_loss) log_rank_0( f"[MLPerf] Validation loss at step {iteration}: {val_loss:.6f} " f"(target: {target_val_loss:.6f})" ) - if val_loss <= target_val_loss: + converged = val_loss <= target_val_loss + if converged: log_rank_0( f"[MLPerf] Convergence reached! val_loss={val_loss:.6f} " f"<= target={target_val_loss:.6f}" @@ -460,18 +635,27 @@ def _capture_wrapper(*a, **kw): try: from megatron.training import get_args as megatron_get_args - megatron_get_args().train_iters = iteration + megatron_args = megatron_get_args() + megatron_args.train_iters = iteration + # Breaking the loop returns into pretrain(), which runs one + # more validation whenever do_valid is set. That evaluation + # is past run_stop, and it draws fresh VAE epsilon and + # flow-matching noise, so it reports a different loss and + # can land above target -- contradicting the evaluation + # that just ended the run. + megatron_args.do_valid = False except Exception: - logger.warning("Could not set args.train_iters for early stop") - else: - if int(os.environ.get("RANK", "0")) == 0: - mlperf_logger._start( - key=mlperf_logger._constants.BLOCK_START, - metadata={"first_epoch_num": 0}, - ) + logger.warning("Could not set args.train_iters/do_valid for early stop") else: + converged = False logger.warning("Could not extract validation loss from evaluate result") + # Training resumes unless this eval ended the run, so the block that + # on_validation_start closed has to be reopened -- including on the + # path where the loss could not be read and training carries on. + if not converged: + mlperf_logger.log_block_start(iteration) + return result _mlperf_evaluate_and_print_results._primus_mlperf_eval_wrapper = True @@ -483,3 +667,39 @@ def _capture_wrapper(*a, **kw): f"[Patch:mlperf_logging] Installed MLPerf logging (gbs={gbs}, " f"target_val_loss={target_val_loss}, log_interval={log_interval})" ) + + +@register_patch( + "megatron.training.mlperf_run_stop", + backend="megatron", + phase="after_train", + description="Close out a run that ended without reaching the quality target", + condition=_mlperf_logging_enabled, + priority=15, +) +def patch_mlperf_terminal_run_stop(ctx: PatchContext): + """Emit run_stop for a run that exhausted its step budget. + + Only convergence emitted run_stop before, so a run that never hit the + target produced a log with run_start and no run_stop. The ruleset requires + exactly one, and a non-converging run is still evidence -- it belongs in + the RCP comparison as a run that did not make it, not as an unparseable + file. log_run_stop is idempotent, so converged runs fall through here. + """ + import megatron.training.training as megatron_training + + mlperf_logger = getattr(megatron_training, "_primus_mlperf_logger", None) + if mlperf_logger is None or mlperf_logger.converged: + return + + iteration = 0 + try: + from megatron.training import get_args as megatron_get_args + + megatron_args = megatron_get_args() + iteration = getattr(megatron_args, "curr_iteration", 0) or getattr(megatron_args, "train_iters", 0) + except Exception: + logger.warning("Could not read the final iteration for the terminal run_stop") + + mlperf_logger.log_run_stop(success=False, global_step=iteration) + log_rank_0(f"[Patch:mlperf_run_stop] Run ended without converging at iteration {iteration}") diff --git a/primus/backends/megatron/patches/mlperf_warmup_patches.py b/primus/backends/megatron/patches/mlperf_warmup_patches.py index 5ac532978..a9fa5a1ee 100644 --- a/primus/backends/megatron/patches/mlperf_warmup_patches.py +++ b/primus/backends/megatron/patches/mlperf_warmup_patches.py @@ -200,66 +200,292 @@ def _reset_single(opt): _log("Reset optimizer step counters") -@register_patch( - "megatron.training.mlperf_warmup", - backend="megatron", - phase="before_train", - description="MLPerf warmup: synthetic data steps before measured training", - condition=_warmup_enabled, - priority=95, -) -def patch_mlperf_warmup(ctx: PatchContext): - """Install warmup hook on train_step at priority 95 (outermost wrapper).""" +def _build_synthetic_iterator(primus_args): + """Build the mock Flux dataloader the warmup steps consume.""" + from torch.utils.data import DataLoader + + from primus.backends.megatron.data.dataloader import MegatronDataloaderWrapper + from primus.backends.megatron.data.synthetic.mock_datasets import ( + PreGeneratedMockFluxSchnellDataset, + ) + + image_size = getattr(primus_args, "image_size", 256) + vae_latent_mode = getattr(primus_args, "vae_latent_mode", "resample") + mbs = getattr(primus_args, "micro_batch_size", 64) + + mock_dataset = PreGeneratedMockFluxSchnellDataset( + num_samples=max(mbs * 4, 256), + image_size=image_size, + vae_latent_mode=vae_latent_mode, + ) + mock_loader = DataLoader(mock_dataset, batch_size=mbs, shuffle=False, drop_last=True) + return MegatronDataloaderWrapper(mock_loader) + + +def _run_warmup_and_restore( + *, + warmup_steps, + train_step_fn, + forward_step_func, + synthetic_iter, + model, + optimizer, + opt_param_scheduler, + config, + forward_backward_func, + iteration=None, +): + """Run synthetic steps, then undo every effect they had on training state. + + The caller supplies ``train_step_fn`` so this works both from inside the + train_step chain and from the pre-data boundary, where the chain has to be + read at call time. + """ import megatron.training.training as mt + from megatron.training import get_args as megatron_get_args - if hasattr(mt.train_step, "_primus_warmup_hook"): - return + megatron_args = megatron_get_args() + models = model if isinstance(model, (list, tuple)) else [model] + transformer_impl = getattr(megatron_args, "transformer_impl", "local") + use_fsdp2_fp8 = getattr(megatron_args, "use_fsdp2_fp8_all_gather", False) - primus_args = get_args(ctx) - warmup_steps = getattr(primus_args, "warmup_train_steps", 2) + # ---- 1. Snapshot model parameters to CPU ---- + _log("Saving model parameters to CPU before warmup") + saved_params = {} + for m in models: + for name, p in m.named_parameters(): + saved_params[name] = p.data.to("cpu", non_blocking=True) + torch.cuda.synchronize() + _log(f"Saved {len(saved_params)} parameter tensors") + + # ---- 2. Neuter optimizer ---- + saved_opt = _neuter_optimizer(optimizer) + + # ---- 3. Suppress training_log and eval during warmup ---- + saved_training_log = mt.training_log + saved_eval = mt.evaluate_and_print_results + mt.training_log = lambda *a, **k: None + mt.evaluate_and_print_results = lambda *a, **k: None + + # ---- 3b. Save LR scheduler state (NeMo never steps the scheduler during warmup) ---- + saved_lr_num_steps = opt_param_scheduler.num_steps + + # ---- 3c. Install Megatron's grad-finalize callback for the warmup steps ---- + # At the pre-data boundary `train()` has not run yet, so it has not assigned + # config.finalize_model_grads_func. That callback is the only caller of + # finish_grad_sync(), so without it a warmup backward dispatches the + # data-parallel reduce-scatter and nothing ever waits on it. The handle is + # still outstanding when the first real step dispatches its own, and + # Megatron asserts "Should not have multiple communication calls + # outstanding at once" before a single iteration completes. + saved_finalize = config.finalize_model_grads_func + if saved_finalize is None: + from megatron.core.distributed import finalize_model_grads + + config.finalize_model_grads_func = finalize_model_grads + _log("Installed finalize_model_grads_func for warmup (unset at this boundary)") + + # ---- 4. Run warmup steps with synthetic data ---- + try: + for step_idx in range(warmup_steps): + _log(f"Warmup step {step_idx + 1}/{warmup_steps}") + train_step_fn( + forward_step_func, + synthetic_iter, + model, + optimizer, + opt_param_scheduler, + config, + forward_backward_func, + iteration=iteration, + ) + finally: + config.finalize_model_grads_func = saved_finalize + _log(f"Completed {warmup_steps} warmup steps") - _lazy_state = { - "initialized": False, - "synthetic_iter": None, - "use_fsdp2_fp8": False, - "transformer_impl": "local", - } + # ---- 5. Restore optimizer ---- + _restore_optimizer(optimizer, saved_opt) + _reset_optimizer_state(optimizer) - _wrapped_chain = mt.train_step - _warmup_done = [False] + # ---- 6. Restore model parameters from CPU ---- + restored = 0 + for m in models: + for name, p in m.named_parameters(): + if name in saved_params: + p.data.copy_(saved_params[name]) + restored += 1 + del saved_params + _log(f"Restored {restored} parameter tensors from CPU snapshot") + + # ---- 7. FP8 reset (spec-aware) ---- + if transformer_impl == "transformer_engine": + te_count = _reset_fp8_te_spec(models) + amax_count = _seed_fp8_amax(models) + _log(f"FP8 TE reset: {te_count} modules, " f"seeded {amax_count} amax tensors") + else: + local_count = _reset_fp8_local_spec(models) + _log(f"FP8 local spec reset: {local_count} modules") - def _lazy_init(): - """One-time initialization on first train_step call, when Megatron args exist.""" - if _lazy_state["initialized"]: - return + # ---- 8. FSDP2 FP8 all-gather recompute ---- + if use_fsdp2_fp8: + try: + from primus.backends.megatron.core.distributed.fsdp2_fp8_all_gather import ( + precompute_fp8_scales_for_fsdp, + ) - from megatron.training import get_args as megatron_get_args + cache_data = getattr(megatron_args, "fp8_precompute_data_cache", True) + use_cpp = getattr(megatron_args, "use_cpp_fp8_quantize", False) + sr = getattr(megatron_args, "fp8_all_gather_stochastic_rounding", False) + precompute_fp8_scales_for_fsdp( + models[0], + cache_data=cache_data, + use_cpp_quantize=use_cpp, + stochastic_rounding=sr, + ) + _log("Recomputed FSDP2 FP8 all-gather scales") + except Exception as e: + _log(f"FSDP2 FP8 recompute failed (non-fatal): {e}") - megatron_args = megatron_get_args() + # ---- 9. Reload model params in optimizer (FSDP2 BF16 master weight) ---- + if hasattr(optimizer, "reload_model_params"): + optimizer.reload_model_params() + _log("Called optimizer.reload_model_params()") - from torch.utils.data import DataLoader + # ---- 10. Post-restore NaN check ---- + nan_params = 0 + for m in models: + for name, p in m.named_parameters(): + if p.data.is_floating_point() and torch.isnan(p.data).any(): + nan_params += 1 + _log(f"Post-restore parameter check: nan_params={nan_params}") + + # ---- 11. Zero gradients ---- + try: + optimizer.zero_grad(set_to_none=True) + except TypeError: + optimizer.zero_grad() + + # ---- 12. Reset counters ---- + megatron_args.consumed_train_samples = 0 + megatron_args.skipped_train_samples = 0 + opt_param_scheduler.num_steps = saved_lr_num_steps + _log( + f"Reset consumed_train_samples=0, skipped_train_samples=0, " + f"lr_scheduler.num_steps={saved_lr_num_steps}" + ) - from primus.backends.megatron.data.dataloader import MegatronDataloaderWrapper - from primus.backends.megatron.data.synthetic.mock_datasets import ( - PreGeneratedMockFluxSchnellDataset, + # ---- 13. Restore training_log and eval ---- + mt.training_log = saved_training_log + mt.evaluate_and_print_results = saved_eval + + # ---- 13b. Invalidate the CudaPrefetchIterator that was built around + # the SYNTHETIC iterator during warmup step 1. + # + # ``patch_grad_zero_and_data_prefetch`` builds a ``CudaPrefetchIterator`` + # the first time its ``_patched_train_step`` runs and caches it in a + # closure-local ``_prefetch_state["iter"]``. Because warmup step 1 + # is the first call into that train_step, the prefetch iterator gets + # bound to ``synthetic_iter``. ``MegatronDataloaderWrapper`` is + # cyclic (never raises ``StopIteration``), so subsequent real + # training steps would silently keep reading from the cycling + # synthetic dataset instead of the actual training dataset -- model + # overfits the mock samples and val_loss on real data stays stuck + # at ~1.38 forever. + # + # Dropping the cached entry forces the next train_step to rebuild + # the prefetch wrapper around its incoming ``data_iterator`` arg + # (the real iterator). + try: + from primus.backends.megatron.patches.delayed_fp8_scaling_patches import ( + reset_prefetch_state, ) - image_size = getattr(primus_args, "image_size", 256) - vae_latent_mode = getattr(primus_args, "vae_latent_mode", "resample") - mbs = getattr(primus_args, "micro_batch_size", 64) + evicted = reset_prefetch_state() + if evicted is None: + _log(" Prefetch reset: no cached iterator to evict") + else: + _log( + f" Prefetch reset: evicted cached {type(evicted).__name__} " + f"(wrapped synthetic warmup iterator) -- next train_step " + f"will rebuild it around the real data_iterator" + ) + except Exception as _e: + _log(f" Prefetch reset failed (non-fatal): {_e}") + + # ---- 14. Synchronize ---- + torch.cuda.synchronize() + if torch.distributed.is_initialized(): + torch.distributed.barrier() + + +def _install_boundary_warmup(primus_args, warmup_steps): + """Run warmup at the pre-data boundary, inside the initialization window. + + Used in MLPerf mode, where warmup must finish before ``run_start`` and + therefore before the data iterators exist. Everything the warmup needs is + captured on the way there by :mod:`mlperf_boundary`. + """ + from primus.backends.megatron.patches import mlperf_boundary + + def _warmup_hook(): + import megatron.training.training as mt + + captured = mlperf_boundary.captured() + model = captured.get("model") + optimizer = captured.get("optimizer") + opt_param_scheduler = captured.get("opt_param_scheduler") + forward_step_func = captured.get("forward_step_func") + missing = [ + name + for name, value in ( + ("model", model), + ("optimizer", optimizer), + ("opt_param_scheduler", opt_param_scheduler), + ("forward_step_func", forward_step_func), + ) + if value is None + ] + if missing: + raise RuntimeError( + "MLPerf warmup runs before the data iterators are built and needs " + "objects captured from Megatron, but these were never captured: " + + ", ".join(missing) + + ". The capture wrappers in mlperf_boundary did not run." + ) - mock_dataset = PreGeneratedMockFluxSchnellDataset( - num_samples=max(mbs * 4, 256), - image_size=image_size, - vae_latent_mode=vae_latent_mode, + models = model if isinstance(model, (list, tuple)) else [model] + # Read train_step now, not at install time: every other before_train + # patch has wrapped it by the time the boundary fires. + _run_warmup_and_restore( + warmup_steps=warmup_steps, + train_step_fn=mt.train_step, + forward_step_func=forward_step_func, + synthetic_iter=_build_synthetic_iterator(primus_args), + model=model, + optimizer=optimizer, + opt_param_scheduler=opt_param_scheduler, + config=mt.get_model_config(models[0]), + forward_backward_func=mt.get_forward_backward_func(), + iteration=0, ) - mock_loader = DataLoader(mock_dataset, batch_size=mbs, shuffle=False, drop_last=True) - _lazy_state["synthetic_iter"] = MegatronDataloaderWrapper(mock_loader) - _lazy_state["use_fsdp2_fp8"] = getattr(megatron_args, "use_fsdp2_fp8_all_gather", False) - _lazy_state["transformer_impl"] = getattr(megatron_args, "transformer_impl", "local") - _lazy_state["initialized"] = True - _log(f"Lazy init complete (warmup_steps={warmup_steps})") + mlperf_boundary.register_pre_run_hook("mlperf_warmup", _warmup_hook, order=10) + mlperf_boundary.install() + log_rank_0( + f"[Patch:mlperf_warmup] Warmup registered at the pre-data boundary " f"(warmup_steps={warmup_steps})" + ) + + +def _install_train_step_warmup(mt, primus_args, warmup_steps): + """Run warmup inside the first train_step, then execute the first real step. + + The non-MLPerf path. Warmup lands after the data iterators exist, which is + fine when no clock is running, and it keeps development recipes on the + behavior they were tuned against. + """ + _wrapped_chain = mt.train_step + _warmup_done = [False] + _synthetic_iter = [None] def _hooked_train_step( forward_step_func, @@ -283,165 +509,22 @@ def _hooked_train_step( iteration=iteration, ) - _lazy_init() - - from megatron.training import get_args as megatron_get_args - - megatron_args = megatron_get_args() - models = model if isinstance(model, (list, tuple)) else [model] - synthetic_iter = _lazy_state["synthetic_iter"] - - # ---- 1. Snapshot model parameters to CPU ---- - _log("Saving model parameters to CPU before warmup") - saved_params = {} - for m in models: - for name, p in m.named_parameters(): - saved_params[name] = p.data.to("cpu", non_blocking=True) - torch.cuda.synchronize() - _log(f"Saved {len(saved_params)} parameter tensors") - - # ---- 2. Neuter optimizer ---- - saved_opt = _neuter_optimizer(optimizer) - - # ---- 3. Suppress training_log and eval during warmup ---- - saved_training_log = mt.training_log - saved_eval = mt.evaluate_and_print_results - mt.training_log = lambda *a, **k: None - mt.evaluate_and_print_results = lambda *a, **k: None - - # ---- 3b. Save LR scheduler state (NeMo never steps the scheduler during warmup) ---- - saved_lr_num_steps = opt_param_scheduler.num_steps - - # ---- 4. Run warmup steps with synthetic data ---- - for step_idx in range(warmup_steps): - _log(f"Warmup step {step_idx + 1}/{warmup_steps}") - _wrapped_chain( - forward_step_func, - synthetic_iter, - model, - optimizer, - opt_param_scheduler, - config, - forward_backward_func, - iteration=iteration, - ) - _log(f"Completed {warmup_steps} warmup steps") - - # ---- 5. Restore optimizer ---- - _restore_optimizer(optimizer, saved_opt) - _reset_optimizer_state(optimizer) - - # ---- 6. Restore model parameters from CPU ---- - restored = 0 - for m in models: - for name, p in m.named_parameters(): - if name in saved_params: - p.data.copy_(saved_params[name]) - restored += 1 - del saved_params - _log(f"Restored {restored} parameter tensors from CPU snapshot") - - # ---- 7. FP8 reset (spec-aware) ---- - if _lazy_state["transformer_impl"] == "transformer_engine": - te_count = _reset_fp8_te_spec(models) - amax_count = _seed_fp8_amax(models) - _log(f"FP8 TE reset: {te_count} modules, " f"seeded {amax_count} amax tensors") - else: - local_count = _reset_fp8_local_spec(models) - _log(f"FP8 local spec reset: {local_count} modules") - - # ---- 8. FSDP2 FP8 all-gather recompute ---- - if _lazy_state["use_fsdp2_fp8"]: - try: - from primus.backends.megatron.core.distributed.fsdp2_fp8_all_gather import ( - precompute_fp8_scales_for_fsdp, - ) - - cache_data = getattr(megatron_args, "fp8_precompute_data_cache", True) - use_cpp = getattr(megatron_args, "use_cpp_fp8_quantize", False) - sr = getattr(megatron_args, "fp8_all_gather_stochastic_rounding", False) - precompute_fp8_scales_for_fsdp( - models[0], - cache_data=cache_data, - use_cpp_quantize=use_cpp, - stochastic_rounding=sr, - ) - _log("Recomputed FSDP2 FP8 all-gather scales") - except Exception as e: - _log(f"FSDP2 FP8 recompute failed (non-fatal): {e}") - - # ---- 9. Reload model params in optimizer (FSDP2 BF16 master weight) ---- - if hasattr(optimizer, "reload_model_params"): - optimizer.reload_model_params() - _log("Called optimizer.reload_model_params()") - - # ---- 10. Post-restore NaN check ---- - nan_params = 0 - for m in models: - for name, p in m.named_parameters(): - if p.data.is_floating_point() and torch.isnan(p.data).any(): - nan_params += 1 - _log(f"Post-restore parameter check: nan_params={nan_params}") - - # ---- 11. Zero gradients ---- - try: - optimizer.zero_grad(set_to_none=True) - except TypeError: - optimizer.zero_grad() - - # ---- 12. Reset counters ---- - megatron_args.consumed_train_samples = 0 - megatron_args.skipped_train_samples = 0 - opt_param_scheduler.num_steps = saved_lr_num_steps - _log( - f"Reset consumed_train_samples=0, skipped_train_samples=0, " - f"lr_scheduler.num_steps={saved_lr_num_steps}" + if _synthetic_iter[0] is None: + _synthetic_iter[0] = _build_synthetic_iterator(primus_args) + + _run_warmup_and_restore( + warmup_steps=warmup_steps, + train_step_fn=_wrapped_chain, + forward_step_func=forward_step_func, + synthetic_iter=_synthetic_iter[0], + model=model, + optimizer=optimizer, + opt_param_scheduler=opt_param_scheduler, + config=config, + forward_backward_func=forward_backward_func, + iteration=iteration, ) - # ---- 13. Restore training_log and eval ---- - mt.training_log = saved_training_log - mt.evaluate_and_print_results = saved_eval - - # ---- 13b. Invalidate the CudaPrefetchIterator that was built around - # the SYNTHETIC iterator during warmup step 1. - # - # ``patch_grad_zero_and_data_prefetch`` builds a ``CudaPrefetchIterator`` - # the first time its ``_patched_train_step`` runs and caches it in a - # closure-local ``_prefetch_state["iter"]``. Because warmup step 1 - # is the first call into that train_step, the prefetch iterator gets - # bound to ``synthetic_iter``. ``MegatronDataloaderWrapper`` is - # cyclic (never raises ``StopIteration``), so subsequent real - # training steps would silently keep reading from the cycling - # synthetic dataset instead of the actual training dataset -- model - # overfits the mock samples and val_loss on real data stays stuck - # at ~1.38 forever. - # - # Dropping the cached entry forces the next train_step to rebuild - # the prefetch wrapper around its incoming ``data_iterator`` arg - # (the real iterator). - try: - from primus.backends.megatron.patches.delayed_fp8_scaling_patches import ( - reset_prefetch_state, - ) - - evicted = reset_prefetch_state() - if evicted is None: - _log(" Prefetch reset: no cached iterator to evict") - else: - _log( - f" Prefetch reset: evicted cached {type(evicted).__name__} " - f"(wrapped synthetic warmup iterator) -- next train_step " - f"will rebuild it around the real data_iterator" - ) - except Exception as _e: - _log(f" Prefetch reset failed (non-fatal): {_e}") - - # ---- 14. Synchronize ---- - torch.cuda.synchronize() - if torch.distributed.is_initialized(): - torch.distributed.barrier() - - # ---- 15. Execute first real step ---- _log("Executing first real train_step with training data") result = _wrapped_chain( forward_step_func, @@ -454,7 +537,6 @@ def _hooked_train_step( iteration=iteration, ) - # ---- 16. Self-remove ---- _warmup_done[0] = True mt.train_step = _wrapped_chain _log("Self-removed warmup hook, train_step = inner wrapped chain") @@ -464,8 +546,28 @@ def _hooked_train_step( _hooked_train_step._primus_warmup_hook = True mt.train_step = _hooked_train_step - _log( - f"Installed MLPerf warmup hook (warmup_steps={warmup_steps}, " - f"deferred init until first train_step)" - ) log_rank_0(f"[Patch:mlperf_warmup] Installed warmup hook " f"(warmup_steps={warmup_steps}, priority=95)") + + +@register_patch( + "megatron.training.mlperf_warmup", + backend="megatron", + phase="before_train", + description="MLPerf warmup: synthetic data steps before measured training", + condition=_warmup_enabled, + priority=95, +) +def patch_mlperf_warmup(ctx: PatchContext): + """Install warmup at priority 95, so it is the outermost wrapper.""" + import megatron.training.training as mt + + if hasattr(mt.train_step, "_primus_warmup_hook"): + return + + primus_args = get_args(ctx) + warmup_steps = getattr(primus_args, "warmup_train_steps", 2) + + if getattr(primus_args, "mlperf_mode", False): + _install_boundary_warmup(primus_args, warmup_steps) + else: + _install_train_step_warmup(mt, primus_args, warmup_steps) diff --git a/primus/backends/megatron/training/diffusion/forward_step.py b/primus/backends/megatron/training/diffusion/forward_step.py index 34ab2f001..8172cdc75 100644 --- a/primus/backends/megatron/training/diffusion/forward_step.py +++ b/primus/backends/megatron/training/diffusion/forward_step.py @@ -37,6 +37,95 @@ logger = logging.getLogger(__name__) +# Validation seeds are offset past the training stream so the two can never +# alias. Training seeds are (seed + 100 * dp_rank) * 10000 + step_count, which +# stays below 2**31 for any plausible seed, world size and step count, so 2**40 +# leaves a wide margin while keeping the eval stream's own per-rank and +# per-microbatch structure intact. +EVAL_RNG_OFFSET = 1 << 40 + +# Microbatch index stride within one evaluation. Must exceed the number of +# microbatches any single evaluation runs so that consecutive evaluations +# cannot collide; asserted where the index is built. +EVAL_RNG_ITERATION_STRIDE = 1 << 20 + +# Where validation timesteps come from. See the block above the injection site +# for why inferring this from a missing data field is not good enough. +DATASET_TIMESTEPS = "dataset" +EQUIDISTANT_TIMESTEPS = "equidistant" +EVAL_TIMESTEP_SOURCES = (DATASET_TIMESTEPS, EQUIDISTANT_TIMESTEPS) + +# MLPerf evaluates at t in {0/8, ..., 7/8}. +NUM_VALIDATION_TIMESTEPS = 8 + +# Warn once per process rather than once per microbatch. +_warned_uncovered_equidistant = False + + +def _warn_if_equidistant_undercovers(batch_size: int) -> None: + """Flag equidistant injection that cannot reach all eight timesteps. + + ``arange(batch_size) % 8`` covers every timestep only when the width is a + multiple of 8; at width 2 it evaluates t=0 and t=1/8 alone. This is a + warning rather than an assertion because the shipped flux_535m recipes run + validation at micro_batch_size 2 against the mock dataset, where the + timestep is supplied per sample and this branch is never reached. + """ + global _warned_uncovered_equidistant + if batch_size % NUM_VALIDATION_TIMESTEPS == 0 or _warned_uncovered_equidistant: + return + _warned_uncovered_equidistant = True + logger.warning( + "eval_timestep_source='equidistant' with validation batch width %d, which is " + "not a multiple of %d: this evaluates only timesteps 0..%d and never sees the " + "rest, so val_loss is not comparable to the MLPerf reference. Raise the " + "validation micro_batch_size to a multiple of %d, or use a dataset carrying " + "per-sample timesteps with eval_timestep_source='dataset'.", + batch_size, + NUM_VALIDATION_TIMESTEPS, + min(batch_size, NUM_VALIDATION_TIMESTEPS) - 1, + NUM_VALIDATION_TIMESTEPS, + ) + + +def resolve_validation_timesteps(batch, eval_timestep_source, batch_size, device, compute_dtype): + """Populate ``batch['timestep']`` and ``batch['timesteps']`` for validation. + + ``timestep`` is the integer index in ``[0, 7]``; ``timesteps`` is that index + divided by 8, which is the sigma the scheduler consumes. Mutates ``batch`` + in place and returns the sigma tensor. + + Raises: + ValueError: on an unknown source, or when ``dataset`` is requested and + the batch carries no timestep field. + """ + if eval_timestep_source not in EVAL_TIMESTEP_SOURCES: + raise ValueError( + f"eval_timestep_source must be one of {list(EVAL_TIMESTEP_SOURCES)}, " + f"got {eval_timestep_source!r}." + ) + + if "timestep" in batch: + val_timesteps = batch["timestep"].float() / NUM_VALIDATION_TIMESTEPS + elif eval_timestep_source == DATASET_TIMESTEPS: + raise ValueError( + "eval_timestep_source='dataset' but the validation batch carries no " + "'timestep' field. The MLCommons val Arrow files provide a per-sample " + "int32 'timestep' column; shards ingested before that column was carried " + "through have sidecars of {'key': ...} only, so every eval against them " + "silently used injected timesteps instead. Re-ingest the val split, or " + "set eval_timestep_source='equidistant' to inject t = index % 8 (which " + "does not reproduce the dataset's image-to-timestep pairing)." + ) + else: + _warn_if_equidistant_undercovers(batch_size) + val_idx = torch.arange(batch_size, device=device) % NUM_VALIDATION_TIMESTEPS + batch["timestep"] = val_idx + val_timesteps = val_idx.to(dtype=compute_dtype) / NUM_VALIDATION_TIMESTEPS + + batch["timesteps"] = val_timesteps + return val_timesteps + def prepare_flux_latents( latents: torch.Tensor, @@ -170,6 +259,8 @@ def flux_forward_step_func( vae_latent_mode="presampled", per_step_rng_reseed=False, step_count=0, + eval_step_index=None, + eval_timestep_source=EQUIDISTANT_TIMESTEPS, ): """ Forward step function for Flux training with distributed data loading. @@ -231,13 +322,23 @@ def flux_forward_step_func( # (noise, timesteps, CFG dropout) from model forward RNG consumption. # Required because TE fused attention advances the default generator even # with dropout=0 when the DPA prologue patch is active. + # + # Validation uses its own stream. The training counter is deliberately + # frozen during eval (see diffusion_trainer) so that eval passes do not + # shift the training seed sequence, but reusing the frozen value reseeds + # every validation microbatch identically, which repeats both the VAE + # reparameterization epsilon and the flow-matching noise across the whole + # evaluation. Both references draw both afresh per batch. if per_step_rng_reseed: from megatron.core import parallel_state as _ps from megatron.training import get_args as _get_args _seed = _get_args().seed _per_rank_seed = _seed + 100 * _ps.get_data_parallel_rank() - _step_seed = (_per_rank_seed * 10000 + step_count) % (2**63) + if eval_step_index is None: + _step_seed = (_per_rank_seed * 10000 + step_count) % (2**63) + else: + _step_seed = (EVAL_RNG_OFFSET + _per_rank_seed * 10000 + eval_step_index) % (2**63) torch.cuda.manual_seed(_step_seed) from megatron.core import tensor_parallel @@ -411,25 +512,34 @@ def flux_forward_step_func( # - val_loss = mean over per-timestep means (equivalent to flat mean given # equal counts). # - # NeMo's official to_webdataset preserves a `timestep` integer per sample - # from the MLCommons Arrow source. Our `primus-cli data diffusion-ingest` - # path (pipelines/ingest.py:33 `ARROW_COLUMNS`) ingests only the 4 tensor - # columns and writes `{"key": ...}` to the json sidecar — so our val shards - # are MISSING the timestep field, which used to make this branch fall - # through to the training path with uniform-random timesteps via the - # `timestep_sampler`. That produced a *different* val_loss estimator than - # the spec's: E_t~U[0,1][MSE] (Monte Carlo over [0,1]) vs the spec's + # The MLCommons Arrow source carries a `timestep` int32 per val sample, and + # `primus-cli data diffusion-ingest` copies it into the json sidecar. + # Datasets ingested before that fix have sidecars of `{"key": ...}` only. + # + # `eval_timestep_source` decides what to do about that, because inferring + # "this is MLPerf validation" from a *missing* data field cannot tell a + # broken ingest apart from a config that never had timesteps to begin with: + # + # "dataset" — the batch must carry `timestep`; a missing field is an + # error naming the likely cause. MLPerf recipes use this. + # "equidistant" — inject t = index % 8 within the batch. Deliberate for + # configs with no per-sample timestep (e.g. flux_535m). + # + # Equidistant injection keeps the *marginal* timestep histogram uniform for + # any batch width that is a multiple of 8, but it does NOT reproduce the + # dataset's own image-to-timestep pairing: the real column assigns each + # image one fixed timestep, and the published val set is ordered so a + # contiguous batch spans only two distinct timesteps, not all eight. At + # widths below 8 it is worse than imprecise — micro_batch_size 2 evaluates + # only t=0 and t=1/8 and never sees the other six. + # + # Falling all the way through to the training path would be worse still: + # uniform-random timesteps via the `timestep_sampler` estimate + # E_t~U[0,1][MSE] (Monte Carlo over [0,1]) rather than the spec's # left-Riemann sum over t∈{0/8..7/8}. The two estimators are not # comparable, so a uniform-random val path can make val_loss converge # spuriously fast relative to the reference convergence point. # - # Fix: when batch is in eval mode (model.training=False, set by - # the evaluation harness via `model_module.eval()`) and lacks a `timestep` - # field, inject equidistant timesteps deterministically by within-batch - # index. With MBS=64, each micro-batch covers each t∈{0..7} exactly 8 - # times. Across 58 micro-batches × 8 DP ranks = 464 micro-batches → exactly - # 3 712 samples per timestep, matching the MLPerf v5.1 spec count. - # # CFG dropout during val: SUPPRESSED. # # Reference-implementation tally for "apply CFG dropout during validation": @@ -443,18 +553,19 @@ def flux_forward_step_func( # ~0.015-0.030 (the 10% unconditional samples pay a ~0.15-0.30 MSE # penalty), which is enough to materially shift the convergence-crossing # step, so we keep it off to match the submission configuration. - is_validation = False - if batch is not None and "timestep" in batch: - is_validation = True - val_timesteps = batch["timestep"].float() / 8.0 - batch["timesteps"] = val_timesteps - elif batch is not None and not model.training: - is_validation = True - batch_size_val = pooled_prompt_embeds.shape[0] - val_idx = torch.arange(batch_size_val, device="cuda") % 8 - batch["timestep"] = val_idx - val_timesteps = val_idx.to(dtype=compute_dtype) / 8.0 - batch["timesteps"] = val_timesteps + # Derived from eval mode alone, independent of where timesteps come from, + # so CFG suppression is identical under both sources. The evaluation + # harness sets this via model_module.eval(). + is_validation = batch is not None and not model.training + + if is_validation: + resolve_validation_timesteps( + batch, + eval_timestep_source, + batch_size=pooled_prompt_embeds.shape[0], + device=pooled_prompt_embeds.device, + compute_dtype=compute_dtype, + ) # Matches NeMo's forward_step which wraps prepare_image_latent_like_reference # in torch.no_grad() — no gradients needed for position IDs, noise sampling, diff --git a/primus/backends/megatron/training/eval_budget.py b/primus/backends/megatron/training/eval_budget.py new file mode 100644 index 000000000..1e83f2235 --- /dev/null +++ b/primus/backends/megatron/training/eval_budget.py @@ -0,0 +1,310 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +One definition of how large an evaluation is, shared by the dataloader +provider and the evaluation loop. + +Previously the two disagreed: the Energon provider sized the validation +LimitDataset with ``get_num_microbatches()`` while the evaluator recomputed +``global_batch_size // (micro_batch_size * data_parallel_size)`` for itself. +Those diverge under batch-size ramp-up, and the evaluator's version silently +floors to zero when the global batch is smaller than one microbatch per rank, +which turns evaluation into a no-op that still reports a loss. + +This module also owns the validation worker count, because worker count is +what actually decides how many samples an evaluation reads. Energon splits +the *data* across ``dp_size * num_workers`` workers but splits the *batch +quota* across ``num_workers`` alone, and those two divisions only agree when +each worker's slice is a whole number of full batches. When they disagree, +short-quota workers silently leave the tail of their slice unread and the +evaluation reports the count it intended rather than the count it achieved. +""" + +import json +import os +from pathlib import Path +from typing import Optional + +__all__ = [ + "DEFAULT_VAL_NUM_WORKERS", + "EvalCoverageError", + "assert_mlperf_timestep_source", + "assert_val_worker_divisibility", + "get_data_parallel_size", + "get_eval_num_microbatches", + "get_val_num_workers", + "read_energon_split_sample_count", + "resolve_eval_iters", +] + +# Energon clamps the worker count to at least one when sharding samples +# (sharder.py, ``max(1, worker_config.num_workers)``) and short-circuits the +# quota split entirely when ``num_workers <= 1`` (limit_dataset.py). At 0 or 1 +# each rank therefore holds a single contiguous slice and takes the whole +# quota from it, so coverage is exact whenever the per-rank slice divides into +# whole microbatches. That makes 0 the only default that cannot silently +# under-read, which matters more for a metric that gates convergence than the +# throughput of a loader that runs a few dozen times per job. +DEFAULT_VAL_NUM_WORKERS = 0 + + +class EvalCoverageError(ValueError): + """Raised when an evaluation would not measure what it reports. + + Covers both reading fewer samples than the configuration claims and + evaluating the samples it does read at timesteps the dataset did not pair + them with. + """ + + +def get_data_parallel_size(args) -> int: + """Data-parallel width, usable before Megatron has computed it. + + Megatron only sets ``data_parallel_size`` while initialising process + groups, but the evaluation budget has to be resolved earlier than that -- + in the ``build_args`` phase, so the Energon provider and the evaluator + both see the corrected ``eval_iters``. Reading the attribute directly + there raises ``AttributeError``, which the patch runner logs and swallows, + leaving ``eval_iters`` at 0; the job then runs to completion having + evaluated nothing. So derive the value the same way Megatron does instead. + """ + dp_size = getattr(args, "data_parallel_size", None) + if dp_size: + return dp_size + + world_size = getattr(args, "world_size", None) or int(os.environ.get("WORLD_SIZE", 0)) + if not world_size: + raise EvalCoverageError( + "Cannot size the evaluation: data_parallel_size is not set yet and " + "neither args.world_size nor the WORLD_SIZE environment variable is " + "available to derive it from." + ) + + divisor = 1 + for name in ( + "tensor_model_parallel_size", + "pipeline_model_parallel_size", + "context_parallel_size", + ): + divisor *= getattr(args, name, 1) or 1 + + if world_size % divisor != 0: + raise EvalCoverageError( + f"world_size ({world_size}) is not divisible by " + f"tensor x pipeline x context parallel size ({divisor}), so " + f"data_parallel_size cannot be derived." + ) + return world_size // divisor + + +def get_eval_num_microbatches(args) -> int: + """Microbatches per evaluation iteration. + + Uses the same global batch as training so that ``eval_iters`` counts in + global batches, matching Megatron's convention. + """ + dp_size = get_data_parallel_size(args) + micro_batch_size = args.micro_batch_size + global_batch_size = args.global_batch_size + + samples_per_microbatch = micro_batch_size * dp_size + if samples_per_microbatch <= 0: + raise EvalCoverageError( + f"micro_batch_size ({micro_batch_size}) * data_parallel_size ({dp_size}) " f"must be positive." + ) + + num_microbatches = global_batch_size // samples_per_microbatch + if num_microbatches <= 0: + raise EvalCoverageError( + f"global_batch_size ({global_batch_size}) is smaller than one microbatch " + f"across data parallelism (micro_batch_size {micro_batch_size} x " + f"data_parallel_size {dp_size} = {samples_per_microbatch}). Evaluation " + f"would run zero microbatches per iteration and report a loss computed " + f"from no samples." + ) + if num_microbatches * samples_per_microbatch != global_batch_size: + raise EvalCoverageError( + f"global_batch_size ({global_batch_size}) is not divisible by " + f"micro_batch_size ({micro_batch_size}) x data_parallel_size ({dp_size}) " + f"= {samples_per_microbatch}. Evaluation cannot cover a whole number of " + f"global batches." + ) + return num_microbatches + + +def get_val_num_workers(args) -> int: + """Dataloader worker count to use for validation. + + Falls back to ``DEFAULT_VAL_NUM_WORKERS`` rather than to the training + ``num_workers``: sharing the training value is exactly what produces the + silent under-read this module exists to prevent. + """ + val_num_workers = getattr(args, "val_num_workers", None) + if val_num_workers is None: + return DEFAULT_VAL_NUM_WORKERS + if val_num_workers < 0: + raise EvalCoverageError(f"val_num_workers must be >= 0, got {val_num_workers}.") + return val_num_workers + + +def assert_val_worker_divisibility(args, eval_samples: int) -> None: + """Fail loudly when the configured shape cannot read every eval sample. + + Coverage is exact when each global worker's slice is a whole number of + full microbatches, i.e. when:: + + eval_samples % (dp_size * max(1, val_num_workers) * micro_batch_size) == 0 + + The ``max(1, ...)`` mirrors Energon's own clamping, and is also what keeps + this from dividing by zero at the default worker count of 0. + """ + dp_size = get_data_parallel_size(args) + micro_batch_size = args.micro_batch_size + val_num_workers = get_val_num_workers(args) + + divisor = dp_size * max(1, val_num_workers) * micro_batch_size + remainder = eval_samples % divisor + if remainder == 0: + return + + per_global_worker = eval_samples // (dp_size * max(1, val_num_workers)) + suggestions = [ + w + for w in range(0, min(64, eval_samples) + 1) + if eval_samples % (dp_size * max(1, w) * micro_batch_size) == 0 + ] + raise EvalCoverageError( + f"Validation would silently read fewer than {eval_samples} samples.\n" + f" eval_samples = {eval_samples}\n" + f" data_parallel_size = {dp_size}\n" + f" micro_batch_size = {micro_batch_size}\n" + f" val_num_workers = {val_num_workers}\n" + f"Energon shards samples across dp_size x max(1, val_num_workers) = " + f"{dp_size * max(1, val_num_workers)} global workers, giving " + f"{per_global_worker} samples each, which is not a whole number of " + f"{micro_batch_size}-sample batches ({remainder} sample(s) over). Workers " + f"whose batch quota is short leave the tail of their slice unread, so the " + f"reported sample count would exceed the count actually evaluated.\n" + f"Valid val_num_workers for this shape: {suggestions}" + ) + + +def assert_mlperf_timestep_source(args) -> None: + """Refuse injected validation timesteps in an MLPerf run. + + ``resolve_validation_timesteps`` takes the dataset's ``timestep`` column + whenever the batch carries one, so a split ingested with that column is + evaluated correctly under either source setting. One combination is left + unsafe: shards ingested before the column was carried through, read under + ``eval_timestep_source='equidistant'``. There the positional fallback + injects ``t = index % 8``, which does not reproduce the pairing of image + to timestep the published split defines, and the run reports a val_loss + indistinguishable in the logs from a correct one. + + The MLPerf recipes set ``dataset`` and close that cell for themselves, but + ``trainer_base.yaml`` defaults to ``equidistant`` for the diffusion + recipes that have no annotated split, so the protection would otherwise + rest on every future submission recipe remembering to override it. A + submission has no legitimate use for injected timesteps, so make the + combination unreachable rather than merely avoidable. + + Belongs at validation-dataloader-build time and not in the ``build_args`` + patch that sizes the eval budget: a failure raised there is logged and + swallowed by the patch runner, which is how ``eval_iters`` silently stayed + at 0. ``eval_timestep_source`` is also a Primus-only key, and those are + merged onto ``args`` only after the ``build_args`` phase has run. + """ + if not getattr(args, "mlperf_mode", False): + return + + # Imported here rather than at module scope to keep this module free of + # the diffusion forward step, which pulls in torch and the Flux model + # utilities; the eval budget is also resolved from a build_args patch. + from primus.backends.megatron.training.diffusion.forward_step import ( + DATASET_TIMESTEPS, + ) + + source = getattr(args, "eval_timestep_source", None) + if source == DATASET_TIMESTEPS: + return + + raise EvalCoverageError( + f"mlperf_mode is set but eval_timestep_source is {source!r}, which lets " + f"validation fall back to injecting t = index % 8 when the shards carry " + f"no per-sample 'timestep' column. That reproduces neither the timesteps " + f"nor the image-to-timestep pairing of the published val split, and it " + f"fails silently: the loss it reports looks exactly like a correct one. " + f"Set eval_timestep_source='{DATASET_TIMESTEPS}' in the recipe, and point " + f"the run at a val split ingested with the timestep column " + f"(primus/configs/data/megatron/diffusion/preprocessing/mlperf_flux1_val.yaml) " + f"so the requirement is met rather than merely asserted." + ) + + +def read_energon_split_sample_count(data_path, split: str = "val") -> Optional[int]: + """Total samples in an Energon split, from the dataset's own index. + + Lets ``full_validation`` mean "all of it" without the count being written + into a recipe by hand, and gives the evaluation loop a third, independent + number to check the configured and observed counts against. + + Returns None when the path is not an Energon dataset (mock data, an + unprepared directory), so callers can fall back rather than fail. + """ + if not data_path: + return None + if isinstance(data_path, (list, tuple)): + if not data_path: + return None + data_path = data_path[0] + + info = Path(str(data_path)) / ".nv-meta" / ".info.json" + try: + shard_counts = json.loads(info.read_text())["shard_counts"] + except (OSError, ValueError, KeyError): + return None + + prefix = f"{split}/" + total = sum(count for shard, count in shard_counts.items() if shard.startswith(prefix)) + return total or None + + +def resolve_eval_iters(args) -> Optional[int]: + """Derive ``eval_iters`` from ``eval_samples`` when the latter is set. + + ``eval_samples`` says what the evaluation is meant to measure -- coverage + of a dataset -- whereas ``eval_iters`` only means that at one particular + global batch size. Returns the derived value, or None when ``eval_samples`` + is unset and ``eval_iters`` should be left alone. + """ + eval_samples = getattr(args, "eval_samples", None) + if eval_samples is None: + return None + + if eval_samples <= 0: + raise EvalCoverageError(f"eval_samples must be positive, got {eval_samples}.") + + global_batch_size = args.global_batch_size + if eval_samples % global_batch_size != 0: + raise EvalCoverageError( + f"eval_samples ({eval_samples}) is not divisible by global_batch_size " + f"({global_batch_size}), so evaluation cannot cover it in whole global " + f"batches. Either adjust global_batch_size or accept a different " + f"eval_samples; {global_batch_size * (eval_samples // global_batch_size)} " + f"and {global_batch_size * (eval_samples // global_batch_size + 1)} are " + f"the nearest reachable counts." + ) + + derived = eval_samples // global_batch_size + + # eval_samples deliberately wins over eval_iters rather than conflicting + # with it. There is no way to tell an eval_iters a recipe chose from one it + # inherited: trainer_base.yaml always supplies a value, and Megatron's + # parser defaults it besides, so treating a mismatch as an error would + # reject every recipe that opts into eval_samples at all. + assert_val_worker_divisibility(args, eval_samples) + return derived diff --git a/primus/backends/megatron/training/eval_session.py b/primus/backends/megatron/training/eval_session.py new file mode 100644 index 000000000..bf5b12411 --- /dev/null +++ b/primus/backends/megatron/training/eval_session.py @@ -0,0 +1,54 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Where one evaluation ends and the next begins. + +Consumers that derive per-microbatch state -- the diffusion trainer's eval RNG +index is the only one today -- need a signal for the start of an evaluation. +The training step is not that signal: two evaluations can run at the same step, +the one inside the training loop and the one ``pretrain`` runs after the loop +exits, and keying off the step alone would let the second continue the first's +counter instead of reproducing it. + +Deliberately free of imports. The producer is the generic evaluator and the +consumer is a diffusion trainer, so anywhere else this could live would either +drag the Megatron evaluation stack into a per-microbatch code path or point the +dependency from generic code at diffusion code. +""" + +_eval_session = 0 + + +def begin_eval_session() -> int: + """Mark the start of one evaluation. Called on every rank. + + Only ``primus_evaluate`` calls this, and the patch that installs it also + requires ``eval_interval > 0`` while ``do_valid`` does not. A configuration + with ``eval_iters > 0`` and ``eval_interval == 0`` therefore still reaches + pretrain's post-training evaluation through stock Megatron ``evaluate``, + where no session opens and consumers fall back to a free-running counter. + No shipped recipe is in that state. + + ``multiple_validation_sets`` would open one session per set rather than per + evaluation, since ``evaluate_and_print_results`` calls ``evaluate`` once per + iterator. Consumers keying on (step, session) would then hand the same index + to the first microbatch of every set. Primus defaults the flag off in + trainer_base.yaml and no recipe turns it on; enabling it means revisiting + this. + """ + global _eval_session + _eval_session += 1 + return _eval_session + + +def current_eval_session() -> int: + """Identify the evaluation in progress. + + The value is not reproducible across runs and carries no meaning beyond + changing when an evaluation starts, which is all a consumer needs to know + to reset per-evaluation state. + """ + return _eval_session diff --git a/primus/backends/megatron/training/evaluator.py b/primus/backends/megatron/training/evaluator.py index 9ea10e438..2df9dc24f 100644 --- a/primus/backends/megatron/training/evaluator.py +++ b/primus/backends/megatron/training/evaluator.py @@ -15,9 +15,143 @@ from megatron.training import ft_integration, get_args, get_timers from megatron.training.utils import is_last_rank +from primus.backends.megatron.training.eval_budget import get_eval_num_microbatches +from primus.backends.megatron.training.eval_session import begin_eval_session from primus.backends.megatron.training.global_vars import get_train_start_time from primus.backends.megatron.training.utils import is_pipeline_stage_containing_loss -from primus.core.utils.module_utils import log_rank_0 +from primus.core.utils.module_utils import debug_rank_0, log_rank_0 + +# The key under which the diffusion validation path reports +# (summed per-sample loss, sample count). Its denominator is the only one that +# is a sample count rather than a microbatch count. +VAL_LOSS_KEY = "loss" + + +def _report_eval(args, message): + """Report evaluation progress or a result on rank 0, at debug under MLPerf mode. + + MLPerf mode keeps the run's output to a single voice: it stubs out + Megatron's print_rank_last and the tensorboard and wandb writers, and + reports the loss itself, both as an mllog eval_accuracy event and as its own + [MLPerf] line. Repeating any of that alongside a submission log is noise, so + under MLPerf mode it drops to debug and stays in debug.log for a post-mortem. + + Only progress and results that confirm things went as configured come + through here. A coverage shortfall raises, and the one mismatch that does + not raise still reports at info, so nothing this quietens can turn a bad run + into a silent one. + """ + if getattr(args, "mlperf_mode", False): + debug_rank_0(message) + else: + log_rank_0(message) + + +def _record_consumed_valid_samples(args, observed_samples, eval_iters, eval_batch_size): + """Account for the samples the evaluation actually read, and say so. + + The previous behaviour added ``eval_batch_size`` per iteration regardless + of how wide the batches really were, so a short final batch on any worker + was counted as a full one and the logged sample count could exceed the + count evaluated. Reporting the intended number is worse than reporting + nothing, because it is indistinguishable from a correct run. + """ + expected = eval_iters * eval_batch_size + + if observed_samples is None: + # No loss-bearing stage on this rank, or a metric shape that carries + # microbatch counts rather than sample counts. Fall back to the + # configured budget rather than skipping the accounting entirely. + args.consumed_valid_samples += expected + return + + args.consumed_valid_samples += observed_samples + + if observed_samples == expected: + _report_eval( + args, + f"[eval] covered {observed_samples} samples " + f"({eval_iters} iterations x {eval_batch_size} per iteration)", + ) + return + + # Context parallelism duplicates the per-sample loss across CP ranks, which + # inflates the reduced denominator; only assert when it cannot. + cp_size = parallel_state.get_context_parallel_world_size() + detail = ( + f"Evaluation read {observed_samples} samples but the configuration " + f"implies {expected} ({eval_iters} iterations x {eval_batch_size}). " + f"Difference: {expected - observed_samples}." + ) + if cp_size == 1: + raise RuntimeError( + f"{detail}\nThis is the silent under-read described in eval_budget: " + f"Energon workers whose batch quota is short leave the tail of their " + f"slice unread. Check val_num_workers against eval_samples." + ) + log_rank_0(f"[eval] {detail} (context_parallel_size={cp_size}, not asserting)") + + +def _reduction_device(numerators): + """Where to build the packed buffer: wherever the accumulators already live. + + Packing onto the accumulators' own device avoids a transfer and keeps the + buffer on the device the process group can reduce over. + """ + for value in numerators.values(): + if isinstance(value, torch.Tensor): + return value.device + return torch.device("cuda") + + +def reduce_eval_losses(numerators, denominators, dp_group): + """Reduce every accumulated (numerator, denominator) pair across data parallelism. + + One reduction, one group, one host sync. + + Every numerator and denominator across every key is packed into a single + fp64 buffer and reduced once. The reduction must produce a value identical + on every rank: the target-eval-loss early stop compares it against a + threshold, and if ranks disagree near the target one can exit train() alone + while the others keep training, desyncing collectives (grad-norm + all-reduce) into an NCCL hang. + + The previous implementation reduced num/den over DP-with-CP and then + reduced the result again over DP-without-CP. That left both sides + multiplied by data_parallel_size, so the ratio was right but the + denominator could not be read as a sample count. + + Returns: + ``(total_loss_dict, observed_samples)``, where ``observed_samples`` is + the globally reduced ``VAL_LOSS_KEY`` denominator -- a true sample + count -- or None when no such key was reported. + """ + keys = sorted(numerators.keys()) + packed = torch.tensor( + [float(value) for key in keys for value in (numerators[key], denominators[key])], + dtype=torch.float64, + device=_reduction_device(numerators), + ) + torch.distributed.all_reduce(packed, op=torch.distributed.ReduceOp.SUM, group=dp_group) + + # Single host sync for every key at once. + reduced = packed.tolist() + total_loss_dict = {} + observed_samples = None + for index, key in enumerate(keys): + numerator, denominator = reduced[2 * index], reduced[2 * index + 1] + # Keep the result as a 0-dim tensor: downstream Megatron code + # (evaluate_and_print_results) and mlperf logging call .item(). + if denominator > 0: + total_loss_dict[key] = torch.tensor( + numerator / denominator, dtype=torch.float32, device=packed.device + ) + else: + total_loss_dict[key] = torch.zeros((), dtype=torch.float32, device=packed.device) + if key == VAL_LOSS_KEY: + observed_samples = int(round(denominator)) + + return total_loss_dict, observed_samples def primus_evaluate( @@ -34,6 +168,10 @@ def primus_evaluate( args = get_args() timers = get_timers() + # Before any forward step, so consumers deriving per-microbatch state can + # tell this evaluation apart from the previous one even at the same step. + begin_eval_session() + timers("evaluate", log_level=0).start(barrier=True) if args.vision_pretraining and args.vision_pretraining_type == "dino": @@ -56,7 +194,9 @@ def primus_evaluate( # make validation batch size independent from training batch size eval_batch_size = args.global_batch_size - eval_num_microbatches = eval_batch_size // (args.micro_batch_size * args.data_parallel_size) + # Shared with the dataloader provider so the loop and the dataset it reads + # from cannot disagree about how large an evaluation is. + eval_num_microbatches = get_eval_num_microbatches(args) forward_backward_func = get_forward_backward_func() if args.enable_cuda_graph and args.cuda_graph_scope == "full_iteration": forward_backward_func = FullCudaGraphWrapper( @@ -69,11 +209,15 @@ def primus_evaluate( with torch.no_grad(): iteration = 0 if verbose: - log_rank_0(f"Evaluating on {eval_iters * eval_batch_size} samples") + _report_eval(args, f"Evaluating on {eval_iters * eval_batch_size} samples") while iteration < eval_iters: iteration += 1 if verbose: - log_rank_0(f"Evaluating iter {iteration}/{eval_iters}") + # One line per iteration, so 58 per evaluation at the MLPerf + # shape and 580 over a ten-evaluation run. Progress is worth + # watching on an ordinary run and worth nothing in a submission + # log, which reports its own eval_start and eval_stop. + _report_eval(args, f"Evaluating iter {iteration}/{eval_iters}") # Don't care about timing during evaluation config.timers = None @@ -123,8 +267,6 @@ def primus_evaluate( total_loss_numerators[key] += numerator total_loss_denominators[key] += denominator - args.consumed_valid_samples += eval_batch_size - if args.exit_duration_in_mins: train_time = (time.time() - get_train_start_time()) / 60.0 done_cuda = torch.tensor( @@ -137,48 +279,26 @@ def primus_evaluate( log_rank_0("Exiting during evaluation, timelimit reached") return None, None, True - # DP all-reduce for tuple-path (validation) metrics so that every - # rank sees the same globally-averaged loss. Scalar/legacy metrics - # are NOT all-reduced, matching upstream Megatron's evaluate(). total_loss_dict = {} + observed_samples = None if is_pipeline_stage_containing_loss(): from megatron.core import mpu - dp_group = mpu.get_data_parallel_group(with_context_parallel=True) - for key in total_loss_numerators.keys(): - num = total_loss_numerators[key] - den = total_loss_denominators[key] - if isinstance(num, torch.Tensor) and isinstance(den, torch.Tensor): - torch.distributed.all_reduce(num, group=dp_group) - torch.distributed.all_reduce(den, group=dp_group) - - for key in total_loss_numerators.keys(): - # Reduce numerator/denominator across data-parallel ranks so the - # validation loss is a TRUE global average, identical on every rank. - # Without this, args._eval_val_loss stays a per-rank local value, and - # the target-eval-loss early stop (mlperf_pretrain_trainer.py) is then - # evaluated inconsistently: near the target one rank's local loss can - # dip <= target and exit train() alone while the others keep training, - # desyncing collectives (grad-norm all-reduce) -> NCCL hang at ~172k. - reduced = torch.tensor( - [float(total_loss_numerators[key]), float(total_loss_denominators[key])], - dtype=torch.float64, - device="cuda", - ) - torch.distributed.all_reduce( - reduced, - op=torch.distributed.ReduceOp.SUM, - group=parallel_state.get_data_parallel_group(), - ) - # Keep the result as a 0-dim tensor: downstream Megatron code - # (evaluate_and_print_results) and mlperf logging call .item() on it. - if reduced[1].item() > 0: - total_loss_dict[key] = (reduced[0] / reduced[1]).to(torch.float32) - else: - total_loss_dict[key] = torch.zeros((), dtype=torch.float32, device="cuda") - if "lm loss" in total_loss_dict: - val = total_loss_dict["lm loss"] - args._eval_val_loss = val.item() if hasattr(val, "item") else float(val) + total_loss_dict, observed_samples = reduce_eval_losses( + total_loss_numerators, + total_loss_denominators, + mpu.get_data_parallel_group(with_context_parallel=True), + ) + + _record_consumed_valid_samples(args, observed_samples, eval_iters, eval_batch_size) + + # Megatron reports the losses with print_rank_last, and torchrun does + # not forward the last rank's stdout, so on a multi-GPU job the number + # the evaluation exists to produce never reaches the console. Repeat it + # through the Primus logger, which does. + if total_loss_dict: + summary = ", ".join(f"{key}={value.item():.6f}" for key, value in sorted(total_loss_dict.items())) + _report_eval(args, f"[eval] {summary}") collected_non_loss_data = None if non_loss_data_func is not None: diff --git a/primus/configs/data/megatron/diffusion/preprocessing/mlperf_flux1_val.yaml b/primus/configs/data/megatron/diffusion/preprocessing/mlperf_flux1_val.yaml new file mode 100644 index 000000000..373378ea2 --- /dev/null +++ b/primus/configs/data/megatron/diffusion/preprocessing/mlperf_flux1_val.yaml @@ -0,0 +1,69 @@ +# MLPerf Flux1 validation-split-only ingest configuration +# +# Same COCO val split as mlperf_flux1.yaml, without the 1.1 TB CC12M train +# split. Exists because the val split has to be re-ingested on its own: the +# MLCommons val Arrow files carry a per-sample int32 `timestep` column that +# earlier ingests dropped, and evaluation against a split missing that column +# fails under `eval_timestep_source: dataset`. Re-fetching only the ~60 GB val +# split is the cheap way to repair an existing dataset. +# +# ============================================================================ +# Usage +# ============================================================================ +# +# primus data diffusion-ingest \ +# --config primus/configs/data/megatron/diffusion/preprocessing/mlperf_flux1_val.yaml \ +# --output-dir /your/path/mlperf_flux1_val +# +# The resulting directory is a self-contained Energon dataset with an empty +# train split, so it can be used directly as `data_path` for an +# evaluation-only run. To repair a combined train+val dataset instead, ingest +# here first, then move `val/` into place and re-run `energon prepare` over +# the combined directory so the shared index covers the new shards. +# +# ============================================================================ +# Datasets +# ============================================================================ + +datasets: + - name: coco + manifest_url: https://training.mlcommons-storage.org/metadata/flux-1-coco-preprocessed.uri + split_name: val + +# ============================================================================ +# Empty Encodings +# ============================================================================ +# Only used for CFG dropout, which validation suppresses. Fetched anyway (a +# couple of MB) so the directory stands alone as a `data_path`: the Flux +# trainer loads these at startup whenever cfg_dropout_prob > 0, before it +# knows the run will only evaluate. + +empty_encodings: + manifest_url: https://training.mlcommons-storage.org/metadata/flux-1-empty-encodings.uri + output_subdir: empty_encodings + +# ============================================================================ +# Output Configuration +# ============================================================================ + +output: + output_dir: ${PRIMUS_DATA_ROOT:/workspace/Primus/data}/mlperf_flux1_val + +# ============================================================================ +# Pipeline Configuration +# ============================================================================ + +pipeline: + max_workers: 4 # concurrent download threads + prefetch_depth: 6 # max Arrow files buffered on disk (~1.2 GB) + max_files: null # null = all 129 val Arrow files + +# ============================================================================ +# Required Training-Time Settings (not consumed by ingest, for reference) +# ============================================================================ +# Evaluating against this split requires: +# eval_timestep_source: dataset # the column this config exists to restore +# eval_samples: 29696 +# vae_latent_mode: resample +# vae_scale: 0.3611 +# vae_shift: 0.1159 diff --git a/primus/configs/modules/megatron/trainer_base.yaml b/primus/configs/modules/megatron/trainer_base.yaml index 65dc2f9e0..7df701679 100755 --- a/primus/configs/modules/megatron/trainer_base.yaml +++ b/primus/configs/modules/megatron/trainer_base.yaml @@ -236,6 +236,27 @@ overlap_moe_expert_parallel_comm: false train_iters: null eval_iters: 32 +# Size the evaluation by samples rather than iterations. When set, eval_iters is +# derived as eval_samples // global_batch_size, so the budget stays correct when +# the batch size changes. Leave null to keep using eval_iters directly. +eval_samples: null +# Dataloader workers for validation only. Defaults to 0 because Energon splits +# validation data across dp_size x num_workers but splits the batch quota across +# num_workers alone; the two only agree when each worker's slice is a whole +# number of batches, and a mismatch silently leaves the tail of each short +# worker's slice unread. See backends/megatron/training/eval_budget.py. +val_num_workers: 0 +# Run gc.collect() and empty_cache() after every evaluation. Off by default: it +# costs a GC pause and an allocator flush inside the measured window. +eval_purge_memory: false +# Where per-sample validation timesteps come from. +# dataset - the batch must carry a 'timestep' field; missing is an error. +# Required for MLPerf, whose val set ships one timestep per image. +# equidistant - inject t = index % 8 within each validation batch. Only correct +# for datasets that have no per-sample timestep, and only covers +# all eight when the validation batch width is a multiple of 8. +# Defaults to equidistant so configs that never had timesteps keep working. +eval_timestep_source: equidistant full_validation: false multiple_validation_sets: false eval_interval: 2000 diff --git a/pyproject.toml b/pyproject.toml index 57bfbaa90..ed585f033 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,7 +48,7 @@ dependencies = [ "mlflow==3.11.1", "pyrsmi", "plotext", - "mlperf-logging @ git+https://github.com/mlcommons/logging.git@6.0.0-rc5", + "mlperf-logging @ git+https://github.com/mlcommons/logging.git@6.0.0-rc6", "mlperf-common @ git+https://github.com/NVIDIA/mlperf-common.git@b86d175a05849d650a8ff69c1e2c37b9f4e61d51", ] diff --git a/requirements.txt b/requirements.txt index 44713cf61..182e12278 100644 --- a/requirements.txt +++ b/requirements.txt @@ -16,7 +16,7 @@ pyrsmi plotext pulp hip-python -mlperf-logging @ git+https://github.com/mlcommons/logging.git@6.0.0-rc5 +mlperf-logging @ git+https://github.com/mlcommons/logging.git@6.0.0-rc6 mlperf-common @ git+https://github.com/NVIDIA/mlperf-common.git@b86d175a05849d650a8ff69c1e2c37b9f4e61d51 megatron-energon==7.3.2 webdataset==1.0.2 diff --git a/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt b/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt index e2402377d..29bf2fd67 100644 --- a/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt +++ b/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt @@ -14,5 +14,5 @@ imageio[ffmpeg] wandb pydantic transformers==4.50.0 -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 diff --git a/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_ingest.py b/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_ingest.py index 65c1860cb..b3e69f976 100644 --- a/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_ingest.py +++ b/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_ingest.py @@ -11,19 +11,128 @@ - Sample offset accumulation - Arrow file cleanup after conversion - Skip-and-log for failed downloads and conversions +- _arrow_to_tar: sidecar metadata and tensor byte fidelity """ import json +import tarfile import tempfile from pathlib import Path from unittest.mock import patch +import pyarrow as pa + from primus.backends.megatron.data.diffusion.preprocessing.pipelines.ingest import ( StreamingIngestPipeline, + _arrow_to_tar, ) _INGEST_MODULE = "primus.backends.megatron.data.diffusion.preprocessing.pipelines.ingest" +# Tensor columns as declared by both published MLPerf Flux splits. The val +# split (flux-1-coco-preprocessed) adds a scalar int32 "timestep"; the train +# split (flux-1-cc12m-preprocessed) does not. Both shapes are covered below so +# the val-only field cannot regress the 4,762-shard train ingest. +_TENSOR_FEATURES = ("t5_encodings", "clip_encodings", "mean", "logvar") + + +def _write_arrow(path, keys=None, timesteps=None, num_rows=None): + """Write a minimal Arrow IPC stream file in the upstream val/train schema. + + Tensor payloads are short unique byte strings rather than real bf16 buffers: + _arrow_to_tar passes them through opaquely, so their content only needs to + be distinguishable enough to detect mis-ordering. Pass keys=None to emit a + table with no __key__ column. + """ + num_rows = len(keys) if keys is not None else num_rows + labels = keys if keys is not None else [str(i) for i in range(num_rows)] + + columns = {} + if keys is not None: + columns["__key__"] = pa.array(keys, pa.string()) + for col in _TENSOR_FEATURES: + columns[col] = pa.array([f"{col}:{label}".encode() for label in labels], pa.binary()) + if timesteps is not None: + columns["timestep"] = pa.array(timesteps, pa.int32()) + + table = pa.table(columns) + with pa.ipc.new_stream(str(path), table.schema) as writer: + writer.write_table(table) + return table + + +def _read_tar(tar_path): + """Return {member_name: bytes} for every entry in a tar.""" + with tarfile.open(str(tar_path), "r") as tar: + return {m.name: tar.extractfile(m).read() for m in tar.getmembers()} + + +class TestArrowToTarTimestep: + """Direct tests of _arrow_to_tar, which the pipeline tests above mock out. + + The val split carries a per-sample timestep that MLPerf validation is + defined in terms of (sigma = t / 8). Dropping it silently downgrades + evaluation to a positional fallback, so these assert it survives ingest. + """ + + def test_val_schema_carries_timestep_into_sidecar(self, tmp_path): + keys = ["227049", "172952", "183786", "502163"] + # Deliberately non-monotonic and with a repeat, so an arange-style or + # row-index-derived value cannot satisfy this assertion by coincidence. + timesteps = [3, 0, 7, 3] + _write_arrow(tmp_path / "in.arrow", keys, timesteps) + + num_rows = _arrow_to_tar(tmp_path / "in.arrow", tmp_path / "out.tar", 0) + + assert num_rows == len(keys) + members = _read_tar(tmp_path / "out.tar") + for key, expected_t in zip(keys, timesteps): + sidecar = json.loads(members[f"{key}.json"]) + assert sidecar == {"key": key, "timestep": expected_t} + + def test_timestep_is_a_plain_int(self, tmp_path): + """int32 must land as a JSON number, not a nested pyarrow scalar repr.""" + _write_arrow(tmp_path / "in.arrow", ["a"], [5]) + + _arrow_to_tar(tmp_path / "in.arrow", tmp_path / "out.tar", 0) + + sidecar = json.loads(_read_tar(tmp_path / "out.tar")["a.json"]) + assert isinstance(sidecar["timestep"], int) + assert not isinstance(sidecar["timestep"], bool) + + def test_train_schema_without_timestep_is_unaffected(self, tmp_path): + """The train split has no timestep column; ingest must not invent one.""" + keys = ["k0", "k1"] + _write_arrow(tmp_path / "in.arrow", keys, timesteps=None) + + _arrow_to_tar(tmp_path / "in.arrow", tmp_path / "out.tar", 0) + + members = _read_tar(tmp_path / "out.tar") + for key in keys: + assert json.loads(members[f"{key}.json"]) == {"key": key} + + def test_tensor_bytes_pass_through_unchanged(self, tmp_path): + """Payloads must survive byte-for-byte and stay matched to their key.""" + keys = ["227049", "172952"] + table = _write_arrow(tmp_path / "in.arrow", keys, [1, 6]) + + _arrow_to_tar(tmp_path / "in.arrow", tmp_path / "out.tar", 0) + + members = _read_tar(tmp_path / "out.tar") + for row, key in enumerate(keys): + for col, ext in zip(_TENSOR_FEATURES, ("t5.bytes", "clip.bytes", "mean.bytes", "logvar.bytes")): + assert members[f"{key}.{ext}"] == table.column(col)[row].as_py() + + def test_missing_key_column_falls_back_to_offset_naming(self, tmp_path): + """Without __key__, names are global-offset based and still get timesteps.""" + _write_arrow(tmp_path / "in.arrow", keys=None, timesteps=[2, 4], num_rows=2) + + _arrow_to_tar(tmp_path / "in.arrow", tmp_path / "out.tar", 100) + + members = _read_tar(tmp_path / "out.tar") + assert json.loads(members["00000100.json"]) == {"key": "00000100", "timestep": 2} + assert json.loads(members["00000101.json"]) == {"key": "00000101", "timestep": 4} + class TestStreamingIngestPipeline: """Tests for the StreamingIngestPipeline with mocked I/O.""" @@ -65,6 +174,38 @@ def fake_download(url, dest, **kwargs): assert mock_download.call_count == 3 assert mock_convert.call_count == 3 + @patch(f"{_INGEST_MODULE}._arrow_to_tar") + @patch(f"{_INGEST_MODULE}.download_with_backoff") + @patch(f"{_INGEST_MODULE}.fetch_manifest") + def test_manifest_is_filtered_to_arrow_files(self, mock_manifest, mock_download, mock_convert): + """The pipeline asks the manifest for Arrow files only. + + MLCommons manifests list dataset_info.json and state.json alongside the + data. Converting one of those fails and reports the whole run as having + failed files; each would also consume a max_files slot, and one sorting + ahead of a data file would shift every shard index after it. + """ + mock_manifest.return_value = ("https://base.url", [("md5_0", "data-00000.arrow")]) + mock_convert.return_value = 100 + + with tempfile.TemporaryDirectory() as tmp: + pipeline = StreamingIngestPipeline( + manifest_url="http://manifest.uri", + input_dir=f"{tmp}/arrows", + output_dir=f"{tmp}/output", + split_name="val", + ) + + def fake_download(url, dest, **kwargs): + dest.parent.mkdir(parents=True, exist_ok=True) + dest.write_bytes(b"fake") + + mock_download.side_effect = fake_download + + pipeline.run() + + assert mock_manifest.call_args.kwargs["suffix_filter"] == ".arrow" + @patch(f"{_INGEST_MODULE}._arrow_to_tar") @patch(f"{_INGEST_MODULE}.download_with_backoff") @patch(f"{_INGEST_MODULE}.fetch_manifest") diff --git a/tests/unit_tests/backends/megatron/diffusion/data/task_encoders/test_task_encoders.py b/tests/unit_tests/backends/megatron/diffusion/data/task_encoders/test_task_encoders.py index 2bac33b3f..a453ced04 100644 --- a/tests/unit_tests/backends/megatron/diffusion/data/task_encoders/test_task_encoders.py +++ b/tests/unit_tests/backends/megatron/diffusion/data/task_encoders/test_task_encoders.py @@ -10,6 +10,7 @@ import pytest +import torch from primus.backends.megatron.data.diffusion.task_encoders import ( DiffusionSample, @@ -168,6 +169,84 @@ def test_batch_with_all_fields(self, sample_latents, sample_prompt_embeds, sampl assert batch["pooled_prompt_embeds"].shape[0] == 2 +class TestValidationTimestepCollation: + """Tests for per-sample timestep validation during batching. + + MLPerf validation is defined per-sample at t in {0..7}; a batch that + silently loses the field falls back to a positional assignment in + forward_step, so these are error paths rather than conveniences. + """ + + @staticmethod + def _sample(key, timestep, latents, prompt_embeds, pooled): + return DiffusionSample( + __key__=key, + __restore_key__=lambda: key, + __subflavors__={"encoding": "preencoded"}, + latents=latents.squeeze(0), + prompt_embeds=prompt_embeds.squeeze(0), + pooled_prompt_embeds=pooled.squeeze(0), + timestep=None if timestep is None else torch.tensor(timestep), + ) + + def _batch(self, timesteps, sample_latents, sample_prompt_embeds, sample_pooled_prompt_embeds): + encoder = EncodedDiffusionTaskEncoder(worker_config=None) + samples = [ + self._sample(f"s{i}", t, sample_latents, sample_prompt_embeds, sample_pooled_prompt_embeds) + for i, t in enumerate(timesteps) + ] + return encoder.batch(samples) + + def test_timesteps_are_collated(self, sample_latents, sample_prompt_embeds, sample_pooled_prompt_embeds): + batch = self._batch([0, 4, 7], sample_latents, sample_prompt_embeds, sample_pooled_prompt_embeds) + + assert "timestep" in batch + assert batch["timestep"].tolist() == [0, 4, 7] + + def test_absent_timesteps_leave_key_out( + self, sample_latents, sample_prompt_embeds, sample_pooled_prompt_embeds + ): + """The training split has no timestep; batching must not invent one.""" + batch = self._batch([None, None], sample_latents, sample_prompt_embeds, sample_pooled_prompt_embeds) + + assert "timestep" not in batch + + def test_partially_missing_timestep_raises_with_keys( + self, sample_latents, sample_prompt_embeds, sample_pooled_prompt_embeds + ): + """Checking samples[0] alone would let this through or crash opaquely.""" + with pytest.raises(ValueError, match="lack a 'timestep' field") as excinfo: + self._batch([3, None, 5], sample_latents, sample_prompt_embeds, sample_pooled_prompt_embeds) + + assert "s1" in str(excinfo.value) + + def test_missing_timestep_on_first_sample_is_still_caught( + self, sample_latents, sample_prompt_embeds, sample_pooled_prompt_embeds + ): + """The inverse ordering: samples[0] absent but later ones present.""" + with pytest.raises(ValueError, match="lack a 'timestep' field"): + self._batch([None, 2, 6], sample_latents, sample_prompt_embeds, sample_pooled_prompt_embeds) + + def test_out_of_range_timestep_raises( + self, sample_latents, sample_prompt_embeds, sample_pooled_prompt_embeds + ): + with pytest.raises(ValueError, match=r"must be in \[0, 7\]") as excinfo: + self._batch([0, 9, 3], sample_latents, sample_prompt_embeds, sample_pooled_prompt_embeds) + + assert "s1" in str(excinfo.value) + + def test_negative_timestep_raises( + self, sample_latents, sample_prompt_embeds, sample_pooled_prompt_embeds + ): + with pytest.raises(ValueError, match=r"must be in \[0, 7\]"): + self._batch([-1, 3], sample_latents, sample_prompt_embeds, sample_pooled_prompt_embeds) + + def test_float_timestep_raises(self, sample_latents, sample_prompt_embeds, sample_pooled_prompt_embeds): + """sigma = t / 8 expects an integer index, not a pre-divided float.""" + with pytest.raises(ValueError, match="must be an integer type"): + self._batch([0.5, 0.25], sample_latents, sample_prompt_embeds, sample_pooled_prompt_embeds) + + class TestRawDiffusionTaskEncoder: """Tests for RawDiffusionTaskEncoder.""" diff --git a/tests/unit_tests/backends/megatron/diffusion/training/test_eval_timestep_source.py b/tests/unit_tests/backends/megatron/diffusion/training/test_eval_timestep_source.py new file mode 100644 index 000000000..bc0b0e59a --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/training/test_eval_timestep_source.py @@ -0,0 +1,148 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Tests for the explicit validation timestep source. + +Before this knob existed, a validation batch with no ``timestep`` field silently +fell back to injecting ``arange(B) % 8``. That is indistinguishable from a +dataset that legitimately has no per-sample timestep, so an ingest that dropped +the column produced a plausible-looking val_loss computed against the wrong +image-to-timestep pairing. These tests pin the two sources apart. +""" + +import logging + +import pytest +import torch + +from primus.backends.megatron.training.diffusion.forward_step import ( + DATASET_TIMESTEPS, + EQUIDISTANT_TIMESTEPS, + NUM_VALIDATION_TIMESTEPS, + resolve_validation_timesteps, +) + + +def _resolve(batch, source, batch_size, compute_dtype=torch.float32): + return resolve_validation_timesteps( + batch, + source, + batch_size=batch_size, + device=torch.device("cpu"), + compute_dtype=compute_dtype, + ) + + +class TestDatasetSource: + def test_uses_the_dataset_values_verbatim(self): + """The dataset's own pairing must survive, not be re-derived by position.""" + # Deliberately not arange % 8: an injection regression would show up here. + timestep = torch.tensor([3, 7, 3, 0, 5, 5]) + batch = {"timestep": timestep} + + sigmas = _resolve(batch, DATASET_TIMESTEPS, batch_size=6) + + assert torch.equal(batch["timestep"], timestep) + torch.testing.assert_close(sigmas, timestep.float() / NUM_VALIDATION_TIMESTEPS, check_dtype=False) + + def test_sigma_is_the_index_over_eight(self): + batch = {"timestep": torch.arange(NUM_VALIDATION_TIMESTEPS)} + + sigmas = _resolve(batch, DATASET_TIMESTEPS, batch_size=NUM_VALIDATION_TIMESTEPS) + + expected = torch.tensor([0.0, 0.125, 0.25, 0.375, 0.5, 0.625, 0.75, 0.875]) + torch.testing.assert_close(sigmas, expected, check_dtype=False) + + def test_missing_field_is_an_error_naming_the_cause(self): + """The whole point of the knob: fail loudly instead of injecting.""" + with pytest.raises(ValueError, match="carries no 'timestep' field"): + _resolve({}, DATASET_TIMESTEPS, batch_size=64) + + def test_missing_field_error_suggests_the_remedy(self): + with pytest.raises(ValueError) as excinfo: + _resolve({}, DATASET_TIMESTEPS, batch_size=64) + + message = str(excinfo.value) + assert "Re-ingest" in message + assert "equidistant" in message + + def test_a_batch_spanning_only_two_timesteps_is_accepted(self): + """The published val set is ordered so a contiguous batch has two sigmas. + + An assertion that every batch covers all eight would reject the real + data, so confirm no such check crept in. + """ + batch = {"timestep": torch.tensor([0, 4] * 32)} + + sigmas = _resolve(batch, DATASET_TIMESTEPS, batch_size=64) + + assert set(sigmas.tolist()) == {0.0, 0.5} + + +class TestEquidistantSource: + def test_injects_index_modulo_eight(self): + batch = {} + + _resolve(batch, EQUIDISTANT_TIMESTEPS, batch_size=16) + + assert torch.equal(batch["timestep"], torch.arange(16) % NUM_VALIDATION_TIMESTEPS) + + def test_does_not_override_a_present_dataset_timestep(self): + """Equidistant is a fallback, not an override.""" + timestep = torch.tensor([6, 1, 6, 1]) + batch = {"timestep": timestep} + + _resolve(batch, EQUIDISTANT_TIMESTEPS, batch_size=4) + + assert torch.equal(batch["timestep"], timestep) + + def test_full_width_batch_covers_every_timestep(self): + batch = {} + + _resolve(batch, EQUIDISTANT_TIMESTEPS, batch_size=64) + + assert set(batch["timestep"].tolist()) == set(range(NUM_VALIDATION_TIMESTEPS)) + + def test_narrow_batch_warns_that_it_misses_timesteps(self, caplog, monkeypatch): + """micro_batch_size 2 evaluates only t=0 and t=1/8. Say so.""" + monkeypatch.setattr( + "primus.backends.megatron.training.diffusion.forward_step._warned_uncovered_equidistant", + False, + ) + batch = {} + + with caplog.at_level(logging.WARNING): + _resolve(batch, EQUIDISTANT_TIMESTEPS, batch_size=2) + + assert "not a multiple of 8" in caplog.text + assert set(batch["timestep"].tolist()) == {0, 1} + + def test_multiple_of_eight_does_not_warn(self, caplog, monkeypatch): + monkeypatch.setattr( + "primus.backends.megatron.training.diffusion.forward_step._warned_uncovered_equidistant", + False, + ) + + with caplog.at_level(logging.WARNING): + _resolve({}, EQUIDISTANT_TIMESTEPS, batch_size=40) + + assert caplog.text == "" + + +class TestSourceValidation: + @pytest.mark.parametrize("source", ["", "Dataset", "arange", None, 8]) + def test_unknown_source_is_rejected(self, source): + with pytest.raises(ValueError, match="eval_timestep_source must be one of"): + _resolve({"timestep": torch.zeros(4, dtype=torch.long)}, source, batch_size=4) + + def test_rejection_happens_even_when_the_batch_would_have_worked(self): + """A typo must not be masked by the batch happening to carry timesteps.""" + batch = {"timestep": torch.arange(8)} + + with pytest.raises(ValueError, match="eval_timestep_source"): + _resolve(batch, "datset", batch_size=8) + + assert "timesteps" not in batch diff --git a/tests/unit_tests/backends/megatron/diffusion/training/test_flux_forward_step_e2e.py b/tests/unit_tests/backends/megatron/diffusion/training/test_flux_forward_step_e2e.py index 88c440b10..d386bcc74 100644 --- a/tests/unit_tests/backends/megatron/diffusion/training/test_flux_forward_step_e2e.py +++ b/tests/unit_tests/backends/megatron/diffusion/training/test_flux_forward_step_e2e.py @@ -157,7 +157,8 @@ def test_resample_path(self, model, scheduler): ) def test_validation_with_timestep_key(self, model, scheduler): - """Batch with 'timestep' key triggers validation mode.""" + """In eval mode, a 'timestep' key is used verbatim as the sigma source.""" + model.eval() batch = self._make_presampled_batch() batch["timestep"] = torch.arange(2) data_iterator = iter([batch]) @@ -181,6 +182,35 @@ def test_validation_with_timestep_key(self, model, scheduler): torch.arange(2).float() / 8.0, ) + def test_timestep_key_alone_does_not_make_a_training_step_validation(self, model, scheduler): + """Validation is decided by eval mode, never by the batch contents. + + This previously went the other way: a 'timestep' key set is_validation + regardless of model mode, which silently suppressed CFG dropout for any + training batch that happened to carry the field. Eval mode is the only + signal now, and primus_evaluate sets it via model_module.eval(). + """ + assert model.training, "fixture should hand back a model in training mode" + batch = self._make_presampled_batch() + batch["timestep"] = torch.arange(2) + data_iterator = iter([batch]) + + with contextlib.ExitStack() as stack: + for p in _patch_parallel_state(): + stack.enter_context(p) + result = flux_forward_step_func( + data_iterator, + model, + scheduler=scheduler, + step_count=1, + ) + + _, _, _, _, _, is_validation = result + assert is_validation is False + # The training path samples its own timesteps and must not have been + # overridden by the stray field. + assert "timesteps" not in batch + def test_validation_equidistant_injection(self, model, scheduler): """model.eval() without timestep key injects equidistant timesteps.""" model.eval() diff --git a/tests/unit_tests/backends/megatron/test_diffusion_trainer_eval_rng.py b/tests/unit_tests/backends/megatron/test_diffusion_trainer_eval_rng.py new file mode 100644 index 000000000..1c87937f4 --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_diffusion_trainer_eval_rng.py @@ -0,0 +1,275 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for the validation RNG stream. + +The training forward-step counter is deliberately frozen during evaluation so +that eval passes do not shift the training seed sequence. Reusing that frozen +value to reseed, however, gives every validation microbatch the same seed, +which repeats one draw of the VAE reparameterization epsilon and one draw of +the flow-matching noise across the entire evaluation. These tests pin the +separate, advancing, resume-deterministic eval stream that replaces it. + +Two properties carry most of the weight. The index keys on the step under +evaluation rather than on the resume point, so an evaluation at a given step +draws the same noise whether or not the run was resumed to get there. And it +restarts per evaluation rather than per step, so two evaluations at the same +step -- the in-loop one that ends a run and pretrain's post-training one -- +reproduce each other instead of chaining. +""" + +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pytest +import torch + +import primus.backends.megatron.training.eval_session as eval_session +from primus.backends.megatron.diffusion_trainer import DiffusionPretrainTrainer +from primus.backends.megatron.training.diffusion.forward_step import ( + EVAL_RNG_ITERATION_STRIDE, + EVAL_RNG_OFFSET, +) +from primus.backends.megatron.training.eval_session import begin_eval_session + +_UNSET = object() + + +@pytest.fixture(autouse=True) +def _isolate_eval_session(): + """The session counter is module state shared by every test in the process.""" + saved = eval_session._eval_session + yield + eval_session._eval_session = saved + + +class _ConcreteDiffusionTrainer(DiffusionPretrainTrainer): + def create_model(self, *args, **kwargs): + return None + + def create_scheduler(self, *args, **kwargs): + return None + + def get_task_encoder(self, *args, **kwargs): + return None + + +def _make_trainer(): + trainer = _ConcreteDiffusionTrainer.__new__(_ConcreteDiffusionTrainer) + trainer._forward_step_count = 0 + trainer._forward_step_count_initialized = False + trainer._eval_session = None + trainer._eval_microbatch_index = 0 + return trainer + + +def _args(iteration, curr_iteration): + """Megatron args as seen from an evaluation. + + ``curr_iteration`` is genuinely absent before the training loop runs its + first step, so _UNSET is a state the real code has to handle, not a + convenience for the test. + """ + if curr_iteration is _UNSET: + return SimpleNamespace(iteration=iteration) + return SimpleNamespace(iteration=iteration, curr_iteration=curr_iteration) + + +def _indices(trainer, count, curr_iteration=_UNSET, iteration=0): + with patch("megatron.training.get_args", return_value=_args(iteration, curr_iteration)): + return [trainer._next_eval_step_index() for _ in range(count)] + + +def _evaluation(trainer, count, curr_iteration=_UNSET, iteration=0): + """One whole evaluation: the session the evaluator opens, then its microbatches.""" + begin_eval_session() + return _indices(trainer, count, curr_iteration=curr_iteration, iteration=iteration) + + +class TestEvalStepIndex: + def test_advances_within_one_evaluation(self): + trainer = _make_trainer() + assert _evaluation(trainer, count=4, curr_iteration=100) == [ + 100 * EVAL_RNG_ITERATION_STRIDE + i for i in range(4) + ] + + def test_restarts_each_evaluation(self): + """Index restarts at zero per evaluation so no counter needs checkpointing.""" + trainer = _make_trainer() + _evaluation(trainer, count=3, curr_iteration=100) + second = _evaluation(trainer, count=3, curr_iteration=200) + + assert second == [200 * EVAL_RNG_ITERATION_STRIDE + i for i in range(3)] + + def test_repeated_evaluation_at_the_same_step_reproduces_the_first(self): + """The regression test for the post-convergence re-evaluation. + + pretrain runs one more validation after the training loop exits, at the + same step as the evaluation that ended the run. Continuing the counter + would give it fresh noise and so a different loss, which is how a run + that had just converged could report a final loss above target. + """ + trainer = _make_trainer() + in_loop = _evaluation(trainer, count=8, curr_iteration=2559) + post_train = _evaluation(trainer, count=8, curr_iteration=2559) + + assert post_train == in_loop + + def test_reproducible_after_resume(self): + """Same step under evaluation, different resume point, same indices.""" + continuous = _evaluation(_make_trainer(), count=8, curr_iteration=511, iteration=0) + resumed = _evaluation(_make_trainer(), count=8, curr_iteration=511, iteration=256) + + assert resumed == continuous + + def test_keys_on_the_step_not_the_resume_point(self): + trainer = _make_trainer() + indices = _evaluation(trainer, count=4, curr_iteration=100, iteration=512) + + assert indices == [100 * EVAL_RNG_ITERATION_STRIDE + i for i in range(4)] + + def test_falls_back_to_the_checkpoint_step_before_the_loop_runs(self): + """--skip-train evaluates without ever setting curr_iteration.""" + trainer = _make_trainer() + indices = _evaluation(trainer, count=4, iteration=512) + + assert indices == [512 * EVAL_RNG_ITERATION_STRIDE + i for i in range(4)] + + def test_consecutive_evaluations_do_not_collide(self): + trainer = _make_trainer() + first = set(_evaluation(trainer, count=64, curr_iteration=100)) + second = set(_evaluation(trainer, count=64, curr_iteration=101)) + + assert not (first & second) + + def test_first_evaluation_resets_without_a_session_signal(self): + """Stock Megatron evaluate opens no session; the first eval still starts at zero.""" + trainer = _make_trainer() + + assert _indices(trainer, count=3, curr_iteration=7) == [ + 7 * EVAL_RNG_ITERATION_STRIDE + i for i in range(3) + ] + + def test_overrun_of_the_stride_is_rejected(self): + trainer = _make_trainer() + trainer._eval_session = eval_session.current_eval_session() + trainer._eval_microbatch_index = EVAL_RNG_ITERATION_STRIDE + + with patch("megatron.training.get_args", return_value=_args(5, 5)): + with pytest.raises(RuntimeError, match="stride"): + trainer._next_eval_step_index() + + +class TestEvalSeedsAreDisjointFromTraining: + """The eval stream must not alias the training stream on any rank. + + Mirrors the derivations in forward_step: training seeds are + (seed + 100 * dp_rank) * 10000 + step_count, eval seeds add EVAL_RNG_OFFSET. + """ + + @staticmethod + def _train_seed(seed, dp_rank, step_count): + return ((seed + 100 * dp_rank) * 10000 + step_count) % (2**63) + + @staticmethod + def _eval_seed(seed, dp_rank, eval_index): + return (EVAL_RNG_OFFSET + (seed + 100 * dp_rank) * 10000 + eval_index) % (2**63) + + def test_no_overlap_across_plausible_run_shapes(self): + seed = 42 + train = {self._train_seed(seed, rank, step) for rank in range(64) for step in range(0, 20000, 97)} + evals = { + self._eval_seed(seed, rank, step * EVAL_RNG_ITERATION_STRIDE + micro) + for rank in range(64) + for step in range(0, 20000, 512) + for micro in range(58) + } + + assert not (train & evals) + + def test_eval_seeds_differ_per_microbatch(self): + seeds = {self._eval_seed(42, 0, i) for i in range(58)} + assert len(seeds) == 58 + + def test_eval_seeds_differ_per_rank(self): + seeds = {self._eval_seed(42, rank, 0) for rank in range(8)} + assert len(seeds) == 8 + + +class TestForwardStepPassesEvalIndex: + """The trainer must route eval and training passes to different streams.""" + + @staticmethod + def _trainer_with_recorder(per_step_rng_reseed=True): + trainer = _make_trainer() + trainer._scheduler = None + trainer.per_step_rng_reseed = per_step_rng_reseed + + class _FakeRuntimeState: + def update_metrics(self, metrics): + pass + + trainer.runtime_state = _FakeRuntimeState() + + recorded = [] + + def recording_func(*args, **kwargs): + recorded.append(dict(kwargs)) + t = torch.zeros(1) + return t, t, t, None, {}, False + + return trainer, recorded, recording_func + + def test_training_pass_has_no_eval_index(self): + trainer, recorded, recording_func = self._trainer_with_recorder() + model = Mock() + model.training = True + + with patch( + "primus.backends.megatron.training.diffusion.forward_step.flux_forward_step_func", + side_effect=recording_func, + ): + trainer.forward_step(data_iterator=None, model=model) + + assert recorded[-1]["eval_step_index"] is None + assert recorded[-1]["step_count"] == 1 + + def test_eval_pass_freezes_training_counter_but_advances_eval_index(self): + trainer, recorded, recording_func = self._trainer_with_recorder() + trainer._forward_step_count = 7 + model = Mock() + model.training = False + + begin_eval_session() + with patch( + "primus.backends.megatron.training.diffusion.forward_step.flux_forward_step_func", + side_effect=recording_func, + ), patch("megatron.training.get_args", return_value=_args(0, 3)): + trainer.forward_step(data_iterator=None, model=model) + trainer.forward_step(data_iterator=None, model=model) + + # Training counter stays put, so the training seed sequence is unshifted. + assert trainer._forward_step_count == 7 + assert [r["step_count"] for r in recorded] == [7, 7] + # ... but each eval microbatch gets its own seed. + eval_indices = [r["eval_step_index"] for r in recorded] + assert eval_indices == [ + 3 * EVAL_RNG_ITERATION_STRIDE, + 3 * EVAL_RNG_ITERATION_STRIDE + 1, + ] + assert len(set(eval_indices)) == 2 + + def test_no_eval_index_when_reseeding_is_off(self): + """Without reseeding the ambient generator already advances per batch.""" + trainer, recorded, recording_func = self._trainer_with_recorder(per_step_rng_reseed=False) + model = Mock() + model.training = False + + with patch( + "primus.backends.megatron.training.diffusion.forward_step.flux_forward_step_func", + side_effect=recording_func, + ): + trainer.forward_step(data_iterator=None, model=model) + + assert recorded[-1]["eval_step_index"] is None diff --git a/tests/unit_tests/backends/megatron/test_energon_provider_eval_guards.py b/tests/unit_tests/backends/megatron/test_energon_provider_eval_guards.py new file mode 100644 index 000000000..50c491ead --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_energon_provider_eval_guards.py @@ -0,0 +1,126 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Tests that the Energon provider's evaluation guards refuse before building. + +The provider wraps validation-dataloader construction in a bare +``except Exception`` that logs, sets ``eval_iters = 0`` and carries on, so a +check that runs inside it does not refuse the run -- it turns the run into one +that evaluates nothing and exits 0. That is the same swallow that left +``eval_iters`` at 0 when the ``build_args`` patch raised. The guards therefore +have to run ahead of that block, and this pins that they do. +""" + +from types import SimpleNamespace + +import pytest + +import primus.backends.megatron.data.energon_dataset_provider as provider_module +from primus.backends.megatron.data.energon_dataset_provider import ( + EnergonDatasetProvider, +) +from primus.backends.megatron.training.eval_budget import EvalCoverageError + +# The MLPerf Flux shape: 58 iterations x 1 microbatch x 64 x DP 8 = 29,696. +DATA_PARALLEL_SIZE = 8 + + +def _args(**overrides): + base = dict( + skip_train=True, + eval_iters=58, + eval_samples=None, + micro_batch_size=64, + global_batch_size=512, + data_parallel_size=DATA_PARALLEL_SIZE, + val_num_workers=0, + num_workers=4, + prefetch_factor=2, + mlperf_mode=True, + eval_timestep_source="dataset", + ) + base.update(overrides) + return SimpleNamespace(**base) + + +@pytest.fixture +def build_dataloaders(monkeypatch): + """Drive create_dataloaders far enough to reach the guards. + + Everything stubbed here is either rank plumbing or an Energon entry point; + the argument handling and the guards themselves are the real code. + """ + import megatron.training + + val_dataset_calls = [] + + def _stub(args, val_datasets_result=None): + monkeypatch.setattr(megatron.training, "get_args", lambda: args) + monkeypatch.setattr(EnergonDatasetProvider, "_is_dataloader_rank", lambda self: True) + monkeypatch.setattr( + EnergonDatasetProvider, "_create_worker_config", lambda self, args, num_workers=None: object() + ) + monkeypatch.setattr(EnergonDatasetProvider, "_get_data_path", lambda self, args: "/dev/null/dataset") + monkeypatch.setattr( + provider_module.parallel_state, "get_data_parallel_world_size", lambda: DATA_PARALLEL_SIZE + ) + + def _get_val_datasets(*call_args, **call_kwargs): + val_dataset_calls.append(call_kwargs) + if isinstance(val_datasets_result, Exception): + raise val_datasets_result + return val_datasets_result or [] + + monkeypatch.setattr(provider_module, "get_val_datasets", _get_val_datasets) + + provider = EnergonDatasetProvider(task_encoder_factory=lambda: object()) + return provider.create_dataloaders(trainer_config=None, train_val_test_num_samples=[0, 0, 0]) + + _stub.val_dataset_calls = val_dataset_calls + return _stub + + +class TestTimestepSourceGuard: + def test_an_mlperf_run_on_injected_timesteps_never_reaches_construction(self, build_dataloaders): + args = _args(eval_timestep_source="equidistant") + + with pytest.raises(EvalCoverageError, match="eval_timestep_source is 'equidistant'"): + build_dataloaders(args) + + assert build_dataloaders.val_dataset_calls == [] + # The refusal has to leave the budget alone. Zeroing it here is what + # the swallowing handler does, and it is what turns a refusal into a + # run that reports nothing and exits 0. + assert args.eval_iters == 58 + + def test_dataset_timesteps_proceed_to_construction(self, build_dataloaders): + _, valid_dataloaders, _ = build_dataloaders(_args()) + + assert len(build_dataloaders.val_dataset_calls) == 1 + assert valid_dataloaders == [] + + +class TestGuardsRunAheadOfTheSwallow: + def test_a_construction_failure_is_still_swallowed(self, build_dataloaders): + """Not a defect to fix here, but the reason guard placement matters. + + A dataset with no validation split has to degrade rather than abort, + which is why the handler exists. It cannot tell that case apart from + a misconfiguration, so anything raised inside it disappears. + """ + args = _args() + + _, valid_dataloaders, _ = build_dataloaders(args, val_datasets_result=RuntimeError("no val split")) + + assert valid_dataloaders is None + assert args.eval_iters == 0 + + def test_the_zero_eval_iters_guard_also_refuses_before_construction(self, build_dataloaders): + """The guard this PR added for a swallowed build_args patch failure.""" + args = _args(eval_samples=29696, eval_iters=0) + + with pytest.raises(EvalCoverageError, match="no evaluation would run"): + build_dataloaders(args) + + assert build_dataloaders.val_dataset_calls == [] diff --git a/tests/unit_tests/backends/megatron/test_eval_budget.py b/tests/unit_tests/backends/megatron/test_eval_budget.py new file mode 100644 index 000000000..8bfa7505b --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_eval_budget.py @@ -0,0 +1,248 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Tests for the shared evaluation budget derivation. + +The numbers here are the MLPerf Flux shape: 29,696 validation samples, +global batch 512, micro batch 64, DP 8. Both shipped recipes under-read that +set today (27,776 at num_workers=16, 28,928 at num_workers=8), so the +divisibility rule is asserted against those exact configurations. +""" + +import json +from types import SimpleNamespace + +import pytest + +from primus.backends.megatron.training.eval_budget import ( + DEFAULT_VAL_NUM_WORKERS, + EvalCoverageError, + assert_mlperf_timestep_source, + assert_val_worker_divisibility, + get_eval_num_microbatches, + get_val_num_workers, + read_energon_split_sample_count, + resolve_eval_iters, +) + +MLPERF_EVAL_SAMPLES = 29696 + + +def _args(**overrides): + base = dict( + data_parallel_size=8, + micro_batch_size=64, + global_batch_size=512, + eval_iters=58, + eval_samples=None, + val_num_workers=0, + ) + base.update(overrides) + return SimpleNamespace(**base) + + +class TestEvalNumMicrobatches: + def test_mlperf_shape(self): + assert get_eval_num_microbatches(_args()) == 1 + + def test_multiple_microbatches(self): + assert get_eval_num_microbatches(_args(global_batch_size=2048)) == 4 + + def test_zero_microbatches_is_an_error_not_a_silent_noop(self): + """global_batch < one microbatch per rank floored to 0 and evaluated nothing.""" + with pytest.raises(EvalCoverageError, match="smaller than one microbatch"): + get_eval_num_microbatches(_args(global_batch_size=256)) + + def test_indivisible_global_batch_is_an_error(self): + with pytest.raises(EvalCoverageError, match="not divisible"): + get_eval_num_microbatches(_args(global_batch_size=768)) + + +class TestValNumWorkers: + def test_defaults_to_zero_not_to_training_num_workers(self): + args = _args(val_num_workers=None) + args.num_workers = 16 + assert get_val_num_workers(args) == DEFAULT_VAL_NUM_WORKERS == 0 + + def test_explicit_value_is_used(self): + assert get_val_num_workers(_args(val_num_workers=2)) == 2 + + def test_negative_rejected(self): + with pytest.raises(EvalCoverageError, match="must be >= 0"): + get_val_num_workers(_args(val_num_workers=-1)) + + +class TestWorkerDivisibility: + @pytest.mark.parametrize("workers", [0, 1, 2, 29, 58]) + def test_accepts_the_shapes_that_cover_the_set(self, workers): + assert_val_worker_divisibility(_args(val_num_workers=workers), MLPERF_EVAL_SAMPLES) + + @pytest.mark.parametrize( + "workers, observed", + [ + (16, 27776), # local_spec recipe + (8, 28928), # te_spec recipe and the MXFP6 convergence config + ], + ) + def test_rejects_the_shipped_recipe_worker_counts(self, workers, observed): + with pytest.raises(EvalCoverageError, match="silently read fewer") as excinfo: + assert_val_worker_divisibility(_args(val_num_workers=workers), MLPERF_EVAL_SAMPLES) + + message = str(excinfo.value) + assert f"val_num_workers = {workers}" in message + # The error must offer a way out, not just refuse. + assert "Valid val_num_workers for this shape: [0, 1, 2, 29, 58]" in message + + def test_default_worker_count_does_not_divide_by_zero(self): + """max(1, val_num_workers) is what keeps the default of 0 usable.""" + assert_val_worker_divisibility(_args(val_num_workers=0), MLPERF_EVAL_SAMPLES) + + +class TestMlperfTimestepSource: + """The one combination that reports a plausible number instead of failing. + + Injected timesteps are correct-looking whatever fraction of the split was + read, so under mlperf_mode the setting is refused rather than trusted to + the recipe. trainer_base.yaml still defaults to equidistant, which the + non-MLPerf diffusion recipes depend on. + """ + + def test_mlperf_run_on_dataset_timesteps_is_accepted(self): + assert_mlperf_timestep_source(_args(mlperf_mode=True, eval_timestep_source="dataset")) + + @pytest.mark.parametrize("source", ["equidistant", "dataset", None]) + def test_a_non_mlperf_run_is_left_alone(self, source): + assert_mlperf_timestep_source(_args(mlperf_mode=False, eval_timestep_source=source)) + + def test_the_trainer_base_default_still_works_outside_mlperf_mode(self): + """Recipes with no annotated split must keep running unchanged.""" + assert_mlperf_timestep_source(_args()) + + def test_mlperf_run_on_injected_timesteps_is_refused(self): + args = _args(mlperf_mode=True, eval_timestep_source="equidistant") + with pytest.raises(EvalCoverageError, match="eval_timestep_source is 'equidistant'") as excinfo: + assert_mlperf_timestep_source(args) + + message = str(excinfo.value) + # Refusing is only half of it: the message has to name the setting to + # change and the config that produces shards the setting can be met on. + assert "eval_timestep_source='dataset'" in message + assert "mlperf_flux1_val.yaml" in message + + def test_an_absent_setting_fails_closed(self): + """An MLPerf recipe that inherits nothing must not be read as 'dataset'.""" + args = _args(mlperf_mode=True) + with pytest.raises(EvalCoverageError, match="eval_timestep_source is None"): + assert_mlperf_timestep_source(args) + + +class TestResolveEvalIters: + def test_null_eval_samples_leaves_eval_iters_alone(self): + assert resolve_eval_iters(_args(eval_samples=None)) is None + + def test_derives_from_samples(self): + args = _args(eval_samples=MLPERF_EVAL_SAMPLES, eval_iters=0) + assert resolve_eval_iters(args) == 58 + + def test_agreeing_eval_iters_is_accepted(self): + args = _args(eval_samples=MLPERF_EVAL_SAMPLES, eval_iters=58) + assert resolve_eval_iters(args) == 58 + + def test_eval_samples_overrides_an_inherited_eval_iters(self): + """eval_samples wins rather than conflicting. + + trainer_base.yaml gives every module an eval_iters, and Megatron's + parser defaults it too, so a recipe opting into eval_samples always + carries some inherited eval_iters as well. Treating that as a conflict + would reject every such recipe at startup. + """ + args = _args(eval_samples=MLPERF_EVAL_SAMPLES, eval_iters=10) + assert resolve_eval_iters(args) == 58 + + def test_override_holds_for_the_trainer_base_default(self): + args = _args(eval_samples=MLPERF_EVAL_SAMPLES, eval_iters=32) + assert resolve_eval_iters(args) == 58 + + def test_indivisible_sample_count_reports_nearest_reachable(self): + args = _args(eval_samples=MLPERF_EVAL_SAMPLES, eval_iters=0, global_batch_size=768) + with pytest.raises(EvalCoverageError, match="not divisible by global_batch_size") as excinfo: + resolve_eval_iters(args) + assert "29184" in str(excinfo.value) and "29952" in str(excinfo.value) + + def test_non_positive_sample_count_rejected(self): + with pytest.raises(EvalCoverageError, match="must be positive"): + resolve_eval_iters(_args(eval_samples=0, eval_iters=0)) + + def test_derivation_also_enforces_coverage(self): + """Deriving a correct iteration count is not enough if workers under-read.""" + args = _args(eval_samples=MLPERF_EVAL_SAMPLES, eval_iters=0, val_num_workers=8) + with pytest.raises(EvalCoverageError, match="silently read fewer"): + resolve_eval_iters(args) + + def test_derivation_tracks_global_batch_size(self): + """The point of eval_samples: the same coverage at a different batch size.""" + args = _args(eval_samples=MLPERF_EVAL_SAMPLES, eval_iters=0, global_batch_size=1024) + assert resolve_eval_iters(args) == 29 + + +class TestReadEnergonSplitSampleCount: + """Reading the split size from the dataset's own index. + + This is what lets full_validation mean "all of it" without the count being + hand-written into a recipe, and gives the evaluation loop an independent + third number to check against. + """ + + @staticmethod + def _dataset(tmp_path, shard_counts): + meta = tmp_path / ".nv-meta" + meta.mkdir() + (meta / ".info.json").write_text( + json.dumps({"energon_version": "7.3.2", "shard_counts": shard_counts}) + ) + return tmp_path + + def test_sums_only_the_requested_split(self, tmp_path): + path = self._dataset( + tmp_path, + { + "train/shard_000000.tar": 231, + "train/shard_000001.tar": 230, + "val/shard_000000.tar": 231, + "val/shard_000001.tar": 230, + }, + ) + + assert read_energon_split_sample_count(path, "val") == 461 + assert read_energon_split_sample_count(path, "train") == 461 + + def test_mlperf_shape(self, tmp_path): + """26 shards of 231 and 103 of 230 is the published val split.""" + counts = {f"val/shard_{i:06d}.tar": (231 if i < 26 else 230) for i in range(129)} + path = self._dataset(tmp_path, counts) + + assert read_energon_split_sample_count(path, "val") == MLPERF_EVAL_SAMPLES + + def test_accepts_list_form_data_path(self, tmp_path): + path = self._dataset(tmp_path, {"val/shard_000000.tar": 7}) + assert read_energon_split_sample_count([str(path)]) == 7 + + @pytest.mark.parametrize("value", [None, [], ""]) + def test_empty_inputs_return_none(self, value): + assert read_energon_split_sample_count(value) is None + + def test_missing_dataset_returns_none_rather_than_raising(self, tmp_path): + """Mock data and unprepared directories must fall back, not fail.""" + assert read_energon_split_sample_count(tmp_path / "nope") is None + + def test_malformed_index_returns_none(self, tmp_path): + meta = tmp_path / ".nv-meta" + meta.mkdir() + (meta / ".info.json").write_text("not json") + + assert read_energon_split_sample_count(tmp_path) is None + + def test_absent_split_returns_none(self, tmp_path): + path = self._dataset(tmp_path, {"train/shard_000000.tar": 231}) + assert read_energon_split_sample_count(path, "val") is None diff --git a/tests/unit_tests/backends/megatron/test_eval_samples_patch.py b/tests/unit_tests/backends/megatron/test_eval_samples_patch.py new file mode 100644 index 000000000..86569eca8 --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_eval_samples_patch.py @@ -0,0 +1,157 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Tests for the build_args patch that derives eval_iters from eval_samples. + +The case these exist for: ``eval_samples`` and ``val_num_workers`` are not +Megatron arguments, and MegatronArgBuilder drops every key Megatron's parser +does not define. The runtime does merge the leftover Primus-only params into +backend_args, but only *after* the build_args phase has run. So a patch in this +phase that reads them straight off args sees them as unset and silently leaves +eval_iters alone -- the exact failure these tests pin, since it turns the +coverage fix into a no-op while still looking like a working config. +""" + +from types import SimpleNamespace +from unittest.mock import patch as mock_patch + +import pytest + +from primus.backends.megatron.patches.args import eval_samples_patches +from primus.backends.megatron.patches.args.eval_samples_patches import ( + patch_eval_samples, +) +from primus.backends.megatron.training.eval_budget import EvalCoverageError +from primus.core.patches import PatchContext + +MLPERF_EVAL_SAMPLES = 29696 + + +@pytest.fixture(autouse=True) +def _silence_patch_logging(): + """The Primus logger is not initialised outside a real run.""" + with mock_patch.object(eval_samples_patches, "log_kv_rank_0", lambda *a, **k: None): + yield + + +def _megatron_args(**overrides): + """backend_args as MegatronArgBuilder produces them: no Primus-only keys.""" + base = dict( + data_parallel_size=8, + micro_batch_size=64, + global_batch_size=512, + eval_iters=32, + full_validation=False, + data_path="/does/not/exist", + ) + base.update(overrides) + return SimpleNamespace(**base) + + +def _ctx(args, **primus_only): + module_config = SimpleNamespace(params=SimpleNamespace(**primus_only)) + return PatchContext( + backend="megatron", + phase="build_args", + extra={"backend_args": args, "module_config": module_config}, + ) + + +class TestPrimusOnlyKeysReachThePatch: + def test_eval_samples_from_module_config_derives_eval_iters(self): + """The regression: eval_samples lives only on the module config here.""" + args = _megatron_args() + + patch_eval_samples(_ctx(args, eval_samples=MLPERF_EVAL_SAMPLES, val_num_workers=0)) + + assert args.eval_iters == 58, "eval_samples on the module config must still derive eval_iters" + + def test_val_num_workers_from_module_config_is_enforced(self): + """A worker count that cannot read the whole set must still be caught.""" + args = _megatron_args() + + with pytest.raises(EvalCoverageError): + patch_eval_samples(_ctx(args, eval_samples=MLPERF_EVAL_SAMPLES, val_num_workers=8)) + + def test_args_value_wins_over_module_config(self): + """If a key already reached args, do not overwrite it from the config.""" + args = _megatron_args(eval_samples=MLPERF_EVAL_SAMPLES, val_num_workers=0) + + patch_eval_samples(_ctx(args, eval_samples=512, val_num_workers=0)) + + assert args.eval_samples == MLPERF_EVAL_SAMPLES + assert args.eval_iters == 58 + + def test_missing_module_config_does_not_crash(self): + args = _megatron_args(eval_iters=58) + ctx = PatchContext( + backend="megatron", + phase="build_args", + extra={"backend_args": args}, + ) + + patch_eval_samples(ctx) + + assert args.eval_iters == 58 + + def test_no_backend_args_is_a_no_op(self): + patch_eval_samples(PatchContext(backend="megatron", phase="build_args", extra={})) + + +class TestEvalItersPath: + def test_plain_eval_iters_still_gets_a_coverage_check(self): + """Configs that never opt into eval_samples are still covered.""" + args = _megatron_args(eval_iters=58) + + with pytest.raises(EvalCoverageError): + patch_eval_samples(_ctx(args, val_num_workers=8)) + + def test_reachable_eval_iters_shape_is_left_alone(self): + args = _megatron_args(eval_iters=58) + + patch_eval_samples(_ctx(args, val_num_workers=0)) + + assert args.eval_iters == 58 + + def test_zero_eval_iters_is_not_checked(self): + """eval_iters 0 means "no validation"; there is nothing to cover.""" + args = _megatron_args(eval_iters=0) + + patch_eval_samples(_ctx(args, val_num_workers=8)) + + assert args.eval_iters == 0 + + +class TestFullValidation: + def test_full_validation_reads_the_split_size(self): + args = _megatron_args(full_validation=True) + + with mock_patch.object( + eval_samples_patches, + "read_energon_split_sample_count", + return_value=MLPERF_EVAL_SAMPLES, + ): + patch_eval_samples(_ctx(args, val_num_workers=0)) + + assert args.eval_samples == MLPERF_EVAL_SAMPLES + assert args.eval_iters == 58 + + def test_full_validation_without_a_readable_index_errors(self): + args = _megatron_args(full_validation=True) + + with mock_patch.object(eval_samples_patches, "read_energon_split_sample_count", return_value=None): + with pytest.raises(ValueError, match="could not be read"): + patch_eval_samples(_ctx(args, val_num_workers=0)) + + def test_explicit_eval_samples_takes_precedence(self): + """full_validation must not re-read the index when a count is given.""" + args = _megatron_args(full_validation=True) + + with mock_patch.object(eval_samples_patches, "read_energon_split_sample_count") as reader: + patch_eval_samples(_ctx(args, eval_samples=MLPERF_EVAL_SAMPLES, val_num_workers=0)) + + reader.assert_not_called() + assert args.eval_iters == 58 diff --git a/tests/unit_tests/backends/megatron/test_evaluator_reduction.py b/tests/unit_tests/backends/megatron/test_evaluator_reduction.py new file mode 100644 index 000000000..f6fe405d6 --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_evaluator_reduction.py @@ -0,0 +1,397 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for the validation reduction and the sample accounting around it. + +These cover the code that produces the number the run actually reports. Two +properties matter and neither is visible from the per-microbatch loss function: +the reported loss must be a ratio of globally summed numerator and denominator +rather than an average of per-rank ratios, and the reduced denominator must +survive as a true sample count so an under-read can be detected instead of +being papered over with the count the configuration intended. +""" + +from contextlib import contextmanager +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch + +import primus.backends.megatron.training.eval_session as eval_session +from primus.backends.megatron.training.evaluator import ( + VAL_LOSS_KEY, + _record_consumed_valid_samples, + _report_eval, + primus_evaluate, + reduce_eval_losses, +) + +GROUP = object() +EVALUATOR = "primus.backends.megatron.training.evaluator" + + +class FakeAllReduce: + """Stand-in for the DP all-reduce that can add other ranks' contributions.""" + + def __init__(self, other_ranks=()): + self.other_ranks = other_ranks + self.buffers = [] + self.groups = [] + + def __call__(self, tensor, op=None, group=None): + self.buffers.append(tensor.clone()) + self.groups.append(group) + for contribution in self.other_ranks: + tensor += torch.tensor(contribution, dtype=tensor.dtype, device=tensor.device) + + @property + def call_count(self): + return len(self.buffers) + + +def _accumulators(pairs): + """Build the (numerator, denominator) dicts the eval loop accumulates.""" + numerators = {key: torch.tensor(num) for key, (num, _) in pairs.items()} + denominators = {key: torch.tensor(float(den)) for key, (_, den) in pairs.items()} + return numerators, denominators + + +class TestReduction: + def test_loss_is_the_ratio_of_the_summed_pair(self): + numerators, denominators = _accumulators({VAL_LOSS_KEY: (150.0, 300.0)}) + with patch("torch.distributed.all_reduce", FakeAllReduce()): + losses, _ = reduce_eval_losses(numerators, denominators, GROUP) + + assert losses[VAL_LOSS_KEY].item() == pytest.approx(0.5) + + def test_ratio_of_sums_not_the_mean_of_per_rank_ratios(self): + """The property the old double-reduce lost. + + Rank 0 reports 100/200 = 0.5 and rank 1 reports 900/100 = 9.0. The + correct sample-weighted answer is 1000/300 = 3.33, not the 4.75 an + average of the two per-rank ratios would give. + """ + numerators, denominators = _accumulators({VAL_LOSS_KEY: (100.0, 200.0)}) + other_rank = FakeAllReduce(other_ranks=[[900.0, 100.0]]) + with patch("torch.distributed.all_reduce", other_rank): + losses, _ = reduce_eval_losses(numerators, denominators, GROUP) + + assert losses[VAL_LOSS_KEY].item() == pytest.approx(1000.0 / 300.0, rel=1e-6) + assert losses[VAL_LOSS_KEY].item() != pytest.approx((0.5 + 9.0) / 2) + + def test_one_reduction_no_matter_how_many_keys(self): + numerators, denominators = _accumulators({"loss": (1.0, 2.0), "aux": (3.0, 4.0), "extra": (5.0, 6.0)}) + fake = FakeAllReduce() + with patch("torch.distributed.all_reduce", fake): + reduce_eval_losses(numerators, denominators, GROUP) + + assert fake.call_count == 1 + + def test_buffer_is_fp64(self): + """bf16 accumulators must not decide the precision of the reported loss.""" + numerators = {VAL_LOSS_KEY: torch.tensor(150.0, dtype=torch.bfloat16)} + denominators = {VAL_LOSS_KEY: torch.tensor(300.0, dtype=torch.bfloat16)} + fake = FakeAllReduce() + with patch("torch.distributed.all_reduce", fake): + reduce_eval_losses(numerators, denominators, GROUP) + + assert fake.buffers[0].dtype == torch.float64 + + def test_keys_are_packed_as_sorted_pairs(self): + """Packing order is what lets one flat buffer be unpacked per key.""" + numerators, denominators = _accumulators({"loss": (1.0, 2.0), "aux": (3.0, 4.0)}) + fake = FakeAllReduce() + with patch("torch.distributed.all_reduce", fake): + reduce_eval_losses(numerators, denominators, GROUP) + + assert fake.buffers[0].tolist() == [3.0, 4.0, 1.0, 2.0] + + def test_each_key_keeps_its_own_ratio(self): + numerators, denominators = _accumulators({"loss": (1.0, 2.0), "aux": (9.0, 3.0)}) + with patch("torch.distributed.all_reduce", FakeAllReduce()): + losses, _ = reduce_eval_losses(numerators, denominators, GROUP) + + assert losses["loss"].item() == pytest.approx(0.5) + assert losses["aux"].item() == pytest.approx(3.0) + + def test_empty_denominator_yields_zero_rather_than_nan(self): + numerators, denominators = _accumulators({VAL_LOSS_KEY: (0.0, 0)}) + with patch("torch.distributed.all_reduce", FakeAllReduce()): + losses, observed = reduce_eval_losses(numerators, denominators, GROUP) + + assert losses[VAL_LOSS_KEY].item() == 0.0 + assert not torch.isnan(losses[VAL_LOSS_KEY]) + assert observed == 0 + + def test_result_is_a_zero_dim_tensor(self): + """Megatron's evaluate_and_print_results and the mlperf logger call .item().""" + numerators, denominators = _accumulators({VAL_LOSS_KEY: (1.0, 2.0)}) + with patch("torch.distributed.all_reduce", FakeAllReduce()): + losses, _ = reduce_eval_losses(numerators, denominators, GROUP) + + assert losses[VAL_LOSS_KEY].shape == () + assert losses[VAL_LOSS_KEY].dtype == torch.float32 + + def test_reduces_over_the_group_it_was_given(self): + numerators, denominators = _accumulators({VAL_LOSS_KEY: (1.0, 2.0)}) + fake = FakeAllReduce() + with patch("torch.distributed.all_reduce", fake): + reduce_eval_losses(numerators, denominators, GROUP) + + assert fake.groups == [GROUP] + + +class TestObservedSampleCount: + def test_is_the_globally_summed_denominator(self): + """29,696 total across 8 ranks of 3,712 -- the MLPerf shape.""" + numerators, denominators = _accumulators({VAL_LOSS_KEY: (2000.0, 3712)}) + seven_more = FakeAllReduce(other_ranks=[[2000.0, 3712.0]] * 7) + with patch("torch.distributed.all_reduce", seven_more): + _, observed = reduce_eval_losses(numerators, denominators, GROUP) + + assert observed == 29696 + + def test_is_an_integer_not_a_float(self): + numerators, denominators = _accumulators({VAL_LOSS_KEY: (1.0, 512)}) + with patch("torch.distributed.all_reduce", FakeAllReduce()): + _, observed = reduce_eval_losses(numerators, denominators, GROUP) + + assert isinstance(observed, int) + + def test_absent_without_the_val_loss_key(self): + """A microbatch-counting metric shape must not masquerade as a count.""" + numerators, denominators = _accumulators({"lm loss": (1.0, 2.0)}) + with patch("torch.distributed.all_reduce", FakeAllReduce()): + _, observed = reduce_eval_losses(numerators, denominators, GROUP) + + assert observed is None + + +class TestConsumedValidSamples: + """The accounting that turns a short read into a failure instead of a log line.""" + + @staticmethod + def _args(consumed=0): + return SimpleNamespace(consumed_valid_samples=consumed) + + @staticmethod + def _record(args, observed, eval_iters=58, eval_batch_size=512, cp_size=1): + with patch( + "primus.backends.megatron.training.evaluator.parallel_state." "get_context_parallel_world_size", + return_value=cp_size, + ): + _record_consumed_valid_samples(args, observed, eval_iters, eval_batch_size) + + def test_full_coverage_accumulates_the_observed_count(self): + args = self._args() + self._record(args, observed=29696) + + assert args.consumed_valid_samples == 29696 + + def test_accumulates_across_evaluations(self): + args = self._args(consumed=29696) + self._record(args, observed=29696) + + assert args.consumed_valid_samples == 2 * 29696 + + def test_under_read_is_an_error_naming_the_shortfall(self): + """The 27,776-sample under-read that started all of this.""" + args = self._args() + with pytest.raises(RuntimeError, match="read 27776 samples"): + self._record(args, observed=27776) + + def test_under_read_error_points_at_the_worker_count(self): + args = self._args() + with pytest.raises(RuntimeError, match="val_num_workers"): + self._record(args, observed=27776) + + def test_over_read_is_also_an_error(self): + """Double-counting is as wrong as under-counting, and just as silent.""" + args = self._args() + with pytest.raises(RuntimeError, match="read 59392 samples"): + self._record(args, observed=2 * 29696) + + def test_context_parallelism_logs_instead_of_raising(self): + """CP duplicates the per-sample loss, so the denominator is legitimately inflated.""" + args = self._args() + with patch("primus.backends.megatron.training.evaluator.log_rank_0") as log: + self._record(args, observed=2 * 29696, cp_size=2) + + assert args.consumed_valid_samples == 2 * 29696 + assert "context_parallel_size=2" in log.call_args[0][0] + + def test_no_observation_falls_back_to_the_configured_budget(self): + """Stages with no loss still have to advance the counter.""" + args = self._args() + self._record(args, observed=None) + + assert args.consumed_valid_samples == 58 * 512 + + def test_fallback_does_not_raise_on_a_mismatch_it_cannot_see(self): + args = self._args() + self._record(args, observed=None, eval_iters=1, eval_batch_size=7) + + assert args.consumed_valid_samples == 7 + + +class TestEvalReporting: + """Which stream the evaluation's own reporting goes to. + + MLPerf mode reports the loss itself, as an mllog event and as its own line, + and silences Megatron's reporting to keep the run to one voice; repeating it + at info level there is noise. What must not follow is a coverage shortfall + going quiet along with it. + """ + + @staticmethod + def _args(mlperf_mode=False): + return SimpleNamespace(consumed_valid_samples=0, mlperf_mode=mlperf_mode) + + @staticmethod + @contextmanager + def _streams(cp_size=1): + """Capture the info and debug streams the evaluator reports through.""" + with patch( + f"{EVALUATOR}.parallel_state.get_context_parallel_world_size", + return_value=cp_size, + ), patch(f"{EVALUATOR}.log_rank_0") as info, patch(f"{EVALUATOR}.debug_rank_0") as debug: + yield info, debug + + def test_reports_at_info_by_default(self): + with self._streams() as (info, debug): + _report_eval(self._args(), "[eval] loss=1.665372") + + assert info.call_args[0][0] == "[eval] loss=1.665372" + assert debug.call_count == 0 + + def test_mlperf_mode_reports_at_debug_instead(self): + with self._streams() as (info, debug): + _report_eval(self._args(mlperf_mode=True), "[eval] loss=1.665372") + + assert debug.call_args[0][0] == "[eval] loss=1.665372" + assert info.call_count == 0 + + def test_an_absent_flag_reads_as_off(self): + """args is built before MLPerf mode is a settled attribute on it.""" + with self._streams() as (info, debug): + _report_eval(SimpleNamespace(), "[eval] loss=1.665372") + + assert info.call_count == 1 + assert debug.call_count == 0 + + def test_mlperf_mode_moves_the_coverage_line_too(self): + args = self._args(mlperf_mode=True) + with self._streams() as (info, debug): + _record_consumed_valid_samples(args, 29696, 58, 512) + + assert "covered 29696 samples" in debug.call_args[0][0] + assert info.call_count == 0 + + def test_mlperf_mode_does_not_quieten_an_under_read(self): + """The whole point of the coverage check survives the quietening.""" + args = self._args(mlperf_mode=True) + with self._streams(): + with pytest.raises(RuntimeError, match="read 27776 samples"): + _record_consumed_valid_samples(args, 27776, 58, 512) + + def test_mlperf_mode_still_reports_a_context_parallel_mismatch(self): + """A mismatch CP can explain is a diagnostic, not a confirmation.""" + args = self._args(mlperf_mode=True) + with self._streams(cp_size=2) as (info, debug): + _record_consumed_valid_samples(args, 2 * 29696, 58, 512) + + assert "context_parallel_size=2" in info.call_args[0][0] + assert debug.call_count == 0 + + +class TestEvaluationSessions: + """Every evaluation must announce itself before it reads anything. + + The diffusion trainer derives its validation RNG index from this signal, so + that two evaluations at the same training step -- the in-loop one that ends + a run and the one pretrain runs afterwards -- reproduce each other rather + than drawing fresh noise. That only holds if the evaluator opens a session, + and the trainer cannot tell that it failed to. + """ + + @pytest.fixture(autouse=True) + def _isolate_eval_session(self): + """Driving primus_evaluate advances process-wide state.""" + saved = eval_session._eval_session + yield + eval_session._eval_session = saved + + @staticmethod + def _args(): + return SimpleNamespace( + vision_pretraining=False, + vision_pretraining_type=None, + global_batch_size=512, + micro_batch_size=64, + seq_length=256, + decoder_seq_length=None, + enable_cuda_graph=False, + cuda_graph_scope=None, + eval_iters=2, + empty_unused_memory_level=0, + exit_duration_in_mins=None, + consumed_valid_samples=0, + mlperf_mode=False, + ) + + @contextmanager + def _evaluator(self, sessions_seen): + """primus_evaluate with everything below the session signal stubbed out.""" + + def forward_backward_func(**kwargs): + sessions_seen.append(eval_session.current_eval_session()) + return [{}] + + with patch(f"{EVALUATOR}.get_args", return_value=self._args()), patch( + f"{EVALUATOR}.get_timers" + ), patch(f"{EVALUATOR}.get_rerun_state_machine"), patch(f"{EVALUATOR}.ft_integration"), patch( + f"{EVALUATOR}.get_eval_num_microbatches", return_value=1 + ), patch( + f"{EVALUATOR}.get_forward_backward_func", return_value=forward_backward_func + ), patch( + f"{EVALUATOR}.is_pipeline_stage_containing_loss", return_value=False + ), patch( + f"{EVALUATOR}._report_eval" + ): + yield + + def _evaluate(self, sessions_seen): + with self._evaluator(sessions_seen): + primus_evaluate( + forward_step_func=lambda *a, **k: None, + data_iterator=None, + model=[MagicMock()], + process_non_loss_data_func=None, + config=SimpleNamespace(timers=None), + ) + + def test_a_session_is_open_before_the_first_forward_step(self): + before = eval_session.current_eval_session() + seen = [] + + self._evaluate(seen) + + assert seen, "the evaluation ran no forward steps" + assert all(session > before for session in seen) + + def test_every_forward_step_of_one_evaluation_sees_the_same_session(self): + seen = [] + + self._evaluate(seen) + + assert len(set(seen)) == 1 + + def test_consecutive_evaluations_open_different_sessions(self): + first, second = [], [] + + self._evaluate(first) + self._evaluate(second) + + assert set(first).isdisjoint(second) diff --git a/tests/unit_tests/backends/megatron/test_mlperf_log_is_checker_valid.py b/tests/unit_tests/backends/megatron/test_mlperf_log_is_checker_valid.py new file mode 100644 index 000000000..95bc5c460 --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_mlperf_log_is_checker_valid.py @@ -0,0 +1,320 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Drive a full run's worth of MLPerf records and hand the file to the checker. + +Every other test in this area asserts against a mock and therefore only says +that the code emitted what the test expected. This one uses the real +``mlperf_logging`` package end to end: real ``mllog`` writes the file, and +``mlperf_logging.compliance_checker`` reads it back under the same ruleset a +submission is judged with. It fails when the emitted log stops being a valid +one, including for reasons nobody wrote an assertion for. + +Skipped when ``mlperf_logging`` is not installed, which is the case on plain +development hosts; the package ships in the training image. +""" + +import types +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +mlperf_logging = pytest.importorskip("mlperf_logging") + +RULESET = "6.0.0" +GLOBAL_BATCH_SIZE = 512 +MICRO_BATCH_SIZE = 64 +EVAL_INTERVAL = 512 +TARGET_VAL_LOSS = 0.586 + + +def _has_ruleset() -> bool: + from pathlib import Path + + root = Path(mlperf_logging.__file__).parent + return (root / "compliance_checker" / f"training_{RULESET}" / "closed_flux1.yaml").exists() + + +pytestmark = pytest.mark.skipif( + not _has_ruleset(), + reason=f"installed mlperf_logging has no training_{RULESET} flux1 ruleset", +) + + +def _install_fake_megatron(monkeypatch): + import sys + + megatron_mod = types.ModuleType("megatron") + training_pkg = types.ModuleType("megatron.training") + training_mod = types.ModuleType("megatron.training.training") + global_vars_mod = types.ModuleType("megatron.training.global_vars") + + training_mod.train_step = lambda *a, **k: ({}, 0, False, False, 0, 0.0, 0, None) + training_mod.training_log = lambda *a, **k: None + training_mod.print_rank_last = lambda msg: None + training_mod.get_tensorboard_writer = lambda: None + training_mod.get_wandb_writer = lambda: None + training_mod.setup_model_and_optimizer = lambda *a, **k: ("model", "optimizer", "scheduler") + training_mod.build_train_valid_test_data_iterators = lambda *a, **k: ("train", "valid", "test") + training_mod.get_model_config = lambda model: SimpleNamespace() + training_mod.get_forward_backward_func = lambda: (lambda *a, **k: None) + training_pkg.pretrain = lambda *a, **k: None + + training_pkg.training = training_mod + megatron_mod.training = training_pkg + + megatron_args = SimpleNamespace( + iteration=0, + curr_iteration=0, + consumed_train_samples=0, + skipped_train_samples=0, + train_iters=16000, + eval_interval=EVAL_INTERVAL, + do_valid=True, + global_batch_size=GLOBAL_BATCH_SIZE, + micro_batch_size=MICRO_BATCH_SIZE, + ) + global_vars_mod.get_args = lambda: megatron_args + training_pkg.get_args = global_vars_mod.get_args + + monkeypatch.setitem(sys.modules, "megatron", megatron_mod) + monkeypatch.setitem(sys.modules, "megatron.training", training_pkg) + monkeypatch.setitem(sys.modules, "megatron.training.training", training_mod) + monkeypatch.setitem(sys.modules, "megatron.training.global_vars", global_vars_mod) + + return training_mod, megatron_args + + +def _make_ctx(): + """A context carrying the values closed_flux1.yaml pins exactly.""" + params = SimpleNamespace( + mlperf_mode=True, + warmup_train_steps=0, + target_val_loss=TARGET_VAL_LOSS, + global_batch_size=GLOBAL_BATCH_SIZE, + micro_batch_size=MICRO_BATCH_SIZE, + data_parallel_size=8, + seed=42, + log_interval=10, + lr=2e-4, + adam_beta1=0.9, + adam_beta2=0.95, + adam_eps=1e-8, + weight_decay=0.1, + clip_grad=1.0, + lr_warmup_iters=1600, + eval_interval=EVAL_INTERVAL, + eval_samples=29696, + train_samples=1099776, + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + context_parallel_size=1, + expert_model_parallel_size=1, + transformer_impl="local", + wall_clock_step_timer=False, + ) + return SimpleNamespace( + extra={"module_config": SimpleNamespace(params=params)}, + backend="megatron", + phase="before_train", + ) + + +@pytest.fixture +def _run_environment(monkeypatch, tmp_path): + from primus.backends.megatron.patches import mlperf_boundary + + mlperf_boundary.reset_for_tests() + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("MLLOG_OUTPUT_FILE", str(tmp_path / "result_0.txt")) + monkeypatch.setenv("MLPERF_CLEAR_CACHES", "true") + monkeypatch.setenv("MLLOG_SUBMISSION_ORG", "AMD") + monkeypatch.setenv("MLLOG_SUBMISSION_DIVISION", "closed") + monkeypatch.setenv("MLLOG_SUBMISSION_PLATFORM", "MI355X") + monkeypatch.setenv("EXP", "examples/megatron/configs/MI355X/diffusion/flux_mlperf.yaml") + monkeypatch.setenv("MLLOG_LOWEST_NUMERICAL_PRECISION_IN_LINEAR", "fp8") + monkeypatch.setenv("MLLOG_LOWEST_NUMERICAL_PRECISION_IN_ATTN", "bfloat16") + monkeypatch.setenv("MLLOG_LOWEST_NUMERICAL_PRECISION_IN_COMM", "bfloat16") + yield tmp_path / "result_0.txt" + mlperf_boundary.reset_for_tests() + _detach_mllog_handlers() + + +def _detach_mllog_handlers(): + """mllog's logger is a process-wide singleton; leave it as it was found.""" + from mlperf_logging import mllog + + logger = mllog.get_mllogger().logger + for handler in list(logger.handlers): + logger.removeHandler(handler) + handler.close() + + +def _drive_a_run(mt, eval_losses): + """Walk the run through the same call sequence Megatron would.""" + for module in ( + "primus.backends.megatron.patches.mlperf_logging_patches", + "primus.backends.megatron.patches.mlperf_warmup_patches", + ): + pytest.MonkeyPatch().setattr(f"{module}.log_rank_0", lambda *a, **k: None) + + from primus.backends.megatron.patches.mlperf_logging_patches import ( + patch_mlperf_logging, + ) + + patch_mlperf_logging(_make_ctx()) + + # Data iterators are built after the model, and this is where the clock + # starts. + mt.build_train_valid_test_data_iterators(None) + + iteration = 0 + for loss in eval_losses: + for _ in range(EVAL_INTERVAL // 10): + iteration += 10 + mt.training_log({"loss": 1.2}, {}, 2e-4, iteration, 1.0, False, False, 0.0, None, 0, None) + + mt.evaluate = lambda *a, _loss=loss, **k: ({"loss": _loss},) + mt.evaluate_and_print_results( + f"iteration {iteration}", + lambda: None, + None, + [MagicMock()], + iteration, + None, + MagicMock(), + ) + + +def _values(log_path, key): + """Every value logged under ``key``, in emission order.""" + import json + + return [ + json.loads(line[line.index("{") :])["value"] + for line in log_path.read_text().splitlines() + if f'"key": "{key}"' in line + ] + + +def _check(log_path): + from mlperf_logging.compliance_checker import mlp_compliance + + checker = mlp_compliance.make_checker(usage="training", ruleset=RULESET, quiet=True, werror=False) + valid, _system_id, _benchmark, _result = mlp_compliance.main( + str(log_path), f"training_{RULESET}/common.yaml", checker + ) + return valid + + +def test_a_converged_run_passes_the_compliance_checker(monkeypatch, _run_environment): + log_path = _run_environment + mt, _ = _install_fake_megatron(monkeypatch) + + # Real Megatron calls evaluate() from inside evaluate_and_print_results; + # the capture wrapper in the patch depends on that. + mt.evaluate_and_print_results = lambda *a, **k: mt.evaluate(*a, **k) + + _drive_a_run(mt, eval_losses=[0.9, 0.7, 0.5]) + + assert log_path.exists(), "rank zero did not write the result file" + assert _check(log_path), log_path.read_text() + + +def test_a_post_convergence_evaluation_cannot_pollute_the_log(monkeypatch, _run_environment): + """The reported symptom, end to end, against the real checker. + + pretrain validates once more after the training loop exits. That evaluation + draws fresh VAE and flow-matching noise, so it reports a different loss and + can land above the target the run just met, leaving the log ending on a + result that contradicts the one that stopped the run. + + Two things prevent it. do_valid is cleared on convergence, so pretrain + skips the evaluation entirely; and the wrapper records nothing once + run_stop has fired, for any caller that evaluates regardless. This asserts + the first and then exercises the second, since the first alone would leave + nothing to test. + """ + log_path = _run_environment + mt, megatron_args = _install_fake_megatron(monkeypatch) + mt.evaluate_and_print_results = lambda *a, **k: mt.evaluate(*a, **k) + + _drive_a_run(mt, eval_losses=[0.9, 0.7, 0.5]) + + assert megatron_args.do_valid is False, "pretrain would still run its post-training eval" + + # Evaluate anyway, at the loss the re-evaluation was observed to produce. + mt.evaluate = lambda *a, **k: ({"loss": 0.91},) + mt.evaluate_and_print_results( + "iteration 1536 on validation set", + lambda: None, + None, + [MagicMock()], + 1536, + None, + MagicMock(), + ) + + assert _check(log_path), log_path.read_text() + assert _values(log_path, "eval_accuracy") == [0.9, 0.7, 0.5] + assert _values(log_path, "run_stop") == ["success"] + + +def test_an_exhausted_run_fails_only_on_quality(monkeypatch, caplog, _run_environment): + """A run that misses the target is rejected, and that is the right answer. + + ``closed_flux1.yaml`` requires at least one ``eval_accuracy`` at or below + 0.586, so no amount of well-formed logging makes a non-converged run pass. + What the terminal ``run_stop`` buys is that the rejection is about the + model and nothing else -- the log is otherwise complete, so the run can be + read by the RCP checker and counted in the campaign instead of being an + unparseable hole in it. + """ + log_path = _run_environment + mt, megatron_args = _install_fake_megatron(monkeypatch) + mt.evaluate_and_print_results = lambda *a, **k: mt.evaluate(*a, **k) + + _drive_a_run(mt, eval_losses=[0.9, 0.8, 0.7]) + megatron_args.curr_iteration = 16000 + + from primus.backends.megatron.patches.mlperf_logging_patches import ( + patch_mlperf_terminal_run_stop, + ) + + patch_mlperf_terminal_run_stop(_make_ctx()) + + assert '"key": "run_stop"' in log_path.read_text() + assert '"status": "aborted"' in log_path.read_text() + + with caplog.at_level("WARNING"): + assert not _check(log_path) + + failures = [record.message for record in caplog.records if "Failed checks" in record.message] + assert failures, "expected the checker to report why it rejected the log" + assert all("eval_accuracy" in message for message in failures), failures + + +def test_a_run_without_a_terminal_record_is_rejected(monkeypatch, caplog, _run_environment): + """Guards the guard: the checker really does fail a log with no run_stop.""" + log_path = _run_environment + mt, _ = _install_fake_megatron(monkeypatch) + mt.evaluate_and_print_results = lambda *a, **k: mt.evaluate(*a, **k) + + # Converges, so quality is not what is missing here. + _drive_a_run(mt, eval_losses=[0.9, 0.5]) + assert '"key": "run_stop"' in log_path.read_text() + + without_run_stop = log_path.with_name("no_run_stop.txt") + without_run_stop.write_text( + "".join(line for line in log_path.read_text().splitlines(keepends=True) if "run_stop" not in line) + ) + + with caplog.at_level("WARNING"): + assert not _check(without_run_stop) + + assert any("run_stop" in record.message for record in caplog.records) diff --git a/tests/unit_tests/backends/megatron/test_mlperf_patches.py b/tests/unit_tests/backends/megatron/test_mlperf_patches.py index 077db6d25..6f257cf8a 100644 --- a/tests/unit_tests/backends/megatron/test_mlperf_patches.py +++ b/tests/unit_tests/backends/megatron/test_mlperf_patches.py @@ -17,6 +17,64 @@ from types import SimpleNamespace from unittest.mock import MagicMock +import pytest + +_MLLOG_CONSTANT_NAMES = { + "INIT_START": "init_start", + "INIT_STOP": "init_stop", + "RUN_START": "run_start", + "RUN_STOP": "run_stop", + "SUBMISSION_BENCHMARK": "submission_benchmark", + "SUBMISSION_ORG": "submission_org", + "SUBMISSION_DIVISION": "submission_division", + "SUBMISSION_PLATFORM": "submission_platform", + "SUBMISSION_STATUS": "submission_status", + "SEED": "seed", + "GLOBAL_BATCH_SIZE": "global_batch_size", + "TRAIN_SAMPLES": "train_samples", + "EVAL_SAMPLES": "eval_samples", + "GRADIENT_ACCUMULATION_STEPS": "gradient_accumulation_steps", + "OPT_NAME": "opt_name", + "OPT_BASE_LR": "opt_base_learning_rate", + "EVAL_ACCURACY": "eval_accuracy", + "EVAL_START": "eval_start", + "EVAL_STOP": "eval_stop", + "EPOCH_START": "epoch_start", + "EPOCH_STOP": "epoch_stop", + "BLOCK_START": "block_start", + "BLOCK_STOP": "block_stop", +} + + +@pytest.fixture(autouse=True) +def _mlperf_submission_environment(monkeypatch, tmp_path): + """Everything the logger now refuses to guess. + + The patch fails closed on each of these, so without them every test in + this module would fail on identity rather than on what it is testing. + """ + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("MLLOG_OUTPUT_FILE", str(tmp_path / "result_0.txt")) + monkeypatch.setenv("MLPERF_CLEAR_CACHES", "true") + monkeypatch.setenv("MLLOG_SUBMISSION_ORG", "AMD") + monkeypatch.setenv("MLLOG_SUBMISSION_DIVISION", "closed") + monkeypatch.setenv("MLLOG_SUBMISSION_PLATFORM", "MI355X") + monkeypatch.setenv("EXP", "examples/megatron/configs/MI355X/diffusion/test.yaml") + monkeypatch.setenv("MLLOG_LOWEST_NUMERICAL_PRECISION_IN_LINEAR", "fp8") + monkeypatch.setenv("MLLOG_LOWEST_NUMERICAL_PRECISION_IN_ATTN", "bfloat16") + monkeypatch.setenv("MLLOG_LOWEST_NUMERICAL_PRECISION_IN_COMM", "bfloat16") + + +@pytest.fixture(autouse=True) +def _clean_boundary(): + """The boundary keeps module-level state, so tests must not inherit it.""" + from primus.backends.megatron.patches import mlperf_boundary + + mlperf_boundary.reset_for_tests() + yield + mlperf_boundary.reset_for_tests() + + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -60,6 +118,14 @@ def fake_get_wandb_writer(): training_mod.get_tensorboard_writer = fake_get_tb_writer training_mod.get_wandb_writer = fake_get_wandb_writer + # The three entry points the MLPerf boundary wraps, plus the two helpers + # the relocated warmup reads off the training module. + training_mod.setup_model_and_optimizer = lambda *a, **k: ("model", "optimizer", "scheduler") + training_mod.build_train_valid_test_data_iterators = lambda *a, **k: ("train", "valid", "test") + training_mod.get_model_config = lambda model: SimpleNamespace() + training_mod.get_forward_backward_func = lambda: (lambda *a, **k: None) + training_pkg.pretrain = lambda *a, **k: None + training_pkg.training = training_mod megatron_mod.training = training_pkg @@ -87,6 +153,90 @@ def fake_get_wandb_writer(): return training_mod, _megatron_args +def _install_fake_grad_finalize(monkeypatch): + """Point megatron.core.distributed.finalize_model_grads at a sentinel. + + Call this *before* ``_install_fake_megatron``: the warmup imports the + callback lazily, and once ``megatron`` is a stub the real package can no + longer be reached to swap the attribute on. + """ + import importlib + import sys + + def finalize_model_grads(*args, **kwargs): + return None + + try: + real = importlib.import_module("megatron.core.distributed") + except ImportError: + core_mod = types.ModuleType("megatron.core") + core_mod.__path__ = [] + distributed_mod = types.ModuleType("megatron.core.distributed") + distributed_mod.finalize_model_grads = finalize_model_grads + core_mod.distributed = distributed_mod + monkeypatch.setitem(sys.modules, "megatron.core", core_mod) + monkeypatch.setitem(sys.modules, "megatron.core.distributed", distributed_mod) + else: + monkeypatch.setattr(real, "finalize_model_grads", finalize_model_grads) + + return finalize_model_grads + + +def _warmup_actors(monkeypatch): + """A model / optimizer / scheduler trio small enough to warm up on CPU.""" + import torch + + # The warmup brackets itself with device syncs; there is no device here. + monkeypatch.setattr(torch.cuda, "synchronize", lambda *a, **k: None) + + model = torch.nn.Linear(2, 2) + optimizer = SimpleNamespace( + param_groups=[{"betas": (0.9, 0.95), "weight_decay": 0.1, "step": 0}], + state={}, + zero_grad=lambda set_to_none=True: None, + ) + return [model], optimizer, SimpleNamespace(num_steps=0) + + +def _install_mock_mllog(monkeypatch, events=None): + """Install a fake mlperf_logging.mllog and return (module, events). + + ``events`` accumulates ``(kind, key, value, metadata)`` in emission order, + which is what most of the assertions below are about. + """ + import sys + + if events is None: + events = [] + + def _record(kind): + def _emit(key, value=None, metadata=None): + events.append((kind, key, value, metadata)) + + return _emit + + mock_mllogger = MagicMock() + mock_mllogger.start = _record("start") + mock_mllogger.end = _record("end") + mock_mllogger.event = _record("event") + + mock_mllog_module = MagicMock() + mock_mllog_module.get_mllogger.return_value = mock_mllogger + mock_mllog_module.constants = SimpleNamespace(**_MLLOG_CONSTANT_NAMES) + + mock_mlperf_pkg = MagicMock() + mock_mlperf_pkg.mllog = mock_mllog_module + + monkeypatch.setitem(sys.modules, "mlperf_logging", mock_mlperf_pkg) + monkeypatch.setitem(sys.modules, "mlperf_logging.mllog", mock_mllog_module) + + return mock_mllog_module, events + + +def _keys(events): + return [entry[1] for entry in events] + + def _make_ctx( mlperf_mode=False, warmup_train_steps=0, @@ -141,8 +291,99 @@ def _make_ctx( ) +def _silence_log_rank_0(monkeypatch): + for module in ( + "primus.backends.megatron.patches.mlperf_logging_patches", + "primus.backends.megatron.patches.mlperf_warmup_patches", + ): + monkeypatch.setattr(f"{module}.log_rank_0", lambda *a, **k: None) + + # ============================================================================ -# Level 1: Patch registration and conditions +# Submission identity and precision disclosure +# ============================================================================ + + +def test_precision_disclosures_are_explicit(monkeypatch): + from primus.backends.megatron.patches.mlperf_logging_patches import ( + _precision_disclosures_from_env, + ) + + monkeypatch.setenv("MLLOG_LOWEST_NUMERICAL_PRECISION_IN_LINEAR", "mxfp6") + monkeypatch.setenv("MLLOG_LOWEST_NUMERICAL_PRECISION_IN_ATTN", "bfloat16") + monkeypatch.setenv("MLLOG_LOWEST_NUMERICAL_PRECISION_IN_COMM", "bfloat16") + + assert _precision_disclosures_from_env() == { + "lowest_numerical_precision_in_linear": "mxfp6", + "lowest_numerical_precision_in_attn": "bfloat16", + "lowest_numerical_precision_in_comm": "bfloat16", + } + + +def test_missing_precision_disclosure_fails_mlperf_startup(monkeypatch): + from primus.backends.megatron.patches.mlperf_logging_patches import ( + _precision_disclosures_from_env, + ) + + for name in ( + "MLLOG_LOWEST_NUMERICAL_PRECISION_IN_LINEAR", + "MLLOG_LOWEST_NUMERICAL_PRECISION_IN_ATTN", + "MLLOG_LOWEST_NUMERICAL_PRECISION_IN_COMM", + ): + monkeypatch.delenv(name, raising=False) + + with pytest.raises(RuntimeError, match="explicit precision disclosures"): + _precision_disclosures_from_env() + + +def test_precision_value_outside_the_checker_enum_warns(monkeypatch, caplog): + """mxfp6 is not yet an accepted disclosure, and a run must say so. + + The value is still emitted -- describing an MXFP6 run as fp8 would be + worse than a log the checker rejects -- but it cannot pass silently. + """ + from primus.backends.megatron.patches.mlperf_logging_patches import ( + _precision_disclosures_from_env, + ) + + monkeypatch.setenv("MLLOG_LOWEST_NUMERICAL_PRECISION_IN_LINEAR", "mxfp6") + + with caplog.at_level("WARNING"): + values = _precision_disclosures_from_env() + + assert values["lowest_numerical_precision_in_linear"] == "mxfp6" + assert "mxfp6" in caplog.text + assert "compliance checker" in caplog.text + + +@pytest.mark.parametrize( + "variable", + ["MLLOG_SUBMISSION_ORG", "MLLOG_SUBMISSION_DIVISION", "MLLOG_SUBMISSION_PLATFORM"], +) +def test_submission_identity_never_defaults(monkeypatch, variable): + from primus.backends.megatron.patches.mlperf_logging_patches import ( + _submission_identity_from_env, + ) + + monkeypatch.delenv(variable, raising=False) + + with pytest.raises(RuntimeError, match=variable): + _submission_identity_from_env() + + +def test_submission_division_must_be_a_real_division(monkeypatch): + from primus.backends.megatron.patches.mlperf_logging_patches import ( + _submission_identity_from_env, + ) + + monkeypatch.setenv("MLLOG_SUBMISSION_DIVISION", "network") + + with pytest.raises(RuntimeError, match="closed"): + _submission_identity_from_env() + + +# ============================================================================ +# Patch registration # ============================================================================ @@ -151,42 +392,8 @@ class TestLoggingPatchMonkeyPatching: def test_installs_wrappers(self, monkeypatch): mt, _ = _install_fake_megatron(monkeypatch) - - monkeypatch.setattr( - "primus.backends.megatron.patches.mlperf_logging_patches.log_rank_0", - lambda *a, **k: None, - ) - - mock_mllog = MagicMock() - mock_mllog.get_mllogger.return_value = MagicMock() - mock_mllog.constants = SimpleNamespace( - INIT_START="init_start", - INIT_STOP="init_stop", - RUN_START="run_start", - RUN_STOP="run_stop", - SUBMISSION_BENCHMARK="submission_benchmark", - SUBMISSION_ORG="submission_org", - SUBMISSION_DIVISION="submission_division", - SUBMISSION_PLATFORM="submission_platform", - SUBMISSION_STATUS="submission_status", - SEED="seed", - GLOBAL_BATCH_SIZE="global_batch_size", - TRAIN_SAMPLES="train_samples", - EVAL_SAMPLES="eval_samples", - GRADIENT_ACCUMULATION_STEPS="gradient_accumulation_steps", - OPT_NAME="opt_name", - OPT_BASE_LR="opt_base_lr", - EVAL_ACCURACY="eval_accuracy", - EVAL_START="eval_start", - EVAL_STOP="eval_stop", - EPOCH_START="epoch_start", - BLOCK_START="block_start", - BLOCK_STOP="block_stop", - ) - mock_mlperf_pkg = MagicMock() - mock_mlperf_pkg.mllog = mock_mllog - monkeypatch.setitem(__import__("sys").modules, "mlperf_logging", mock_mlperf_pkg) - monkeypatch.setitem(__import__("sys").modules, "mlperf_logging.mllog", mock_mllog) + _silence_log_rank_0(monkeypatch) + _install_mock_mllog(monkeypatch) from primus.backends.megatron.patches.mlperf_logging_patches import ( patch_mlperf_logging, @@ -196,8 +403,7 @@ def test_installs_wrappers(self, monkeypatch): original_eval = mt.evaluate_and_print_results original_prl = mt.print_rank_last - ctx = _make_ctx(mlperf_mode=True) - patch_mlperf_logging(ctx) + patch_mlperf_logging(_make_ctx(mlperf_mode=True)) assert mt.training_log is not original_tl assert mt.evaluate_and_print_results is not original_eval @@ -206,42 +412,8 @@ def test_installs_wrappers(self, monkeypatch): def test_idempotent(self, monkeypatch): mt, _ = _install_fake_megatron(monkeypatch) - - monkeypatch.setattr( - "primus.backends.megatron.patches.mlperf_logging_patches.log_rank_0", - lambda *a, **k: None, - ) - - mock_mllog = MagicMock() - mock_mllog.get_mllogger.return_value = MagicMock() - mock_mllog.constants = SimpleNamespace( - INIT_START="init_start", - INIT_STOP="init_stop", - RUN_START="run_start", - RUN_STOP="run_stop", - SUBMISSION_BENCHMARK="submission_benchmark", - SUBMISSION_ORG="submission_org", - SUBMISSION_DIVISION="submission_division", - SUBMISSION_PLATFORM="submission_platform", - SUBMISSION_STATUS="submission_status", - SEED="seed", - GLOBAL_BATCH_SIZE="global_batch_size", - TRAIN_SAMPLES="train_samples", - EVAL_SAMPLES="eval_samples", - GRADIENT_ACCUMULATION_STEPS="gradient_accumulation_steps", - OPT_NAME="opt_name", - OPT_BASE_LR="opt_base_lr", - EVAL_ACCURACY="eval_accuracy", - EVAL_START="eval_start", - EVAL_STOP="eval_stop", - EPOCH_START="epoch_start", - BLOCK_START="block_start", - BLOCK_STOP="block_stop", - ) - mock_mlperf_pkg = MagicMock() - mock_mlperf_pkg.mllog = mock_mllog - monkeypatch.setitem(__import__("sys").modules, "mlperf_logging", mock_mlperf_pkg) - monkeypatch.setitem(__import__("sys").modules, "mlperf_logging.mllog", mock_mllog) + _silence_log_rank_0(monkeypatch) + _install_mock_mllog(monkeypatch) from primus.backends.megatron.patches.mlperf_logging_patches import ( patch_mlperf_logging, @@ -258,194 +430,513 @@ def test_idempotent(self, monkeypatch): assert mt.training_log is first_tl assert mt.evaluate_and_print_results is first_eval + def test_rank_zero_writes_the_log_to_a_file(self, monkeypatch, tmp_path): + """The submitted artifact is a file, not whatever landed on stdout.""" + _install_fake_megatron(monkeypatch) + _silence_log_rank_0(monkeypatch) + mock_mllog, _ = _install_mock_mllog(monkeypatch) + + output = tmp_path / "result_3.txt" + monkeypatch.setenv("MLLOG_OUTPUT_FILE", str(output)) + + from primus.backends.megatron.patches.mlperf_logging_patches import ( + patch_mlperf_logging, + ) + + patch_mlperf_logging(_make_ctx(mlperf_mode=True)) -class TestTrainingLogFirstCall: - """Verify INIT_STOP + RUN_START fire on first post-warmup training_log.""" + mock_mllog.config.assert_called_once() + assert mock_mllog.config.call_args.kwargs["filename"] == str(output) - def test_first_call_emits_init_stop_run_start(self, monkeypatch): - mt, args = _install_fake_megatron(monkeypatch) + def test_missing_output_file_fails_startup(self, monkeypatch): + _install_fake_megatron(monkeypatch) + _silence_log_rank_0(monkeypatch) + _install_mock_mllog(monkeypatch) + monkeypatch.delenv("MLLOG_OUTPUT_FILE", raising=False) - monkeypatch.setattr( - "primus.backends.megatron.patches.mlperf_logging_patches.log_rank_0", - lambda *a, **k: None, + from primus.backends.megatron.patches.mlperf_logging_patches import ( + patch_mlperf_logging, ) - emitted_events = [] + with pytest.raises(RuntimeError, match="MLLOG_OUTPUT_FILE"): + patch_mlperf_logging(_make_ctx(mlperf_mode=True)) + + +# ============================================================================ +# The measured-time boundary +# ============================================================================ + - mock_mllogger = MagicMock() - mock_constants = SimpleNamespace( - INIT_START="init_start", - INIT_STOP="init_stop", - RUN_START="run_start", - RUN_STOP="run_stop", - SUBMISSION_BENCHMARK="submission_benchmark", - SUBMISSION_ORG="submission_org", - SUBMISSION_DIVISION="submission_division", - SUBMISSION_PLATFORM="submission_platform", - SUBMISSION_STATUS="submission_status", - SEED="seed", - GLOBAL_BATCH_SIZE="global_batch_size", - TRAIN_SAMPLES="train_samples", - EVAL_SAMPLES="eval_samples", - GRADIENT_ACCUMULATION_STEPS="gradient_accumulation_steps", - OPT_NAME="opt_name", - OPT_BASE_LR="opt_base_lr", - EVAL_ACCURACY="eval_accuracy", - EVAL_START="eval_start", - EVAL_STOP="eval_stop", - EPOCH_START="epoch_start", - BLOCK_START="block_start", - BLOCK_STOP="block_stop", +class TestMeasuredTimeBoundary: + """run_start must precede every read of the real dataset.""" + + def test_transition_fires_before_the_data_iterators_are_built(self, monkeypatch): + mt, _ = _install_fake_megatron(monkeypatch) + _silence_log_rank_0(monkeypatch) + _, events = _install_mock_mllog(monkeypatch) + + order = [] + original_build = mt.build_train_valid_test_data_iterators + mt.build_train_valid_test_data_iterators = lambda *a, **k: ( + order.append("build_data_iterators"), + original_build(*a, **k), + )[1] + + from primus.backends.megatron.patches.mlperf_logging_patches import ( + patch_mlperf_logging, ) - def track_start(key, value=None, metadata=None): - emitted_events.append(("start", key)) + patch_mlperf_logging(_make_ctx(mlperf_mode=True)) + events.clear() - def track_end(key, value=None, metadata=None): - emitted_events.append(("end", key)) + mt.build_train_valid_test_data_iterators(None) - def track_event(key, value=None, metadata=None): - emitted_events.append(("event", key)) + keys = _keys(events) + assert "init_stop" in keys and "run_start" in keys + assert keys.index("init_stop") < keys.index("run_start") + # The wrapper records nothing itself, so the only way the dataset call + # can be ordered against the log is that it has not happened yet. + assert order == ["build_data_iterators"] - mock_mllogger.start = track_start - mock_mllogger.end = track_end - mock_mllogger.event = track_event + def test_transition_fires_once(self, monkeypatch): + """Virtual pipelining builds iterators per stage; the clock starts once.""" + mt, _ = _install_fake_megatron(monkeypatch) + _silence_log_rank_0(monkeypatch) + _, events = _install_mock_mllog(monkeypatch) - mock_mllog_module = MagicMock() - mock_mllog_module.get_mllogger.return_value = mock_mllogger - mock_mllog_module.constants = mock_constants + from primus.backends.megatron.patches.mlperf_logging_patches import ( + patch_mlperf_logging, + ) - # Wire the top-level mlperf_logging mock so that - # `from mlperf_logging import mllog` resolves correctly - mock_mlperf_pkg = MagicMock() - mock_mlperf_pkg.mllog = mock_mllog_module + patch_mlperf_logging(_make_ctx(mlperf_mode=True)) + events.clear() - import sys as _sys + mt.build_train_valid_test_data_iterators(None) + mt.build_train_valid_test_data_iterators(None) - monkeypatch.setitem(_sys.modules, "mlperf_logging", mock_mlperf_pkg) - monkeypatch.setitem(_sys.modules, "mlperf_logging.mllog", mock_mllog_module) - monkeypatch.setenv("RANK", "0") + assert _keys(events).count("run_start") == 1 + def test_pre_run_hooks_finish_before_the_clock_starts(self, monkeypatch): + """Warmup is initialization, so it belongs on the init side of run_start.""" + mt, _ = _install_fake_megatron(monkeypatch) + _silence_log_rank_0(monkeypatch) + _, events = _install_mock_mllog(monkeypatch) + + from primus.backends.megatron.patches import mlperf_boundary from primus.backends.megatron.patches.mlperf_logging_patches import ( patch_mlperf_logging, ) - ctx = _make_ctx(mlperf_mode=True, log_interval=1) - patch_mlperf_logging(ctx) + patch_mlperf_logging(_make_ctx(mlperf_mode=True)) + events.clear() + + mlperf_boundary.register_pre_run_hook( + "fake_warmup", lambda: events.append(("hook", "warmup", None, None)), order=10 + ) + + mt.build_train_valid_test_data_iterators(None) + + keys = _keys(events) + assert keys.index("warmup") < keys.index("run_start") + + def test_model_optimizer_and_forward_step_are_captured(self, monkeypatch): + """Warmup needs these, and nothing hands them to a before_train patch.""" + mt, _ = _install_fake_megatron(monkeypatch) + _silence_log_rank_0(monkeypatch) + _install_mock_mllog(monkeypatch) + + from primus.backends.megatron.patches import mlperf_boundary + from primus.backends.megatron.patches.mlperf_logging_patches import ( + patch_mlperf_logging, + ) + + patch_mlperf_logging(_make_ctx(mlperf_mode=True)) + + import megatron.training as megatron_training_pkg + + def forward_step_func(*a, **k): + return None + + megatron_training_pkg.pretrain(None, None, None, forward_step_func) + mt.setup_model_and_optimizer(None, None) + + captured = mlperf_boundary.captured() + assert captured["forward_step_func"] is forward_step_func + assert captured["model"] == "model" + assert captured["optimizer"] == "optimizer" + assert captured["opt_param_scheduler"] == "scheduler" + + def test_training_log_backstop_does_not_restart_the_clock(self, monkeypatch): + """The old first-training_log trigger stays, but must now be inert.""" + mt, _ = _install_fake_megatron(monkeypatch) + _silence_log_rank_0(monkeypatch) + _, events = _install_mock_mllog(monkeypatch) + + from primus.backends.megatron.patches.mlperf_logging_patches import ( + patch_mlperf_logging, + ) - emitted_events.clear() + patch_mlperf_logging(_make_ctx(mlperf_mode=True, log_interval=1)) + mt.build_train_valid_test_data_iterators(None) + events.clear() - # First call should emit INIT_STOP + RUN_START mt.training_log({"loss": 0.5}, {}, 1e-4, 1, 1.0, False, False, 0.0, None, 0, None) - event_keys = [e[1] for e in emitted_events] - assert "init_stop" in event_keys, f"INIT_STOP not emitted. Events: {emitted_events}" - assert "run_start" in event_keys, f"RUN_START not emitted. Events: {emitted_events}" + keys = _keys(events) + assert "init_stop" not in keys + assert "run_start" not in keys + + def test_warmup_registers_at_the_boundary_in_mlperf_mode(self, monkeypatch): + mt, _ = _install_fake_megatron(monkeypatch) + _silence_log_rank_0(monkeypatch) + + from primus.backends.megatron.patches import mlperf_boundary + from primus.backends.megatron.patches.mlperf_warmup_patches import ( + patch_mlperf_warmup, + ) + + original_train_step = mt.train_step + patch_mlperf_warmup(_make_ctx(mlperf_mode=True, warmup_train_steps=2)) + + assert mt.train_step is original_train_step, "warmup must not wrap train_step in MLPerf mode" + assert [name for _order, name, _fn in mlperf_boundary._HOOKS] == ["mlperf_warmup"] + + def test_warmup_stays_on_train_step_outside_mlperf_mode(self, monkeypatch): + """Development recipes keep the behavior they were tuned against.""" + mt, _ = _install_fake_megatron(monkeypatch) + _silence_log_rank_0(monkeypatch) + + from primus.backends.megatron.patches import mlperf_boundary + from primus.backends.megatron.patches.mlperf_warmup_patches import ( + patch_mlperf_warmup, + ) + + patch_mlperf_warmup(_make_ctx(mlperf_mode=False, warmup_train_steps=2)) + + assert getattr(mt.train_step, "_primus_warmup_hook", False) is True + assert mlperf_boundary._HOOKS == [] + + def test_warmup_steps_run_with_a_grad_finalize_callback(self, monkeypatch): + """Every warmup backward must have something that waits on its grad sync. + + The boundary warmup runs before Megatron's ``train()``, which is where + ``config.finalize_model_grads_func`` is assigned. That callback is the + only caller of ``finish_grad_sync()``, so if it is missing a warmup + backward dispatches the data-parallel reduce-scatter and nothing ever + waits on it. The first real step then asserts "Should not have multiple + communication calls outstanding at once" and training never starts. + """ + sentinel = _install_fake_grad_finalize(monkeypatch) + _, megatron_args = _install_fake_megatron(monkeypatch) + # The local-spec FP8 reset pulls in real Megatron enums, which the fake + # megatron above does not provide and this test is not about. + megatron_args.transformer_impl = "transformer_engine" + _silence_log_rank_0(monkeypatch) + model, optimizer, scheduler = _warmup_actors(monkeypatch) + config = SimpleNamespace(finalize_model_grads_func=None) + + seen = [] + + def _train_step(fwd, data_iter, mdl, opt, sched, cfg, fwdbwd, iteration=None): + seen.append(cfg.finalize_model_grads_func) + + from primus.backends.megatron.patches.mlperf_warmup_patches import ( + _run_warmup_and_restore, + ) + + _run_warmup_and_restore( + warmup_steps=2, + train_step_fn=_train_step, + forward_step_func=lambda *a, **k: None, + synthetic_iter=iter(()), + model=model, + optimizer=optimizer, + opt_param_scheduler=scheduler, + config=config, + forward_backward_func=lambda *a, **k: None, + iteration=0, + ) + + assert seen == [sentinel, sentinel] + assert config.finalize_model_grads_func is None, "warmup must leave the config as it found it" - init_stop_idx = event_keys.index("init_stop") - run_start_idx = event_keys.index("run_start") - assert init_stop_idx < run_start_idx + def test_warmup_keeps_a_finalize_callback_the_caller_already_set(self, monkeypatch): + _install_fake_grad_finalize(monkeypatch) + _, megatron_args = _install_fake_megatron(monkeypatch) + megatron_args.transformer_impl = "transformer_engine" + _silence_log_rank_0(monkeypatch) + model, optimizer, scheduler = _warmup_actors(monkeypatch) - # Second call should NOT emit INIT_STOP/RUN_START again - emitted_events.clear() - mt.training_log({"loss": 0.4}, {}, 1e-4, 2, 1.0, False, False, 0.0, None, 0, None) + def preexisting(*args, **kwargs): + return None - event_keys_2 = [e[1] for e in emitted_events] - assert "init_stop" not in event_keys_2 - assert "run_start" not in event_keys_2 + config = SimpleNamespace(finalize_model_grads_func=preexisting) + + seen = [] + + def _train_step(fwd, data_iter, mdl, opt, sched, cfg, fwdbwd, iteration=None): + seen.append(cfg.finalize_model_grads_func) + + from primus.backends.megatron.patches.mlperf_warmup_patches import ( + _run_warmup_and_restore, + ) + + _run_warmup_and_restore( + warmup_steps=1, + train_step_fn=_train_step, + forward_step_func=lambda *a, **k: None, + synthetic_iter=iter(()), + model=model, + optimizer=optimizer, + opt_param_scheduler=scheduler, + config=config, + forward_backward_func=lambda *a, **k: None, + iteration=0, + ) + + assert seen == [preexisting] + assert config.finalize_model_grads_func is preexisting # ============================================================================ -# Level 3: Component tests +# Run lifecycle # ============================================================================ +def _patch_with_eval_loss(monkeypatch, loss): + mt, megatron_args = _install_fake_megatron(monkeypatch) + _silence_log_rank_0(monkeypatch) + _, events = _install_mock_mllog(monkeypatch) + + mt.evaluate = lambda *a, **k: ({"loss": loss},) + mt.evaluate_and_print_results = lambda *a, **k: mt.evaluate(*a, **k) + + from primus.backends.megatron.patches.mlperf_logging_patches import ( + patch_mlperf_logging, + ) + + megatron_args.train_iters = 5000 + patch_mlperf_logging(_make_ctx(mlperf_mode=True, target_val_loss=0.586)) + mt.build_train_valid_test_data_iterators(None) + events.clear() + return mt, megatron_args, events + + +def _run_eval(mt, iteration): + mt.evaluate_and_print_results( + f"iteration {iteration}", + lambda: None, + None, + [MagicMock()], + iteration, + None, + MagicMock(), + ) + + class TestConvergenceDetection: """Verify convergence detection in evaluate_and_print_results wrapper.""" def test_convergence_sets_train_iters(self, monkeypatch): + mt, megatron_args, _ = _patch_with_eval_loss(monkeypatch, 0.500) + _run_eval(mt, 512) + assert megatron_args.train_iters == 512 + + def test_convergence_clears_do_valid(self, monkeypatch): + """Breaking the loop returns into pretrain, which validates once more. + + That evaluation is past run_stop and draws fresh noise, so it reports a + different loss and can land above the target the run just met. do_valid + is the flag pretrain branches on, so clearing it is what stops it. + """ + mt, megatron_args, _ = _patch_with_eval_loss(monkeypatch, 0.500) + assert megatron_args.do_valid is True + + _run_eval(mt, 512) + + assert megatron_args.do_valid is False + + def test_a_missed_target_leaves_do_valid_alone(self, monkeypatch): + """A run still in progress must keep evaluating.""" + mt, megatron_args, _ = _patch_with_eval_loss(monkeypatch, 0.900) + _run_eval(mt, 512) + + assert megatron_args.do_valid is True + + def test_convergence_emits_a_successful_run_stop(self, monkeypatch): + mt, _, events = _patch_with_eval_loss(monkeypatch, 0.500) + _run_eval(mt, 512) + + stops = [entry for entry in events if entry[1] == "run_stop"] + assert len(stops) == 1 + assert stops[0][2] == "success" + assert stops[0][3]["status"] == "success" + + def test_a_missed_target_reopens_the_block(self, monkeypatch): + """block_stop fires on entry to eval; training resuming needs a new block.""" + mt, _, events = _patch_with_eval_loss(monkeypatch, 0.900) + _run_eval(mt, 512) + + keys = _keys(events) + assert keys.index("block_stop") < keys.index("eval_start") + assert keys.index("eval_stop") < keys.index("block_start") + assert "run_stop" not in keys + + def test_an_unreadable_loss_still_reopens_the_block(self, monkeypatch): + mt, _, events = _patch_with_eval_loss(monkeypatch, 0.900) + mt.evaluate = lambda *a, **k: ({},) + events.clear() + + _run_eval(mt, 512) + + assert "block_start" in _keys(events) + + +class TestNothingIsLoggedAfterRunStop: + """The submission log ends at run_stop, whoever evaluates afterwards. + + Clearing do_valid on convergence means pretrain's post-training validation + never runs, so in practice nothing reaches the wrapper past run_stop. This + guard is what makes that a property of the wrapper rather than of the order + the callers happen to run in, which is why every test here calls the + wrapper directly. + """ + + def test_an_evaluation_after_run_stop_emits_no_records(self, monkeypatch): + mt, _, events = _patch_with_eval_loss(monkeypatch, 0.500) + _run_eval(mt, 512) + assert "run_stop" in _keys(events) + events.clear() + + _run_eval(mt, 512) + + assert _keys(events) == [] + + def test_a_second_run_stop_is_never_emitted(self, monkeypatch): + """run_stop is EXACTLY_ONE, and the repeat evaluation also converges.""" + mt, _, events = _patch_with_eval_loss(monkeypatch, 0.500) + _run_eval(mt, 512) + + _run_eval(mt, 512) + + assert _keys(events).count("run_stop") == 1 + + def test_the_evaluation_itself_still_runs(self, monkeypatch): + """The guard suppresses records, not the caller's evaluation.""" + mt, _, _ = _patch_with_eval_loss(monkeypatch, 0.500) + _run_eval(mt, 512) + + calls = [] + mt.evaluate = lambda *a, **k: (calls.append(1), {"loss": 0.500})[1] + + _run_eval(mt, 512) + + assert calls == [1] + + def test_an_aborted_run_also_closes_the_log(self, monkeypatch): + """converged would miss this; the guard keys on run_stop having fired.""" + mt, _, events = _patch_with_eval_loss(monkeypatch, 0.900) + mt._primus_mlperf_logger.log_run_stop(success=False, global_step=5000) + events.clear() + + _run_eval(mt, 5000) + + assert _keys(events) == [] + + +class TestTerminalRunStop: + """A run that never converges still has to produce a parseable log.""" + + def _patch(self, monkeypatch): mt, megatron_args = _install_fake_megatron(monkeypatch) + _silence_log_rank_0(monkeypatch) + _, events = _install_mock_mllog(monkeypatch) + + from primus.backends.megatron.patches.mlperf_logging_patches import ( + patch_mlperf_logging, + ) + + patch_mlperf_logging(_make_ctx(mlperf_mode=True)) + mt.build_train_valid_test_data_iterators(None) + events.clear() + return mt, megatron_args, events - monkeypatch.setattr( - "primus.backends.megatron.patches.mlperf_logging_patches.log_rank_0", - lambda *a, **k: None, - ) - - # Mock mlperf_logging — wire .mllog attribute on the top-level package - mock_mllogger = MagicMock() - mock_constants = SimpleNamespace( - INIT_START="init_start", - INIT_STOP="init_stop", - RUN_START="run_start", - RUN_STOP="run_stop", - SUBMISSION_BENCHMARK="submission_benchmark", - SUBMISSION_ORG="submission_org", - SUBMISSION_DIVISION="submission_division", - SUBMISSION_PLATFORM="submission_platform", - SUBMISSION_STATUS="submission_status", - SEED="seed", - GLOBAL_BATCH_SIZE="global_batch_size", - TRAIN_SAMPLES="train_samples", - EVAL_SAMPLES="eval_samples", - GRADIENT_ACCUMULATION_STEPS="gradient_accumulation_steps", - OPT_NAME="opt_name", - OPT_BASE_LR="opt_base_lr", - EVAL_ACCURACY="eval_accuracy", - EVAL_START="eval_start", - EVAL_STOP="eval_stop", - EPOCH_START="epoch_start", - BLOCK_START="block_start", - BLOCK_STOP="block_stop", - ) - mock_mllogger.start = MagicMock() - mock_mllogger.end = MagicMock() - mock_mllogger.event = MagicMock() - - mock_mllog_module = MagicMock() - mock_mllog_module.get_mllogger.return_value = mock_mllogger - mock_mllog_module.constants = mock_constants - - mock_mlperf_pkg = MagicMock() - mock_mlperf_pkg.mllog = mock_mllog_module - - import sys as _sys - - monkeypatch.setitem(_sys.modules, "mlperf_logging", mock_mlperf_pkg) - monkeypatch.setitem(_sys.modules, "mlperf_logging.mllog", mock_mllog_module) - monkeypatch.setenv("RANK", "0") - - # Set evaluate to return a loss below target BEFORE patching. - # Also make evaluate_and_print_results call evaluate() internally, - # mirroring real Megatron behavior so _captured_loss gets populated. - mt.evaluate = lambda *a, **k: ({"loss": 0.500},) - - def fake_eval_and_print(*a, **k): - mt.evaluate(*a, **k) - - mt.evaluate_and_print_results = fake_eval_and_print + def test_exhausted_run_emits_an_aborted_run_stop(self, monkeypatch): + mt, megatron_args, events = self._patch(monkeypatch) + megatron_args.curr_iteration = 5000 + + from primus.backends.megatron.patches.mlperf_logging_patches import ( + patch_mlperf_terminal_run_stop, + ) + + patch_mlperf_terminal_run_stop(_make_ctx(mlperf_mode=True)) + + stops = [entry for entry in events if entry[1] == "run_stop"] + assert len(stops) == 1 + assert stops[0][2] == "aborted" + assert stops[0][3]["status"] == "aborted" + assert stops[0][3]["samples_count"] == 5000 * 512 + + def test_a_converged_run_is_not_stopped_twice(self, monkeypatch): + """run_stop is EXACTLY_ONE in the ruleset.""" + mt, megatron_args, events = self._patch(monkeypatch) + + mlperf_logger = mt._primus_mlperf_logger + mlperf_logger.log_run_stop(success=True, global_step=512) + + from primus.backends.megatron.patches.mlperf_logging_patches import ( + patch_mlperf_terminal_run_stop, + ) + + patch_mlperf_terminal_run_stop(_make_ctx(mlperf_mode=True)) + + assert _keys(events).count("run_stop") == 1 + + +class TestHyperparameterRecords: + """The values the closed_flux1 ruleset pins exactly.""" + + def test_evaluation_frequency_is_reported_in_samples(self, monkeypatch): + """closed_flux1.yaml requires evaluation_frequency == 262144.""" + _install_fake_megatron(monkeypatch) + _silence_log_rank_0(monkeypatch) + _, events = _install_mock_mllog(monkeypatch) from primus.backends.megatron.patches.mlperf_logging_patches import ( patch_mlperf_logging, ) - target_val_loss = 0.586 - megatron_args.train_iters = 5000 + patch_mlperf_logging(_make_ctx(mlperf_mode=True, global_batch_size=512, eval_interval=512)) - ctx = _make_ctx(mlperf_mode=True, target_val_loss=target_val_loss) - patch_mlperf_logging(ctx) + frequency = [entry for entry in events if entry[1] == "evaluation_frequency"] + assert len(frequency) == 1 + assert frequency[0][2] == 262144 + + def test_config_filename_names_a_real_recipe(self, monkeypatch): + _install_fake_megatron(monkeypatch) + _silence_log_rank_0(monkeypatch) + _install_mock_mllog(monkeypatch) + monkeypatch.delenv("EXP", raising=False) - # Call eval at iteration 512 - mt.evaluate_and_print_results( - "iteration 512", - lambda: None, - None, - [MagicMock()], - 512, - None, - MagicMock(), + from primus.backends.megatron.patches.mlperf_logging_patches import ( + patch_mlperf_logging, ) - assert megatron_args.train_iters == 512 + with pytest.raises(RuntimeError, match="EXP"): + patch_mlperf_logging(_make_ctx(mlperf_mode=True)) + + def test_cache_clear_is_reported_not_assumed(self, monkeypatch): + _install_fake_megatron(monkeypatch) + _silence_log_rank_0(monkeypatch) + _install_mock_mllog(monkeypatch) + monkeypatch.delenv("MLPERF_CLEAR_CACHES", raising=False) + + from primus.backends.megatron.patches.mlperf_logging_patches import ( + patch_mlperf_logging, + ) + + with pytest.raises(RuntimeError, match="MLPERF_CLEAR_CACHES"): + patch_mlperf_logging(_make_ctx(mlperf_mode=True)) # ============================================================================