From 2c5981cc95a5a6fae3b65cdd28c4136f39318b81 Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Thu, 27 Aug 2026 09:25:37 -0500 Subject: [PATCH 01/15] fix(flux): make MLPerf validation measure what it claims to Flux evaluation reported a val_loss over 29,696 samples at the dataset's own timesteps. It was doing neither, and both defects were invisible in the logs because every number printed was the configured one. Coverage. Energon shards validation data across dp_size * num_workers but splits the batch quota across num_workers alone, so the two divisions agree only when each worker's slice is a whole number of microbatches. They did not agree in any shipped recipe: the tail of every short worker's slice went unread, giving 27,776 samples at num_workers 16 and 28,928 at 8, with the same samples missing on every eval. This adds val_num_workers (defaulting to 0), builds validation its own WorkerConfig, and asserts eval_samples % (dp_size * max(1, val_num_workers) * micro_batch_size) == 0 so an unreachable shape fails at startup with the list of counts that work. The budget is now expressed as eval_samples rather than eval_iters, which stays correct across batch-size changes, and full_validation reads the true split size from the dataset index instead of being a no-op. Timesteps. The MLCommons val Arrow files carry a per-sample int32 timestep; ingest read only the four tensor columns and dropped it, so every eval ever run against these shards silently fell back to injecting arange(B) % 8. That is not a near miss: the published val set pairs each image with one fixed timestep and orders them so a contiguous batch spans two sigmas, not eight. Ingest now carries the column into the sidecar, collate validates it on every sample rather than samples[0], and eval_timestep_source makes the choice explicit so a missing field is an error for MLPerf instead of a silent substitution. is_validation now derives from eval mode alone rather than from the batch carrying a timestep. The old signal was data-driven, so a training batch that happened to carry the field silently suppressed CFG dropout. Noise. per_step_rng_reseed froze the training step counter during eval and reseeded every microbatch identically, replaying one draw of VAE epsilon and flow noise across the whole validation set. Both references redraw per batch. Eval now advances its own counter in a seed range disjoint from training, derived from (iteration, microbatch) so it survives resume without checkpointing a counter, and leaves the training stream untouched. Reporting and cost. consumed_valid_samples counted a flat global batch per iteration regardless of the real width; it now uses the observed loss denominator. Three all-reduces and two host round-trips collapse to one packed fp64 reduction, the missing MLPerf evaluation_frequency event is logged, and the unconditional per-eval gc.collect()/empty_cache() moves behind eval_purge_memory, off by default. Verified against /mnt/dcgpuval/mlperf_flux1: a harness rebuilding the provider's exact dataloader reproduces 28,928, 27,776 and 5,120 on the unfixed shapes and 29,696 with 3,712 per timestep after the fix. A full re-ingest of the val split yields exactly 3,712 per timestep, and driving both splits through the reader shows the same batch of 64 images carrying sigmas {0, 0.5} from the real column against all eight from the fallback. The reduction is split out of primus_evaluate as reduce_eval_losses so the step that produces the reported number is testable without a process group. Its tests pin that the loss is a ratio of summed numerator and denominator rather than a mean of per-rank ratios, that the denominator survives as a sample count, and that a short or doubled read raises instead of being reported as the count the configuration intended. Deployment note: the recipes set eval_timestep_source: dataset, so evaluation fails loudly until a val split ingested with the timestep column is in place. --- ...chnell_resample_local_spec_fp8_mlperf.yaml | 12 +- ...n_schnell_resample_te_spec_fp8_mlperf.yaml | 12 +- .../preprocessing/pipelines/ingest.py | 14 +- .../data/diffusion/task_encoders/image.py | 51 +++- .../megatron/data/energon_dataset_provider.py | 49 +++- primus/backends/megatron/diffusion_trainer.py | 54 +++- .../megatron/flux_pretrain_trainer.py | 14 ++ .../megatron/patches/args/__init__.py | 2 + .../patches/args/eval_samples_patches.py | 102 ++++++++ .../patches/mlperf_logging_patches.py | 63 +++-- .../training/diffusion/forward_step.py | 167 ++++++++++--- .../backends/megatron/training/eval_budget.py | 210 ++++++++++++++++ .../backends/megatron/training/evaluator.py | 158 ++++++++---- .../modules/megatron/trainer_base.yaml | 21 ++ .../data/preprocessing/test_ingest.py | 109 +++++++++ .../data/task_encoders/test_task_encoders.py | 79 ++++++ .../training/test_eval_timestep_source.py | 148 +++++++++++ .../training/test_flux_forward_step_e2e.py | 32 ++- .../test_diffusion_trainer_eval_rng.py | 202 +++++++++++++++ .../backends/megatron/test_eval_budget.py | 209 ++++++++++++++++ .../megatron/test_eval_samples_patch.py | 157 ++++++++++++ .../megatron/test_evaluator_reduction.py | 231 ++++++++++++++++++ 22 files changed, 1977 insertions(+), 119 deletions(-) create mode 100644 primus/backends/megatron/patches/args/eval_samples_patches.py create mode 100644 primus/backends/megatron/training/eval_budget.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/training/test_eval_timestep_source.py create mode 100644 tests/unit_tests/backends/megatron/test_diffusion_trainer_eval_rng.py create mode 100644 tests/unit_tests/backends/megatron/test_eval_budget.py create mode 100644 tests/unit_tests/backends/megatron/test_eval_samples_patch.py create mode 100644 tests/unit_tests/backends/megatron/test_evaluator_reduction.py 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..06252438a 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 @@ -75,7 +75,17 @@ modules: # Training iterations train_iters: 5000 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 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..60ccd9a9e 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 @@ -88,7 +88,17 @@ modules: train_iters: 5000 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 diff --git a/primus/backends/megatron/data/diffusion/preprocessing/pipelines/ingest.py b/primus/backends/megatron/data/diffusion/preprocessing/pipelines/ingest.py index 5c68bc572..ba2c19d94 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)) 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..157f83f70 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,11 @@ 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 ( + assert_val_worker_divisibility, + get_eval_num_microbatches, + get_val_num_workers, +) from primus.core.utils.module_utils import log_rank_0 @@ -93,8 +97,11 @@ 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) @@ -124,22 +131,38 @@ def create_dataloaders( # 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) + 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 +170,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 +232,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..939220c98 100644 --- a/primus/backends/megatron/diffusion_trainer.py +++ b/primus/backends/megatron/diffusion_trainer.py @@ -50,6 +50,8 @@ def __init__(self, *args, **kwargs): self._compiled_loss_fn = None self._forward_step_count = 0 self._forward_step_count_initialized = False + self._eval_rng_iteration = 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 +227,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 +235,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 +265,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 +312,38 @@ 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 ``(iteration, microbatch index within this evaluation)`` + rather than from a free-running counter, so it is reproducible after a + checkpoint resume without having to checkpoint the counter itself: the + iteration comes from the checkpoint and the index restarts at zero for + each evaluation. + """ + from megatron.training import get_args + + from primus.backends.megatron.training.diffusion.forward_step import ( + EVAL_RNG_ITERATION_STRIDE, + ) + + iteration = getattr(get_args(), "iteration", 0) + if iteration != self._eval_rng_iteration: + self._eval_rng_iteration = iteration + 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"per-iteration stride {EVAL_RNG_ITERATION_STRIDE} that keeps " + f"consecutive evaluations' RNG streams disjoint." + ) + + return iteration * 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_logging_patches.py b/primus/backends/megatron/patches/mlperf_logging_patches.py index b53173d8a..06a704174 100644 --- a/primus/backends/megatron/patches/mlperf_logging_patches.py +++ b/primus/backends/megatron/patches/mlperf_logging_patches.py @@ -148,9 +148,22 @@ def log_hyperparams(self, args): key=self._constants.TRAIN_SAMPLES, value=getattr(args, "train_samples", 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. Required by the MLPerf logging + # rules and previously absent; the constant name differs across + # mlperf_logging releases, so fall back to the literal key. + self._event( + key=getattr(self._constants, "EVAL_FREQUENCY", "eval_frequency"), + value=getattr(args, "eval_interval", 0) * self.gbs, ) gas = max(self.gbs // self.mbs, 1) self._event(key=self._constants.GRADIENT_ACCUMULATION_STEPS, value=gas) @@ -323,6 +336,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, @@ -404,47 +418,32 @@ 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} " 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..94989a5e3 --- /dev/null +++ b/primus/backends/megatron/training/eval_budget.py @@ -0,0 +1,210 @@ +############################################################################### +# 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 +from pathlib import Path +from typing import Optional + +__all__ = [ + "DEFAULT_VAL_NUM_WORKERS", + "EvalCoverageError", + "assert_val_worker_divisibility", + "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 read the samples it claims to.""" + + +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 = args.data_parallel_size + 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 = args.data_parallel_size + 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 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/evaluator.py b/primus/backends/megatron/training/evaluator.py index 9ea10e438..6525188b9 100644 --- a/primus/backends/megatron/training/evaluator.py +++ b/primus/backends/megatron/training/evaluator.py @@ -15,10 +15,116 @@ 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.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 +# 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 _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: + # 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( forward_step_func, @@ -56,7 +162,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( @@ -123,8 +231,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 +243,18 @@ 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) collected_non_loss_data = None if non_loss_data_func is not None: 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/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..62a3a9bb8 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.""" 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..b5afe6086 --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_diffusion_trainer_eval_rng.py @@ -0,0 +1,202 @@ +# 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. +""" + +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pytest +import torch + +from primus.backends.megatron.diffusion_trainer import DiffusionPretrainTrainer +from primus.backends.megatron.training.diffusion.forward_step import ( + EVAL_RNG_ITERATION_STRIDE, + EVAL_RNG_OFFSET, +) + + +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_rng_iteration = None + trainer._eval_microbatch_index = 0 + return trainer + + +def _indices(trainer, iteration, count): + with patch("megatron.training.get_args", return_value=SimpleNamespace(iteration=iteration)): + return [trainer._next_eval_step_index() for _ in range(count)] + + +class TestEvalStepIndex: + def test_advances_within_one_evaluation(self): + trainer = _make_trainer() + assert _indices(trainer, iteration=100, count=4) == [ + 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() + _indices(trainer, iteration=100, count=3) + second = _indices(trainer, iteration=200, count=3) + + assert second == [200 * EVAL_RNG_ITERATION_STRIDE + i for i in range(3)] + + def test_consecutive_evaluations_do_not_collide(self): + trainer = _make_trainer() + first = set(_indices(trainer, iteration=100, count=64)) + second = set(_indices(trainer, iteration=101, count=64)) + + assert not (first & second) + + def test_reproducible_after_resume(self): + """A fresh trainer at the same iteration must produce the same indices.""" + before = _indices(_make_trainer(), iteration=512, count=8) + after_resume = _indices(_make_trainer(), iteration=512, count=8) + + assert before == after_resume + + def test_overrun_of_the_stride_is_rejected(self): + trainer = _make_trainer() + trainer._eval_rng_iteration = 5 + trainer._eval_microbatch_index = EVAL_RNG_ITERATION_STRIDE + + with patch("megatron.training.get_args", return_value=SimpleNamespace(iteration=5)): + with pytest.raises(RuntimeError, match="per-iteration 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, iteration * EVAL_RNG_ITERATION_STRIDE + micro) + for rank in range(64) + for iteration 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 + + with patch( + "primus.backends.megatron.training.diffusion.forward_step.flux_forward_step_func", + side_effect=recording_func, + ), patch("megatron.training.get_args", return_value=SimpleNamespace(iteration=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_eval_budget.py b/tests/unit_tests/backends/megatron/test_eval_budget.py new file mode 100644 index 000000000..299e4bac3 --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_eval_budget.py @@ -0,0 +1,209 @@ +# 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_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 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..68f762a9c --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_evaluator_reduction.py @@ -0,0 +1,231 @@ +# 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 types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from primus.backends.megatron.training.evaluator import ( + VAL_LOSS_KEY, + _record_consumed_valid_samples, + reduce_eval_losses, +) + +GROUP = object() + + +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 From 9fce2cd77b5ef9ff72f4875ee60e26dd4394efbc Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Thu, 27 Aug 2026 14:17:00 -0500 Subject: [PATCH 02/15] fix(diffusion data): keep metadata files and bookcorpus out of Flux ingest MLCommons manifests list dataset_info.json and state.json beside the Arrow data. The ingest pipeline fetched the manifest unfiltered, so conversion failed on both and reported every run as having failed files. They also consumed max_files slots, and one sorting ahead of a data file would have shifted every shard index after it. parse_md5_manifest already accepted a suffix_filter; fetch_manifest just never passed one through. Separately, prepare_dataset_if_needed read an absent train_data_path as proof that a run wants a tokenised bookcorpus corpus. Every diffusion recipe sets dataloader_type: external and brings its own Energon pipeline, so that built a dataset the run never opens and demanded HF_TOKEN for a tokenizer it never loads. Adds the val-only ingest config the MLPerf validation split is built from. --- examples/megatron/prepare.py | 12 ++++ .../data/diffusion/preprocessing/download.py | 14 +++- .../preprocessing/pipelines/ingest.py | 7 +- .../preprocessing/mlperf_flux1_val.yaml | 69 +++++++++++++++++++ .../data/preprocessing/test_ingest.py | 32 +++++++++ 5 files changed, 131 insertions(+), 3 deletions(-) create mode 100644 primus/configs/data/megatron/diffusion/preprocessing/mlperf_flux1_val.yaml 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/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 ba2c19d94..60000f5c8 100644 --- a/primus/backends/megatron/data/diffusion/preprocessing/pipelines/ingest.py +++ b/primus/backends/megatron/data/diffusion/preprocessing/pipelines/ingest.py @@ -140,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/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/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_ingest.py b/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_ingest.py index 62a3a9bb8..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 @@ -174,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") From 331692ee3e2c0b1844fe31a1ff2e5722c8c1127d Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Thu, 27 Aug 2026 14:17:15 -0500 Subject: [PATCH 03/15] fix(flux): size the eval budget before process groups exist The eval_samples patch derives eval_iters during build_args, before Megatron initialises process groups, so args.data_parallel_size did not exist yet. The patch runner logged the AttributeError and carried on, leaving eval_iters at 0: the job evaluated nothing and exited 0. Derive the width from world_size and the parallel sizes instead, and refuse to build validation dataloaders when eval_samples is set but eval_iters is 0, so a swallowed patch failure cannot pass for a clean run. Under skip_train the provider no longer builds a train dataset, which an evaluation-only dataset has no split for. Evaluation reporting now goes through one helper that drops to debug under MLPerf mode, where the submission log reports the loss itself and Megatron's own reporting is already suppressed. Only confirmations are quietened: a coverage shortfall still raises, and the mismatch that context parallelism explains still reports at info. --- .../megatron/data/energon_dataset_provider.py | 60 ++++++++++----- .../backends/megatron/training/eval_budget.py | 46 +++++++++++- .../backends/megatron/training/evaluator.py | 48 +++++++++++- .../megatron/test_evaluator_reduction.py | 73 +++++++++++++++++++ 4 files changed, 203 insertions(+), 24 deletions(-) diff --git a/primus/backends/megatron/data/energon_dataset_provider.py b/primus/backends/megatron/data/energon_dataset_provider.py index 157f83f70..743853657 100644 --- a/primus/backends/megatron/data/energon_dataset_provider.py +++ b/primus/backends/megatron/data/energon_dataset_provider.py @@ -32,6 +32,7 @@ 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_val_worker_divisibility, get_eval_num_microbatches, get_val_num_workers, @@ -106,27 +107,48 @@ def create_dataloaders( # 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 diff --git a/primus/backends/megatron/training/eval_budget.py b/primus/backends/megatron/training/eval_budget.py index 94989a5e3..d1a468496 100644 --- a/primus/backends/megatron/training/eval_budget.py +++ b/primus/backends/megatron/training/eval_budget.py @@ -25,6 +25,7 @@ """ import json +import os from pathlib import Path from typing import Optional @@ -32,6 +33,7 @@ "DEFAULT_VAL_NUM_WORKERS", "EvalCoverageError", "assert_val_worker_divisibility", + "get_data_parallel_size", "get_eval_num_microbatches", "get_val_num_workers", "read_energon_split_sample_count", @@ -53,13 +55,53 @@ class EvalCoverageError(ValueError): """Raised when an evaluation would not read the samples it claims to.""" +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 = args.data_parallel_size + dp_size = get_data_parallel_size(args) micro_batch_size = args.micro_batch_size global_batch_size = args.global_batch_size @@ -114,7 +156,7 @@ def assert_val_worker_divisibility(args, eval_samples: int) -> None: 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 = args.data_parallel_size + dp_size = get_data_parallel_size(args) micro_batch_size = args.micro_batch_size val_num_workers = get_val_num_workers(args) diff --git a/primus/backends/megatron/training/evaluator.py b/primus/backends/megatron/training/evaluator.py index 6525188b9..4b8f25673 100644 --- a/primus/backends/megatron/training/evaluator.py +++ b/primus/backends/megatron/training/evaluator.py @@ -18,7 +18,7 @@ from primus.backends.megatron.training.eval_budget import get_eval_num_microbatches 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 @@ -26,6 +26,26 @@ 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. @@ -46,6 +66,14 @@ def _record_consumed_valid_samples(args, observed_samples, eval_iters, eval_batc 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 + if observed_samples != expected: # Context parallelism duplicates the per-sample loss across CP ranks, # which inflates the reduced denominator; only assert when it cannot. @@ -177,11 +205,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 @@ -256,6 +288,16 @@ def primus_evaluate( _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: collected_non_loss_data = non_loss_data_func(model) diff --git a/tests/unit_tests/backends/megatron/test_evaluator_reduction.py b/tests/unit_tests/backends/megatron/test_evaluator_reduction.py index 68f762a9c..e2935cfe8 100644 --- a/tests/unit_tests/backends/megatron/test_evaluator_reduction.py +++ b/tests/unit_tests/backends/megatron/test_evaluator_reduction.py @@ -12,6 +12,7 @@ being papered over with the count the configuration intended. """ +from contextlib import contextmanager from types import SimpleNamespace from unittest.mock import patch @@ -21,10 +22,12 @@ from primus.backends.megatron.training.evaluator import ( VAL_LOSS_KEY, _record_consumed_valid_samples, + _report_eval, reduce_eval_losses, ) GROUP = object() +EVALUATOR = "primus.backends.megatron.training.evaluator" class FakeAllReduce: @@ -229,3 +232,73 @@ def test_fallback_does_not_raise_on_a_mismatch_it_cannot_see(self): 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 From d58d07165369e9c03eb95a958e77b9f8bfb7e079 Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Thu, 27 Aug 2026 23:22:57 -0500 Subject: [PATCH 04/15] style(flux): collapse the eval loss summary onto one line Black's formatting of the line, which the pre-commit gate enforces and which the preceding commits had not been run through. --- primus/backends/megatron/training/evaluator.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/primus/backends/megatron/training/evaluator.py b/primus/backends/megatron/training/evaluator.py index 4b8f25673..a25e7df8c 100644 --- a/primus/backends/megatron/training/evaluator.py +++ b/primus/backends/megatron/training/evaluator.py @@ -293,9 +293,7 @@ def primus_evaluate( # 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()) - ) + 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 From 41cbd16041f4ed8ab7a50a6b3c871b5e11921415 Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Thu, 27 Aug 2026 23:37:34 -0500 Subject: [PATCH 05/15] fix(flux): drop the always-true guard on the coverage mismatch path The equal case returns early just above, so the `!= expected` test could never be false. Outdenting its body leaves behaviour identical and stops the guard from reading as though a third case existed. --- .../backends/megatron/training/evaluator.py | 29 +++++++++---------- 1 file changed, 14 insertions(+), 15 deletions(-) diff --git a/primus/backends/megatron/training/evaluator.py b/primus/backends/megatron/training/evaluator.py index a25e7df8c..d7cbf7bf8 100644 --- a/primus/backends/megatron/training/evaluator.py +++ b/primus/backends/megatron/training/evaluator.py @@ -74,22 +74,21 @@ def _record_consumed_valid_samples(args, observed_samples, eval_iters, eval_batc ) return - if observed_samples != expected: - # 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}." + # 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." ) - 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)") + log_rank_0(f"[eval] {detail} (context_parallel_size={cp_size}, not asserting)") def _reduction_device(numerators): From b7c5ba86d157db16d63dcb20ce50b9138b8ea0e6 Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Fri, 28 Aug 2026 08:01:17 -0500 Subject: [PATCH 06/15] fix(flux): emit the MLPerf v6.1 keys the compliance checker actually requires Run the Flux MLLOG output through the official training_6.1.0 compliance checker and it fails before it reaches anything about the run itself. Several of these are v6.1 renames that were still being emitted under their v5 spelling, so the keys were present but invisible to the checker. Precision disclosures. v6.1 requires EXACTLY_ONE of lowest_numerical_precision_in_linear, _in_attn and _in_comm, and the logger emitted none. They are read from the environment and startup fails if any is unset, rather than defaulting: the value is a submission claim about what the run did, so inferring it from the config is how a stale claim gets logged after a precision change. The checker validates them against its own enum, which is what makes a wrong value cheap to catch and a missing one expensive. Epoch metadata. BLOCK_START/BLOCK_STOP took first_epoch_num and EVAL_START took epoch_num. v6.1 keys these on samples_count, and Flux is sample-based -- there is one epoch, so the old metadata carried no information either. Gradient accumulation. gbs // mbs counts the accumulation steps of the whole job rather than of one rank, so it over-reported by the data-parallel size on every multi-GPU run. Also emits evaluation_frequency under its v6.1 name rather than the EVAL_FREQUENCY constant, which is absent from the pinned mlperf_logging and silently fell back to the v5 literal; adds the required cache_clear, parallelism, micro_batch_size, config_filename, warmup-steps and clip-norm keys; and makes run_start idempotent, since an eval that precedes the first training batch otherwise emits a second one and the checker rejects the duplicate. --- .../patches/mlperf_logging_patches.py | 64 ++++++++++++++++--- .../backends/megatron/test_mlperf_patches.py | 42 ++++++++++++ 2 files changed, 98 insertions(+), 8 deletions(-) diff --git a/primus/backends/megatron/patches/mlperf_logging_patches.py b/primus/backends/megatron/patches/mlperf_logging_patches.py index 06a704174..7b5dc8d8f 100644 --- a/primus/backends/megatron/patches/mlperf_logging_patches.py +++ b/primus/backends/megatron/patches/mlperf_logging_patches.py @@ -25,6 +25,26 @@ 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", +} + + +def _precision_disclosures_from_env() -> dict[str, str]: + """Return mandatory v6.1 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) + ) + return values + def _mlperf_logging_enabled(ctx: PatchContext) -> bool: args = get_args(ctx) @@ -94,6 +114,7 @@ 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.profiler = os.getenv("PROFILER", "") self.profiler_warmup_steps = int(os.getenv("PROF_WARMUP_STEPS", "0")) @@ -122,6 +143,8 @@ def _end(self, key, value=None, metadata=None): def log_init(self, seed: int): if int(os.environ.get("RANK", "0")) == 0: + clear_caches = os.environ.get("MLPERF_CLEAR_CACHES", "false").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( @@ -144,9 +167,20 @@ def log_hyperparams(self, args): if int(os.environ.get("RANK", "0")) != 0: 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), + ("config_filename", os.environ.get("EXP", "unknown")), + ): + 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 @@ -162,10 +196,11 @@ def log_hyperparams(self, args): # rules and previously absent; the constant name differs across # mlperf_logging releases, so fall back to the literal key. self._event( - key=getattr(self._constants, "EVAL_FREQUENCY", "eval_frequency"), + 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( @@ -188,13 +223,22 @@ 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 int(os.environ.get("RANK", "0")) == 0 and not self._run_started: 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._start(key=self._constants.BLOCK_START, metadata={"samples_count": 0}) + self._run_started = True def on_train_batch_end(self, global_step: int, loss: float, lr: float): self.timer.mark_training_start() @@ -216,6 +260,7 @@ 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() @@ -232,9 +277,12 @@ 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() @@ -466,7 +514,7 @@ def _capture_wrapper(*a, **kw): if int(os.environ.get("RANK", "0")) == 0: mlperf_logger._start( key=mlperf_logger._constants.BLOCK_START, - metadata={"first_epoch_num": 0}, + metadata={"samples_count": iteration * gbs}, ) else: logger.warning("Could not extract validation loss from evaluate result") diff --git a/tests/unit_tests/backends/megatron/test_mlperf_patches.py b/tests/unit_tests/backends/megatron/test_mlperf_patches.py index 077db6d25..b8ca0d8e2 100644 --- a/tests/unit_tests/backends/megatron/test_mlperf_patches.py +++ b/tests/unit_tests/backends/megatron/test_mlperf_patches.py @@ -17,6 +17,16 @@ from types import SimpleNamespace from unittest.mock import MagicMock +import pytest + + +@pytest.fixture(autouse=True) +def _explicit_precision_environment(monkeypatch): + 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") + + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -141,6 +151,38 @@ def _make_ctx( ) +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() + + # ============================================================================ # Level 1: Patch registration and conditions # ============================================================================ From e01c0bcc83c76507cf601ee4c8b1acbb871dea27 Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Mon, 31 Aug 2026 00:06:40 -0500 Subject: [PATCH 07/15] chore(mlperf): move to the mlperf-logging release that ships the flux1 ruleset 6.0.0-rc5 stops at training_5.0.0, so there was no way to check a v6 Flux log against anything: no training_6.0.0 config directory, no closed_flux1.yaml, no rcps_flux1.json. Work on the logging contract was therefore being done against a guess at what the checker wanted. 6.0.0-rc6 adds training_6.0.0 with closed_flux1.yaml and the reference convergence points, which is what makes the rest of this branch verifiable rather than plausible. --- examples/mlperf/flux1/requirements.txt | 2 +- pyproject.toml | 2 +- requirements.txt | 2 +- .../hooks/train/pretrain/diffusion/requirements-diffusion.txt | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) 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/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 From 8e1c22aff7eed7e896849bcc4e27c085fc0125e3 Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Mon, 31 Aug 2026 00:06:55 -0500 Subject: [PATCH 08/15] fix(flux): start the MLPerf clock before Megatron opens the dataset run_start has to precede every read of the real dataset, and it did not. It fired from the first training_log call, which is after the first train_step, which is after build_train_valid_test_data_iterators has already constructed the Energon iterators. Everything between those two points -- reading the .nv-meta shard metadata, spawning loader workers, filling the prefetch queue -- happened inside the initialization window and was excluded from the measured time. The warmup steps landed on the wrong side too: they ran on synthetic data inside the first train_step, so compile and FP8 bring-up were being charged to the run. There is no lifecycle hook at the point where the boundary belongs. Primus' phases (setup / build_args / before_train / after_train) all fire before pretrain(), and pretrain() runs setup_model_and_optimizer, then the data builder, then train() with nothing in between. Megatron-LM is a pinned submodule, so the call site cannot be edited either without forking it. mlperf_boundary creates the seam from outside by wrapping three entry points during before_train: pretrain to capture forward_step_func, which is a pretrain() argument and is reachable nowhere else beforehand; setup_model_and_optimizer to capture the model, optimizer and scheduler that warmup needs; and build_train_valid_test_data_iterators, whose first call is the boundary. Firing there runs the pre-run hooks, synchronizes, emits the transition on rank zero, and holds every rank on a second barrier so none of them can start opening shards while rank zero is still writing run_start. The once-guard covers virtual pipelining, which calls the builder per stage. The warmup body is now shared between the boundary and the original train_step hook rather than duplicated. Only MLPerf mode moves to the boundary: development recipes have no clock to protect and keep the behaviour they were tuned against, and the relocated path needs objects that only exist once the capture wrappers have run, so it fails loudly rather than silently skipping warmup if they have not. The old first-training_log trigger stays as a backstop and is now a no-op once the boundary has fired. --- .../megatron/patches/mlperf_boundary.py | 160 ++++++ .../megatron/patches/mlperf_warmup_patches.py | 498 ++++++++++-------- 2 files changed, 451 insertions(+), 207 deletions(-) create mode 100644 primus/backends/megatron/patches/mlperf_boundary.py 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_warmup_patches.py b/primus/backends/megatron/patches/mlperf_warmup_patches.py index 5ac532978..da0db9d25 100644 --- a/primus/backends/megatron/patches/mlperf_warmup_patches.py +++ b/primus/backends/megatron/patches/mlperf_warmup_patches.py @@ -200,66 +200,274 @@ 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 + + # ---- 4. Run warmup steps with synthetic data ---- + 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, + ) + _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 +491,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 +519,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 +528,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) From 781eaa13350a71147a31c8ec7ff330ceda1972a6 Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Mon, 31 Aug 2026 00:07:12 -0500 Subject: [PATCH 09/15] fix(flux): stop the MLPerf logger inventing what the submission claims Four things in the log were being decided by defaults rather than by the run. The result never reached a file. mllog was left on its default handler, so the records went to stdout interleaved with every other rank's output and with framework noise, and the artifact a checker reads would have been whatever survived a later filtering step. Rank zero now writes MLLOG_OUTPUT_FILE directly. Submission identity defaulted. org fell back to "AMD", division to "closed", platform to "MI355X" and config_filename to "unknown". Those four values decide which division the log is judged in and which recipe a reviewer goes looking for, and a default is exactly how a wrong one gets submitted without anyone noticing. cache_clear defaulted to false in the same way, which let a run that never dropped caches make a claim about the machine it ran on. All of them are now required, and startup fails naming the variable. A run that missed the target produced no run_stop at all -- only convergence emitted one. run_stop is EXACTLY_ONE in the ruleset, so those logs did not parse, and a non-converging run is still part of the campaign the RCP checker compares. It is emitted from after_train now, idempotently, so a converged run does not get a second contradictory record. The block reopened only on the path where the validation loss could be read. When it could not, block_stop had already fired on entry to eval and training carried on outside any block. Precision disclosures are passed through unchanged even when the checker will reject them: training_6.0.0/common.yaml accepts a fixed vocabulary and mxfp6 is not in it, so an MXFP6 run produces a log that has to wait for the format to be approved upstream. Emitting fp8 instead would make it pass by describing the run falsely. It warns instead. The tests now drive the real mlperf_logging end to end -- real mllog writes the file, the real compliance checker reads it back under training_6.0.0 -- so a converged run is verified against the ruleset a submission is judged with rather than against a mock's idea of it. They skip where the package is absent. That is also how two assumptions turned out to be wrong: the checker requires at least one eval_accuracy at or below the target, so no amount of well-formed logging makes an exhausted run pass, and epoch_start needs no matching epoch_stop. --- .../patches/mlperf_logging_patches.py | 220 +++++- .../test_mlperf_log_is_checker_valid.py | 270 +++++++ .../backends/megatron/test_mlperf_patches.py | 670 ++++++++++++------ 3 files changed, 912 insertions(+), 248 deletions(-) create mode 100644 tests/unit_tests/backends/megatron/test_mlperf_log_is_checker_valid.py diff --git a/primus/backends/megatron/patches/mlperf_logging_patches.py b/primus/backends/megatron/patches/mlperf_logging_patches.py index 7b5dc8d8f..6d874eeee 100644 --- a/primus/backends/megatron/patches/mlperf_logging_patches.py +++ b/primus/backends/megatron/patches/mlperf_logging_patches.py @@ -31,9 +31,52 @@ "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 mandatory v6.1 precision disclosures without guessing policy names.""" + """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() @@ -43,6 +86,25 @@ def _precision_disclosures_from_env() -> dict[str, str]: 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 @@ -115,6 +177,20 @@ def __init__( 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")) @@ -142,29 +218,31 @@ 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: - clear_caches = os.environ.get("MLPERF_CLEAR_CACHES", "false").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=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(): @@ -175,7 +253,9 @@ def log_hyperparams(self, args): ("context_parallelism", getattr(args, "context_parallel_size", 1)), ("expert_parallelism", getattr(args, "expert_model_parallel_size", 1)), ("micro_batch_size", self.mbs), - ("config_filename", os.environ.get("EXP", "unknown")), + # 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( @@ -233,12 +313,21 @@ def log_hyperparams(self, args): ) def log_init_stop_run_start(self): - if int(os.environ.get("RANK", "0")) == 0 and not self._run_started: + 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={"samples_count": 0}) - self._run_started = True + 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() @@ -246,7 +335,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( @@ -264,7 +353,7 @@ def on_validation_start(self, global_step: int): 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( @@ -287,7 +376,7 @@ def on_validation_start(self, global_step: int): 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, @@ -327,9 +416,14 @@ def converged(self): return self._converged 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, @@ -396,6 +490,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 @@ -498,7 +606,8 @@ def _capture_wrapper(*a, **kw): 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}" @@ -510,15 +619,16 @@ def _capture_wrapper(*a, **kw): megatron_get_args().train_iters = iteration 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={"samples_count": iteration * gbs}, - ) 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 @@ -530,3 +640,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/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..b0ab6ae2a --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_mlperf_log_is_checker_valid.py @@ -0,0 +1,270 @@ +############################################################################### +# 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 _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_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 b8ca0d8e2..601ffe622 100644 --- a/tests/unit_tests/backends/megatron/test_mlperf_patches.py +++ b/tests/unit_tests/backends/megatron/test_mlperf_patches.py @@ -19,14 +19,62 @@ 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 _explicit_precision_environment(monkeypatch): - monkeypatch.setenv("MLLOG_LOWEST_NUMERICAL_PRECISION_IN_LINEAR", "mxfp6") +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 # --------------------------------------------------------------------------- @@ -70,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 @@ -97,6 +153,45 @@ def fake_get_wandb_writer(): return training_mod, _megatron_args +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, @@ -151,6 +246,19 @@ 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) + + +# ============================================================================ +# 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, @@ -183,8 +291,54 @@ def test_missing_precision_disclosure_fails_mlperf_startup(monkeypatch): _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() + + # ============================================================================ -# Level 1: Patch registration and conditions +# Patch registration # ============================================================================ @@ -193,42 +347,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, @@ -238,8 +358,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 @@ -248,42 +367,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, @@ -300,195 +385,358 @@ 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)) + - 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", +# ============================================================================ +# The measured-time boundary +# ============================================================================ + + +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) + + from primus.backends.megatron.patches.mlperf_logging_patches import ( + patch_mlperf_logging, + ) - mock_mllog_module = MagicMock() - mock_mllog_module.get_mllogger.return_value = mock_mllogger - mock_mllog_module.constants = mock_constants + patch_mlperf_logging(_make_ctx(mlperf_mode=True)) + events.clear() - # 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 + mt.build_train_valid_test_data_iterators(None) + mt.build_train_valid_test_data_iterators(None) - import sys as _sys + assert _keys(events).count("run_start") == 1 - monkeypatch.setitem(_sys.modules, "mlperf_logging", mock_mlperf_pkg) - monkeypatch.setitem(_sys.modules, "mlperf_logging.mllog", mock_mllog_module) - monkeypatch.setenv("RANK", "0") + 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) - emitted_events.clear() + 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, + ) + + 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)) - init_stop_idx = event_keys.index("init_stop") - run_start_idx = event_keys.index("run_start") - assert init_stop_idx < run_start_idx + 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"] - # 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 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) - 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 + 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 == [] # ============================================================================ -# Level 3: Component tests +# Run lifecycle # ============================================================================ class TestConvergenceDetection: """Verify convergence detection in evaluate_and_print_results wrapper.""" - def test_convergence_sets_train_iters(self, monkeypatch): + def _patch_with_eval_loss(self, monkeypatch, loss): mt, megatron_args = _install_fake_megatron(monkeypatch) + _silence_log_rank_0(monkeypatch) + _, events = _install_mock_mllog(monkeypatch) - 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 + 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, ) - target_val_loss = 0.586 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 - ctx = _make_ctx(mlperf_mode=True, target_val_loss=target_val_loss) - patch_mlperf_logging(ctx) - - # Call eval at iteration 512 + def _run_eval(self, mt, iteration): mt.evaluate_and_print_results( - "iteration 512", + f"iteration {iteration}", lambda: None, None, [MagicMock()], - 512, + iteration, None, MagicMock(), ) + def test_convergence_sets_train_iters(self, monkeypatch): + mt, megatron_args, _ = self._patch_with_eval_loss(monkeypatch, 0.500) + self._run_eval(mt, 512) assert megatron_args.train_iters == 512 + def test_convergence_emits_a_successful_run_stop(self, monkeypatch): + mt, _, events = self._patch_with_eval_loss(monkeypatch, 0.500) + self._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 = self._patch_with_eval_loss(monkeypatch, 0.900) + self._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 = self._patch_with_eval_loss(monkeypatch, 0.900) + mt.evaluate = lambda *a, **k: ({},) + events.clear() + + self._run_eval(mt, 512) + + assert "block_start" in _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 + + 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, + ) + + patch_mlperf_logging(_make_ctx(mlperf_mode=True, global_batch_size=512, eval_interval=512)) + + 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) + + from primus.backends.megatron.patches.mlperf_logging_patches import ( + patch_mlperf_logging, + ) + + 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)) + # ============================================================================ # Level 4: Helper function unit tests (CPU-only, no GPU required) From 6426afd22298093afeb1366a2c131514fc17eded Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Mon, 31 Aug 2026 00:07:29 -0500 Subject: [PATCH 10/15] fix(flux): give the MLPerf recipes a step budget that can reach the target train_iters was 5000, which at GBS 512 is 2560000 samples. The reference convergence points for flux1 at that batch size span 7077888 to 7602176 samples across 20 NVIDIA BF16 runs (rcp_checker training_6.0.0/rcps_flux1.json), so the budget ran out around a third of the way to the earliest of them. Every run was going to exhaust below target no matter how well it trained, and none of them could ever have counted as a result. 16000 steps is 8192000 samples, clearing the slowest reference run with room for a run that converges later than the reference spread. It is a cap and not a target: the run ends when eval_accuracy reaches 0.586 and how many samples that took is the result. lr_decay_iters moves with it -- inert while lr_decay_style is constant, but 4000 next to a 16000-step budget reads as a decay schedule that ends early. The seed also comes from outside now. A campaign is ten runs with ten different seeds, and the value in the log has to be the one the launcher chose rather than the one the file was written with. --- ...chnell_resample_local_spec_fp8_mlperf.yaml | 20 ++++++++++++++++--- ...n_schnell_resample_te_spec_fp8_mlperf.yaml | 19 +++++++++++++++--- 2 files changed, 33 insertions(+), 6 deletions(-) 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 06252438a..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,7 +73,17 @@ 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 # Cover the whole MLPerf validation set. eval_iters is derived from this # (29696 / 512 = 58); setting both is rejected. The previous eval_iters of @@ -114,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 # ========================================== @@ -195,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 60ccd9a9e..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,7 +86,16 @@ 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 # Cover the whole MLPerf validation set. eval_iters is derived from this # (29696 / 512 = 58); setting both is rejected. The previous eval_iters of @@ -130,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 # ========================================== @@ -193,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). From 11d47e7b074ae0f7ccaa2499b4c244b63e15d4af Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Mon, 31 Aug 2026 00:07:29 -0500 Subject: [PATCH 11/15] feat(mlperf): add the Megatron launcher for flux1 submission runs The logging patch now fails closed on everything it would otherwise guess, so something has to supply the answers. This is that something, which makes it part of the submission rather than a convenience wrapper: it sets the submission identity, the precision disclosures, the result path and the recipe name, and it drops the page cache and reports cache_clear=false when it lacks the privileges to do so, instead of claiming a cold start it did not get. Each run is checked against the ruleset as soon as it finishes. Checking at the end of the campaign instead would mean finding out after ten runs and several days of compute that none of the logs parse. The campaign driver does ten runs with distinct seeds, then an RCP comparison over the set. Ten is what rcp_checker requires for flux1, not a round number. A run that ends without converging still produces a result file and still counts as one of the ten -- it is part of the distribution being compared, and dropping the runs that went badly is how the comparison stops meaning anything. --- examples/mlperf/flux1/megatron/README.md | 69 ++++++++++ .../mlperf/flux1/megatron/run_and_time.sh | 120 ++++++++++++++++++ .../mlperf/flux1/megatron/run_campaign.sh | 57 +++++++++ 3 files changed, 246 insertions(+) create mode 100644 examples/mlperf/flux1/megatron/README.md create mode 100755 examples/mlperf/flux1/megatron/run_and_time.sh create mode 100755 examples/mlperf/flux1/megatron/run_campaign.sh 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" From 67d03094d024fc6793777479f2ce391063086d00 Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Mon, 31 Aug 2026 14:37:29 -0500 Subject: [PATCH 12/15] fix(flux): stop a converged run evaluating once more after run_stop pretrain() runs a final validation pass whenever args.do_valid is set, and the early stop only clears args.train_iters, so reaching the target loss was followed by one more evaluation. That evaluation draws its own VAE epsilon and flow-matching noise, so it can land above the target and contradict the evaluation that just ended the run, leaving a result file whose last eval_accuracy is worse than the one at run_stop. Clear do_valid alongside train_iters so the pass never runs, and refuse to emit any MLLOG record once run_stop is out, so an evaluation reached by some other path cannot append to a closed log either. The second guard keys on run_stopped rather than converged, which also covers aborted runs. Verified on 8x MI355X: a run converging at step 2 of a 4-step budget now ends at run_stop with a single eval_accuracy and nothing after it. --- .../patches/mlperf_logging_patches.py | 29 +++- .../test_mlperf_log_is_checker_valid.py | 50 ++++++ .../backends/megatron/test_mlperf_patches.py | 143 +++++++++++++----- 3 files changed, 186 insertions(+), 36 deletions(-) diff --git a/primus/backends/megatron/patches/mlperf_logging_patches.py b/primus/backends/megatron/patches/mlperf_logging_patches.py index 6d874eeee..a808ae63e 100644 --- a/primus/backends/megatron/patches/mlperf_logging_patches.py +++ b/primus/backends/megatron/patches/mlperf_logging_patches.py @@ -415,6 +415,15 @@ 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. @@ -556,6 +565,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. @@ -616,9 +633,17 @@ 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") + 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") 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 index b0ab6ae2a..95bc5c460 100644 --- 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 @@ -191,6 +191,17 @@ def _drive_a_run(mt, eval_losses): ) +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 @@ -215,6 +226,45 @@ def test_a_converged_run_passes_the_compliance_checker(monkeypatch, _run_environ 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. diff --git a/tests/unit_tests/backends/megatron/test_mlperf_patches.py b/tests/unit_tests/backends/megatron/test_mlperf_patches.py index 601ffe622..a9eb04fdf 100644 --- a/tests/unit_tests/backends/megatron/test_mlperf_patches.py +++ b/tests/unit_tests/backends/megatron/test_mlperf_patches.py @@ -577,46 +577,69 @@ def test_warmup_stays_on_train_step_outside_mlperf_mode(self, monkeypatch): # ============================================================================ -class TestConvergenceDetection: - """Verify convergence detection in evaluate_and_print_results wrapper.""" +def _patch_with_eval_loss(monkeypatch, loss): + mt, megatron_args = _install_fake_megatron(monkeypatch) + _silence_log_rank_0(monkeypatch) + _, events = _install_mock_mllog(monkeypatch) - def _patch_with_eval_loss(self, 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) - 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, + ) - 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(), + ) - 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(self, 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, _ = self._patch_with_eval_loss(monkeypatch, 0.500) - self._run_eval(mt, 512) + 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 = self._patch_with_eval_loss(monkeypatch, 0.500) - self._run_eval(mt, 512) + 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 @@ -625,8 +648,8 @@ def test_convergence_emits_a_successful_run_stop(self, monkeypatch): 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 = self._patch_with_eval_loss(monkeypatch, 0.900) - self._run_eval(mt, 512) + 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") @@ -634,15 +657,67 @@ def test_a_missed_target_reopens_the_block(self, monkeypatch): assert "run_stop" not in keys def test_an_unreadable_loss_still_reopens_the_block(self, monkeypatch): - mt, _, events = self._patch_with_eval_loss(monkeypatch, 0.900) + mt, _, events = _patch_with_eval_loss(monkeypatch, 0.900) mt.evaluate = lambda *a, **k: ({},) events.clear() - self._run_eval(mt, 512) + _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.""" From ff91c9757db33b571dfe0cd4dbf8d04460a4b9de Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Mon, 31 Aug 2026 14:37:42 -0500 Subject: [PATCH 13/15] fix(flux): key the evaluation RNG on the step being evaluated The per-microbatch evaluation RNG index was built from args.iteration and a microbatch counter that only reset when that value changed. Neither is what it was meant to be. args.iteration is the resume point: Megatron assigns it at setup and at checkpoint load and never inside the training loop, so an evaluation at step k drew one set of noise when reached continuously and a different set when reached after a resume. And because it never changes mid-run, the counter never reset, so two evaluations at the same step continued one stream rather than reproducing each other. Key on args.curr_iteration, the step the loop last completed, and reset the counter per evaluation rather than per step. The evaluator signals the boundary through a module-global counter in eval_session, which is import-free so the per-microbatch path does not pull in the Megatron evaluation stack. Measured on 8x MI355X with two evaluations forced at the same step: they reported 1.688877 and 1.690349 before, and 1.685579 twice after. A drift of that size is enough to carry a follow-up evaluation across a 0.586 target. Note this changes which noise every evaluation draws, so convergence evidence collected before it does not carry over. --- primus/backends/megatron/diffusion_trainer.py | 50 ++++++-- .../megatron/training/eval_session.py | 54 +++++++++ .../backends/megatron/training/evaluator.py | 5 + .../test_diffusion_trainer_eval_rng.py | 111 +++++++++++++++--- .../megatron/test_evaluator_reduction.py | 95 ++++++++++++++- 5 files changed, 282 insertions(+), 33 deletions(-) create mode 100644 primus/backends/megatron/training/eval_session.py diff --git a/primus/backends/megatron/diffusion_trainer.py b/primus/backends/megatron/diffusion_trainer.py index 939220c98..ffa71efc1 100644 --- a/primus/backends/megatron/diffusion_trainer.py +++ b/primus/backends/megatron/diffusion_trainer.py @@ -50,7 +50,9 @@ def __init__(self, *args, **kwargs): self._compiled_loss_fn = None self._forward_step_count = 0 self._forward_step_count_initialized = False - self._eval_rng_iteration = None + # 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 @@ -315,21 +317,43 @@ def diffusion_loss_func(output_tensor, non_loss_data=False): def _next_eval_step_index(self) -> int: """Index identifying this validation microbatch within the run. - Built from ``(iteration, microbatch index within this evaluation)`` - rather than from a free-running counter, so it is reproducible after a - checkpoint resume without having to checkpoint the counter itself: the - iteration comes from the checkpoint and the index restarts at zero for - each evaluation. + 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, ) - - iteration = getattr(get_args(), "iteration", 0) - if iteration != self._eval_rng_iteration: - self._eval_rng_iteration = iteration + 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 @@ -338,11 +362,11 @@ def _next_eval_step_index(self) -> int: if index >= EVAL_RNG_ITERATION_STRIDE: raise RuntimeError( f"Evaluation ran {index + 1} microbatches, at or beyond the " - f"per-iteration stride {EVAL_RNG_ITERATION_STRIDE} that keeps " - f"consecutive evaluations' RNG streams disjoint." + f"stride {EVAL_RNG_ITERATION_STRIDE} that keeps consecutive " + f"training steps' evaluation RNG streams disjoint." ) - return iteration * EVAL_RNG_ITERATION_STRIDE + index + return step * EVAL_RNG_ITERATION_STRIDE + index def get_forward_step(self): """ 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 d7cbf7bf8..2df9dc24f 100644 --- a/primus/backends/megatron/training/evaluator.py +++ b/primus/backends/megatron/training/evaluator.py @@ -16,6 +16,7 @@ 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 debug_rank_0, log_rank_0 @@ -167,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": 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 index b5afe6086..1c87937f4 100644 --- a/tests/unit_tests/backends/megatron/test_diffusion_trainer_eval_rng.py +++ b/tests/unit_tests/backends/megatron/test_diffusion_trainer_eval_rng.py @@ -10,6 +10,13 @@ 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 @@ -18,11 +25,23 @@ 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): @@ -40,52 +59,105 @@ def _make_trainer(): trainer = _ConcreteDiffusionTrainer.__new__(_ConcreteDiffusionTrainer) trainer._forward_step_count = 0 trainer._forward_step_count_initialized = False - trainer._eval_rng_iteration = None + trainer._eval_session = None trainer._eval_microbatch_index = 0 return trainer -def _indices(trainer, iteration, count): - with patch("megatron.training.get_args", return_value=SimpleNamespace(iteration=iteration)): +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 _indices(trainer, iteration=100, count=4) == [ + 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() - _indices(trainer, iteration=100, count=3) - second = _indices(trainer, iteration=200, count=3) + _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(_indices(trainer, iteration=100, count=64)) - second = set(_indices(trainer, iteration=101, count=64)) + first = set(_evaluation(trainer, count=64, curr_iteration=100)) + second = set(_evaluation(trainer, count=64, curr_iteration=101)) assert not (first & second) - def test_reproducible_after_resume(self): - """A fresh trainer at the same iteration must produce the same indices.""" - before = _indices(_make_trainer(), iteration=512, count=8) - after_resume = _indices(_make_trainer(), iteration=512, count=8) + 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 before == after_resume + 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_rng_iteration = 5 + trainer._eval_session = eval_session.current_eval_session() trainer._eval_microbatch_index = EVAL_RNG_ITERATION_STRIDE - with patch("megatron.training.get_args", return_value=SimpleNamespace(iteration=5)): - with pytest.raises(RuntimeError, match="per-iteration stride"): + with patch("megatron.training.get_args", return_value=_args(5, 5)): + with pytest.raises(RuntimeError, match="stride"): trainer._next_eval_step_index() @@ -108,9 +180,9 @@ 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, iteration * EVAL_RNG_ITERATION_STRIDE + micro) + self._eval_seed(seed, rank, step * EVAL_RNG_ITERATION_STRIDE + micro) for rank in range(64) - for iteration in range(0, 20000, 512) + for step in range(0, 20000, 512) for micro in range(58) } @@ -169,10 +241,11 @@ def test_eval_pass_freezes_training_counter_but_advances_eval_index(self): 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=SimpleNamespace(iteration=3)): + ), 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) diff --git a/tests/unit_tests/backends/megatron/test_evaluator_reduction.py b/tests/unit_tests/backends/megatron/test_evaluator_reduction.py index e2935cfe8..f6fe405d6 100644 --- a/tests/unit_tests/backends/megatron/test_evaluator_reduction.py +++ b/tests/unit_tests/backends/megatron/test_evaluator_reduction.py @@ -14,15 +14,17 @@ from contextlib import contextmanager from types import SimpleNamespace -from unittest.mock import patch +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, ) @@ -302,3 +304,94 @@ def test_mlperf_mode_still_reports_a_context_parallel_mismatch(self): 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) From 72251fe43dd9f897c2140272783963e0648d4913 Mon Sep 17 00:00:00 2001 From: GP Huang Date: Tue, 1 Sep 2026 11:07:12 +0300 Subject: [PATCH 14/15] fix(flux): wait on the warmup's gradient reduce-scatter (#1069) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Every Flux run with `mlperf_mode: true` and `warmup_train_steps > 0` dies during the first real training step, before a single iteration completes. This makes the warmup wait on the gradient reduce-scatter it dispatches. The MLPerf warmup runs at the pre-data boundary, which is before Megatron's `train()` executes. `train()` is where `config.finalize_model_grads_func` gets assigned, and that callback is the only caller of `finish_grad_sync()`. So with `overlap_grad_reduce: true` 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: ``` File ".../megatron/core/distributed/param_and_grad_buffer.py", line 530, in start_grad_sync self.grad_reduce_handle is None AssertionError: Should not have multiple communication calls outstanding at once ``` ### The fix Install `finalize_model_grads` for the duration of the warmup loop and restore whatever was there before, so the boundary warmup does not change what `train()` later sees. The `None` guard leaves the in-`train_step` warmup path used by development recipes untouched, since by then `train()` has already installed the callback — which is why only MLPerf mode was affected. ### Why the escape hatch does not cover this `start_grad_sync` no-ops when a handle is outstanding *and* `is_first_batch` is still set. Warmup step 1 populates the ready counts without dispatching; the `reset()` at the head of warmup step 2 records the golden counts and clears `is_first_batch`; warmup step 2 then dispatches for real. By the first training step the flag is long gone, so the assertion fires. ## Test plan - [x] Two unit tests in `tests/unit_tests/backends/megatron/test_mlperf_patches.py`: the warmup steps see a grad-finalize callback, and a callback the caller already set is preserved rather than replaced. Both fail without the production change. - [x] Full `test_mlperf_patches.py` passes (46 tests) on an MI355X dev container. - [x] Bisected on 8x MI355X with the Flux 12B MXFP6 MLPerf recipe. The FP8 MLPerf recipe fails identically, so this is not precision-specific; `mlperf_mode: false` passes; `warmup_train_steps: 0` passes. - [x] With this change the full MLPerf recipe at `warmup_train_steps: 2` trains 20 iterations clean, and a full convergence run is in flight. --------- Co-authored-by: guangphu --- .../megatron/patches/mlperf_warmup_patches.py | 42 ++++-- .../backends/megatron/test_mlperf_patches.py | 126 ++++++++++++++++++ 2 files changed, 156 insertions(+), 12 deletions(-) diff --git a/primus/backends/megatron/patches/mlperf_warmup_patches.py b/primus/backends/megatron/patches/mlperf_warmup_patches.py index da0db9d25..a9fa5a1ee 100644 --- a/primus/backends/megatron/patches/mlperf_warmup_patches.py +++ b/primus/backends/megatron/patches/mlperf_warmup_patches.py @@ -270,19 +270,37 @@ def _run_warmup_and_restore( # ---- 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 ---- - 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, - ) + 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") # ---- 5. Restore optimizer ---- diff --git a/tests/unit_tests/backends/megatron/test_mlperf_patches.py b/tests/unit_tests/backends/megatron/test_mlperf_patches.py index a9eb04fdf..6f257cf8a 100644 --- a/tests/unit_tests/backends/megatron/test_mlperf_patches.py +++ b/tests/unit_tests/backends/megatron/test_mlperf_patches.py @@ -153,6 +153,51 @@ 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). @@ -571,6 +616,87 @@ def test_warmup_stays_on_train_step_outside_mlperf_mode(self, monkeypatch): 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" + + 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) + + def preexisting(*args, **kwargs): + return None + + 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 + # ============================================================================ # Run lifecycle From fe20330fccfdf3521aa1f8f586bbdc82a7430dd7 Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Tue, 1 Sep 2026 07:21:30 -0500 Subject: [PATCH 15/15] fix(flux): refuse injected validation timesteps under mlperf_mode resolve_validation_timesteps prefers the dataset's timestep column whenever the batch carries one, so a re-ingested split is evaluated correctly under either setting of eval_timestep_source. That leaves exactly one unsafe combination: shards ingested before the column was carried through, read under equidistant. There the positional fallback injects t = index % 8, which reproduces neither the timesteps nor the image-to-timestep pairing the published split defines, and it fails silently -- the loss it reports is indistinguishable in the logs from a correct one. The MLPerf recipes all 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 rested on every future submission recipe remembering to override it. A recipe copied from a non-MLPerf diffusion config would land back on injected timesteps. There are also two complete copies of the val split on shared storage, one annotated and one not, so the combination is reachable today by pointing a run at the wrong path. The check goes beside the existing eval_iters guard in the Energon provider rather than 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, merged onto args only after build_args has run. Placement within the provider matters for the same reason: validation dataloader construction is wrapped in a bare except that disables evaluation and carries on, so the guard has to precede it. Adds the first tests to drive the provider's guards through create_dataloaders, pinning that an MLPerf run on injected timesteps raises before any dataset is built and leaves eval_iters alone, and that a genuine construction failure is still swallowed as before. Also corrects the comment on evaluation_frequency, which described a getattr fallback the key no longer uses. No mlperf_logging release defines a constant for that name; it appears only in the checker's own rulesets, so the literal is the spelling rather than a fallback. --- .../megatron/data/energon_dataset_provider.py | 4 + .../patches/mlperf_logging_patches.py | 8 +- .../backends/megatron/training/eval_budget.py | 60 ++++++++- .../test_energon_provider_eval_guards.py | 126 ++++++++++++++++++ .../backends/megatron/test_eval_budget.py | 39 ++++++ 5 files changed, 233 insertions(+), 4 deletions(-) create mode 100644 tests/unit_tests/backends/megatron/test_energon_provider_eval_guards.py diff --git a/primus/backends/megatron/data/energon_dataset_provider.py b/primus/backends/megatron/data/energon_dataset_provider.py index 743853657..cf3bdf704 100644 --- a/primus/backends/megatron/data/energon_dataset_provider.py +++ b/primus/backends/megatron/data/energon_dataset_provider.py @@ -33,6 +33,7 @@ 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, @@ -163,6 +164,9 @@ def create_dataloaders( * (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} " diff --git a/primus/backends/megatron/patches/mlperf_logging_patches.py b/primus/backends/megatron/patches/mlperf_logging_patches.py index a808ae63e..f29415cad 100644 --- a/primus/backends/megatron/patches/mlperf_logging_patches.py +++ b/primus/backends/megatron/patches/mlperf_logging_patches.py @@ -272,9 +272,11 @@ def log_hyperparams(self, args): key=self._constants.EVAL_SAMPLES, value=getattr(args, "eval_samples", None) or eval_iters * self.gbs, ) - # How often evaluation runs, in samples. Required by the MLPerf logging - # rules and previously absent; the constant name differs across - # mlperf_logging releases, so fall back to the literal key. + # 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, diff --git a/primus/backends/megatron/training/eval_budget.py b/primus/backends/megatron/training/eval_budget.py index d1a468496..1e83f2235 100644 --- a/primus/backends/megatron/training/eval_budget.py +++ b/primus/backends/megatron/training/eval_budget.py @@ -32,6 +32,7 @@ __all__ = [ "DEFAULT_VAL_NUM_WORKERS", "EvalCoverageError", + "assert_mlperf_timestep_source", "assert_val_worker_divisibility", "get_data_parallel_size", "get_eval_num_microbatches", @@ -52,7 +53,12 @@ class EvalCoverageError(ValueError): - """Raised when an evaluation would not read the samples it claims to.""" + """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: @@ -187,6 +193,58 @@ def assert_val_worker_divisibility(args, eval_samples: int) -> None: ) +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. 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 index 299e4bac3..8bfa7505b 100644 --- a/tests/unit_tests/backends/megatron/test_eval_budget.py +++ b/tests/unit_tests/backends/megatron/test_eval_budget.py @@ -18,6 +18,7 @@ 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, @@ -98,6 +99,44 @@ def test_default_worker_count_does_not_divide_by_zero(self): 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