From cbf4ad52dd7b5867e512661b4ad331f952586b63 Mon Sep 17 00:00:00 2001 From: Guangpu Huang Date: Tue, 11 Aug 2026 07:07:24 +0000 Subject: [PATCH 01/14] fix(megatron): preserve exact diffusion resume state Restore rank-local Energon state and reset warmup RNG accounting so resumed diffusion runs continue the same data and seed sequence. Add opt-in runtime fingerprints and precision census for validation. --- primus/backends/megatron/data/dataloader.py | 87 ++++- .../data/diffusion/task_encoders/image.py | 20 +- .../megatron/data/energon_dataset_provider.py | 70 +++- primus/backends/megatron/diffusion_trainer.py | 15 + .../megatron/flux_pretrain_trainer.py | 41 ++ .../megatron/patches/mlperf_warmup_patches.py | 51 ++- .../training/diffusion/forward_step.py | 46 +++ .../megatron/test_dataloader_checkpoint.py | 362 ++++++++++++++++++ .../megatron/test_diffusion_audit_markers.py | 109 ++++++ 9 files changed, 784 insertions(+), 17 deletions(-) create mode 100644 tests/unit_tests/backends/megatron/test_dataloader_checkpoint.py create mode 100644 tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py diff --git a/primus/backends/megatron/data/dataloader.py b/primus/backends/megatron/data/dataloader.py index 4a53ceb13..ac8eae8dc 100644 --- a/primus/backends/megatron/data/dataloader.py +++ b/primus/backends/megatron/data/dataloader.py @@ -23,10 +23,13 @@ """ import logging -from typing import Any, Iterator +from pathlib import Path +from typing import Any, Callable, Iterator, Optional logger = logging.getLogger(__name__) +DATALOADER_STATE_KEY = "dataloader_state_dict" + def cyclic_iter(iterator: Iterator) -> Iterator: """ @@ -129,26 +132,102 @@ def save_state(self) -> Any: ) return None - def restore_state(self, state: Any): + def restore_state(self, state: Any, *, strict: bool = False) -> bool: """ Restore dataloader state from checkpoint (if supported). Uses duck typing - checks for restore_state_rank() method. - No-op if not supported (e.g., for synthetic/mock data). + No-op if not supported (e.g., for synthetic/mock data), unless strict=True. Args: state: Dataloader state dictionary + strict: Raise when the wrapped dataloader cannot restore state. + + Returns: + True when state was restored, otherwise False. """ if hasattr(self._dataloader, "restore_state_rank"): self._dataloader.restore_state_rank(state) # Recreate iterator after restore self._iter = iter(cyclic_iter(self._dataloader)) logger.info("Restored dataloader state from checkpoint") + return True else: + if strict: + raise RuntimeError( + f"{type(self._dataloader).__name__} does not support restore_state_rank()" + ) logger.debug( f"{type(self._dataloader).__name__} does not support restore_state_rank() " f"(OK for synthetic data)" ) + return False + + +def restore_dataloader_state_from_checkpoint( + dataloader: MegatronDataloaderWrapper, + checkpoint_root: str, + iteration: int, + data_parallel_rank: int, + *, + checkpoint_name_fn: Optional[Callable[..., str]] = None, + load_fn: Optional[Callable[[Path], Any]] = None, +) -> Path: + """Restore one data-parallel rank's dataloader checkpoint. + + ``maybe_save_dataloader_state`` in Megatron writes one + ``train_dataloader_dprankNNN.pt`` file per data-parallel rank under the + model checkpoint iteration directory. This helper resolves that exact + path and fails closed when the file or payload is invalid. + """ + if iteration < 0: + raise ValueError(f"iteration must be non-negative, got {iteration}") + if data_parallel_rank < 0: + raise ValueError(f"data_parallel_rank must be non-negative, got {data_parallel_rank}") + + if checkpoint_name_fn is None: + from megatron.training.checkpointing import get_checkpoint_name + + checkpoint_name_fn = get_checkpoint_name + + checkpoint_path = Path( + checkpoint_name_fn( + checkpoint_root, + iteration, + tensor_rank=0, + pipeline_rank=0, + basename=f"train_dataloader_dprank{data_parallel_rank:03d}.pt", + ) + ) + if not checkpoint_path.is_file(): + raise FileNotFoundError( + "Energon dataloader checkpoint is missing for " + f"iteration={iteration}, data_parallel_rank={data_parallel_rank}: " + f"{checkpoint_path}" + ) + + if load_fn is None: + import torch + + # This file is produced by the same trusted Megatron checkpoint run. + def _torch_load(path: Path) -> Any: + return torch.load(path, map_location="cpu", weights_only=False) + + load_fn = _torch_load + + payload = load_fn(checkpoint_path) + if not isinstance(payload, dict): + raise TypeError( + f"Expected a dictionary in dataloader checkpoint {checkpoint_path}, " + f"got {type(payload).__name__}" + ) + if DATALOADER_STATE_KEY not in payload: + raise KeyError( + f"Dataloader checkpoint {checkpoint_path} has no {DATALOADER_STATE_KEY!r} entry" + ) + + dataloader.restore_state(payload[DATALOADER_STATE_KEY], strict=True) + return checkpoint_path # Backwards compatibility alias (DEPRECATED) @@ -157,7 +236,9 @@ def restore_state(self, state: Any): __all__ = [ + "DATALOADER_STATE_KEY", "MegatronDataloaderWrapper", "cyclic_iter", "EnergonDataloader", # Deprecated, use MegatronDataloaderWrapper + "restore_dataloader_state_from_checkpoint", ] diff --git a/primus/backends/megatron/data/diffusion/task_encoders/image.py b/primus/backends/megatron/data/diffusion/task_encoders/image.py index 837bbcfdb..8a78daadb 100644 --- a/primus/backends/megatron/data/diffusion/task_encoders/image.py +++ b/primus/backends/megatron/data/diffusion/task_encoders/image.py @@ -12,8 +12,11 @@ Position IDs are generated in the model code, not here. """ +import hashlib import io +import json import logging +import os from dataclasses import dataclass from pathlib import Path from typing import Any, Dict, List, Optional @@ -33,6 +36,17 @@ logger = logging.getLogger(__name__) +def _sample_key_fingerprint(samples: List["DiffusionSample"]) -> str: + keys = [] + for sample in samples: + key = getattr(sample, "__key__", None) + if key is None: + raise RuntimeError("Cannot audit data continuity: sample has no __key__") + keys.append(str(key)) + payload = json.dumps(keys, ensure_ascii=False, separators=(",", ":")).encode() + return hashlib.sha256(payload).hexdigest() + + # ============================================================================ # Sample Definition (with proper Sample inheritance) # ============================================================================ @@ -304,7 +318,7 @@ def __init__(self, worker_config: Optional[WorkerConfig] = None): self.worker_config = worker_config logger.info("Initialized EncodedDiffusionTaskEncoder (preencoded / preencoded_numpy modes)") - def batch(self, samples: List[DiffusionSample]) -> Dict[str, torch.Tensor]: + def batch(self, samples: List[DiffusionSample]) -> Dict[str, Any]: """ Batch pre-encoded samples. @@ -331,6 +345,10 @@ def batch(self, samples: List[DiffusionSample]) -> Dict[str, torch.Tensor]: if samples[0].timestep is not None: batch["timestep"] = torch.stack([s.timestep for s in samples]) + if os.getenv("PRIMUS_AUDIT_BATCH_FINGERPRINTS") == "1": + batch["_audit_sample_key_sha256"] = _sample_key_fingerprint(samples) + batch["_audit_sample_count"] = len(samples) + return batch diff --git a/primus/backends/megatron/data/energon_dataset_provider.py b/primus/backends/megatron/data/energon_dataset_provider.py index 4ac1728dc..cc4e77195 100644 --- a/primus/backends/megatron/data/energon_dataset_provider.py +++ b/primus/backends/megatron/data/energon_dataset_provider.py @@ -30,7 +30,10 @@ get_val_datasets, ) -from primus.backends.megatron.data.dataloader import MegatronDataloaderWrapper +from primus.backends.megatron.data.dataloader import ( + MegatronDataloaderWrapper, + restore_dataloader_state_from_checkpoint, +) from primus.backends.megatron.data.dataset_provider import DatasetProvider from primus.core.utils.module_utils import log_rank_0 @@ -119,6 +122,7 @@ def create_dataloaders( train_dataset, worker_config=worker_config, prefetch_factor=prefetch_factor ) train_dataloader = MegatronDataloaderWrapper(train_dataloader) + self._restore_train_dataloader_state(args, train_dataloader) log_rank_0("Created training dataloader") # Create validation dataloaders if evaluation is enabled @@ -174,6 +178,70 @@ def create_dataloaders( return train_dataloader, valid_dataloaders, test_dataloader + def _restore_train_dataloader_state( + self, args: Any, train_dataloader: MegatronDataloaderWrapper + ) -> Optional[str]: + """Restore Energon position when resuming from a model checkpoint. + + Megatron saves Energon state separately from the model checkpoint when + ``dataloader_save`` is configured. ``require_dataloader_restore`` makes + successful rank-local restoration mandatory; this is intended for exact + training continuation rather than finetuning or optional weight loads. + """ + require_restore = getattr(args, "require_dataloader_restore", False) + if not isinstance(require_restore, bool): + raise TypeError( + "require_dataloader_restore must be a boolean, " + f"got {require_restore!r}" + ) + + load_path = getattr(args, "load", None) + finetune = bool(getattr(args, "finetune", False)) + iteration = getattr(args, "iteration", 0) + resumed = ( + bool(load_path) + and not finetune + and isinstance(iteration, int) + and not isinstance(iteration, bool) + and iteration > 0 + ) + + if not resumed: + if require_restore: + raise RuntimeError( + "Exact Energon continuation requires a successfully loaded " + "non-finetune checkpoint with iteration > 0; " + f"got load={load_path!r}, finetune={finetune!r}, " + f"iteration={iteration!r}" + ) + if load_path and not finetune: + log_rank_0( + "WARNING: Checkpoint load did not resume a positive iteration; " + "the Energon dataloader position will not be restored" + ) + return None + + dataloader_save = getattr(args, "dataloader_save", None) + if not dataloader_save: + message = ( + "Resuming an Energon run without dataloader_save; " + "the dataloader position cannot be restored" + ) + if require_restore: + raise RuntimeError(message) + log_rank_0(f"WARNING: {message}") + return None + + data_parallel_rank = parallel_state.get_data_parallel_rank() + checkpoint_path = restore_dataloader_state_from_checkpoint( + train_dataloader, + str(dataloader_save), + iteration, + data_parallel_rank, + ) + log_rank_0(f"Restored Energon dataloader state from: {checkpoint_path}") + return str(checkpoint_path) + @property def is_distributed(self) -> bool: """ diff --git a/primus/backends/megatron/diffusion_trainer.py b/primus/backends/megatron/diffusion_trainer.py index dd91f3285..90d5eb2cf 100644 --- a/primus/backends/megatron/diffusion_trainer.py +++ b/primus/backends/megatron/diffusion_trainer.py @@ -324,6 +324,21 @@ def diffusion_forward_step(data_iterator, model): return self.forward_step(data_iterator, model, return_schedule_plan=False) + def reset_forward_step_count(iteration=0): + from megatron.core.num_microbatches_calculator import ( + get_num_microbatches, + ) + + self._forward_step_count = int(iteration) * get_num_microbatches() + self._forward_step_count_initialized = True + log_rank_0( + "[DiffusionPretrainTrainer] Reset forward-step RNG counter " + f"to {self._forward_step_count}" + ) + + diffusion_forward_step._primus_reset_forward_step_count = ( + reset_forward_step_count + ) return diffusion_forward_step def get_datasets_provider(self): diff --git a/primus/backends/megatron/flux_pretrain_trainer.py b/primus/backends/megatron/flux_pretrain_trainer.py index 36da05ce3..7d04971cd 100644 --- a/primus/backends/megatron/flux_pretrain_trainer.py +++ b/primus/backends/megatron/flux_pretrain_trainer.py @@ -13,7 +13,9 @@ - Custom forward step function """ +import json import os +from collections import Counter import numpy as np import torch @@ -28,6 +30,44 @@ ) from primus.core.utils.module_utils import log_rank_0 +PRECISION_LINEAR_CLASSES = ( + "MXFP4ColumnParallelLinear", + "MXFP4RowParallelLinear", + "Float8ColumnParallelLinear", + "Float8RowParallelLinear", +) + + +def _precision_linear_class_census(model) -> dict[str, int]: + observed = Counter(type(module).__name__ for module in model.modules()) + return {name: observed.get(name, 0) for name in PRECISION_LINEAR_CLASSES} + + +def _emit_precision_linear_class_census(model) -> None: + """Emit the actually instantiated precision-linear classes on every rank.""" + if os.getenv("PRIMUS_AUDIT_LINEAR_CLASS_CENSUS") != "1": + return + + from megatron.core import parallel_state + + counts = _precision_linear_class_census(model) + if sum(counts.values()) <= 0: + raise RuntimeError("No MXFP4 or Float8 linear modules were instantiated") + global_rank = ( + torch.distributed.get_rank() + if torch.distributed.is_initialized() + else int(os.getenv("RANK", "-1")) + ) + payload = { + "global_rank": global_rank, + "data_parallel_rank": parallel_state.get_data_parallel_rank(), + "classes": counts, + } + print( + "PRIMUS_LINEAR_CLASS_CENSUS=" + json.dumps(payload, sort_keys=True), + flush=True, + ) + def _restore_chimera_rng_state(args) -> None: """Restore canonical RNG state after chimera model init. @@ -367,6 +407,7 @@ def create_model(self, pre_process=True, post_process=True): # Create Flux model (backend=None lets model select based on config.transformer_impl) model = Flux(config=config, backend=backend) + _emit_precision_linear_class_census(model) if self.nemo_chimera_init: _restore_chimera_rng_state(args) diff --git a/primus/backends/megatron/patches/mlperf_warmup_patches.py b/primus/backends/megatron/patches/mlperf_warmup_patches.py index 5ac532978..f7a816b0a 100644 --- a/primus/backends/megatron/patches/mlperf_warmup_patches.py +++ b/primus/backends/megatron/patches/mlperf_warmup_patches.py @@ -20,6 +20,7 @@ """ import logging +import os import torch import torch.distributed @@ -313,20 +314,46 @@ def _hooked_train_step( 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, - ) + previous_warmup_marker = os.environ.get( + "PRIMUS_SYNTHETIC_WARMUP_ACTIVE" + ) + os.environ["PRIMUS_SYNTHETIC_WARMUP_ACTIVE"] = "1" + try: + 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, + ) + finally: + if previous_warmup_marker is None: + os.environ.pop("PRIMUS_SYNTHETIC_WARMUP_ACTIVE", None) + else: + os.environ["PRIMUS_SYNTHETIC_WARMUP_ACTIVE"] = ( + previous_warmup_marker + ) _log(f"Completed {warmup_steps} warmup steps") + reset_forward_counter = getattr( + forward_step_func, + "_primus_reset_forward_step_count", + None, + ) + if callable(reset_forward_counter): + reset_forward_counter(0) + _log("Reset diffusion forward-step RNG counter after warmup") + elif os.getenv("PRIMUS_AUDIT_BATCH_FINGERPRINTS") == "1": + raise RuntimeError( + "Batch continuity audit requires a forward-step counter reset " + "after synthetic warmup" + ) + # ---- 5. Restore optimizer ---- _restore_optimizer(optimizer, saved_opt) _reset_optimizer_state(optimizer) diff --git a/primus/backends/megatron/training/diffusion/forward_step.py b/primus/backends/megatron/training/diffusion/forward_step.py index 1a112133d..a62200b6d 100644 --- a/primus/backends/megatron/training/diffusion/forward_step.py +++ b/primus/backends/megatron/training/diffusion/forward_step.py @@ -17,7 +17,9 @@ Architecture follows functional composition for clarity and testability. """ +import json import logging +import os from typing import Optional, Tuple import torch @@ -38,6 +40,47 @@ logger = logging.getLogger(__name__) +def _emit_batch_fingerprint(batch: dict, step_count: int) -> None: + """Emit a rank-local sample-key digest for explicit continuity audits.""" + if os.getenv("PRIMUS_AUDIT_BATCH_FINGERPRINTS") != "1": + return + if os.getenv("PRIMUS_SYNTHETIC_WARMUP_ACTIVE") == "1": + return + + fingerprint = batch.get("_audit_sample_key_sha256") + sample_count = batch.get("_audit_sample_count") + if ( + not isinstance(fingerprint, str) + or len(fingerprint) != 64 + or any(character not in "0123456789abcdef" for character in fingerprint) + or not isinstance(sample_count, int) + or sample_count <= 0 + ): + raise RuntimeError( + "Batch-fingerprint audit was requested but the Energon batch has no " + "valid sample-key fingerprint" + ) + + from megatron.core import parallel_state + + global_rank = ( + torch.distributed.get_rank() + if torch.distributed.is_initialized() + else int(os.getenv("RANK", "-1")) + ) + payload = { + "global_rank": global_rank, + "data_parallel_rank": parallel_state.get_data_parallel_rank(), + "step": int(step_count), + "sample_count": sample_count, + "sample_keys_sha256": fingerprint, + } + print( + "PRIMUS_BATCH_FINGERPRINT=" + json.dumps(payload, sort_keys=True), + flush=True, + ) + + def prepare_flux_latents( latents: torch.Tensor, scheduler, @@ -385,6 +428,9 @@ def flux_forward_step_func( if not pooled_prompt_embeds.is_cuda: pooled_prompt_embeds = pooled_prompt_embeds.cuda(non_blocking=True) + if batch is not None: + _emit_batch_fingerprint(batch, step_count) + # Obtain latents based on vae_latent_mode if vae_latent_mode == "resample": # Resample mode: reconstruct latents from posterior parameters each step diff --git a/tests/unit_tests/backends/megatron/test_dataloader_checkpoint.py b/tests/unit_tests/backends/megatron/test_dataloader_checkpoint.py new file mode 100644 index 000000000..e17543a53 --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_dataloader_checkpoint.py @@ -0,0 +1,362 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""CPU-only tests for rank-local Energon dataloader checkpoint restore.""" + +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +from primus.backends.megatron.data import energon_dataset_provider as provider_module +from primus.backends.megatron.data.dataloader import ( + DATALOADER_STATE_KEY, + MegatronDataloaderWrapper, + restore_dataloader_state_from_checkpoint, +) +from primus.backends.megatron.data.energon_dataset_provider import ( + EnergonDatasetProvider, +) + + +class _StatefulLoader: + def __init__(self): + self.position = 0 + self.restored_states = [] + + def __iter__(self): + return iter([self.position]) + + def restore_state_rank(self, state): + self.restored_states.append(state) + self.position = state["position"] + + +class _StatelessLoader: + def __iter__(self): + return iter([0]) + + +class _SequentialStatefulLoader: + def __init__(self): + self.position = 0 + + def __iter__(self): + while True: + sample_id = self.position + self.position += 1 + yield sample_id + + def save_state_rank(self): + return {"position": self.position} + + def restore_state_rank(self, state): + self.position = state["position"] + + +def test_restore_dataloader_state_uses_rank_specific_checkpoint(tmp_path): + checkpoint_path = tmp_path / "train_dataloader_dprank003.pt" + checkpoint_path.touch() + calls = [] + + def checkpoint_name(root, iteration, **kwargs): + calls.append((root, iteration, kwargs)) + return str(checkpoint_path) + + loader = _StatefulLoader() + wrapper = MegatronDataloaderWrapper(loader) + + restored_path = restore_dataloader_state_from_checkpoint( + wrapper, + str(tmp_path), + iteration=17, + data_parallel_rank=3, + checkpoint_name_fn=checkpoint_name, + load_fn=lambda path: {"dataloader_state_dict": {"position": 41}}, + ) + + assert restored_path == checkpoint_path + assert loader.restored_states == [{"position": 41}] + assert next(wrapper) == 41 + assert calls == [ + ( + str(tmp_path), + 17, + { + "tensor_rank": 0, + "pipeline_rank": 0, + "basename": "train_dataloader_dprank003.pt", + }, + ) + ] + + +def test_restore_dataloader_state_fails_when_rank_file_is_missing(tmp_path): + missing_path = tmp_path / "missing.pt" + wrapper = MegatronDataloaderWrapper(_StatefulLoader()) + + with pytest.raises(FileNotFoundError, match="data_parallel_rank=2"): + restore_dataloader_state_from_checkpoint( + wrapper, + str(tmp_path), + iteration=8, + data_parallel_rank=2, + checkpoint_name_fn=lambda *args, **kwargs: str(missing_path), + ) + + +def test_restore_dataloader_state_fails_when_payload_key_is_missing(tmp_path): + checkpoint_path = tmp_path / "state.pt" + checkpoint_path.touch() + wrapper = MegatronDataloaderWrapper(_StatefulLoader()) + + with pytest.raises(KeyError, match="dataloader_state_dict"): + restore_dataloader_state_from_checkpoint( + wrapper, + str(tmp_path), + iteration=8, + data_parallel_rank=0, + checkpoint_name_fn=lambda *args, **kwargs: str(checkpoint_path), + load_fn=lambda path: {"wrong_key": {}}, + ) + + +def test_restore_dataloader_state_requires_restore_capability(tmp_path): + checkpoint_path = tmp_path / "state.pt" + checkpoint_path.touch() + wrapper = MegatronDataloaderWrapper(_StatelessLoader()) + + with pytest.raises(RuntimeError, match="does not support restore_state_rank"): + restore_dataloader_state_from_checkpoint( + wrapper, + str(tmp_path), + iteration=8, + data_parallel_rank=0, + checkpoint_name_fn=lambda *args, **kwargs: str(checkpoint_path), + load_fn=lambda path: {"dataloader_state_dict": {}}, + ) + + +def test_restore_dataloader_state_real_path_and_torch_round_trip( + tmp_path, monkeypatch +): + from megatron.training import checkpointing + + monkeypatch.setattr( + checkpointing.mpu, "get_pipeline_model_parallel_world_size", lambda: 1 + ) + monkeypatch.setattr( + checkpointing.mpu, "get_expert_model_parallel_world_size", lambda: 1 + ) + monkeypatch.setattr( + checkpointing.mpu, "get_expert_model_parallel_rank", lambda: 0 + ) + + uninterrupted_loader = _SequentialStatefulLoader() + uninterrupted = MegatronDataloaderWrapper(uninterrupted_loader) + assert [next(uninterrupted) for _ in range(5)] == list(range(5)) + state = uninterrupted.save_state() + expected_tail = [next(uninterrupted) for _ in range(5)] + + checkpoint_path = Path( + checkpointing.get_checkpoint_name( + str(tmp_path), + 5, + pipeline_parallel=False, + tensor_rank=0, + pipeline_rank=0, + expert_parallel=False, + expert_rank=0, + basename="train_dataloader_dprank000.pt", + ) + ) + checkpoint_path.parent.mkdir(parents=True) + torch.save({DATALOADER_STATE_KEY: state}, checkpoint_path) + + resumed = MegatronDataloaderWrapper(_SequentialStatefulLoader()) + restored_path = restore_dataloader_state_from_checkpoint( + resumed, + str(tmp_path), + iteration=5, + data_parallel_rank=0, + ) + + assert restored_path == checkpoint_path + assert [next(resumed) for _ in range(5)] == expected_tail == list(range(5, 10)) + + +def test_provider_restores_before_first_resumed_batch(monkeypatch, tmp_path): + provider = EnergonDatasetProvider(lambda: None) + loader = MegatronDataloaderWrapper(_SequentialStatefulLoader()) + calls = [] + + def restore(dataloader, checkpoint_root, iteration, data_parallel_rank): + calls.append((checkpoint_root, iteration, data_parallel_rank)) + dataloader.restore_state({"position": 5}, strict=True) + return tmp_path / "train_dataloader_dprank003.pt" + + monkeypatch.setattr( + provider_module.parallel_state, "get_data_parallel_rank", lambda: 3 + ) + monkeypatch.setattr( + provider_module, "restore_dataloader_state_from_checkpoint", restore + ) + args = SimpleNamespace( + load="/checkpoints", + dataloader_save="/dataloader-state", + iteration=5, + finetune=False, + require_dataloader_restore=True, + ) + + restored_path = provider._restore_train_dataloader_state(args, loader) + + assert restored_path == str(tmp_path / "train_dataloader_dprank003.pt") + assert calls == [("/dataloader-state", 5, 3)] + assert next(loader) == 5 + + +@pytest.mark.parametrize( + ("args", "error"), + [ + ( + SimpleNamespace( + load=None, + iteration=0, + finetune=False, + require_dataloader_restore=False, + ), + None, + ), + ( + SimpleNamespace( + load="/checkpoints", + iteration=5, + finetune=True, + require_dataloader_restore=False, + ), + None, + ), + ( + SimpleNamespace( + load="/checkpoints", + iteration=0, + finetune=False, + require_dataloader_restore=False, + ), + None, + ), + ( + SimpleNamespace( + load=None, + iteration=0, + finetune=False, + require_dataloader_restore=True, + ), + RuntimeError, + ), + ( + SimpleNamespace( + load="/checkpoints", + iteration=5, + finetune=True, + require_dataloader_restore=True, + ), + RuntimeError, + ), + ( + SimpleNamespace( + load="/checkpoints", + iteration=0, + finetune=False, + require_dataloader_restore=True, + ), + RuntimeError, + ), + ], +) +def test_provider_classifies_resume_state(args, error): + provider = EnergonDatasetProvider(lambda: None) + loader = MegatronDataloaderWrapper(_StatefulLoader()) + + if error is None: + assert provider._restore_train_dataloader_state(args, loader) is None + else: + with pytest.raises(error, match="successfully loaded"): + provider._restore_train_dataloader_state(args, loader) + + +@pytest.mark.parametrize("required", [False, True]) +def test_provider_handles_missing_dataloader_save(required): + provider = EnergonDatasetProvider(lambda: None) + loader = MegatronDataloaderWrapper(_StatefulLoader()) + args = SimpleNamespace( + load="/checkpoints", + iteration=5, + finetune=False, + require_dataloader_restore=required, + ) + + if required: + with pytest.raises(RuntimeError, match="without dataloader_save"): + provider._restore_train_dataloader_state(args, loader) + else: + assert provider._restore_train_dataloader_state(args, loader) is None + + +def test_provider_propagates_missing_rank_state(monkeypatch): + provider = EnergonDatasetProvider(lambda: None) + loader = MegatronDataloaderWrapper(_StatefulLoader()) + args = SimpleNamespace( + load="/checkpoints", + dataloader_save="/dataloader-state", + iteration=5, + finetune=False, + require_dataloader_restore=True, + ) + + monkeypatch.setattr( + provider_module.parallel_state, "get_data_parallel_rank", lambda: 7 + ) + + def missing(*args, **kwargs): + raise FileNotFoundError("train_dataloader_dprank007.pt") + + monkeypatch.setattr( + provider_module, "restore_dataloader_state_from_checkpoint", missing + ) + with pytest.raises(FileNotFoundError, match="dprank007"): + provider._restore_train_dataloader_state(args, loader) + + +def test_provider_rejects_non_boolean_required_flag(): + provider = EnergonDatasetProvider(lambda: None) + args = SimpleNamespace(require_dataloader_restore="true") + + with pytest.raises(TypeError, match="must be a boolean"): + provider._restore_train_dataloader_state( + args, MegatronDataloaderWrapper(_StatefulLoader()) + ) + + +@pytest.mark.parametrize( + ("iteration", "data_parallel_rank", "message"), + [ + (-1, 0, "iteration must be non-negative"), + (0, -1, "data_parallel_rank must be non-negative"), + ], +) +def test_restore_dataloader_state_rejects_invalid_coordinates( + tmp_path: Path, iteration: int, data_parallel_rank: int, message: str +): + with pytest.raises(ValueError, match=message): + restore_dataloader_state_from_checkpoint( + MegatronDataloaderWrapper(_StatefulLoader()), + str(tmp_path), + iteration=iteration, + data_parallel_rank=data_parallel_rank, + ) diff --git a/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py b/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py new file mode 100644 index 000000000..76fd2cb24 --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py @@ -0,0 +1,109 @@ +import json +from types import SimpleNamespace + +import pytest +import torch.nn as nn + +from primus.backends.megatron.data.diffusion.task_encoders.image import ( + _sample_key_fingerprint, +) +from primus.backends.megatron.flux_pretrain_trainer import ( + FluxPretrainTrainer, + _precision_linear_class_census, +) +from primus.backends.megatron.training.diffusion.forward_step import ( + _emit_batch_fingerprint, +) + + +def test_sample_key_fingerprint_is_order_sensitive(): + first = [ + SimpleNamespace(**{"__key__": "sample-a"}), + SimpleNamespace(**{"__key__": "sample-b"}), + ] + second = list(reversed(first)) + + assert _sample_key_fingerprint(first) != _sample_key_fingerprint(second) + assert _sample_key_fingerprint(first) == _sample_key_fingerprint(first) + + +def test_sample_key_fingerprint_requires_energon_key(): + with pytest.raises(RuntimeError, match="sample has no __key__"): + _sample_key_fingerprint([SimpleNamespace()]) + + +def test_precision_linear_class_census_reports_exact_classes(): + mxfp4_column = type("MXFP4ColumnParallelLinear", (nn.Module,), {})() + mxfp4_row = type("MXFP4RowParallelLinear", (nn.Module,), {})() + float8_column = type("Float8ColumnParallelLinear", (nn.Module,), {})() + model = nn.ModuleList([mxfp4_column, mxfp4_row, float8_column]) + + assert _precision_linear_class_census(model) == { + "MXFP4ColumnParallelLinear": 1, + "MXFP4RowParallelLinear": 1, + "Float8ColumnParallelLinear": 1, + "Float8RowParallelLinear": 0, + } + + +def test_emit_batch_fingerprint_is_fail_closed(monkeypatch): + monkeypatch.setenv("PRIMUS_AUDIT_BATCH_FINGERPRINTS", "1") + with pytest.raises(RuntimeError, match="no valid sample-key fingerprint"): + _emit_batch_fingerprint({}, step_count=1) + + +def test_emit_batch_fingerprint_skips_synthetic_warmup( + monkeypatch, capsys +): + monkeypatch.setenv("PRIMUS_AUDIT_BATCH_FINGERPRINTS", "1") + monkeypatch.setenv("PRIMUS_SYNTHETIC_WARMUP_ACTIVE", "1") + + _emit_batch_fingerprint({}, step_count=1) + + assert capsys.readouterr().out == "" + + +def test_diffusion_forward_step_exposes_counter_reset(monkeypatch): + from megatron.core import num_microbatches_calculator + + trainer = FluxPretrainTrainer.__new__(FluxPretrainTrainer) + trainer._forward_step_count = 2 + trainer._forward_step_count_initialized = True + monkeypatch.setattr( + num_microbatches_calculator, + "get_num_microbatches", + lambda: 1, + ) + + forward_step = trainer.get_forward_step() + forward_step._primus_reset_forward_step_count(0) + + assert trainer._forward_step_count == 0 + assert trainer._forward_step_count_initialized is True + + +def test_emit_batch_fingerprint_logs_rank_local_payload( + monkeypatch, capsys +): + from megatron.core import parallel_state + + monkeypatch.setenv("PRIMUS_AUDIT_BATCH_FINGERPRINTS", "1") + monkeypatch.setenv("RANK", "3") + monkeypatch.setattr(parallel_state, "get_data_parallel_rank", lambda: 3) + _emit_batch_fingerprint( + { + "_audit_sample_key_sha256": "a" * 64, + "_audit_sample_count": 64, + }, + step_count=6, + ) + + line = capsys.readouterr().out.strip() + payload = json.loads(line.split("=", 1)[1]) + assert payload == { + "data_parallel_rank": 3, + "global_rank": 3, + "sample_count": 64, + "sample_keys_sha256": "a" * 64, + "step": 6, + } From 93f4b1c1ddba329b5cebec279717efb4a3829f5c Mon Sep 17 00:00:00 2001 From: GP Huang <13152353+gphuang@users.noreply.github.com> Date: Tue, 11 Aug 2026 02:25:15 -0500 Subject: [PATCH 02/14] fix(megatron): satisfy resume patch lint checks Apply the repository's pinned isort and Black formatting so the resume-state PR can pass its pre-commit gate. --- primus/backends/megatron/data/dataloader.py | 8 ++--- .../megatron/data/energon_dataset_provider.py | 5 +-- primus/backends/megatron/diffusion_trainer.py | 11 ++---- .../megatron/flux_pretrain_trainer.py | 4 +-- .../megatron/patches/mlperf_warmup_patches.py | 11 ++---- .../training/diffusion/forward_step.py | 4 +-- .../megatron/test_dataloader_checkpoint.py | 36 +++++-------------- .../megatron/test_diffusion_audit_markers.py | 8 ++--- 8 files changed, 22 insertions(+), 65 deletions(-) diff --git a/primus/backends/megatron/data/dataloader.py b/primus/backends/megatron/data/dataloader.py index ac8eae8dc..2f9f81cd6 100644 --- a/primus/backends/megatron/data/dataloader.py +++ b/primus/backends/megatron/data/dataloader.py @@ -154,9 +154,7 @@ def restore_state(self, state: Any, *, strict: bool = False) -> bool: return True else: if strict: - raise RuntimeError( - f"{type(self._dataloader).__name__} does not support restore_state_rank()" - ) + raise RuntimeError(f"{type(self._dataloader).__name__} does not support restore_state_rank()") logger.debug( f"{type(self._dataloader).__name__} does not support restore_state_rank() " f"(OK for synthetic data)" @@ -222,9 +220,7 @@ def _torch_load(path: Path) -> Any: f"got {type(payload).__name__}" ) if DATALOADER_STATE_KEY not in payload: - raise KeyError( - f"Dataloader checkpoint {checkpoint_path} has no {DATALOADER_STATE_KEY!r} entry" - ) + raise KeyError(f"Dataloader checkpoint {checkpoint_path} has no {DATALOADER_STATE_KEY!r} entry") dataloader.restore_state(payload[DATALOADER_STATE_KEY], strict=True) return checkpoint_path diff --git a/primus/backends/megatron/data/energon_dataset_provider.py b/primus/backends/megatron/data/energon_dataset_provider.py index cc4e77195..c37d99653 100644 --- a/primus/backends/megatron/data/energon_dataset_provider.py +++ b/primus/backends/megatron/data/energon_dataset_provider.py @@ -190,10 +190,7 @@ def _restore_train_dataloader_state( """ require_restore = getattr(args, "require_dataloader_restore", False) if not isinstance(require_restore, bool): - raise TypeError( - "require_dataloader_restore must be a boolean, " - f"got {require_restore!r}" - ) + raise TypeError("require_dataloader_restore must be a boolean, " f"got {require_restore!r}") load_path = getattr(args, "load", None) finetune = bool(getattr(args, "finetune", False)) diff --git a/primus/backends/megatron/diffusion_trainer.py b/primus/backends/megatron/diffusion_trainer.py index 90d5eb2cf..7e521db1e 100644 --- a/primus/backends/megatron/diffusion_trainer.py +++ b/primus/backends/megatron/diffusion_trainer.py @@ -325,20 +325,15 @@ def diffusion_forward_step(data_iterator, model): return self.forward_step(data_iterator, model, return_schedule_plan=False) def reset_forward_step_count(iteration=0): - from megatron.core.num_microbatches_calculator import ( - get_num_microbatches, - ) + from megatron.core.num_microbatches_calculator import get_num_microbatches self._forward_step_count = int(iteration) * get_num_microbatches() self._forward_step_count_initialized = True log_rank_0( - "[DiffusionPretrainTrainer] Reset forward-step RNG counter " - f"to {self._forward_step_count}" + "[DiffusionPretrainTrainer] Reset forward-step RNG counter " f"to {self._forward_step_count}" ) - diffusion_forward_step._primus_reset_forward_step_count = ( - reset_forward_step_count - ) + diffusion_forward_step._primus_reset_forward_step_count = reset_forward_step_count return diffusion_forward_step def get_datasets_provider(self): diff --git a/primus/backends/megatron/flux_pretrain_trainer.py b/primus/backends/megatron/flux_pretrain_trainer.py index 7d04971cd..e0dca5506 100644 --- a/primus/backends/megatron/flux_pretrain_trainer.py +++ b/primus/backends/megatron/flux_pretrain_trainer.py @@ -54,9 +54,7 @@ def _emit_precision_linear_class_census(model) -> None: if sum(counts.values()) <= 0: raise RuntimeError("No MXFP4 or Float8 linear modules were instantiated") global_rank = ( - torch.distributed.get_rank() - if torch.distributed.is_initialized() - else int(os.getenv("RANK", "-1")) + torch.distributed.get_rank() if torch.distributed.is_initialized() else int(os.getenv("RANK", "-1")) ) payload = { "global_rank": global_rank, diff --git a/primus/backends/megatron/patches/mlperf_warmup_patches.py b/primus/backends/megatron/patches/mlperf_warmup_patches.py index f7a816b0a..2ae5f878d 100644 --- a/primus/backends/megatron/patches/mlperf_warmup_patches.py +++ b/primus/backends/megatron/patches/mlperf_warmup_patches.py @@ -314,9 +314,7 @@ def _hooked_train_step( saved_lr_num_steps = opt_param_scheduler.num_steps # ---- 4. Run warmup steps with synthetic data ---- - previous_warmup_marker = os.environ.get( - "PRIMUS_SYNTHETIC_WARMUP_ACTIVE" - ) + previous_warmup_marker = os.environ.get("PRIMUS_SYNTHETIC_WARMUP_ACTIVE") os.environ["PRIMUS_SYNTHETIC_WARMUP_ACTIVE"] = "1" try: for step_idx in range(warmup_steps): @@ -335,9 +333,7 @@ def _hooked_train_step( if previous_warmup_marker is None: os.environ.pop("PRIMUS_SYNTHETIC_WARMUP_ACTIVE", None) else: - os.environ["PRIMUS_SYNTHETIC_WARMUP_ACTIVE"] = ( - previous_warmup_marker - ) + os.environ["PRIMUS_SYNTHETIC_WARMUP_ACTIVE"] = previous_warmup_marker _log(f"Completed {warmup_steps} warmup steps") reset_forward_counter = getattr( @@ -350,8 +346,7 @@ def _hooked_train_step( _log("Reset diffusion forward-step RNG counter after warmup") elif os.getenv("PRIMUS_AUDIT_BATCH_FINGERPRINTS") == "1": raise RuntimeError( - "Batch continuity audit requires a forward-step counter reset " - "after synthetic warmup" + "Batch continuity audit requires a forward-step counter reset " "after synthetic warmup" ) # ---- 5. Restore optimizer ---- diff --git a/primus/backends/megatron/training/diffusion/forward_step.py b/primus/backends/megatron/training/diffusion/forward_step.py index a62200b6d..38128a0d2 100644 --- a/primus/backends/megatron/training/diffusion/forward_step.py +++ b/primus/backends/megatron/training/diffusion/forward_step.py @@ -64,9 +64,7 @@ def _emit_batch_fingerprint(batch: dict, step_count: int) -> None: from megatron.core import parallel_state global_rank = ( - torch.distributed.get_rank() - if torch.distributed.is_initialized() - else int(os.getenv("RANK", "-1")) + torch.distributed.get_rank() if torch.distributed.is_initialized() else int(os.getenv("RANK", "-1")) ) payload = { "global_rank": global_rank, diff --git a/tests/unit_tests/backends/megatron/test_dataloader_checkpoint.py b/tests/unit_tests/backends/megatron/test_dataloader_checkpoint.py index e17543a53..6b2da0a29 100644 --- a/tests/unit_tests/backends/megatron/test_dataloader_checkpoint.py +++ b/tests/unit_tests/backends/megatron/test_dataloader_checkpoint.py @@ -141,20 +141,12 @@ def test_restore_dataloader_state_requires_restore_capability(tmp_path): ) -def test_restore_dataloader_state_real_path_and_torch_round_trip( - tmp_path, monkeypatch -): +def test_restore_dataloader_state_real_path_and_torch_round_trip(tmp_path, monkeypatch): from megatron.training import checkpointing - monkeypatch.setattr( - checkpointing.mpu, "get_pipeline_model_parallel_world_size", lambda: 1 - ) - monkeypatch.setattr( - checkpointing.mpu, "get_expert_model_parallel_world_size", lambda: 1 - ) - monkeypatch.setattr( - checkpointing.mpu, "get_expert_model_parallel_rank", lambda: 0 - ) + monkeypatch.setattr(checkpointing.mpu, "get_pipeline_model_parallel_world_size", lambda: 1) + monkeypatch.setattr(checkpointing.mpu, "get_expert_model_parallel_world_size", lambda: 1) + monkeypatch.setattr(checkpointing.mpu, "get_expert_model_parallel_rank", lambda: 0) uninterrupted_loader = _SequentialStatefulLoader() uninterrupted = MegatronDataloaderWrapper(uninterrupted_loader) @@ -199,12 +191,8 @@ def restore(dataloader, checkpoint_root, iteration, data_parallel_rank): dataloader.restore_state({"position": 5}, strict=True) return tmp_path / "train_dataloader_dprank003.pt" - monkeypatch.setattr( - provider_module.parallel_state, "get_data_parallel_rank", lambda: 3 - ) - monkeypatch.setattr( - provider_module, "restore_dataloader_state_from_checkpoint", restore - ) + monkeypatch.setattr(provider_module.parallel_state, "get_data_parallel_rank", lambda: 3) + monkeypatch.setattr(provider_module, "restore_dataloader_state_from_checkpoint", restore) args = SimpleNamespace( load="/checkpoints", dataloader_save="/dataloader-state", @@ -319,16 +307,12 @@ def test_provider_propagates_missing_rank_state(monkeypatch): require_dataloader_restore=True, ) - monkeypatch.setattr( - provider_module.parallel_state, "get_data_parallel_rank", lambda: 7 - ) + monkeypatch.setattr(provider_module.parallel_state, "get_data_parallel_rank", lambda: 7) def missing(*args, **kwargs): raise FileNotFoundError("train_dataloader_dprank007.pt") - monkeypatch.setattr( - provider_module, "restore_dataloader_state_from_checkpoint", missing - ) + monkeypatch.setattr(provider_module, "restore_dataloader_state_from_checkpoint", missing) with pytest.raises(FileNotFoundError, match="dprank007"): provider._restore_train_dataloader_state(args, loader) @@ -338,9 +322,7 @@ def test_provider_rejects_non_boolean_required_flag(): args = SimpleNamespace(require_dataloader_restore="true") with pytest.raises(TypeError, match="must be a boolean"): - provider._restore_train_dataloader_state( - args, MegatronDataloaderWrapper(_StatefulLoader()) - ) + provider._restore_train_dataloader_state(args, MegatronDataloaderWrapper(_StatefulLoader())) @pytest.mark.parametrize( diff --git a/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py b/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py index 76fd2cb24..2e431dd74 100644 --- a/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py +++ b/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py @@ -52,9 +52,7 @@ def test_emit_batch_fingerprint_is_fail_closed(monkeypatch): _emit_batch_fingerprint({}, step_count=1) -def test_emit_batch_fingerprint_skips_synthetic_warmup( - monkeypatch, capsys -): +def test_emit_batch_fingerprint_skips_synthetic_warmup(monkeypatch, capsys): monkeypatch.setenv("PRIMUS_AUDIT_BATCH_FINGERPRINTS", "1") monkeypatch.setenv("PRIMUS_SYNTHETIC_WARMUP_ACTIVE", "1") @@ -82,9 +80,7 @@ def test_diffusion_forward_step_exposes_counter_reset(monkeypatch): assert trainer._forward_step_count_initialized is True -def test_emit_batch_fingerprint_logs_rank_local_payload( - monkeypatch, capsys -): +def test_emit_batch_fingerprint_logs_rank_local_payload(monkeypatch, capsys): from megatron.core import parallel_state monkeypatch.setenv("PRIMUS_AUDIT_BATCH_FINGERPRINTS", "1") From 298eaad9102c027e9f4d4a010e99c8e580d01120 Mon Sep 17 00:00:00 2001 From: GP Huang <13152353+gphuang@users.noreply.github.com> Date: Tue, 11 Aug 2026 03:04:25 -0500 Subject: [PATCH 03/14] fix(megatron): harden diffusion resume review gaps Reject empty or topology-incompatible Energon state, avoid mutating resumed runs during synthetic warmup, and keep validation batches out of continuity fingerprints. --- primus/backends/megatron/data/dataloader.py | 59 ++++++++++++++- .../megatron/data/energon_dataset_provider.py | 13 +++- .../megatron/patches/mlperf_warmup_patches.py | 27 ++++++- .../training/diffusion/forward_step.py | 14 +++- .../megatron/test_dataloader_checkpoint.py | 75 +++++++++++++++++++ .../megatron/test_diffusion_audit_markers.py | 22 ++++++ 6 files changed, 200 insertions(+), 10 deletions(-) diff --git a/primus/backends/megatron/data/dataloader.py b/primus/backends/megatron/data/dataloader.py index 2f9f81cd6..cb6c043e1 100644 --- a/primus/backends/megatron/data/dataloader.py +++ b/primus/backends/megatron/data/dataloader.py @@ -29,6 +29,8 @@ logger = logging.getLogger(__name__) DATALOADER_STATE_KEY = "dataloader_state_dict" +DATALOADER_STATE_FORMAT_VERSION = 1 +DATALOADER_STATE_PAYLOAD_KEY = "state" def cyclic_iter(iterator: Iterator) -> Iterator: @@ -94,14 +96,32 @@ class MegatronDataloaderWrapper: Megatron-LM examples/multimodal/dataloader_provider.py:EnergonDataloader """ - def __init__(self, dataloader): + def __init__( + self, + dataloader, + *, + data_parallel_rank: Optional[int] = None, + data_parallel_world_size: Optional[int] = None, + ): """ Initialize wrapper for any iterable. Args: dataloader: Any iterable (PyTorch DataLoader, Energon loader, etc.) """ + if (data_parallel_rank is None) != (data_parallel_world_size is None): + raise ValueError("data_parallel_rank and data_parallel_world_size must be provided together") + if data_parallel_rank is not None: + if data_parallel_rank < 0: + raise ValueError("data_parallel_rank must be non-negative") + if data_parallel_world_size <= 0: + raise ValueError("data_parallel_world_size must be positive") + if data_parallel_rank >= data_parallel_world_size: + raise ValueError("data_parallel_rank must be smaller than data_parallel_world_size") + self._dataloader = dataloader + self._data_parallel_rank = data_parallel_rank + self._data_parallel_world_size = data_parallel_world_size self._iter = iter(cyclic_iter(dataloader)) logger.debug(f"Initialized MegatronDataloaderWrapper for {type(dataloader).__name__}") @@ -124,7 +144,17 @@ def save_state(self) -> Any: Dataloader state dictionary, or None if not supported """ if hasattr(self._dataloader, "save_state_rank"): - return self._dataloader.save_state_rank() + state = self._dataloader.save_state_rank() + if self._data_parallel_rank is None: + return state + if state is None: + raise RuntimeError("Stateful dataloader returned no state for exact checkpoint continuation") + return { + "format_version": DATALOADER_STATE_FORMAT_VERSION, + "data_parallel_rank": self._data_parallel_rank, + "data_parallel_world_size": self._data_parallel_world_size, + DATALOADER_STATE_PAYLOAD_KEY: state, + } else: logger.debug( f"{type(self._dataloader).__name__} does not support save_state_rank() " @@ -146,6 +176,29 @@ def restore_state(self, state: Any, *, strict: bool = False) -> bool: Returns: True when state was restored, otherwise False. """ + if state is None and strict: + raise RuntimeError("Dataloader checkpoint contains an empty state") + + if self._data_parallel_rank is not None: + if not isinstance(state, dict): + raise RuntimeError("Dataloader checkpoint has no topology-aware state envelope") + expected_metadata = { + "format_version": DATALOADER_STATE_FORMAT_VERSION, + "data_parallel_rank": self._data_parallel_rank, + "data_parallel_world_size": self._data_parallel_world_size, + } + observed_metadata = {key: state.get(key) for key in expected_metadata} + if observed_metadata != expected_metadata: + raise RuntimeError( + "Dataloader checkpoint topology does not match the current run: " + f"saved={observed_metadata}, expected={expected_metadata}" + ) + if DATALOADER_STATE_PAYLOAD_KEY not in state: + raise RuntimeError("Dataloader checkpoint state envelope has no payload") + state = state[DATALOADER_STATE_PAYLOAD_KEY] + if state is None: + raise RuntimeError("Dataloader checkpoint state payload is empty") + if hasattr(self._dataloader, "restore_state_rank"): self._dataloader.restore_state_rank(state) # Recreate iterator after restore @@ -232,7 +285,9 @@ def _torch_load(path: Path) -> Any: __all__ = [ + "DATALOADER_STATE_FORMAT_VERSION", "DATALOADER_STATE_KEY", + "DATALOADER_STATE_PAYLOAD_KEY", "MegatronDataloaderWrapper", "cyclic_iter", "EnergonDataloader", # Deprecated, use MegatronDataloaderWrapper diff --git a/primus/backends/megatron/data/energon_dataset_provider.py b/primus/backends/megatron/data/energon_dataset_provider.py index c37d99653..68bbe454e 100644 --- a/primus/backends/megatron/data/energon_dataset_provider.py +++ b/primus/backends/megatron/data/energon_dataset_provider.py @@ -121,7 +121,11 @@ def create_dataloaders( train_dataloader = get_savable_loader( train_dataset, worker_config=worker_config, prefetch_factor=prefetch_factor ) - train_dataloader = MegatronDataloaderWrapper(train_dataloader) + train_dataloader = MegatronDataloaderWrapper( + train_dataloader, + data_parallel_rank=parallel_state.get_data_parallel_rank(), + data_parallel_world_size=parallel_state.get_data_parallel_world_size(), + ) self._restore_train_dataloader_state(args, train_dataloader) log_rank_0("Created training dataloader") @@ -218,6 +222,13 @@ def _restore_train_dataloader_state( ) return None + if not require_restore: + log_rank_0( + "WARNING: Resuming without require_dataloader_restore=True; " + "the Energon dataloader position will not be restored" + ) + return None + dataloader_save = getattr(args, "dataloader_save", None) if not dataloader_save: message = ( diff --git a/primus/backends/megatron/patches/mlperf_warmup_patches.py b/primus/backends/megatron/patches/mlperf_warmup_patches.py index 2ae5f878d..e9fbcb2c4 100644 --- a/primus/backends/megatron/patches/mlperf_warmup_patches.py +++ b/primus/backends/megatron/patches/mlperf_warmup_patches.py @@ -40,6 +40,12 @@ def _warmup_enabled(ctx: PatchContext) -> bool: return args is not None and getattr(args, "warmup_train_steps", 0) > 0 +def _is_resumed_training(args) -> bool: + """Return whether Megatron has restored a positive training iteration.""" + iteration = getattr(args, "iteration", 0) + return isinstance(iteration, int) and not isinstance(iteration, bool) and iteration > 0 + + def _reset_fp8_te_spec(models): """Reset FP8 state for TransformerEngine spec modules. @@ -284,11 +290,26 @@ def _hooked_train_step( iteration=iteration, ) - _lazy_init() - from megatron.training import get_args as megatron_get_args megatron_args = megatron_get_args() + if _is_resumed_training(megatron_args): + _warmup_done[0] = True + mt.train_step = _wrapped_chain + _log("Skipping synthetic warmup for resumed training at " f"iteration={megatron_args.iteration}") + return _wrapped_chain( + forward_step_func, + data_iterator, + model, + optimizer, + opt_param_scheduler, + config, + forward_backward_func, + iteration=iteration, + ) + + _lazy_init() + models = model if isinstance(model, (list, tuple)) else [model] synthetic_iter = _lazy_state["synthetic_iter"] @@ -342,7 +363,7 @@ def _hooked_train_step( None, ) if callable(reset_forward_counter): - reset_forward_counter(0) + reset_forward_counter(megatron_args.iteration) _log("Reset diffusion forward-step RNG counter after warmup") elif os.getenv("PRIMUS_AUDIT_BATCH_FINGERPRINTS") == "1": raise RuntimeError( diff --git a/primus/backends/megatron/training/diffusion/forward_step.py b/primus/backends/megatron/training/diffusion/forward_step.py index d6d49c56b..ec70cabc5 100644 --- a/primus/backends/megatron/training/diffusion/forward_step.py +++ b/primus/backends/megatron/training/diffusion/forward_step.py @@ -40,10 +40,12 @@ logger = logging.getLogger(__name__) -def _emit_batch_fingerprint(batch: dict, step_count: int) -> None: +def _emit_batch_fingerprint(batch: dict, step_count: int, *, is_training: bool = True) -> None: """Emit a rank-local sample-key digest for explicit continuity audits.""" if os.getenv("PRIMUS_AUDIT_BATCH_FINGERPRINTS") != "1": return + if not is_training: + return if os.getenv("PRIMUS_SYNTHETIC_WARMUP_ACTIVE") == "1": return @@ -426,9 +428,6 @@ def flux_forward_step_func( if not pooled_prompt_embeds.is_cuda: pooled_prompt_embeds = pooled_prompt_embeds.cuda(non_blocking=True) - if batch is not None: - _emit_batch_fingerprint(batch, step_count) - # Obtain latents based on vae_latent_mode if vae_latent_mode == "resample": # Resample mode: reconstruct latents from posterior parameters each step @@ -500,6 +499,13 @@ def flux_forward_step_func( val_timesteps = val_idx.to(dtype=compute_dtype) / 8.0 batch["timesteps"] = val_timesteps + if batch is not None: + _emit_batch_fingerprint( + batch, + step_count, + is_training=not is_validation, + ) + # Matches NeMo's forward_step which wraps prepare_image_latent_like_reference # in torch.no_grad() — no gradients needed for position IDs, noise sampling, # timestep sampling, or latent packing. diff --git a/tests/unit_tests/backends/megatron/test_dataloader_checkpoint.py b/tests/unit_tests/backends/megatron/test_dataloader_checkpoint.py index 6b2da0a29..5656f35d9 100644 --- a/tests/unit_tests/backends/megatron/test_dataloader_checkpoint.py +++ b/tests/unit_tests/backends/megatron/test_dataloader_checkpoint.py @@ -14,7 +14,9 @@ from primus.backends.megatron.data import energon_dataset_provider as provider_module from primus.backends.megatron.data.dataloader import ( + DATALOADER_STATE_FORMAT_VERSION, DATALOADER_STATE_KEY, + DATALOADER_STATE_PAYLOAD_KEY, MegatronDataloaderWrapper, restore_dataloader_state_from_checkpoint, ) @@ -125,6 +127,53 @@ def test_restore_dataloader_state_fails_when_payload_key_is_missing(tmp_path): ) +def test_restore_dataloader_state_rejects_empty_state(tmp_path): + checkpoint_path = tmp_path / "state.pt" + checkpoint_path.touch() + wrapper = MegatronDataloaderWrapper(_StatefulLoader()) + + with pytest.raises(RuntimeError, match="empty state"): + restore_dataloader_state_from_checkpoint( + wrapper, + str(tmp_path), + iteration=8, + data_parallel_rank=0, + checkpoint_name_fn=lambda *args, **kwargs: str(checkpoint_path), + load_fn=lambda path: {DATALOADER_STATE_KEY: None}, + ) + + +def test_restore_dataloader_state_rejects_changed_data_parallel_topology(): + source = MegatronDataloaderWrapper( + _SequentialStatefulLoader(), + data_parallel_rank=3, + data_parallel_world_size=8, + ) + state = source.save_state() + + assert state == { + "format_version": DATALOADER_STATE_FORMAT_VERSION, + "data_parallel_rank": 3, + "data_parallel_world_size": 8, + DATALOADER_STATE_PAYLOAD_KEY: {"position": 0}, + } + + matching = MegatronDataloaderWrapper( + _SequentialStatefulLoader(), + data_parallel_rank=3, + data_parallel_world_size=8, + ) + assert matching.restore_state(state, strict=True) + + changed_world_size = MegatronDataloaderWrapper( + _SequentialStatefulLoader(), + data_parallel_rank=3, + data_parallel_world_size=4, + ) + with pytest.raises(RuntimeError, match="topology does not match"): + changed_world_size.restore_state(state, strict=True) + + def test_restore_dataloader_state_requires_restore_capability(tmp_path): checkpoint_path = tmp_path / "state.pt" checkpoint_path.touch() @@ -208,6 +257,32 @@ def restore(dataloader, checkpoint_root, iteration, data_parallel_rank): assert next(loader) == 5 +def test_provider_skips_restore_when_exact_continuation_is_not_required( + monkeypatch, +): + provider = EnergonDatasetProvider(lambda: None) + loader = MegatronDataloaderWrapper(_SequentialStatefulLoader()) + args = SimpleNamespace( + load="/checkpoints", + dataloader_save="/dataloader-state", + iteration=5, + finetune=False, + require_dataloader_restore=False, + ) + + def unexpected_restore(*args, **kwargs): + raise AssertionError("optional resume must not partially restore rank state") + + monkeypatch.setattr( + provider_module, + "restore_dataloader_state_from_checkpoint", + unexpected_restore, + ) + + assert provider._restore_train_dataloader_state(args, loader) is None + assert next(loader) == 0 + + @pytest.mark.parametrize( ("args", "error"), [ diff --git a/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py b/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py index 2e431dd74..792b25c56 100644 --- a/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py +++ b/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py @@ -11,6 +11,7 @@ FluxPretrainTrainer, _precision_linear_class_census, ) +from primus.backends.megatron.patches.mlperf_warmup_patches import _is_resumed_training from primus.backends.megatron.training.diffusion.forward_step import ( _emit_batch_fingerprint, ) @@ -61,6 +62,27 @@ def test_emit_batch_fingerprint_skips_synthetic_warmup(monkeypatch, capsys): assert capsys.readouterr().out == "" +def test_emit_batch_fingerprint_skips_validation(monkeypatch, capsys): + monkeypatch.setenv("PRIMUS_AUDIT_BATCH_FINGERPRINTS", "1") + + _emit_batch_fingerprint({}, step_count=5, is_training=False) + + assert capsys.readouterr().out == "" + + +@pytest.mark.parametrize( + ("iteration", "expected"), + [ + (0, False), + (5, True), + (None, False), + (True, False), + ], +) +def test_warmup_resume_detection(iteration, expected): + assert _is_resumed_training(SimpleNamespace(iteration=iteration)) is expected + + def test_diffusion_forward_step_exposes_counter_reset(monkeypatch): from megatron.core import num_microbatches_calculator From 7a7e4f07b9d51f2c722a6585b2b8548d2f7673c1 Mon Sep 17 00:00:00 2001 From: Guangpu Huang Date: Tue, 11 Aug 2026 06:30:17 -0500 Subject: [PATCH 04/14] fix(megatron): emit diffusion audit markers through logging The batch-fingerprint and linear-class-census markers were written with a bare print, but the training process reaches the run log through a path that keeps only fd 2, so every marker was discarded. The audit then reported zero markers on a healthy run and no continuity check could ever pass. Route both through the module logger, which the same run demonstrably preserves. --- primus/backends/megatron/flux_pretrain_trainer.py | 11 +++++++---- .../megatron/training/diffusion/forward_step.py | 8 ++++---- 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/primus/backends/megatron/flux_pretrain_trainer.py b/primus/backends/megatron/flux_pretrain_trainer.py index e0dca5506..476659b62 100644 --- a/primus/backends/megatron/flux_pretrain_trainer.py +++ b/primus/backends/megatron/flux_pretrain_trainer.py @@ -14,6 +14,7 @@ """ import json +import logging import os from collections import Counter @@ -30,6 +31,8 @@ ) from primus.core.utils.module_utils import log_rank_0 +logger = logging.getLogger(__name__) + PRECISION_LINEAR_CLASSES = ( "MXFP4ColumnParallelLinear", "MXFP4RowParallelLinear", @@ -61,10 +64,10 @@ def _emit_precision_linear_class_census(model) -> None: "data_parallel_rank": parallel_state.get_data_parallel_rank(), "classes": counts, } - print( - "PRIMUS_LINEAR_CLASS_CENSUS=" + json.dumps(payload, sort_keys=True), - flush=True, - ) + # Emit through logging, not print: the launcher pipes the training process + # through a filter that drops fd 1, so a bare print never reaches the run + # log and the audit silently reports zero markers. + logger.info("PRIMUS_LINEAR_CLASS_CENSUS=%s", json.dumps(payload, sort_keys=True)) def _restore_chimera_rng_state(args) -> None: diff --git a/primus/backends/megatron/training/diffusion/forward_step.py b/primus/backends/megatron/training/diffusion/forward_step.py index ec70cabc5..4dab50c32 100644 --- a/primus/backends/megatron/training/diffusion/forward_step.py +++ b/primus/backends/megatron/training/diffusion/forward_step.py @@ -75,10 +75,10 @@ def _emit_batch_fingerprint(batch: dict, step_count: int, *, is_training: bool = "sample_count": sample_count, "sample_keys_sha256": fingerprint, } - print( - "PRIMUS_BATCH_FINGERPRINT=" + json.dumps(payload, sort_keys=True), - flush=True, - ) + # Emit through logging, not print: the launcher pipes the training process + # through a filter that drops fd 1, so a bare print never reaches the run + # log and the audit silently reports zero markers. + logger.info("PRIMUS_BATCH_FINGERPRINT=%s", json.dumps(payload, sort_keys=True)) def prepare_flux_latents( From 7e4810a5425b60939b2a72c7da289fa414a53273 Mon Sep 17 00:00:00 2001 From: Guangpu Huang Date: Tue, 11 Aug 2026 06:52:19 -0500 Subject: [PATCH 05/14] docs(megatron): state only what the probe established about fd 1 The comment asserted that the launcher filters fd 1, which the probe never showed. It showed that fd 1 does not reach the run log on this path while logging and fd 2 do, which is all the fix relies on. --- primus/backends/megatron/flux_pretrain_trainer.py | 6 +++--- primus/backends/megatron/training/diffusion/forward_step.py | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/primus/backends/megatron/flux_pretrain_trainer.py b/primus/backends/megatron/flux_pretrain_trainer.py index 476659b62..ae86a7d64 100644 --- a/primus/backends/megatron/flux_pretrain_trainer.py +++ b/primus/backends/megatron/flux_pretrain_trainer.py @@ -64,9 +64,9 @@ def _emit_precision_linear_class_census(model) -> None: "data_parallel_rank": parallel_state.get_data_parallel_rank(), "classes": counts, } - # Emit through logging, not print: the launcher pipes the training process - # through a filter that drops fd 1, so a bare print never reaches the run - # log and the audit silently reports zero markers. + # Emit through logging, not print: fd 1 does not survive to the run log on + # this launch path, so a bare print leaves the audit reporting zero markers + # on a healthy run. Logging and fd 2 both survive. logger.info("PRIMUS_LINEAR_CLASS_CENSUS=%s", json.dumps(payload, sort_keys=True)) diff --git a/primus/backends/megatron/training/diffusion/forward_step.py b/primus/backends/megatron/training/diffusion/forward_step.py index 4dab50c32..ac9aa38a7 100644 --- a/primus/backends/megatron/training/diffusion/forward_step.py +++ b/primus/backends/megatron/training/diffusion/forward_step.py @@ -75,9 +75,9 @@ def _emit_batch_fingerprint(batch: dict, step_count: int, *, is_training: bool = "sample_count": sample_count, "sample_keys_sha256": fingerprint, } - # Emit through logging, not print: the launcher pipes the training process - # through a filter that drops fd 1, so a bare print never reaches the run - # log and the audit silently reports zero markers. + # Emit through logging, not print: fd 1 does not survive to the run log on + # this launch path, so a bare print leaves the audit reporting zero markers + # on a healthy run. Logging and fd 2 both survive. logger.info("PRIMUS_BATCH_FINGERPRINT=%s", json.dumps(payload, sort_keys=True)) From f9c5ddbe1a3122984a269f92affe8f236f6132c9 Mon Sep 17 00:00:00 2001 From: Guangpu Huang Date: Tue, 11 Aug 2026 08:54:24 -0500 Subject: [PATCH 06/14] fix(megatron): remove unreachable dataloader restore fallback The optional warning path in the missing dataloader_save branch was dead code because this block is reachable only when require_dataloader_restore is true. Raise unconditionally to keep exact-resume behavior explicit and remove the unreachable branch flagged by code-quality review. --- primus/backends/megatron/data/energon_dataset_provider.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/primus/backends/megatron/data/energon_dataset_provider.py b/primus/backends/megatron/data/energon_dataset_provider.py index 68bbe454e..b206b54f2 100644 --- a/primus/backends/megatron/data/energon_dataset_provider.py +++ b/primus/backends/megatron/data/energon_dataset_provider.py @@ -235,10 +235,7 @@ def _restore_train_dataloader_state( "Resuming an Energon run without dataloader_save; " "the dataloader position cannot be restored" ) - if require_restore: - raise RuntimeError(message) - log_rank_0(f"WARNING: {message}") - return None + raise RuntimeError(message) data_parallel_rank = parallel_state.get_data_parallel_rank() checkpoint_path = restore_dataloader_state_from_checkpoint( From c7c692a4a93fb868c7e5fbf472f0f12f2e2ba4bc Mon Sep 17 00:00:00 2001 From: GP Huang Date: Tue, 11 Aug 2026 16:55:43 +0300 Subject: [PATCH 07/14] fix(megatron): checkpoint heterogeneous Flux layers independently (#971) ## Summary Stacked on PR #970 for [Issue \#220](https://github.com/AMD-AGI/tiger-training-internal/issues/220). Flux combines 19 joint and 38 single transformer blocks whose parameter schemas differ. The homogeneous distributed-checkpoint path collapses them into one layer-stacked namespace, so checkpoint validation sees both 6-chunk and 3-chunk adaLN tensors under the same key and aborts before writing a checkpoint. - Require Megatron's supported non-homogeneous, per-layer checkpoint namespace for every `FluxConfig`. - Fail closed if a caller tries to select the invalid homogeneous layout. - Add regression coverage for distinct adaLN keys and a real `torch_dist` save/load round trip. ## Test plan - [x] `black --check` on all changed files. - [x] Config regressions: default enables heterogeneous checkpointing; explicit disable is rejected (`2 passed` in the pinned v26.5 image). - [ ] Tiny joint+single Flux `torch_dist` save/load round trip (GPU window pending; n15-09 currently has a foreign eight-GPU tenant). - [ ] Issue 220 Option 6 `stage1` checkpoint save, then resumed `control` continuity smoke after review. --- .../core/models/diffusion/flux/config.py | 12 +++ .../core/models/diffusion/flux/model.py | 8 +- .../megatron/diffusion/test_flux_config.py | 17 +++- .../diffusion/test_flux_dist_checkpoint.py | 87 +++++++++++++++++++ 4 files changed, 116 insertions(+), 8 deletions(-) create mode 100644 tests/unit_tests/backends/megatron/diffusion/test_flux_dist_checkpoint.py diff --git a/primus/backends/megatron/core/models/diffusion/flux/config.py b/primus/backends/megatron/core/models/diffusion/flux/config.py index 83446615b..b3c5e70fa 100644 --- a/primus/backends/megatron/core/models/diffusion/flux/config.py +++ b/primus/backends/megatron/core/models/diffusion/flux/config.py @@ -78,6 +78,8 @@ class FluxConfig(BaseDiffusionConfig): apply_rope_fusion: Whether to apply RoPE fusion optimization (default: False) add_qkv_bias: Whether to add bias to QKV projections (default: True) single_block_bias: Whether to add bias to single block linear layers (default: True) + hetereogenous_dist_checkpoint: Keep joint and single blocks in distinct + per-layer distributed-checkpoint namespaces (required, always True) activation_func: Activation function (default: openai_gelu_no_jit) use_te_rng_tracker: Whether to use Transformer Engine RNG tracker (default: False) """ @@ -91,6 +93,10 @@ class FluxConfig(BaseDiffusionConfig): # Architecture: Number of layers num_joint_layers: int = 19 # Default: Flux 12B num_single_layers: int = 38 # Default: Flux 12B + # Megatron preserves this misspelling in its public TransformerConfig API. + # Flux's joint and single blocks have different parameter sets and shapes, + # so they cannot share the homogeneous layer-stacked checkpoint namespace. + hetereogenous_dist_checkpoint: bool = True # Architecture: Dimensions (Flux standard: 3072) hidden_size: int = 3072 @@ -214,6 +220,12 @@ def __post_init__(self): super().__post_init__() # BaseDiffusionConfig (mapping) -> TransformerConfig (validation) + if not self.hetereogenous_dist_checkpoint: + raise ValueError( + "Flux requires hetereogenous_dist_checkpoint=True because its joint " + "and single transformer blocks have different checkpoint schemas" + ) + # Xavier uniform for Megatron parallel linear layers (common Flux reference default). self.init_method = nn.init.xavier_uniform_ self.output_layer_init_method = nn.init.xavier_uniform_ diff --git a/primus/backends/megatron/core/models/diffusion/flux/model.py b/primus/backends/megatron/core/models/diffusion/flux/model.py index b451045fa..800462c68 100644 --- a/primus/backends/megatron/core/models/diffusion/flux/model.py +++ b/primus/backends/megatron/core/models/diffusion/flux/model.py @@ -604,8 +604,8 @@ def sharded_state_dict( indexed sharded keys (``transformer.layers..*``) are required here: the homogeneous layer-stacked path would leave unclaimed slots for the params absent in single blocks and raise a CheckpointingException at - save time. TransformerBlock.sharded_state_dict provides this - automatically, so no config toggle is needed. + save time. ``FluxConfig.hetereogenous_dist_checkpoint`` enables + TransformerBlock's supported per-layer checkpoint namespace. Args: prefix: Prefix for state dict keys (e.g., 'module.') @@ -617,8 +617,8 @@ def sharded_state_dict( """ sharded_state_dict = {} - # Delegate transformer layers to TransformerBlock - # Handles heterogeneous layers automatically + # Delegate transformer layers to TransformerBlock. FluxConfig requires + # its heterogeneous per-layer checkpoint namespace. sharded_state_dict.update( self.transformer.sharded_state_dict(f"{prefix}transformer.", sharded_offsets, metadata) ) diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_config.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_config.py index 2a4348832..f286490a4 100644 --- a/tests/unit_tests/backends/megatron/diffusion/test_flux_config.py +++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_config.py @@ -9,10 +9,6 @@ import pytest -from tests.utils import skip_if_no_cuda - -skip_if_no_cuda() - from primus.backends.megatron.core.models.diffusion.common import BaseDiffusionConfig from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig from tests.utils import PrimusUT @@ -40,6 +36,19 @@ def test_base_config_validation_invalid_channels(self): class TestFluxConfig(PrimusUT): """Tests for FluxConfig class.""" + def test_distributed_checkpoint_is_non_homogeneous(self): + """Flux must use per-layer checkpoint keys for its two block families.""" + config = FluxConfig.flux_535m() + self.assertTrue(config.hetereogenous_dist_checkpoint) + + def test_homogeneous_distributed_checkpoint_is_rejected(self): + """A homogeneous layer stack cannot represent Flux's parameter schemas.""" + with self.assertRaisesRegex( + ValueError, + "Flux requires hetereogenous_dist_checkpoint=True", + ): + FluxConfig.flux_535m(hetereogenous_dist_checkpoint=False) + def test_validation_positive_joint_layers(self): """Test validation fails for non-positive num_joint_layers.""" with self.assertRaises(ValueError) as cm: diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_dist_checkpoint.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_dist_checkpoint.py new file mode 100644 index 000000000..80d0b36e6 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_dist_checkpoint.py @@ -0,0 +1,87 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +"""Distributed-checkpoint regression tests for Flux's heterogeneous layers.""" + +import pytest +import torch +from megatron.core.dist_checkpointing import load, save + +from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig +from primus.backends.megatron.core.models.diffusion.flux.model import Flux +from tests.utils import PrimusUT + + +def _tiny_flux_config() -> FluxConfig: + """Build the smallest useful joint+single Flux model for checkpoint tests.""" + return FluxConfig( + num_joint_layers=1, + num_single_layers=1, + hidden_size=16, + num_attention_heads=2, + ffn_hidden_size=64, + in_channels=4, + context_dim=16, + vec_in_dim=16, + model_channels=16, + axes_dim=(2, 2, 4), + transformer_impl="local", + params_dtype=torch.float32, + ) + + +def _runtime_device() -> torch.device: + """Use CUDA when available, otherwise run on CPU for CI coverage.""" + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +class TestFluxDistCheckpoint(PrimusUT): + """Verify that joint and single blocks have independent checkpoint schemas.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + """Initialize single-rank model parallelism.""" + + def test_adaln_uses_distinct_per_layer_sharded_keys(self): + """The 6-chunk and 3-chunk adaLN weights must not share a global key.""" + model = Flux(_tiny_flux_config()).to(_runtime_device()) + sharded_state_dict = model.sharded_state_dict() + + double_local_key = "transformer.layers.0.adaln.adaLN_modulation.1.weight" + single_local_key = "transformer.layers.1.adaln.adaLN_modulation.1.weight" + double_weight = sharded_state_dict[double_local_key] + single_weight = sharded_state_dict[single_local_key] + + assert double_weight.key == double_local_key + assert single_weight.key == single_local_key + assert double_weight.key != single_weight.key + assert double_weight.prepend_axis_num == 0 + assert single_weight.prepend_axis_num == 0 + assert double_weight.global_shape == (6 * model.hidden_size, model.hidden_size) + assert single_weight.global_shape == (3 * model.hidden_size, model.hidden_size) + + def test_torch_dist_save_load_round_trip(self, tmp_path): + """Save and restore both adaLN shapes through the real torch_dist backend.""" + source = Flux(_tiny_flux_config()).to(_runtime_device()) + double_weight = source.transformer.layers[0].adaln.adaLN_modulation[-1].weight + single_weight = source.transformer.layers[1].adaln.adaLN_modulation[-1].weight + with torch.no_grad(): + double_weight.fill_(1.25) + single_weight.fill_(-2.5) + + checkpoint_dir = tmp_path / "flux_torch_dist" + checkpoint_dir.mkdir() + save({"model": source.sharded_state_dict()}, checkpoint_dir) + + target = Flux(_tiny_flux_config()).to(_runtime_device()) + loaded = load({"model": target.sharded_state_dict()}, checkpoint_dir) + target.load_state_dict(loaded["model"]) + + torch.testing.assert_close( + target.transformer.layers[0].adaln.adaLN_modulation[-1].weight, + double_weight, + ) + torch.testing.assert_close( + target.transformer.layers[1].adaln.adaLN_modulation[-1].weight, + single_weight, + ) From 1a272bc4666230b72bce6e29f1012a6fb8eda268 Mon Sep 17 00:00:00 2001 From: GP Huang <13152353+gphuang@users.noreply.github.com> Date: Tue, 11 Aug 2026 09:16:25 -0500 Subject: [PATCH 08/14] test(megatron): align audit-marker tests with logger output Capture batch-fingerprint markers from the forward-step logger instead of stdout so tests validate the current fail-closed logging path. Add an explicit CUDA-only Flux torch_dist round-trip test to make GPU runtime coverage first-class while keeping the existing CPU/CUDA fallback test. --- .../diffusion/test_flux_dist_checkpoint.py | 53 +++++++++++-------- .../megatron/test_diffusion_audit_markers.py | 24 ++++++--- 2 files changed, 48 insertions(+), 29 deletions(-) diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_dist_checkpoint.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_dist_checkpoint.py index 80d0b36e6..a7fe92de2 100644 --- a/tests/unit_tests/backends/megatron/diffusion/test_flux_dist_checkpoint.py +++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_dist_checkpoint.py @@ -38,6 +38,32 @@ def _runtime_device() -> torch.device: class TestFluxDistCheckpoint(PrimusUT): """Verify that joint and single blocks have independent checkpoint schemas.""" + @staticmethod + def _assert_torch_dist_round_trip(tmp_path, device: torch.device) -> None: + source = Flux(_tiny_flux_config()).to(device) + double_weight = source.transformer.layers[0].adaln.adaLN_modulation[-1].weight + single_weight = source.transformer.layers[1].adaln.adaLN_modulation[-1].weight + with torch.no_grad(): + double_weight.fill_(1.25) + single_weight.fill_(-2.5) + + checkpoint_dir = tmp_path / "flux_torch_dist" + checkpoint_dir.mkdir() + save({"model": source.sharded_state_dict()}, checkpoint_dir) + + target = Flux(_tiny_flux_config()).to(device) + loaded = load({"model": target.sharded_state_dict()}, checkpoint_dir) + target.load_state_dict(loaded["model"]) + + torch.testing.assert_close( + target.transformer.layers[0].adaln.adaLN_modulation[-1].weight, + double_weight, + ) + torch.testing.assert_close( + target.transformer.layers[1].adaln.adaLN_modulation[-1].weight, + single_weight, + ) + @pytest.fixture(autouse=True) def setup_parallel(self, init_parallel_state): """Initialize single-rank model parallelism.""" @@ -62,26 +88,9 @@ def test_adaln_uses_distinct_per_layer_sharded_keys(self): def test_torch_dist_save_load_round_trip(self, tmp_path): """Save and restore both adaLN shapes through the real torch_dist backend.""" - source = Flux(_tiny_flux_config()).to(_runtime_device()) - double_weight = source.transformer.layers[0].adaln.adaLN_modulation[-1].weight - single_weight = source.transformer.layers[1].adaln.adaLN_modulation[-1].weight - with torch.no_grad(): - double_weight.fill_(1.25) - single_weight.fill_(-2.5) + self._assert_torch_dist_round_trip(tmp_path, _runtime_device()) - checkpoint_dir = tmp_path / "flux_torch_dist" - checkpoint_dir.mkdir() - save({"model": source.sharded_state_dict()}, checkpoint_dir) - - target = Flux(_tiny_flux_config()).to(_runtime_device()) - loaded = load({"model": target.sharded_state_dict()}, checkpoint_dir) - target.load_state_dict(loaded["model"]) - - torch.testing.assert_close( - target.transformer.layers[0].adaln.adaLN_modulation[-1].weight, - double_weight, - ) - torch.testing.assert_close( - target.transformer.layers[1].adaln.adaLN_modulation[-1].weight, - single_weight, - ) + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_torch_dist_save_load_round_trip_cuda(self, tmp_path): + """Explicit GPU runtime coverage for the Flux torch_dist round trip.""" + self._assert_torch_dist_round_trip(tmp_path, torch.device("cuda")) diff --git a/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py b/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py index 792b25c56..681f15640 100644 --- a/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py +++ b/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py @@ -16,6 +16,8 @@ _emit_batch_fingerprint, ) +_FORWARD_STEP_LOGGER = "primus.backends.megatron.training.diffusion.forward_step" + def test_sample_key_fingerprint_is_order_sensitive(): first = [ @@ -53,21 +55,23 @@ def test_emit_batch_fingerprint_is_fail_closed(monkeypatch): _emit_batch_fingerprint({}, step_count=1) -def test_emit_batch_fingerprint_skips_synthetic_warmup(monkeypatch, capsys): +def test_emit_batch_fingerprint_skips_synthetic_warmup(monkeypatch, caplog): monkeypatch.setenv("PRIMUS_AUDIT_BATCH_FINGERPRINTS", "1") monkeypatch.setenv("PRIMUS_SYNTHETIC_WARMUP_ACTIVE", "1") + caplog.set_level("INFO", logger=_FORWARD_STEP_LOGGER) _emit_batch_fingerprint({}, step_count=1) - assert capsys.readouterr().out == "" + assert not any(record.message.startswith("PRIMUS_BATCH_FINGERPRINT=") for record in caplog.records) -def test_emit_batch_fingerprint_skips_validation(monkeypatch, capsys): +def test_emit_batch_fingerprint_skips_validation(monkeypatch, caplog): monkeypatch.setenv("PRIMUS_AUDIT_BATCH_FINGERPRINTS", "1") + caplog.set_level("INFO", logger=_FORWARD_STEP_LOGGER) _emit_batch_fingerprint({}, step_count=5, is_training=False) - assert capsys.readouterr().out == "" + assert not any(record.message.startswith("PRIMUS_BATCH_FINGERPRINT=") for record in caplog.records) @pytest.mark.parametrize( @@ -102,12 +106,13 @@ def test_diffusion_forward_step_exposes_counter_reset(monkeypatch): assert trainer._forward_step_count_initialized is True -def test_emit_batch_fingerprint_logs_rank_local_payload(monkeypatch, capsys): +def test_emit_batch_fingerprint_logs_rank_local_payload(monkeypatch, caplog): from megatron.core import parallel_state monkeypatch.setenv("PRIMUS_AUDIT_BATCH_FINGERPRINTS", "1") monkeypatch.setenv("RANK", "3") monkeypatch.setattr(parallel_state, "get_data_parallel_rank", lambda: 3) + caplog.set_level("INFO", logger=_FORWARD_STEP_LOGGER) _emit_batch_fingerprint( { "_audit_sample_key_sha256": "a" * 64, @@ -116,8 +121,13 @@ def test_emit_batch_fingerprint_logs_rank_local_payload(monkeypatch, capsys): step_count=6, ) - line = capsys.readouterr().out.strip() - payload = json.loads(line.split("=", 1)[1]) + marker_lines = [ + record.message + for record in caplog.records + if record.name == _FORWARD_STEP_LOGGER and record.message.startswith("PRIMUS_BATCH_FINGERPRINT=") + ] + assert len(marker_lines) == 1 + payload = json.loads(marker_lines[0].split("=", 1)[1]) assert payload == { "data_parallel_rank": 3, "global_rank": 3, From d1fe53ff47a75c47e8fa738fb1cd53b72357d67d Mon Sep 17 00:00:00 2001 From: GP Huang <13152353+gphuang@users.noreply.github.com> Date: Tue, 11 Aug 2026 09:37:08 -0500 Subject: [PATCH 09/14] fix(test): avoid pytest fixture injection in unittest case Use TemporaryDirectory-backed paths inside TestFluxDistCheckpoint methods instead of tmp_path parameters, which unittest-style test cases cannot receive from pytest. This keeps the GPU torch_dist round-trip coverage runnable in CI and local container validation. --- .../megatron/diffusion/test_flux_dist_checkpoint.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_dist_checkpoint.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_dist_checkpoint.py index a7fe92de2..1b461f286 100644 --- a/tests/unit_tests/backends/megatron/diffusion/test_flux_dist_checkpoint.py +++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_dist_checkpoint.py @@ -3,6 +3,9 @@ """Distributed-checkpoint regression tests for Flux's heterogeneous layers.""" +import tempfile +from pathlib import Path + import pytest import torch from megatron.core.dist_checkpointing import load, save @@ -86,11 +89,13 @@ def test_adaln_uses_distinct_per_layer_sharded_keys(self): assert double_weight.global_shape == (6 * model.hidden_size, model.hidden_size) assert single_weight.global_shape == (3 * model.hidden_size, model.hidden_size) - def test_torch_dist_save_load_round_trip(self, tmp_path): + def test_torch_dist_save_load_round_trip(self): """Save and restore both adaLN shapes through the real torch_dist backend.""" - self._assert_torch_dist_round_trip(tmp_path, _runtime_device()) + with tempfile.TemporaryDirectory() as tmpdir: + self._assert_torch_dist_round_trip(Path(tmpdir), _runtime_device()) @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") - def test_torch_dist_save_load_round_trip_cuda(self, tmp_path): + def test_torch_dist_save_load_round_trip_cuda(self): """Explicit GPU runtime coverage for the Flux torch_dist round trip.""" - self._assert_torch_dist_round_trip(tmp_path, torch.device("cuda")) + with tempfile.TemporaryDirectory() as tmpdir: + self._assert_torch_dist_round_trip(Path(tmpdir), torch.device("cuda")) From 8019c30dcc818f54925030ebba2a15ef9434a6b4 Mon Sep 17 00:00:00 2001 From: GP Huang Date: Wed, 12 Aug 2026 17:03:37 +0300 Subject: [PATCH 10/14] fix(megatron): allow BF16 precision census (#978) ## Summary Allow the precision-class audit to report an all-zero MXFP4/FP8 census for native BF16 models. This keeps the audit fail-closed in the experiment validator while avoiding an unconditional model-startup failure for Option 7. ## Test plan - [x] Added a unit test for the native BF16 zero-count payload - [x] Python compilation passes - [ ] Targeted unit test passes in the pinned runtime container - [ ] Option 7 BF16-resume smoke passes --------- Co-authored-by: Guangpu Huang --- .../megatron/flux_pretrain_trainer.py | 2 -- .../megatron/test_diffusion_audit_markers.py | 24 +++++++++++++++++++ 2 files changed, 24 insertions(+), 2 deletions(-) diff --git a/primus/backends/megatron/flux_pretrain_trainer.py b/primus/backends/megatron/flux_pretrain_trainer.py index ae86a7d64..38925a45e 100644 --- a/primus/backends/megatron/flux_pretrain_trainer.py +++ b/primus/backends/megatron/flux_pretrain_trainer.py @@ -54,8 +54,6 @@ def _emit_precision_linear_class_census(model) -> None: from megatron.core import parallel_state counts = _precision_linear_class_census(model) - if sum(counts.values()) <= 0: - raise RuntimeError("No MXFP4 or Float8 linear modules were instantiated") global_rank = ( torch.distributed.get_rank() if torch.distributed.is_initialized() else int(os.getenv("RANK", "-1")) ) diff --git a/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py b/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py index 681f15640..17ef0ecb2 100644 --- a/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py +++ b/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py @@ -9,6 +9,7 @@ ) from primus.backends.megatron.flux_pretrain_trainer import ( FluxPretrainTrainer, + _emit_precision_linear_class_census, _precision_linear_class_census, ) from primus.backends.megatron.patches.mlperf_warmup_patches import _is_resumed_training @@ -49,6 +50,29 @@ def test_precision_linear_class_census_reports_exact_classes(): } +def test_precision_linear_class_census_emits_zero_counts_for_bf16(monkeypatch, capsys): + from megatron.core import parallel_state + + monkeypatch.setenv("PRIMUS_AUDIT_LINEAR_CLASS_CENSUS", "1") + monkeypatch.setenv("RANK", "3") + monkeypatch.setattr(parallel_state, "get_data_parallel_rank", lambda: 3) + + _emit_precision_linear_class_census(nn.Linear(2, 2)) + + line = capsys.readouterr().out.strip() + payload = json.loads(line.split("=", 1)[1]) + assert payload == { + "data_parallel_rank": 3, + "global_rank": 3, + "classes": { + "MXFP4ColumnParallelLinear": 0, + "MXFP4RowParallelLinear": 0, + "Float8ColumnParallelLinear": 0, + "Float8RowParallelLinear": 0, + }, + } + + def test_emit_batch_fingerprint_is_fail_closed(monkeypatch): monkeypatch.setenv("PRIMUS_AUDIT_BATCH_FINGERPRINTS", "1") with pytest.raises(RuntimeError, match="no valid sample-key fingerprint"): From 1e3c9accb9ff3485c2be433aadb1385f87091205 Mon Sep 17 00:00:00 2001 From: GP Huang Date: Thu, 13 Aug 2026 19:26:23 +0300 Subject: [PATCH 11/14] feat(megatron): select MXFP4 forward precision (#986) ## Summary Tracks [Issue 220](https://github.com/AMD-AGI/tiger-training-internal/issues/220). Add independent MXFP4 forward precision selection so late recovery runs can use FP8 or BF16 forward GEMMs while keeping MXFP4 data- and weight-gradient GEMMs. - Add fail-closed `mxfp4_forward_precision` config plumbing and runtime mode census. - Reuse existing FP8 and BF16 operators; no new kernel implementation. - Cover direct autograd, compile, Flux integration, config validation, and audit modes. ## Test plan - [x] Python compile check - [x] `git diff --check` - [ ] Targeted CPU-safe unit tests in the v26.5 container - [ ] MI355X direct linear and Flux 535M forward/backward tests - [ ] Option 6/7 resume smoke with exact mode census --------- Co-authored-by: Guangpu Huang --- .../diffusion-models/mxfp4_training.md | 9 +- ...gon_schnell_resample_local_spec_mxfp4.yaml | 1 + .../extensions/primus_turbo_mxfp4_local.py | 120 +++++++++++++++--- .../core/models/diffusion/common/config.py | 18 +++ .../megatron/flux_pretrain_trainer.py | 30 +++++ .../distributed/test_flux_mxfp4_local_spec.py | 26 ++++ .../megatron/diffusion/test_flux_config.py | 24 ++++ .../megatron/test_diffusion_audit_markers.py | 26 +++- .../megatron/test_primus_turbo_mxfp4_local.py | 91 ++++++++++++- 9 files changed, 321 insertions(+), 24 deletions(-) diff --git a/docs/04-technical-guides/diffusion-models/mxfp4_training.md b/docs/04-technical-guides/diffusion-models/mxfp4_training.md index d45f64f52..58f66e7c3 100644 --- a/docs/04-technical-guides/diffusion-models/mxfp4_training.md +++ b/docs/04-technical-guides/diffusion-models/mxfp4_training.md @@ -8,7 +8,10 @@ MXFP4 stores activations and weights in 4-bit microscale floating-point with one - Uses a **local spec** (`PrimusTurboMXFP4LocalSpecProvider`) with **no Transformer Engine dependency**—MXFP4 linear layers are self-contained autograd `Function`s that call Primus-Turbo's `gemm_fp4_impl` directly, so the path is `torch.compile`-friendly with minimal graph breaks. - Keeps **attention, optimizer state / main params, and inter-rank communication in BF16**. Only the MMA inputs of the column- and row-parallel linears are quantized. -- Supports two backward modes via `mxfp4_backward_precision`: pure **MXFP4** (default) or **FP8** hybrid (E5M2 backward with tensorwise scaling on HipBLASLt). +- Selects forward and backward precision independently. Forward supports + **MXFP4** (default), tensorwise **FP8 E4M3**, or **BF16** through + `mxfp4_forward_precision`; backward supports MXFP4 or tensorwise FP8 E5M2 + through `mxfp4_backward_precision`. - Dispatches the FP4 GEMM through Primus-Turbo's pluggable backend layer, which can route to either AITER (recommended for MI355X) or HipBLASLt. ## Table of contents @@ -67,6 +70,7 @@ The relevant overrides in [`examples/megatron/configs/MI355X/diffusion/flux_12b_ # MXFP4 precision fp4: "mxfp4" fp4_recipe: "mxfp4" # default is "nvfp4" in trainer_base.yaml; must override +mxfp4_forward_precision: "mxfp4" # "mxfp4", "fp8", or "bf16" mxfp4_backward_precision: "mxfp4" # "mxfp4" (pure) or "fp8" (hybrid) # Local spec + Primus-Turbo @@ -86,6 +90,7 @@ gradient_accumulation_fusion: false |------|--------|-------| | `fp4` | `"mxfp4"` | Top-level switch to enable FP4. | | `fp4_recipe` | `"mxfp4"` for this guide | Default in [`primus/configs/modules/megatron/trainer_base.yaml`](../../../primus/configs/modules/megatron/trainer_base.yaml) is `nvfp4`; the MXFP4 config overrides it. | +| `mxfp4_forward_precision` | `"mxfp4"`, `"fp8"`, or `"bf16"` | Selects only the forward linear GEMMs. FP8 uses dynamic tensorwise E4M3 on HipBLASLt; BF16 uses the native linear operation. MXFP4 tensors are still prepared for backward when `mxfp4_backward_precision: "mxfp4"`. | | `mxfp4_backward_precision` | `"mxfp4"` or `"fp8"` | Exhaustive set (checked by branch in [`primus_turbo_mxfp4_local.py`](../../../primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py)). `"fp8"` uses E5M2 with tensorwise HipBLASLt for backward. | | `mxfp4_gradient_stochastic_rounding` | `true` / `false` | Optional. Enables SR on FP4 gradient quantization. | @@ -198,7 +203,7 @@ Formal A/B benchmarks vs BF16 and FP8 (delayed and tensorwise) are pending and w - MXFP4 spec provider: [`primus/backends/megatron/core/extensions/primus_turbo_local_spec.py`](../../../primus/backends/megatron/core/extensions/primus_turbo_local_spec.py) (`PrimusTurboMXFP4LocalSpecProvider`). - MXFP4 linear-layer autograd / fwd-bwd: [`primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py`](../../../primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py). - Config schema defaults: [`primus/configs/modules/megatron/trainer_base.yaml`](../../../primus/configs/modules/megatron/trainer_base.yaml). -- Dataclass field `mxfp4_backward_precision`: [`primus/backends/megatron/core/models/diffusion/common/config.py`](../../../primus/backends/megatron/core/models/diffusion/common/config.py). +- Dataclass fields `mxfp4_forward_precision` and `mxfp4_backward_precision`: [`primus/backends/megatron/core/models/diffusion/common/config.py`](../../../primus/backends/megatron/core/models/diffusion/common/config.py). - FP4 backend selection (Primus-Turbo): `primus_turbo/common/constants.py`, `primus_turbo/pytorch/core/backend.py`, `primus_turbo/pytorch/kernels/gemm/gemm_fp4_impl.py`. - AITER tuned-config loader: `aiter/jit/core.py` (`AITER_CONFIG_GEMM_A4W4`). - AITER A4W4 dispatch + hit/miss logging: `aiter/ops/gemm_op_a4w4.py`. diff --git a/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml index 10283f997..35afd136f 100644 --- a/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml +++ b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml @@ -141,6 +141,7 @@ modules: fp4: "mxfp4" fp4_recipe: "mxfp4" + mxfp4_forward_precision: "mxfp4" # "mxfp4", "fp8", or "bf16" mxfp4_backward_precision: "mxfp4" # "mxfp4" (pure) or "fp8" (hybrid) empty_unused_memory_level: 0 diff --git a/primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py b/primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py index 73f5ed68b..e4eab45dc 100644 --- a/primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py +++ b/primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py @@ -12,7 +12,8 @@ Key properties: - Uses setup_context pattern with primitive-only args so torch.compile can trace through without graph breaks. -- Two backward modes: pure MXFP4 or hybrid (FP4 fwd / FP8 bwd). +- Independently selects MXFP4, FP8, or BF16 forward precision and + MXFP4 or FP8 backward precision. - gemm_fp4_impl and gemm_fp8_impl are torch.library.custom_op with register_fake. - Zero TransformerEngine dependencies. - Requires tensor_model_parallel_size=1, no GAF, no sequence_parallel. @@ -229,6 +230,14 @@ def _quantize_mxfp4_dual_backward(ctx, *grad_outputs): _FP4_DTYPE = torch.float4_e2m1fn_x2 _GRAN_VALUE = ScalingGranularity.MX_BLOCKWISE.value _DEFAULT_BACKEND = BackendType.HIPBLASLT.value +_FORWARD_PRECISION_MXFP4 = 0 +_FORWARD_PRECISION_FP8 = 1 +_FORWARD_PRECISION_BF16 = 2 +_FORWARD_PRECISION_VALUES = { + "mxfp4": _FORWARD_PRECISION_MXFP4, + "fp8": _FORWARD_PRECISION_FP8, + "bf16": _FORWARD_PRECISION_BF16, +} def _quantize_input_dual(input_2d, preshuffle): @@ -291,9 +300,13 @@ def _quantize_grad_dual(grad_2d, preshuffle, use_sr=True): class MXFP4LinearFunction(torch.autograd.Function): """MXFP4 linear (Y = X @ W^T) with MX block-of-32 scaling. - Two modes via backward_is_fp8 bool primitive: - - Pure MXFP4: forward + backward both use FP4 quantization + gemm_fp4_impl - - Hybrid: forward uses FP4, backward re-quantizes saved BF16 to FP8 tensorwise + Forward and backward precision are independent: + - Forward: MXFP4, dynamic tensorwise FP8 E4M3, or BF16. + - Backward: MXFP4 or dynamic tensorwise FP8 E5M2. + + MXFP4 backward always uses the transposed MXFP4 tensors prepared during + forward. This keeps its memory and compute path identical when only the + forward GEMM moves to higher precision. Uses setup_context pattern with primitive-only args for torch.compile. """ @@ -308,6 +321,10 @@ def forward( fp8_gran_value, fp8_backend_value, use_gradient_sr, + forward_precision_value=_FORWARD_PRECISION_MXFP4, + fp8_fwd_dtype=None, + fp8_fwd_gran_value=0, + fp8_fwd_backend_value=0, ): out_dtype = input.dtype orig_shape = input.shape @@ -316,19 +333,39 @@ def forward( a_fp4, a_scale, a_t_fp4, a_t_scale = _quantize_input_dual(input_2d, preshuffle) b_fp4, b_scale, b_t_fp4, b_t_scale = _quantize_weight_dual(weight, preshuffle) - output = gemm_fp4_impl( - a_fp4, - a_scale, - False, - b_fp4, - b_scale, - True, - out_dtype, - False, - granularity=_GRAN_VALUE, - default_backend=_DEFAULT_BACKEND, - preshuffled=preshuffle, - ) + if forward_precision_value == _FORWARD_PRECISION_MXFP4: + output = gemm_fp4_impl( + a_fp4, + a_scale, + False, + b_fp4, + b_scale, + True, + out_dtype, + False, + granularity=_GRAN_VALUE, + default_backend=_DEFAULT_BACKEND, + preshuffled=preshuffle, + ) + elif forward_precision_value == _FORWARD_PRECISION_FP8: + a_fp8, a_scale_inv = _quantize_fp8_tw(input_2d, fp8_fwd_dtype) + b_fp8, b_scale_inv = _quantize_fp8_tw(weight, fp8_fwd_dtype) + output = gemm_fp8_impl( + a_fp8, + a_scale_inv, + False, + b_fp8, + b_scale_inv, + True, + out_dtype, + False, + granularity=fp8_fwd_gran_value, + default_backend=fp8_fwd_backend_value, + ) + elif forward_precision_value == _FORWARD_PRECISION_BF16: + output = torch.nn.functional.linear(input_2d, weight) + else: + raise ValueError(f"Unsupported MXFP4 forward precision value: {forward_precision_value}") output = output.reshape(*orig_shape[:-1], output.shape[-1]) if backward_is_fp8: @@ -348,6 +385,7 @@ def forward( @staticmethod def setup_context(ctx, inputs, output): + ctx.num_inputs = len(inputs) ( _, _, @@ -357,7 +395,7 @@ def setup_context(ctx, inputs, output): fp8_gran_value, fp8_backend_value, use_gradient_sr, - ) = inputs + ) = inputs[:8] ctx.preshuffle = preshuffle ctx.backward_is_fp8 = backward_is_fp8 @@ -458,7 +496,7 @@ def backward(ctx, grad_output, *_): preshuffled=preshuffle, ) - return grad_input, grad_weight, None, None, None, None, None, None + return (grad_input, grad_weight) + (None,) * (ctx.num_inputs - 2) # --------------------------------------------------------------------------- @@ -492,9 +530,27 @@ def __init__(self, *args, **kwargs): self._preshuffle = _enable_preshuffle() _assert_preshuffle_contract(self.config, self._preshuffle) + self._forward_precision = getattr(self.config, "mxfp4_forward_precision", "mxfp4") + if self._forward_precision not in _FORWARD_PRECISION_VALUES: + raise ValueError( + "mxfp4_forward_precision must be one of " + f"{tuple(_FORWARD_PRECISION_VALUES)}, got {self._forward_precision!r}" + ) + self._forward_precision_value = _FORWARD_PRECISION_VALUES[self._forward_precision] self._backward_is_fp8 = getattr(self.config, "mxfp4_backward_precision", "mxfp4") == "fp8" self._use_gradient_sr = getattr(self.config, "mxfp4_gradient_stochastic_rounding", False) + if self._forward_precision == "fp8": + from primus_turbo.pytorch.core.low_precision import float8_e4m3 + + self._fp8_fwd_dtype = float8_e4m3 + self._fp8_fwd_gran_value = ScalingGranularity.TENSORWISE.value + self._fp8_fwd_backend_value = BackendType.HIPBLASLT.value + else: + self._fp8_fwd_dtype = None + self._fp8_fwd_gran_value = 0 + self._fp8_fwd_backend_value = 0 + if self._backward_is_fp8: from primus_turbo.pytorch.core.low_precision import float8_e5m2 @@ -518,6 +574,10 @@ def _forward_impl(self, input, weight, *args, **kwargs): self._fp8_gran_value, self._fp8_backend_value, self._use_gradient_sr, + self._forward_precision_value, + self._fp8_fwd_dtype, + self._fp8_fwd_gran_value, + self._fp8_fwd_backend_value, ) output = result[0] @@ -552,9 +612,27 @@ def __init__(self, *args, **kwargs): self._preshuffle = _enable_preshuffle() _assert_preshuffle_contract(self.config, self._preshuffle) + self._forward_precision = getattr(self.config, "mxfp4_forward_precision", "mxfp4") + if self._forward_precision not in _FORWARD_PRECISION_VALUES: + raise ValueError( + "mxfp4_forward_precision must be one of " + f"{tuple(_FORWARD_PRECISION_VALUES)}, got {self._forward_precision!r}" + ) + self._forward_precision_value = _FORWARD_PRECISION_VALUES[self._forward_precision] self._backward_is_fp8 = getattr(self.config, "mxfp4_backward_precision", "mxfp4") == "fp8" self._use_gradient_sr = getattr(self.config, "mxfp4_gradient_stochastic_rounding", False) + if self._forward_precision == "fp8": + from primus_turbo.pytorch.core.low_precision import float8_e4m3 + + self._fp8_fwd_dtype = float8_e4m3 + self._fp8_fwd_gran_value = ScalingGranularity.TENSORWISE.value + self._fp8_fwd_backend_value = BackendType.HIPBLASLT.value + else: + self._fp8_fwd_dtype = None + self._fp8_fwd_gran_value = 0 + self._fp8_fwd_backend_value = 0 + if self._backward_is_fp8: from primus_turbo.pytorch.core.low_precision import float8_e5m2 @@ -578,6 +656,10 @@ def _forward_impl(self, input, weight, *args, **kwargs): self._fp8_gran_value, self._fp8_backend_value, self._use_gradient_sr, + self._forward_precision_value, + self._fp8_fwd_dtype, + self._fp8_fwd_gran_value, + self._fp8_fwd_backend_value, ) output = result[0] diff --git a/primus/backends/megatron/core/models/diffusion/common/config.py b/primus/backends/megatron/core/models/diffusion/common/config.py index de4074744..a5a19f07f 100644 --- a/primus/backends/megatron/core/models/diffusion/common/config.py +++ b/primus/backends/megatron/core/models/diffusion/common/config.py @@ -32,6 +32,8 @@ class BaseDiffusionConfig(TransformerConfig): fp8_scaling_strategy: FP8 scaling strategy for local spec provider (default: 'dynamic') fp8_force_nt_layout: FP8 backward GEMM layout (default: False) fp8_reduce_amax: Whether to allreduce amax across ranks (default: False) + mxfp4_forward_precision: MXFP4 forward precision, 'mxfp4', 'fp8', or 'bf16' + (default: 'mxfp4') mxfp4_backward_precision: MXFP4 backward precision, 'mxfp4' or 'fp8' (default: 'mxfp4') mxfp4_gradient_stochastic_rounding: Stochastic rounding on gradients (default: False) sensitive_layers_enabled: Enable sensitive layer configuration (default: False) @@ -71,6 +73,11 @@ class BaseDiffusionConfig(TransformerConfig): # Whether to allreduce amax across DP/TP ranks for delayed FP8 scaling fp8_reduce_amax: bool = False + # MXFP4 forward precision: "mxfp4" (pure), "fp8", or "bf16". + # The latter two retain MXFP4 backward unless the backward knob below is + # independently changed. + mxfp4_forward_precision: str = "mxfp4" + # MXFP4 backward precision: "mxfp4" (pure) or "fp8" (hybrid) mxfp4_backward_precision: str = "mxfp4" @@ -168,6 +175,17 @@ def validate(self): Raises: ValueError: If configuration is invalid """ + if self.mxfp4_forward_precision not in {"mxfp4", "fp8", "bf16"}: + raise ValueError( + "mxfp4_forward_precision must be 'mxfp4', 'fp8', or 'bf16', " + f"got {self.mxfp4_forward_precision!r}" + ) + + if self.mxfp4_backward_precision not in {"mxfp4", "fp8"}: + raise ValueError( + "mxfp4_backward_precision must be 'mxfp4' or 'fp8', " f"got {self.mxfp4_backward_precision!r}" + ) + if self.in_channels <= 0: raise ValueError(f"in_channels must be positive, got {self.in_channels}") diff --git a/primus/backends/megatron/flux_pretrain_trainer.py b/primus/backends/megatron/flux_pretrain_trainer.py index 38925a45e..9d710adc7 100644 --- a/primus/backends/megatron/flux_pretrain_trainer.py +++ b/primus/backends/megatron/flux_pretrain_trainer.py @@ -46,6 +46,21 @@ def _precision_linear_class_census(model) -> dict[str, int]: return {name: observed.get(name, 0) for name in PRECISION_LINEAR_CLASSES} +def _mxfp4_gemm_mode_census(model) -> dict[str, int]: + """Count the effective forward/backward modes of instantiated MXFP4 linears.""" + observed = Counter() + for module in model.modules(): + if type(module).__name__ not in { + "MXFP4ColumnParallelLinear", + "MXFP4RowParallelLinear", + }: + continue + forward_precision = getattr(module, "_forward_precision", "mxfp4") + backward_precision = "fp8" if getattr(module, "_backward_is_fp8", False) else "mxfp4" + observed[f"{forward_precision}_forward_{backward_precision}_backward"] += 1 + return dict(sorted(observed.items())) + + def _emit_precision_linear_class_census(model) -> None: """Emit the actually instantiated precision-linear classes on every rank.""" if os.getenv("PRIMUS_AUDIT_LINEAR_CLASS_CENSUS") != "1": @@ -66,6 +81,19 @@ def _emit_precision_linear_class_census(model) -> None: # this launch path, so a bare print leaves the audit reporting zero markers # on a healthy run. Logging and fd 2 both survive. logger.info("PRIMUS_LINEAR_CLASS_CENSUS=%s", json.dumps(payload, sort_keys=True)) + modes = _mxfp4_gemm_mode_census(model) + if modes: + logger.info( + "PRIMUS_MXFP4_GEMM_MODE_CENSUS=%s", + json.dumps( + { + "global_rank": global_rank, + "data_parallel_rank": parallel_state.get_data_parallel_rank(), + "modes": modes, + }, + sort_keys=True, + ), + ) def _restore_chimera_rng_state(args) -> None: @@ -544,6 +572,7 @@ def _build_flux_config_from_yaml(self): { "fp4": fp4_enabled, "fp4_recipe": fp4_recipe, + "mxfp4_forward_precision": getattr(params, "mxfp4_forward_precision", "mxfp4"), "mxfp4_backward_precision": getattr(params, "mxfp4_backward_precision", "mxfp4"), } ) @@ -741,6 +770,7 @@ def _log_flux_config(self, config, args): "fp8_force_nt_layout", "fp4", "fp4_recipe", + "mxfp4_forward_precision", "mxfp4_backward_precision", "mxfp4_gradient_stochastic_rounding", "sensitive_layers_enabled", diff --git a/tests/integration_tests/backends/megatron/diffusion/distributed/test_flux_mxfp4_local_spec.py b/tests/integration_tests/backends/megatron/diffusion/distributed/test_flux_mxfp4_local_spec.py index 6926817fd..6ac0f69e5 100644 --- a/tests/integration_tests/backends/megatron/diffusion/distributed/test_flux_mxfp4_local_spec.py +++ b/tests/integration_tests/backends/megatron/diffusion/distributed/test_flux_mxfp4_local_spec.py @@ -201,3 +201,29 @@ def test_flux_535m_mxfp4_hybrid_forward_backward(self): break assert has_mxfp4_grad, "No MXFP4 linear has a weight gradient in hybrid mode" + + @requires_mxfp4 + def test_flux_535m_high_precision_forward_mxfp4_backward(self): + """Higher-precision forward retains MXFP4 modules and finite gradients.""" + for forward_precision in ("fp8", "bf16"): + with self.subTest(forward_precision=forward_precision): + config = self._make_mxfp4_config(mxfp4_forward_precision=forward_precision) + model = Flux(config).cuda().to(torch.bfloat16) + model.train() + + precision_modules = [ + module + for module in model.modules() + if isinstance(module, (MXFP4ColumnParallelLinear, MXFP4RowParallelLinear)) + ] + assert precision_modules + assert all(module._forward_precision == forward_precision for module in precision_modules) + assert all(not module._backward_is_fp8 for module in precision_modules) + + output = model(*self._make_inputs(batch_size=2)) + assert torch.isfinite(output).all() + output.sum().backward() + assert any( + module.weight.grad is not None and torch.isfinite(module.weight.grad).all() + for module in precision_modules + ) diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_config.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_config.py index f286490a4..d66d10a19 100644 --- a/tests/unit_tests/backends/megatron/diffusion/test_flux_config.py +++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_config.py @@ -32,6 +32,30 @@ def test_base_config_validation_invalid_channels(self): config.validate() self.assertIn("in_channels must be positive", str(cm.exception)) + def test_mxfp4_forward_precision_values(self): + for forward_precision in ("mxfp4", "fp8", "bf16"): + with self.subTest(forward_precision=forward_precision): + config = FluxConfig.flux_535m( + mxfp4_forward_precision=forward_precision, + ) + self.assertEqual(config.mxfp4_forward_precision, forward_precision) + + def test_invalid_mxfp4_forward_precision_is_rejected(self): + with self.assertRaisesRegex(ValueError, "mxfp4_forward_precision"): + BaseDiffusionConfig( + num_attention_heads=8, + num_layers=1, + mxfp4_forward_precision="fp16", + ) + + def test_invalid_mxfp4_backward_precision_is_rejected(self): + with self.assertRaisesRegex(ValueError, "mxfp4_backward_precision"): + BaseDiffusionConfig( + num_attention_heads=8, + num_layers=1, + mxfp4_backward_precision="bf16", + ) + class TestFluxConfig(PrimusUT): """Tests for FluxConfig class.""" diff --git a/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py b/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py index 17ef0ecb2..30396c5b5 100644 --- a/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py +++ b/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py @@ -10,6 +10,7 @@ from primus.backends.megatron.flux_pretrain_trainer import ( FluxPretrainTrainer, _emit_precision_linear_class_census, + _mxfp4_gemm_mode_census, _precision_linear_class_census, ) from primus.backends.megatron.patches.mlperf_warmup_patches import _is_resumed_training @@ -18,6 +19,7 @@ ) _FORWARD_STEP_LOGGER = "primus.backends.megatron.training.diffusion.forward_step" +_FLUX_TRAINER_LOGGER = "primus.backends.megatron.flux_pretrain_trainer" def test_sample_key_fingerprint_is_order_sensitive(): @@ -50,16 +52,36 @@ def test_precision_linear_class_census_reports_exact_classes(): } -def test_precision_linear_class_census_emits_zero_counts_for_bf16(monkeypatch, capsys): +def test_mxfp4_gemm_mode_census_reports_forward_backward_split(): + mxfp4_column = type("MXFP4ColumnParallelLinear", (nn.Module,), {})() + mxfp4_column._forward_precision = "fp8" + mxfp4_column._backward_is_fp8 = False + mxfp4_row = type("MXFP4RowParallelLinear", (nn.Module,), {})() + mxfp4_row._forward_precision = "bf16" + mxfp4_row._backward_is_fp8 = False + model = nn.ModuleList([mxfp4_column, mxfp4_row, nn.Linear(2, 2)]) + + assert _mxfp4_gemm_mode_census(model) == { + "bf16_forward_mxfp4_backward": 1, + "fp8_forward_mxfp4_backward": 1, + } + + +def test_precision_linear_class_census_emits_zero_counts_for_bf16(monkeypatch, caplog): from megatron.core import parallel_state monkeypatch.setenv("PRIMUS_AUDIT_LINEAR_CLASS_CENSUS", "1") monkeypatch.setenv("RANK", "3") monkeypatch.setattr(parallel_state, "get_data_parallel_rank", lambda: 3) + caplog.set_level("INFO", logger=_FLUX_TRAINER_LOGGER) _emit_precision_linear_class_census(nn.Linear(2, 2)) - line = capsys.readouterr().out.strip() + line = next( + record.message + for record in caplog.records + if record.message.startswith("PRIMUS_LINEAR_CLASS_CENSUS=") + ) payload = json.loads(line.split("=", 1)[1]) assert payload == { "data_parallel_rank": 3, diff --git a/tests/unit_tests/backends/megatron/test_primus_turbo_mxfp4_local.py b/tests/unit_tests/backends/megatron/test_primus_turbo_mxfp4_local.py index a0fc099b9..c1ae48e65 100644 --- a/tests/unit_tests/backends/megatron/test_primus_turbo_mxfp4_local.py +++ b/tests/unit_tests/backends/megatron/test_primus_turbo_mxfp4_local.py @@ -6,7 +6,7 @@ Tests cross-validation against Primus-Turbo's FP4GemmMXFunction reference, torch.compile graph-break validation, Megatron linear backward flow, -2-step training loop, hybrid (FP4 fwd / FP8 bwd) mode, and init guards. +2-step training loop, independent forward/backward precision, and init guards. """ import functools @@ -188,6 +188,55 @@ def test_backward_matches_reference_fp4gemm(self): f"(max abs diff vs Primus-Turbo PR #383 reference: {(x.grad - x_ref.grad).abs().max().item():.6e})" ) + @requires_mxfp4 + def test_high_precision_forward_with_mxfp4_backward(self): + from primus_turbo.pytorch.core.backend import BackendType + from primus_turbo.pytorch.core.low_precision import ( + ScalingGranularity, + float8_e4m3, + ) + + from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import ( + _FORWARD_PRECISION_VALUES, + MXFP4LinearFunction, + _enable_preshuffle, + ) + + for forward_precision in ("fp8", "bf16"): + with self.subTest(forward_precision=forward_precision): + torch.manual_seed(42) + x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda", requires_grad=True) + w = torch.randn(512, 256, dtype=torch.bfloat16, device="cuda", requires_grad=True) + preshuffle = _enable_preshuffle() + + output = MXFP4LinearFunction.apply( + x, + w, + preshuffle, + False, + None, + 0, + 0, + False, + _FORWARD_PRECISION_VALUES[forward_precision], + float8_e4m3 if forward_precision == "fp8" else None, + ScalingGranularity.TENSORWISE.value if forward_precision == "fp8" else 0, + BackendType.HIPBLASLT.value if forward_precision == "fp8" else 0, + )[0] + output.square().mean().backward() + + reference = torch.nn.functional.linear(x.detach(), w.detach()) + signal = (reference.float() ** 2).mean() + noise = ((output.float() - reference.float()) ** 2).mean() + if forward_precision == "bf16": + assert torch.equal(output, reference) + else: + snr_db = 10 * torch.log10(signal / noise).item() + assert snr_db > 10, f"FP8 forward SNR {snr_db:.1f} dB is below 10 dB" + + assert torch.isfinite(x.grad).all() + assert torch.isfinite(w.grad).all() + # --------------------------------------------------------------------------- # torch.compile graph-break validation @@ -260,6 +309,46 @@ def test_no_graph_break_hybrid(self): f"Reasons: {explanation.break_reasons}" ) + @requires_mxfp4 + def test_no_graph_break_high_precision_forward(self): + from primus_turbo.pytorch.core.backend import BackendType + from primus_turbo.pytorch.core.low_precision import ( + ScalingGranularity, + float8_e4m3, + ) + + from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import ( + _FORWARD_PRECISION_VALUES, + MXFP4LinearFunction, + _enable_preshuffle, + ) + + for forward_precision in ("fp8", "bf16"): + with self.subTest(forward_precision=forward_precision): + torch._dynamo.reset() + x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + w = torch.randn(512, 256, dtype=torch.bfloat16, device="cuda") + + explanation = torch._dynamo.explain(MXFP4LinearFunction.apply)( + x, + w, + _enable_preshuffle(), + False, + None, + 0, + 0, + False, + _FORWARD_PRECISION_VALUES[forward_precision], + float8_e4m3 if forward_precision == "fp8" else None, + ScalingGranularity.TENSORWISE.value if forward_precision == "fp8" else 0, + BackendType.HIPBLASLT.value if forward_precision == "fp8" else 0, + ) + + assert explanation.graph_break_count == 0, ( + f"Expected 0 graph breaks, got {explanation.graph_break_count}. " + f"Reasons: {explanation.break_reasons}" + ) + @requires_mxfp4 def test_compiled_forward_matches_eager(self): from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import ( From f031983a7ce09fb48ce32c8321d55544d4c839cb Mon Sep 17 00:00:00 2001 From: GP Huang Date: Thu, 13 Aug 2026 19:44:03 +0300 Subject: [PATCH 12/14] test(megatron): cover compiled MXFP4 backward (#987) ## Summary Closes the post-merge review gap from PR #986. Adds regression coverage for FP8/BF16 forward with MXFP4 backward through `torch.compile`, including exact eager/compiled output and gradient comparisons. It also covers YAML-to-`FluxConfig` precision propagation and runtime mode-census emission. ## Test plan - [x] 35 config and audit tests passed - [x] Compiled-backward checks passed for FP8 and BF16 forward modes - [x] Flux 535M forward-only integration test passed - [x] `black` formatting verified - [x] GPU memory returned to idle Co-authored-by: Guangpu Huang --- .../training/test_flux_model_creation.py | 17 ++++++ .../megatron/test_diffusion_audit_markers.py | 25 +++++++++ .../megatron/test_primus_turbo_mxfp4_local.py | 55 +++++++++++++++++++ 3 files changed, 97 insertions(+) diff --git a/tests/unit_tests/backends/megatron/diffusion/training/test_flux_model_creation.py b/tests/unit_tests/backends/megatron/diffusion/training/test_flux_model_creation.py index 9a965ccfc..ba59e87f8 100644 --- a/tests/unit_tests/backends/megatron/diffusion/training/test_flux_model_creation.py +++ b/tests/unit_tests/backends/megatron/diffusion/training/test_flux_model_creation.py @@ -245,6 +245,23 @@ def test_build_flux_config_from_yaml_fp8_settings(self, monkeypatch: pytest.Monk assert config.fp8_dot_product_attention is False assert config.fp8_multi_head_attention is False + def test_build_flux_config_from_yaml_mxfp4_precision_settings(self, monkeypatch: pytest.MonkeyPatch): + backend_args = SimpleNamespace( + mock_data=True, + fp4="mxfp4", + fp4_recipe="mxfp4", + mxfp4_forward_precision="bf16", + mxfp4_backward_precision="mxfp4", + ) + + trainer = _build_flux_trainer(monkeypatch, backend_args) + config = trainer._build_flux_config_from_yaml() + + assert config.fp4 == "mxfp4" + assert config.fp4_recipe == "mxfp4" + assert config.mxfp4_forward_precision == "bf16" + assert config.mxfp4_backward_precision == "mxfp4" + def test_create_model_does_not_overwrite_existing_torch_compile_attrs( self, monkeypatch: pytest.MonkeyPatch ): diff --git a/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py b/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py index 30396c5b5..6bfb815d1 100644 --- a/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py +++ b/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py @@ -95,6 +95,31 @@ def test_precision_linear_class_census_emits_zero_counts_for_bf16(monkeypatch, c } +def test_precision_linear_class_census_emits_mxfp4_gemm_modes(monkeypatch, caplog): + from megatron.core import parallel_state + + monkeypatch.setenv("PRIMUS_AUDIT_LINEAR_CLASS_CENSUS", "1") + monkeypatch.setenv("RANK", "5") + monkeypatch.setattr(parallel_state, "get_data_parallel_rank", lambda: 5) + caplog.set_level("INFO", logger=_FLUX_TRAINER_LOGGER) + + mxfp4_column = type("MXFP4ColumnParallelLinear", (nn.Module,), {})() + mxfp4_column._forward_precision = "fp8" + mxfp4_column._backward_is_fp8 = False + _emit_precision_linear_class_census(nn.ModuleList([mxfp4_column])) + + line = next( + record.message + for record in caplog.records + if record.message.startswith("PRIMUS_MXFP4_GEMM_MODE_CENSUS=") + ) + assert json.loads(line.split("=", 1)[1]) == { + "data_parallel_rank": 5, + "global_rank": 5, + "modes": {"fp8_forward_mxfp4_backward": 1}, + } + + def test_emit_batch_fingerprint_is_fail_closed(monkeypatch): monkeypatch.setenv("PRIMUS_AUDIT_BATCH_FINGERPRINTS", "1") with pytest.raises(RuntimeError, match="no valid sample-key fingerprint"): diff --git a/tests/unit_tests/backends/megatron/test_primus_turbo_mxfp4_local.py b/tests/unit_tests/backends/megatron/test_primus_turbo_mxfp4_local.py index c1ae48e65..d8336cf5b 100644 --- a/tests/unit_tests/backends/megatron/test_primus_turbo_mxfp4_local.py +++ b/tests/unit_tests/backends/megatron/test_primus_turbo_mxfp4_local.py @@ -349,6 +349,61 @@ def test_no_graph_break_high_precision_forward(self): f"Reasons: {explanation.break_reasons}" ) + @requires_mxfp4 + def test_compiled_backward_matches_eager_high_precision_forward(self): + from primus_turbo.pytorch.core.backend import BackendType + from primus_turbo.pytorch.core.low_precision import ( + ScalingGranularity, + float8_e4m3, + ) + + from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import ( + _FORWARD_PRECISION_VALUES, + MXFP4LinearFunction, + _enable_preshuffle, + ) + + for forward_precision in ("fp8", "bf16"): + with self.subTest(forward_precision=forward_precision): + torch._dynamo.reset() + torch.manual_seed(42) + x_eager = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda", requires_grad=True) + w_eager = torch.randn(512, 256, dtype=torch.bfloat16, device="cuda", requires_grad=True) + x_compiled = x_eager.detach().clone().requires_grad_(True) + w_compiled = w_eager.detach().clone().requires_grad_(True) + upstream = torch.randn(128, 512, dtype=torch.bfloat16, device="cuda") + preshuffle = _enable_preshuffle() + forward_precision_value = _FORWARD_PRECISION_VALUES[forward_precision] + fp8_dtype = float8_e4m3 if forward_precision == "fp8" else None + fp8_granularity = ScalingGranularity.TENSORWISE.value if forward_precision == "fp8" else 0 + fp8_backend = BackendType.HIPBLASLT.value if forward_precision == "fp8" else 0 + + def forward(x, weight): + return MXFP4LinearFunction.apply( + x, + weight, + preshuffle, + False, + None, + 0, + 0, + False, + forward_precision_value, + fp8_dtype, + fp8_granularity, + fp8_backend, + )[0] + + eager_output = forward(x_eager, w_eager) + eager_output.backward(upstream) + + compiled_output = torch.compile(forward, fullgraph=True)(x_compiled, w_compiled) + compiled_output.backward(upstream) + + assert torch.equal(compiled_output, eager_output) + assert torch.equal(x_compiled.grad, x_eager.grad) + assert torch.equal(w_compiled.grad, w_eager.grad) + @requires_mxfp4 def test_compiled_forward_matches_eager(self): from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import ( From eb95743321157066be097ca50b3d74ba6861dd33 Mon Sep 17 00:00:00 2001 From: GP Huang Date: Thu, 13 Aug 2026 20:22:19 +0300 Subject: [PATCH 13/14] feat(megatron): add sampled Flux weight-state audit (#985) ## Summary Tracks [Issue \#220](https://github.com/AMD-AGI/tiger-training-internal/issues/220), Task 9 Option 9. Add an opt-in Flux audit that records deterministic rank-zero parameter samples after selected completed training iterations. This lets early MXFP4, FP8, and BF16 forks compare model drift without extra full optimizer checkpoints or changes to Torch RNG state. - Sample after the model forward so overlapped distributed-optimizer gathers have refreshed every parameter buffer. - Record per-parameter shape, dtype, finite counts, sampled moments, and SHA256. - Publish one strict JSON file per completed iteration using atomic, no-overwrite semantics; identical restarts are idempotent and conflicting output fails closed. - Skip validation and synthetic warmup, and explicitly record completed-iteration, next-iteration, forward-counter, and microbatch coordinates. ## Test plan - [x] Changed modules compile with `py_compile` - [x] `git diff --check` - [ ] CI formatting and dependency checks - [ ] Targeted audit-marker tests in the pinned v26.5 container - [ ] Real Megatron wrapper smoke verifies post-gather step-5,120 and step-8,192 snapshots --------- Co-authored-by: Guangpu Huang <13152353+gphuang@users.noreply.github.com> --- .../training/diffusion/forward_step.py | 537 +++++++++++++- .../training/test_flux_forward_step_e2e.py | 5 +- .../megatron/test_diffusion_audit_markers.py | 688 +++++++++++++++++- 3 files changed, 1190 insertions(+), 40 deletions(-) diff --git a/primus/backends/megatron/training/diffusion/forward_step.py b/primus/backends/megatron/training/diffusion/forward_step.py index ac9aa38a7..ddc8384ee 100644 --- a/primus/backends/megatron/training/diffusion/forward_step.py +++ b/primus/backends/megatron/training/diffusion/forward_step.py @@ -17,9 +17,13 @@ Architecture follows functional composition for clarity and testability. """ +import hashlib import json import logging import os +import stat +import tempfile +from pathlib import Path from typing import Optional, Tuple import torch @@ -38,6 +42,8 @@ ) logger = logging.getLogger(__name__) +_EMITTED_MODEL_WEIGHT_ITERATIONS: set[int] = set() +_MODEL_WEIGHT_AUDIT_CONTEXT_UNSET = object() def _emit_batch_fingerprint(batch: dict, step_count: int, *, is_training: bool = True) -> None: @@ -81,6 +87,427 @@ def _emit_batch_fingerprint(batch: dict, step_count: int, *, is_training: bool = logger.info("PRIMUS_BATCH_FINGERPRINT=%s", json.dumps(payload, sort_keys=True)) +def _parse_model_weight_steps(value: str) -> set[int]: + """Parse completed training iterations requested for weight auditing.""" + steps = set() + for token in value.split(","): + token = token.strip() + if not token or not token.isdigit() or int(token) <= 0: + raise RuntimeError( + "PRIMUS_AUDIT_MODEL_WEIGHT_STEPS must be a comma-separated " "list of positive integers" + ) + step = int(token) + if step in steps: + raise RuntimeError("PRIMUS_AUDIT_MODEL_WEIGHT_STEPS contains duplicate step " f"{step}") + steps.add(step) + if not steps: + raise RuntimeError("PRIMUS_AUDIT_MODEL_WEIGHT_STEPS is empty") + return steps + + +def _sample_model_weights(model, sample_size: int) -> dict: + """Return deterministic, low-overhead per-parameter weight samples.""" + if sample_size <= 0: + raise ValueError("sample_size must be positive") + + metadata = [] + sampled_tensors = [] + for name, parameter in sorted(model.named_parameters(), key=lambda item: item[0]): + tensor = parameter.detach().reshape(-1) + if tensor.numel() == 0: + continue + count = min(sample_size, tensor.numel()) + indices = torch.arange(count, device=tensor.device, dtype=torch.int64) * tensor.numel() // count + sampled = tensor.index_select(0, indices).to(dtype=torch.float32) + metadata.append( + { + "name": name, + "shape": list(parameter.shape), + "dtype": str(parameter.dtype), + "numel": parameter.numel(), + "sample_count": count, + "requires_grad": parameter.requires_grad, + } + ) + sampled_tensors.append(sampled) + + if not sampled_tensors: + raise RuntimeError("model-weight audit found no parameters") + devices = {tensor.device for tensor in sampled_tensors} + if len(devices) != 1: + raise RuntimeError( + "model-weight audit requires all sampled parameters on one device, " + f"found {sorted(map(str, devices))}" + ) + + combined = torch.cat(sampled_tensors).cpu() + parameters = [] + offset = 0 + total_sum = 0.0 + total_sum_squares = 0.0 + total_absmax = 0.0 + total_nonfinite_count = 0 + all_finite = True + for item in metadata: + count = item["sample_count"] + sample = combined[offset : offset + count] + offset += count + finite_mask = torch.isfinite(sample) + nonfinite_count = int((~finite_mask).sum().item()) + finite = nonfinite_count == 0 + sample_sum = float(sample.double().sum().item()) if finite else None + sample_sum_squares = float(sample.double().square().sum().item()) if finite else None + sample_absmax = float(sample.abs().max().item()) if finite else None + item.update( + { + "sample_finite": finite, + "sample_nonfinite_count": nonfinite_count, + "sample_sum": sample_sum, + "sample_sum_squares": sample_sum_squares, + "sample_absmax": sample_absmax, + "sample_sha256": hashlib.sha256(sample.contiguous().numpy().tobytes()).hexdigest(), + } + ) + parameters.append(item) + all_finite = all_finite and finite + total_nonfinite_count += nonfinite_count + if finite: + total_sum += sample_sum + total_sum_squares += sample_sum_squares + total_absmax = max(total_absmax, sample_absmax) + + return { + "parameter_count": len(parameters), + "parameter_numel": sum(item["numel"] for item in parameters), + "sample_count": combined.numel(), + "sample_finite": all_finite, + "sample_nonfinite_count": total_nonfinite_count, + "sample_sum": total_sum if all_finite else None, + "sample_sum_squares": total_sum_squares if all_finite else None, + "sample_absmax": total_absmax if all_finite else None, + "parameters": parameters, + } + + +def _model_weight_iteration_coordinate() -> tuple[int, int, int, int]: + """Return canonical completed/next iterations and current run metadata. + + Megatron restores ``args.iteration`` from the checkpoint before entering + the training loop. The pinned Megatron training loop then records its + active completed-iteration coordinate in ``args.curr_iteration`` before + every ``train_step`` while leaving ``args.iteration`` at the restored + baseline. Prefer that active coordinate when present and retain the + restored value as the resume-safe fallback. + + Forward-call counts are deliberately excluded: gradient accumulation can + change, and Megatron may replay a forward before any optimizer update. + Neither event is allowed to shift the Megatron training-loop coordinate. + """ + from megatron.core.num_microbatches_calculator import get_num_microbatches + from megatron.training import get_args + + args = get_args() + restored_iteration = getattr(args, "iteration", None) + if ( + isinstance(restored_iteration, bool) + or not isinstance(restored_iteration, int) + or restored_iteration < 0 + ): + raise RuntimeError("model-weight audit requires args.iteration to be a nonnegative integer") + + completed_iteration = getattr(args, "curr_iteration", restored_iteration) + if ( + isinstance(completed_iteration, bool) + or not isinstance(completed_iteration, int) + or completed_iteration < 0 + ): + raise RuntimeError( + "model-weight audit requires args.curr_iteration to be a nonnegative integer when present" + ) + if completed_iteration < restored_iteration: + raise RuntimeError( + "model-weight audit requires args.curr_iteration to be at least the restored args.iteration" + ) + + train_iters = getattr(args, "train_iters", None) + if isinstance(train_iters, bool) or not isinstance(train_iters, int) or train_iters <= 0: + raise RuntimeError("model-weight audit requires args.train_iters to be a positive integer") + + num_microbatches = get_num_microbatches() + if isinstance(num_microbatches, bool) or not isinstance(num_microbatches, int) or num_microbatches <= 0: + raise RuntimeError( + "model-weight audit requires get_num_microbatches() to return a " "positive integer" + ) + return completed_iteration, completed_iteration + 1, num_microbatches, train_iters + + +def _reject_json_constant(value: str): + raise ValueError(f"non-standard JSON constant {value}") + + +def _encode_strict_json(payload: dict) -> bytes: + """Encode one deterministic JSON object with no non-finite constants.""" + if not isinstance(payload, dict): + raise TypeError("model-weight audit payload must be a JSON object") + return ( + json.dumps( + payload, + sort_keys=True, + allow_nan=False, + separators=(",", ":"), + ) + + "\n" + ).encode() + + +def _read_strict_json(path: Path) -> dict: + """Read a regular, non-symlink JSON object and reject non-finite values.""" + flags = os.O_RDONLY + if hasattr(os, "O_NOFOLLOW"): + flags |= os.O_NOFOLLOW + if hasattr(os, "O_NONBLOCK"): + flags |= os.O_NONBLOCK + descriptor = os.open(path, flags) + try: + if not stat.S_ISREG(os.fstat(descriptor).st_mode): + raise RuntimeError(f"model-weight audit output is not a regular file: {path}") + handle = os.fdopen(descriptor) + descriptor = None + with handle: + payload = json.load(handle, parse_constant=_reject_json_constant) + finally: + if descriptor is not None: + os.close(descriptor) + if not isinstance(payload, dict): + raise RuntimeError(f"model-weight audit output is not a JSON object: {path}") + # json.load(parse_constant=...) rejects NaN/Infinity tokens. Re-encoding + # additionally rejects valid JSON numbers that overflow to Python infinity + # (for example 1e9999). + _encode_strict_json(payload) + return payload + + +def _model_weight_summary_identity(payload: dict) -> tuple[dict, bool]: + """Return restart identity while validating replay-variant provenance.""" + provenance_fields = {"forward_step_count", "num_microbatches"} + present_fields = provenance_fields.intersection(payload) + if present_fields and present_fields != provenance_fields: + missing = sorted(provenance_fields - present_fields) + raise RuntimeError( + "model-weight audit output has incomplete replay provenance; " f"missing fields: {missing}" + ) + + has_provenance = bool(present_fields) + identity = dict(payload) + if has_provenance: + for field in sorted(provenance_fields): + value = identity.pop(field) + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise RuntimeError(f"model-weight audit output field {field} must be a positive integer") + return identity, has_provenance + + +def _write_model_weight_summary_once(output: Path, payload: dict) -> None: + """Publish one strict JSON record atomically without overwriting.""" + encoded = _encode_strict_json(payload) + identity, has_provenance = _model_weight_summary_identity(payload) + encoded_identity = _encode_strict_json(identity) + descriptor, temporary_value = tempfile.mkstemp( + dir=output.parent, + prefix=f".{output.name}.", + suffix=".tmp", + ) + temporary = Path(temporary_value) + try: + os.fchmod(descriptor, 0o600) + handle = os.fdopen(descriptor, "wb") + descriptor = None + with handle: + handle.write(encoded) + handle.flush() + os.fsync(handle.fileno()) + try: + os.link(temporary, output, follow_symlinks=False) + except FileExistsError: + existing = _read_strict_json(output) + existing_identity, existing_has_provenance = _model_weight_summary_identity(existing) + if ( + existing_has_provenance != has_provenance + or _encode_strict_json(existing_identity) != encoded_identity + ): + raise RuntimeError( + "model-weight audit output already exists with different " f"content: {output}" + ) + else: + directory_flags = os.O_RDONLY + if hasattr(os, "O_DIRECTORY"): + directory_flags |= os.O_DIRECTORY + directory_descriptor = os.open(output.parent, directory_flags) + try: + os.fsync(directory_descriptor) + finally: + os.close(directory_descriptor) + finally: + if descriptor is not None: + os.close(descriptor) + temporary.unlink(missing_ok=True) + + +def _model_weight_audit_context(step_count: int, *, is_training: bool): + """Validate audit configuration on every rank before rank-zero selection.""" + raw_steps = os.getenv("PRIMUS_AUDIT_MODEL_WEIGHT_STEPS") + if raw_steps is None: + return None + if not is_training or os.getenv("PRIMUS_SYNTHETIC_WARMUP_ACTIVE") == "1": + return None + + completed_iterations = _parse_model_weight_steps(raw_steps) + ( + completed_training_iteration, + next_training_iteration, + num_microbatches, + train_iters, + ) = _model_weight_iteration_coordinate() + terminal_iterations = sorted(step for step in completed_iterations if step >= train_iters) + if terminal_iterations: + first_terminal = terminal_iterations[0] + raise RuntimeError( + "model-weight audit cannot safely sample completed iteration " + f"{first_terminal} with train_iters={train_iters}: overlap-param-gather " + "buffers are refreshed by the next model forward, so selected completed " + f"iteration {first_terminal} requires training through at least {first_terminal + 1}" + ) + + raw_sample_size = os.getenv("PRIMUS_AUDIT_MODEL_WEIGHT_SAMPLE_SIZE", "256") + if not raw_sample_size.isdigit() or not 1 <= int(raw_sample_size) <= 4096: + raise RuntimeError("PRIMUS_AUDIT_MODEL_WEIGHT_SAMPLE_SIZE must be an integer in [1, 4096]") + sample_size = int(raw_sample_size) + + output_value = os.getenv("PRIMUS_AUDIT_MODEL_WEIGHT_PATH") + if not output_value: + raise RuntimeError("PRIMUS_AUDIT_MODEL_WEIGHT_PATH is required when weight auditing is enabled") + output_directory = Path(output_value) + if not output_directory.is_absolute(): + raise RuntimeError("PRIMUS_AUDIT_MODEL_WEIGHT_PATH must be absolute") + + if isinstance(step_count, bool) or not isinstance(step_count, int) or step_count <= 0: + raise RuntimeError("model-weight audit requires step_count to be a positive integer") + forward_step_count = step_count + + # Every deterministic check above runs on every training rank. Branching + # earlier can leave peers entering model collectives after rank zero fails. + global_rank = ( + torch.distributed.get_rank() if torch.distributed.is_initialized() else int(os.getenv("RANK", "-1")) + ) + if global_rank != 0: + return None + + return { + "completed_iterations": completed_iterations, + "global_rank": global_rank, + "completed_training_iteration": completed_training_iteration, + "next_training_iteration": next_training_iteration, + "num_microbatches": num_microbatches, + "sample_size": sample_size, + "output_directory": output_directory, + "forward_step_count": forward_step_count, + } + + +def _emit_model_weight_summary( + model, + step_count: int, + *, + is_training: bool = True, + audit_context=_MODEL_WEIGHT_AUDIT_CONTEXT_UNSET, +) -> None: + """Write one rank-0 sampled weight summary at requested training steps.""" + context = ( + _model_weight_audit_context(step_count, is_training=is_training) + if audit_context is _MODEL_WEIGHT_AUDIT_CONTEXT_UNSET + else audit_context + ) + if context is None: + return + completed_iterations = context["completed_iterations"] + global_rank = context["global_rank"] + completed_training_iteration = context["completed_training_iteration"] + next_training_iteration = context["next_training_iteration"] + num_microbatches = context["num_microbatches"] + sample_size = context["sample_size"] + output_directory = context["output_directory"] + forward_step_count = context["forward_step_count"] + if ( + completed_training_iteration not in completed_iterations + or completed_training_iteration in _EMITTED_MODEL_WEIGHT_ITERATIONS + ): + return + + output_directory.mkdir(parents=True, exist_ok=True) + output = output_directory / f"completed_iteration_{completed_training_iteration:07d}.json" + + payload = { + "version": 1, + "global_rank": global_rank, + "completed_training_iteration": completed_training_iteration, + "next_training_iteration": next_training_iteration, + "forward_step_count": forward_step_count, + # Audit convention: zero means the first eligible audit forward for + # this completed iteration. It is not arithmetic on cumulative calls. + "microbatch_index": 0, + "num_microbatches": num_microbatches, + "sample_size_per_parameter": sample_size, + **_sample_model_weights(model, sample_size), + } + try: + _write_model_weight_summary_once(output, payload) + except BaseException: + _EMITTED_MODEL_WEIGHT_ITERATIONS.discard(completed_training_iteration) + raise + + _EMITTED_MODEL_WEIGHT_ITERATIONS.add(completed_training_iteration) + logger.info( + "PRIMUS_MODEL_WEIGHT_SUMMARY=%s", + json.dumps( + { + "path": str(output), + "sample_count": payload["sample_count"], + "sample_finite": payload["sample_finite"], + "completed_training_iteration": completed_training_iteration, + "next_training_iteration": next_training_iteration, + "forward_step_count": forward_step_count, + }, + sort_keys=True, + allow_nan=False, + ), + ) + + +def _emit_model_weight_summary_after_forward( + model_output, + model, + step_count: int, + *, + is_training: bool = True, + audit_context=_MODEL_WEIGHT_AUDIT_CONTEXT_UNSET, +): + """Audit an already-evaluated model forward and return its output unchanged. + + Using this helper as a wrapper around ``model(...)`` makes Python finish the + model call and all forward pre-hooks before auditing. This is required for + current distributed-optimizer buffers when ``overlap_param_gather`` is on. + The real call site carries its pre-forward validated context into this + wrapper, which still runs before backward and the optimizer update. + """ + _emit_model_weight_summary( + model, + step_count, + is_training=is_training, + audit_context=audit_context, + ) + return model_output + + def prepare_flux_latents( latents: torch.Tensor, scheduler, @@ -198,6 +625,37 @@ def prepare_flux_latents( _eager_prepare_flux_latents = prepare_flux_latents +def _is_validation_forward(model, _batch=None) -> bool: + """Classify from rank-consistent model state, not rank-local batch data.""" + return not model.training + + +def _pregenerated_diffusion_inputs( + batch, + *, + tp_size, + is_validation, + compute_dtype, + tensor_parallel, +): + """Return deterministic inputs without splitting TP collective participation.""" + if tp_size == 1: + return batch.get("noise"), batch.get("timesteps") + if not is_validation: + return None, None + + # Evaluation mode is rank-consistent, so every TP rank enters this + # collective even though only TP rank zero owns ``batch``. + batch_timesteps = tensor_parallel.broadcast_data( + ["timesteps"], + batch, + compute_dtype, + ).get("timesteps") + if not batch_timesteps.is_cuda: + batch_timesteps = batch_timesteps.cuda(non_blocking=True) + return None, batch_timesteps + + def flux_forward_step_func( data_iterator, model, @@ -257,9 +715,10 @@ def flux_forward_step_func( to isolate training random ops from model forward RNG consumption (default: False). step_count: Monotonically increasing counter identifying this forward - call. Used to derive a unique per-step RNG seed. Managed by the - caller (DiffusionPretrainTrainer) and reconstructed from checkpoint - state on resume as iteration * num_microbatches. + call. Used only to derive a per-step RNG seed and as audit + provenance; training-loop coordinates come from Megatron iteration + state. Managed by the caller (DiffusionPretrainTrainer) and + reconstructed on resume as iteration * num_microbatches. Returns: Tuple of (noise_pred, clean_latents, noise, loss_mask, metrics_dict, is_validation) @@ -268,7 +727,7 @@ def flux_forward_step_func( - noise: Sampled noise [B, C, H, W] - loss_mask: Optional mask for variable-length sequences [B] or None - metrics_dict: Dictionary with training metrics - - is_validation: True when batch contains "timestep" key (MLPerf validation mode) + - is_validation: True when the model is in evaluation mode """ # Reseed default CUDA generator per step to isolate training random ops # (noise, timesteps, CFG dropout) from model forward RNG consumption. @@ -486,19 +945,28 @@ 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 + # Evaluation mode is identical on every tensor-parallel rank, including + # ranks where ``batch`` is intentionally None. Batch-local classification + # can split ranks before model collectives when TP > 1. + is_validation = _is_validation_forward(model, batch) + if batch is not None and is_validation and "timestep" in batch: + val_timesteps = batch["timestep"].to(dtype=compute_dtype) / 8.0 batch["timesteps"] = val_timesteps - elif batch is not None and not model.training: - is_validation = True + elif batch is not None and is_validation: 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 + # Validate every rank before noise preparation and the expensive model + # forward, then carry rank zero's parsed context through publication. + # Direct emitter calls build and validate the same context themselves. + model_weight_audit_context = _model_weight_audit_context( + step_count, + is_training=not is_validation, + ) + if batch is not None: _emit_batch_fingerprint( batch, @@ -537,23 +1005,13 @@ def flux_forward_step_func( ) # Extract pre-generated noise/timesteps from batch (deterministic tests) - batch_noise = None - batch_timesteps = None - if batch is not None: - if tp_size == 1: - batch_noise = batch.get("noise") - batch_timesteps = batch.get("timesteps") - else: - if "noise" in batch: - batch_noise = tensor_parallel.broadcast_data(["noise"], batch, compute_dtype).get("noise") - if not batch_noise.is_cuda: - batch_noise = batch_noise.cuda(non_blocking=True) - if "timesteps" in batch: - batch_timesteps = tensor_parallel.broadcast_data(["timesteps"], batch, compute_dtype).get( - "timesteps" - ) - if not batch_timesteps.is_cuda: - batch_timesteps = batch_timesteps.cuda(non_blocking=True) + batch_noise, batch_timesteps = _pregenerated_diffusion_inputs( + batch, + tp_size=tp_size, + is_validation=is_validation, + compute_dtype=compute_dtype, + tensor_parallel=tensor_parallel, + ) # Prepare latents (noise, packing, scheduling). # Eager wrapper — see _eager_prepare_flux_latents NOTE for why compile @@ -619,14 +1077,23 @@ def flux_forward_step_func( timesteps_norm = sigma_1d.to(dtype=packed_noisy_latents.dtype) with torch.amp.autocast("cuda", enabled=True, dtype=compute_dtype): - noise_pred = model( - img=packed_noisy_latents, - txt=prompt_embeds, - y=pooled_prompt_embeds, - timesteps=timesteps_norm, - img_ids=img_ids, - txt_ids=text_ids, - guidance=guidance_vec, + # The model call is evaluated before the wrapper. With + # overlap_param_gather, this guarantees every forward pre-hook has + # refreshed its distributed-optimizer parameter buffer before sampling. + noise_pred = _emit_model_weight_summary_after_forward( + model( + img=packed_noisy_latents, + txt=prompt_embeds, + y=pooled_prompt_embeds, + timesteps=timesteps_norm, + img_ids=img_ids, + txt_ids=text_ids, + guidance=guidance_vec, + ), + model, + step_count, + is_training=not is_validation, + audit_context=model_weight_audit_context, ) # Unpack latents from sequence format 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..b4dcde733 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.""" + """Evaluation mode marks a pre-sampled timestep batch as validation.""" + model.eval() batch = self._make_presampled_batch() batch["timestep"] = torch.arange(2) data_iterator = iter([batch]) @@ -176,6 +177,8 @@ def test_validation_with_timestep_key(self, model, scheduler): assert is_validation is True # The forward step writes derived timesteps (timestep / 8.0) into the batch. assert "timesteps" in batch + expected_dtype = torch.bfloat16 if model.config.bf16 else model.config.params_dtype + assert batch["timesteps"].dtype == expected_dtype assert torch.equal( batch["timesteps"].float().cpu(), torch.arange(2).float() / 8.0, diff --git a/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py b/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py index 6bfb815d1..0aa028d18 100644 --- a/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py +++ b/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py @@ -1,7 +1,11 @@ import json +import os +import threading +from concurrent.futures import ThreadPoolExecutor from types import SimpleNamespace import pytest +import torch import torch.nn as nn from primus.backends.megatron.data.diffusion.task_encoders.image import ( @@ -15,11 +19,64 @@ ) from primus.backends.megatron.patches.mlperf_warmup_patches import _is_resumed_training from primus.backends.megatron.training.diffusion.forward_step import ( + _EMITTED_MODEL_WEIGHT_ITERATIONS, _emit_batch_fingerprint, + _emit_model_weight_summary, + _emit_model_weight_summary_after_forward, + _encode_strict_json, + _is_validation_forward, + _model_weight_audit_context, + _model_weight_iteration_coordinate, + _parse_model_weight_steps, + _pregenerated_diffusion_inputs, + _read_strict_json, + _sample_model_weights, + _write_model_weight_summary_once, ) _FORWARD_STEP_LOGGER = "primus.backends.megatron.training.diffusion.forward_step" _FLUX_TRAINER_LOGGER = "primus.backends.megatron.flux_pretrain_trainer" +_UNSET = object() + + +@pytest.fixture(autouse=True) +def _clear_emitted_model_weight_iterations(): + _EMITTED_MODEL_WEIGHT_ITERATIONS.clear() + yield + _EMITTED_MODEL_WEIGHT_ITERATIONS.clear() + + +def _set_training_state( + monkeypatch, + *, + iteration, + train_iters=20_000, + curr_iteration=_UNSET, + num_microbatches=1, +): + from megatron import training as megatron_training + from megatron.core import num_microbatches_calculator + + args = SimpleNamespace(iteration=iteration, train_iters=train_iters) + if curr_iteration is not _UNSET: + args.curr_iteration = curr_iteration + microbatch_state = {"value": num_microbatches} + monkeypatch.setattr(megatron_training, "get_args", lambda: args) + monkeypatch.setattr( + num_microbatches_calculator, + "get_num_microbatches", + lambda: microbatch_state["value"], + ) + return args, microbatch_state + + +def _enable_weight_audit(monkeypatch, output, *, steps, sample_size=3, rank=0): + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False) + monkeypatch.delenv("PRIMUS_SYNTHETIC_WARMUP_ACTIVE", raising=False) + monkeypatch.setenv("RANK", str(rank)) + monkeypatch.setenv("PRIMUS_AUDIT_MODEL_WEIGHT_STEPS", steps) + monkeypatch.setenv("PRIMUS_AUDIT_MODEL_WEIGHT_SAMPLE_SIZE", str(sample_size)) + monkeypatch.setenv("PRIMUS_AUDIT_MODEL_WEIGHT_PATH", str(output)) def test_sample_key_fingerprint_is_order_sensitive(): @@ -77,12 +134,13 @@ def test_precision_linear_class_census_emits_zero_counts_for_bf16(monkeypatch, c _emit_precision_linear_class_census(nn.Linear(2, 2)) - line = next( + marker_lines = [ record.message for record in caplog.records - if record.message.startswith("PRIMUS_LINEAR_CLASS_CENSUS=") - ) - payload = json.loads(line.split("=", 1)[1]) + if record.name == _FLUX_TRAINER_LOGGER and record.message.startswith("PRIMUS_LINEAR_CLASS_CENSUS=") + ] + assert len(marker_lines) == 1 + payload = json.loads(marker_lines[0].split("=", 1)[1]) assert payload == { "data_parallel_rank": 3, "global_rank": 3, @@ -206,3 +264,625 @@ def test_emit_batch_fingerprint_logs_rank_local_payload(monkeypatch, caplog): "sample_keys_sha256": "a" * 64, "step": 6, } + + +def test_parse_model_weight_steps_is_fail_closed(): + assert _parse_model_weight_steps("5120, 8192") == {5120, 8192} + with pytest.raises(RuntimeError, match="duplicate step 5120"): + _parse_model_weight_steps("5120,5120") + with pytest.raises(RuntimeError, match="positive integers"): + _parse_model_weight_steps("5120,not-a-step") + for invalid in ("", "0", "-1", "1,"): + with pytest.raises(RuntimeError): + _parse_model_weight_steps(invalid) + + +def test_sample_model_weights_is_deterministic_and_sensitive(): + model = nn.Sequential(nn.Linear(4, 3, bias=True), nn.Linear(3, 2, bias=False)) + with torch.no_grad(): + for index, parameter in enumerate(model.parameters()): + parameter.copy_( + torch.arange(parameter.numel(), dtype=parameter.dtype).reshape_as(parameter) + index + ) + + first = _sample_model_weights(model, sample_size=4) + second = _sample_model_weights(model, sample_size=4) + assert first == second + assert first["sample_finite"] + assert first["parameter_count"] == 3 + + with torch.no_grad(): + model[0].weight.flatten()[0].add_(1) + changed = _sample_model_weights(model, sample_size=4) + first_by_name = {item["name"]: item for item in first["parameters"]} + changed_by_name = {item["name"]: item for item in changed["parameters"]} + assert changed_by_name["0.weight"]["sample_sha256"] != first_by_name["0.weight"]["sample_sha256"] + + +def test_sample_model_weights_serializes_nonfinite_values_as_strict_json(): + model = nn.Linear(4, 2, bias=False) + with torch.no_grad(): + model.weight.flatten()[0] = torch.nan + model.weight.flatten()[2] = torch.inf + + summary = _sample_model_weights(model, sample_size=4) + + assert summary["sample_finite"] is False + assert summary["sample_nonfinite_count"] == 2 + assert summary["sample_sum"] is None + assert summary["sample_sum_squares"] is None + assert summary["sample_absmax"] is None + json.dumps(summary, allow_nan=False) + + +def test_model_weight_iteration_coordinate_uses_canonical_megatron_iteration(monkeypatch): + args, microbatches = _set_training_state( + monkeypatch, + iteration=5120, + num_microbatches=4, + ) + + assert _model_weight_iteration_coordinate() == (5120, 5121, 4, 20_000) + + microbatches["value"] = 8 + assert _model_weight_iteration_coordinate() == (5120, 5121, 8, 20_000) + + # The pinned Megatron loop keeps args.iteration at the restored checkpoint + # and advances args.curr_iteration immediately before each train_step. + args.curr_iteration = 8192 + assert _model_weight_iteration_coordinate() == (8192, 8193, 8, 20_000) + + +@pytest.mark.parametrize("iteration", [None, True, -1, 1.5]) +def test_model_weight_iteration_coordinate_rejects_invalid_restored_iteration(monkeypatch, iteration): + _set_training_state(monkeypatch, iteration=iteration) + + with pytest.raises(RuntimeError, match="args.iteration"): + _model_weight_iteration_coordinate() + + +@pytest.mark.parametrize("num_microbatches", [None, True, 0, -1, 1.5]) +def test_model_weight_iteration_coordinate_rejects_invalid_microbatch_metadata(monkeypatch, num_microbatches): + _set_training_state( + monkeypatch, + iteration=0, + num_microbatches=num_microbatches, + ) + + with pytest.raises(RuntimeError, match="positive integer"): + _model_weight_iteration_coordinate() + + +@pytest.mark.parametrize("curr_iteration", [None, True, -1, 1.5]) +def test_model_weight_iteration_coordinate_rejects_invalid_active_iteration(monkeypatch, curr_iteration): + _set_training_state( + monkeypatch, + iteration=0, + curr_iteration=curr_iteration, + ) + + with pytest.raises(RuntimeError, match="args.curr_iteration"): + _model_weight_iteration_coordinate() + + +def test_emit_model_weight_summary_uses_resume_coordinate_and_suppresses_replays( + monkeypatch, tmp_path, caplog +): + output = tmp_path / "weights" + model = nn.Linear(4, 2) + _, microbatches = _set_training_state( + monkeypatch, + iteration=5120, + num_microbatches=4, + ) + _enable_weight_audit(monkeypatch, output, steps="5120,8192") + caplog.set_level("INFO", logger=_FORWARD_STEP_LOGGER) + + # step_count is deliberately unrelated to the training-loop coordinate. + _emit_model_weight_summary(model, step_count=999_999) + + record_path = output / "completed_iteration_0005120.json" + record = json.loads(record_path.read_text()) + assert record["completed_training_iteration"] == 5120 + assert record["next_training_iteration"] == 5121 + assert record["forward_step_count"] == 999_999 + assert record["microbatch_index"] == 0 + assert record["num_microbatches"] == 4 + assert record["sample_size_per_parameter"] == 3 + assert record["sample_finite"] + assert any(record.message.startswith("PRIMUS_MODEL_WEIGHT_SUMMARY=") for record in caplog.records) + + original = record_path.read_bytes() + microbatches["value"] = 8 + _emit_model_weight_summary(model, step_count=1) + assert record_path.read_bytes() == original + assert list(output.glob("*.json")) == [record_path] + assert not list(output.glob("*.tmp")) + + +def test_emit_model_weight_summary_uses_active_iteration_after_resume(monkeypatch, tmp_path): + output = tmp_path / "weights" + model = nn.Linear(2, 2) + _set_training_state( + monkeypatch, + iteration=5120, + curr_iteration=8192, + num_microbatches=4, + ) + _enable_weight_audit(monkeypatch, output, steps="8192") + + _emit_model_weight_summary(model, step_count=13) + + record = json.loads((output / "completed_iteration_0008192.json").read_text()) + assert record["completed_training_iteration"] == 8192 + assert record["next_training_iteration"] == 8193 + assert record["forward_step_count"] == 13 + + +def test_emit_model_weight_summary_rejects_terminal_requested_iteration(monkeypatch, tmp_path): + output = tmp_path / "weights" + model = nn.Linear(2, 2) + _set_training_state( + monkeypatch, + iteration=5120, + train_iters=8192, + ) + _enable_weight_audit(monkeypatch, output, steps="5120,8192") + + with pytest.raises(RuntimeError, match="iteration 8192 requires training through at least 8193"): + _emit_model_weight_summary(model, step_count=5121) + assert not output.exists() + + +@pytest.mark.parametrize("train_iters", [None, True, 0, -1, 1.5]) +def test_emit_model_weight_summary_rejects_invalid_train_iters(monkeypatch, tmp_path, train_iters): + output = tmp_path / "weights" + _set_training_state( + monkeypatch, + iteration=0, + train_iters=train_iters, + ) + _enable_weight_audit(monkeypatch, output, steps="1") + + with pytest.raises(RuntimeError, match="args.train_iters"): + _emit_model_weight_summary(nn.Linear(2, 2), step_count=1) + assert not output.exists() + + +def test_emit_model_weight_summary_restart_allows_changed_replay_provenance(monkeypatch, tmp_path): + output = tmp_path / "weights" + model = nn.Linear(2, 2) + _, microbatches = _set_training_state(monkeypatch, iteration=5120) + _enable_weight_audit(monkeypatch, output, steps="5120") + + _emit_model_weight_summary(model, step_count=5121) + _EMITTED_MODEL_WEIGHT_ITERATIONS.clear() + original = (output / "completed_iteration_0005120.json").read_bytes() + microbatches["value"] = 4 + _emit_model_weight_summary(model, step_count=7) + + assert (output / "completed_iteration_0005120.json").read_bytes() == original + assert not list(output.glob("*.tmp")) + + +def test_emit_model_weight_summary_rejects_conflicting_restart(monkeypatch, tmp_path): + output = tmp_path / "weights" + model = nn.Linear(2, 2) + _set_training_state(monkeypatch, iteration=5120) + _enable_weight_audit(monkeypatch, output, steps="5120") + + _emit_model_weight_summary(model, step_count=5121) + _EMITTED_MODEL_WEIGHT_ITERATIONS.clear() + with torch.no_grad(): + model.weight.flatten()[0].add_(1) + + with pytest.raises(RuntimeError, match="already exists with different content"): + _emit_model_weight_summary(model, step_count=5121) + assert not list(output.glob("*.tmp")) + + +def test_emit_model_weight_summary_requires_rank_zero(monkeypatch, tmp_path): + class TraversalForbidden: + def named_parameters(self): + raise AssertionError("rank one must not traverse parameters") + + output = tmp_path / "weights" + _set_training_state(monkeypatch, iteration=5120) + _enable_weight_audit(monkeypatch, output, steps="5120", rank=1) + + _emit_model_weight_summary(TraversalForbidden(), step_count=5121) + + assert not output.exists() + + +def test_rank_one_rejects_terminal_selection_before_model_traversal(monkeypatch, tmp_path): + class TraversalForbidden: + def named_parameters(self): + raise AssertionError("terminal validation must precede model traversal") + + output = tmp_path / "weights" + _set_training_state( + monkeypatch, + iteration=5120, + train_iters=8192, + ) + _enable_weight_audit(monkeypatch, output, steps="8192", rank=1) + + with pytest.raises(RuntimeError, match="iteration 8192 requires training through at least 8193"): + _emit_model_weight_summary(TraversalForbidden(), step_count=5121) + assert not output.exists() + + +def test_validation_forward_classification_is_batch_independent_across_tp_ranks(): + model = nn.Linear(2, 2) + model.eval() + + # Loader and non-loader tensor-parallel ranks see a batch and None, + # respectively, but both classify the model's eval forward identically. + assert _is_validation_forward(model, {"timestep": torch.tensor([0])}) is True + assert _is_validation_forward(model, None) is True + + model.train() + assert _is_validation_forward(model, {"timestep": torch.tensor([0])}) is False + assert _is_validation_forward(model, None) is False + + +@pytest.mark.parametrize("owns_batch", [True, False]) +def test_tp_validation_timesteps_broadcast_on_every_rank(owns_batch): + timestep = SimpleNamespace(is_cuda=True) + + class FakeTensorParallel: + def __init__(self): + self.calls = [] + + def broadcast_data(self, keys, batch, dtype): + self.calls.append((keys, batch, dtype)) + return {"timesteps": timestep} + + tensor_parallel = FakeTensorParallel() + batch = {"timesteps": timestep} if owns_batch else None + noise, observed_timesteps = _pregenerated_diffusion_inputs( + batch, + tp_size=2, + is_validation=True, + compute_dtype=torch.bfloat16, + tensor_parallel=tensor_parallel, + ) + + assert noise is None + assert observed_timesteps is timestep + assert tensor_parallel.calls == [ + (["timesteps"], batch, torch.bfloat16), + ] + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("sample", "0", "SAMPLE_SIZE"), + ("sample", "4097", "SAMPLE_SIZE"), + ("path", None, "PATH is required"), + ("path", "relative/weights", "PATH must be absolute"), + ("step_count", 0, "step_count"), + ("step_count", True, "step_count"), + ], +) +def test_rank_one_rejects_bad_audit_configuration_before_model_traversal( + monkeypatch, tmp_path, field, value, message +): + class TraversalForbidden: + def named_parameters(self): + raise AssertionError("configuration validation must precede model traversal") + + output = tmp_path / "weights" + _set_training_state(monkeypatch, iteration=1, train_iters=2) + _enable_weight_audit(monkeypatch, output, steps="1", rank=1) + step_count = 1 + if field == "sample": + monkeypatch.setenv("PRIMUS_AUDIT_MODEL_WEIGHT_SAMPLE_SIZE", value) + elif field == "path": + if value is None: + monkeypatch.delenv("PRIMUS_AUDIT_MODEL_WEIGHT_PATH") + else: + monkeypatch.setenv("PRIMUS_AUDIT_MODEL_WEIGHT_PATH", value) + else: + step_count = value + + with pytest.raises(RuntimeError, match=message): + _emit_model_weight_summary(TraversalForbidden(), step_count=step_count) + assert not output.exists() + + +def test_rank_zero_reuses_prevalidated_sample_size_and_output_path(monkeypatch, tmp_path): + output = tmp_path / "weights" + model = nn.Linear(2, 2) + _set_training_state(monkeypatch, iteration=1, train_iters=2) + _enable_weight_audit(monkeypatch, output, steps="1", sample_size=2) + context = _model_weight_audit_context(17, is_training=True) + + monkeypatch.setenv("PRIMUS_AUDIT_MODEL_WEIGHT_SAMPLE_SIZE", "invalid-after-preflight") + monkeypatch.setenv("PRIMUS_AUDIT_MODEL_WEIGHT_PATH", "relative/after-preflight") + _emit_model_weight_summary( + model, + step_count=17, + audit_context=context, + ) + + record = json.loads((output / "completed_iteration_0000001.json").read_text()) + assert record["sample_size_per_parameter"] == 2 + assert record["forward_step_count"] == 17 + + +def test_emit_model_weight_summary_env_unset_does_not_traverse_model(monkeypatch): + class TraversalForbidden: + def named_parameters(self): + raise AssertionError("model traversal must stay behind the opt-in gate") + + monkeypatch.delenv("PRIMUS_AUDIT_MODEL_WEIGHT_STEPS", raising=False) + + def distributed_state_forbidden(): + raise AssertionError("distributed state must stay behind the opt-in gate") + + monkeypatch.setattr(torch.distributed, "is_initialized", distributed_state_forbidden) + _emit_model_weight_summary(TraversalForbidden(), step_count=1) + + +def test_emit_model_weight_summary_skips_synthetic_warmup_before_traversal(monkeypatch, tmp_path): + class TraversalForbidden: + def named_parameters(self): + raise AssertionError("synthetic warmup must not traverse parameters") + + output = tmp_path / "weights" + _enable_weight_audit(monkeypatch, output, steps="1") + monkeypatch.setenv("PRIMUS_SYNTHETIC_WARMUP_ACTIVE", "1") + + _emit_model_weight_summary(TraversalForbidden(), step_count=1) + assert not output.exists() + + +def test_emit_model_weight_summary_skips_validation_before_traversal(monkeypatch, tmp_path): + class TraversalForbidden: + def named_parameters(self): + raise AssertionError("validation must not traverse parameters") + + output = tmp_path / "weights" + _enable_weight_audit(monkeypatch, output, steps="1") + + _emit_model_weight_summary(TraversalForbidden(), step_count=1, is_training=False) + assert not output.exists() + + +def test_emit_model_weight_summary_does_not_mutate_cpu_rng(monkeypatch, tmp_path): + output = tmp_path / "weights" + model = nn.Linear(4, 2) + _set_training_state(monkeypatch, iteration=1, train_iters=2) + _enable_weight_audit(monkeypatch, output, steps="1") + + def cuda_seed_forbidden(*_args, **_kwargs): + raise AssertionError("weight auditing must not reseed CUDA RNG") + + monkeypatch.setattr(torch.cuda, "manual_seed", cuda_seed_forbidden) + torch.manual_seed(1234) + before = torch.random.get_rng_state().clone() + + _emit_model_weight_summary(model, step_count=987) + + assert torch.equal(torch.random.get_rng_state(), before) + + +def test_model_weight_summary_observes_forward_pre_hook_weight(monkeypatch, tmp_path): + output = tmp_path / "weights" + model = nn.Linear(1, 1, bias=False) + with torch.no_grad(): + model.weight.fill_(1) + + def refresh_parameter(module, _inputs): + with torch.no_grad(): + module.weight.fill_(7) + + model.register_forward_pre_hook(refresh_parameter) + _set_training_state(monkeypatch, iteration=5, train_iters=6) + _enable_weight_audit(monkeypatch, output, steps="5", sample_size=1) + context = _model_weight_audit_context(123, is_training=True) + + model_output = _emit_model_weight_summary_after_forward( + model(torch.ones(1, 1)), + model, + step_count=123, + audit_context=context, + ) + + record = json.loads((output / "completed_iteration_0000005.json").read_text()) + assert model_output.item() == 7 + assert record["sample_sum"] == 7 + assert record["parameters"][0]["sample_sum"] == 7 + + +def _publish_and_capture(output, payload): + try: + _write_model_weight_summary_once(output, payload) + except BaseException as error: + return error + return None + + +def test_atomic_concurrent_identical_publication_is_idempotent(monkeypatch, tmp_path): + output = tmp_path / "summary.json" + payload = {"version": 1, "value": 7} + barrier = threading.Barrier(2) + real_link = os.link + + def racing_link(*args, **kwargs): + barrier.wait(timeout=5) + return real_link(*args, **kwargs) + + monkeypatch.setattr(os, "link", racing_link) + with ThreadPoolExecutor(max_workers=2) as executor: + results = list(executor.map(lambda _: _publish_and_capture(output, payload), range(2))) + + assert results == [None, None] + assert _read_strict_json(output) == payload + assert not list(tmp_path.glob("*.tmp")) + + +def test_atomic_concurrent_different_publication_fails_closed(monkeypatch, tmp_path): + output = tmp_path / "summary.json" + payloads = [{"version": 1, "value": 7}, {"version": 1, "value": 8}] + barrier = threading.Barrier(2) + real_link = os.link + + def racing_link(*args, **kwargs): + barrier.wait(timeout=5) + return real_link(*args, **kwargs) + + monkeypatch.setattr(os, "link", racing_link) + with ThreadPoolExecutor(max_workers=2) as executor: + results = list(executor.map(lambda payload: _publish_and_capture(output, payload), payloads)) + + assert sum(result is None for result in results) == 1 + errors = [result for result in results if result is not None] + assert len(errors) == 1 + assert isinstance(errors[0], RuntimeError) + assert "different content" in str(errors[0]) + assert _read_strict_json(output) in payloads + assert not list(tmp_path.glob("*.tmp")) + + +def test_restart_canonicalizes_key_order_and_whitespace(tmp_path): + output = tmp_path / "summary.json" + original = b'{\n "value": 7,\n "version": 1\n}\n' + output.write_bytes(original) + + _write_model_weight_summary_once(output, {"version": 1, "value": 7}) + + assert output.read_bytes() == original + assert not list(tmp_path.glob("*.tmp")) + + +@pytest.mark.parametrize( + ("existing", "payload"), + [ + ({"version": True, "global_rank": False}, {"version": 1, "global_rank": 0}), + ({"version": 1.0, "global_rank": 0}, {"version": 1, "global_rank": 0}), + ], +) +def test_restart_rejects_json_type_conflicts(tmp_path, existing, payload): + output = tmp_path / "summary.json" + output.write_bytes(_encode_strict_json(existing)) + + with pytest.raises(RuntimeError, match="different content"): + _write_model_weight_summary_once(output, payload) + assert not list(tmp_path.glob("*.tmp")) + + +@pytest.mark.parametrize( + "existing", + [ + {"version": 1, "forward_step_count": False, "num_microbatches": 1}, + {"version": 1, "forward_step_count": 7}, + ], +) +def test_restart_rejects_malformed_replay_provenance(tmp_path, existing): + output = tmp_path / "summary.json" + output.write_bytes(_encode_strict_json(existing)) + payload = { + "version": 1, + "forward_step_count": 7, + "num_microbatches": 1, + } + + with pytest.raises(RuntimeError, match="replay provenance|positive integer"): + _write_model_weight_summary_once(output, payload) + assert not list(tmp_path.glob("*.tmp")) + + +@pytest.mark.parametrize("encoded", [b'{"value":NaN}\n', b'{"value":1e9999}\n']) +def test_read_strict_json_rejects_nonfinite_payloads(tmp_path, encoded): + output = tmp_path / "summary.json" + output.write_bytes(encoded) + + with pytest.raises(ValueError): + _read_strict_json(output) + + +def test_read_strict_json_closes_descriptor_on_parse_failure(monkeypatch, tmp_path): + output = tmp_path / "summary.json" + output.write_text("{not-json}\n") + real_open = os.open + opened_descriptors = [] + + def capture_open(*args, **kwargs): + descriptor = real_open(*args, **kwargs) + opened_descriptors.append(descriptor) + return descriptor + + monkeypatch.setattr(os, "open", capture_open) + with pytest.raises(json.JSONDecodeError): + _read_strict_json(output) + + assert len(opened_descriptors) == 1 + with pytest.raises(OSError): + os.fstat(opened_descriptors[0]) + + +@pytest.mark.skipif(not hasattr(os, "O_NOFOLLOW"), reason="O_NOFOLLOW is unavailable") +def test_publication_rejects_symlink_target(tmp_path): + target = tmp_path / "target.json" + target.write_bytes(_encode_strict_json({"version": 1})) + output = tmp_path / "summary.json" + output.symlink_to(target) + + with pytest.raises(OSError): + _write_model_weight_summary_once(output, {"version": 1}) + assert output.is_symlink() + assert not list(tmp_path.glob("*.tmp")) + + +def test_publication_rejects_nonregular_target(tmp_path): + output = tmp_path / "summary.json" + output.mkdir() + + with pytest.raises((OSError, RuntimeError), match="regular file|directory"): + _write_model_weight_summary_once(output, {"version": 1}) + assert not list(tmp_path.glob("*.tmp")) + + +def test_file_fsync_failure_closes_descriptor_and_removes_temporary(monkeypatch, tmp_path): + output = tmp_path / "summary.json" + fsync_descriptors = [] + + def fail_file_fsync(descriptor): + fsync_descriptors.append(descriptor) + raise OSError("injected file fsync failure") + + monkeypatch.setattr(os, "fsync", fail_file_fsync) + with pytest.raises(OSError, match="file fsync"): + _write_model_weight_summary_once(output, {"version": 1}) + + assert len(fsync_descriptors) == 1 + with pytest.raises(OSError): + os.fstat(fsync_descriptors[0]) + assert not output.exists() + assert not list(tmp_path.glob("*.tmp")) + + +def test_directory_fsync_failure_closes_descriptor_and_removes_temporary(monkeypatch, tmp_path): + output = tmp_path / "summary.json" + real_fsync = os.fsync + fsync_descriptors = [] + + def fail_directory_fsync(descriptor): + fsync_descriptors.append(descriptor) + if len(fsync_descriptors) == 2: + raise OSError("injected directory fsync failure") + return real_fsync(descriptor) + + monkeypatch.setattr(os, "fsync", fail_directory_fsync) + with pytest.raises(OSError, match="directory fsync"): + _write_model_weight_summary_once(output, {"version": 1}) + + assert len(fsync_descriptors) == 2 + with pytest.raises(OSError): + os.fstat(fsync_descriptors[1]) + assert output.exists() + assert not list(tmp_path.glob("*.tmp")) From 67e52597d301bc7db347db2d0d9604bc299ba174 Mon Sep 17 00:00:00 2001 From: Guangpu Huang Date: Fri, 14 Aug 2026 06:41:06 +0000 Subject: [PATCH 14/14] fix(test): narrow audit helper exception handling Use Exception instead of BaseException in the concurrent publication test helper so the test still captures expected runtime failures without swallowing system-level termination signals. --- .../backends/megatron/test_diffusion_audit_markers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py b/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py index 0aa028d18..119d51512 100644 --- a/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py +++ b/tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py @@ -701,7 +701,7 @@ def refresh_parameter(module, _inputs): def _publish_and_capture(output, payload): try: _write_model_weight_summary_once(output, payload) - except BaseException as error: + except Exception as error: return error return None