From 14d1138c036f9a0f150d91a90606b6a861a5ee7a Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:20:42 -0700 Subject: [PATCH 01/62] Remove parameter sweep support The benchmark driver no longer expands list-valued config parameters into a cross product of runs: each `scaffold benchmark` invocation now runs exactly one worker in the resolved run directory, so the per-combination `param_set_i` subdirectories are gone and restart/checkpoint paths stay in one place. Config validation now always rejects list values for scalar keys, naming each offending key and its value ("problem_scale: parameter sweeps are no longer supported; got list [6, 7]") instead of the previous TypeErrors from `math.floor(list)` / `int - list` deep inside Config; load_config's config type "sweep" is renamed "benchmark". This supersedes findings R11, R12 and R16 of the round-2 review, which all stem from the half-finished sweep path. --- README.md | 6 +- ScaFFold/benchmark.py | 103 ++++++++------------------------- ScaFFold/cli.py | 11 ++-- ScaFFold/utils/config_utils.py | 37 +++++++----- tests/test_config.py | 91 +++++++++++++++++++++-------- 5 files changed, 119 insertions(+), 129 deletions(-) diff --git a/README.md b/README.md index f2827e8..3206b02 100644 --- a/README.md +++ b/README.md @@ -61,9 +61,9 @@ The model is trained from a random initialization until convergence, which is de ScaFFold benchmark training always uses PyTorch distributed execution with DistConv spatial parallelism. For a singleton run, launch one distributed rank rather than disabling distributed execution. -`benchmark` creates a folder for the benchmark run(s) at `base_run_dir` set in the config file. For reproducibility, we store a copy of the benchmark run config yml. Within each run subfolder, `benchmark` creates a yml config for that specific run. +Each `benchmark` invocation performs exactly one benchmark run, in a run folder created under `base_run_dir` set in the config file. Every run parameter must be single-valued; a list (e.g. `problem_scale: [6, 7]`) is rejected by name, since parameter sweeps are not supported. To compare parameter settings, launch one benchmark run per setting. For reproducibility, the run folder holds a copy of the benchmark config yml plus the fully merged `config.yaml` for that run. -After each run completes, statistics from the run are stored in `train_stats.csv`. Additionally, users can inspect plots of the training and validation losses over time in ` Date: Fri, 31 Jul 2026 15:21:01 -0700 Subject: [PATCH 02/62] Skip the final checkpoint save when a resume has nothing to train A resume whose checkpoint already covers config.epochs leaves the epoch loop before any epoch body runs, so val_loss_avg was never bound and the final-save block crashed with UnboundLocalError on every rank. train() now logs that there was nothing to resume and returns cleanly, leaving the existing checkpoint (which already covers those epochs) alone. Round-2 review: R01. --- ScaFFold/utils/trainer.py | 22 +++++++++- tests/test_resume.py | 84 ++++++++++++++++++++++++++++++++++++--- 2 files changed, 99 insertions(+), 7 deletions(-) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 108df71..29e7222 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -776,6 +776,11 @@ def train(self, profiler=None): # other ranks would make them disagree about whether to call # save_checkpoint on exit, deadlocking its internal collective. last_checkpoint_epoch = None + # Whether this invocation completed at least one NEW epoch. A resume + # whose checkpoint already covers config.epochs (or an --epochs lowered + # below the checkpointed epoch) leaves the loop at the max-epoch check + # before any epoch body runs, so none of the per-epoch metrics exist. + completed_new_epoch = False with open(self.outfile_path, "a", newline="") as outfile: start = time.time() while dice_score_train < self.config.target_dice: @@ -1030,6 +1035,7 @@ def train(self, profiler=None): dice_score_train = val_score epoch += 1 + completed_new_epoch = True # This check must exist otherwise the condition dice_score_train < self.config.target_dice will evaluate to False and incorrectly exit the training if math.isnan(dice_score_train): @@ -1039,12 +1045,26 @@ def train(self, profiler=None): completed_epochs = epoch - 1 + if not completed_new_epoch: + # The loop exited without running a single new epoch: the resumed + # checkpoint already covers every epoch this run was asked for. + # There is nothing new to save (the existing checkpoint already + # records epoch `completed_epochs`) and none of the per-epoch + # metrics the final save would write were ever computed, so skip + # it and return normally -- the caller's post-processing still has + # the CSV the original run left behind. + self.log.warning( + "Nothing to resume: the loaded checkpoint already covers epoch " + "%s, so no new epoch was trained and no checkpoint was written. " + "Increase 'epochs' (or lower 'target_dice') to train further.", + completed_epochs, + ) # Save a final checkpoint when the run exits (convergence or max epochs) # at an epoch that was not a checkpoint interval, so the converged # weights that produced the reported metrics are not lost. Skipped when # checkpointing is disabled, when no epoch completed, or when the last # completed epoch was already checkpointed inside the loop. - if ( + elif ( self.config.checkpoint_interval > 0 and completed_epochs >= 1 and last_checkpoint_epoch != completed_epochs diff --git a/tests/test_resume.py b/tests/test_resume.py index 209a6eb..93acdaf 100644 --- a/tests/test_resume.py +++ b/tests/test_resume.py @@ -28,6 +28,7 @@ from __future__ import annotations +import logging from pathlib import Path import numpy as np @@ -142,19 +143,27 @@ def cfg(): ) -def _stub_trainer(run_dir, *, train_from_scratch, log): +def _stub_trainer(run_dir, *, train_from_scratch, log, **config_overrides): """A PyTorchTrainer carrying only what cleanup_or_resume/train touch. Built via ``object.__new__`` (as the checkpointing tests do) so no dataset, model, or process group is needed. A real CheckpointManager over a tiny linear model backs the resume path so load/save round-trip faithfully. + + ``config_overrides`` set additional config fields (``epochs``, + ``target_dice``, ``checkpoint_interval``, ...) for tests that drive + ``train()``'s loop-entry and final-save control flow. """ t = object.__new__(PyTorchTrainer) - t.config = SimpleNamespace( - train_from_scratch=train_from_scratch, - run_dir=str(run_dir), - checkpoint_interval=-1, - ) + config_fields = { + "train_from_scratch": train_from_scratch, + "run_dir": str(run_dir), + "checkpoint_interval": -1, + "epochs": -1, + "target_dice": 0.95, + } + config_fields.update(config_overrides) + t.config = SimpleNamespace(**config_fields) t.world_rank = 0 t.global_step = 0 t.total_optimizer_steps = 0 @@ -346,6 +355,69 @@ def test_step_counters_roundtrip(tmp_path): assert t2.total_optimizer_steps == 37 +def test_restart_of_completed_run_trains_and_saves_nothing(tmp_path, caplog): + """Restarting a run whose checkpoint already covers ``epochs`` exits cleanly. + + The epoch loop breaks at the max-epoch check before any new epoch runs, so + no epoch metric exists to checkpoint: the final-save block must be skipped + (the last checkpoint already covers the completed epochs) rather than + saving with an unbound ``val_loss_avg``. ``train()`` has to return normally + so the worker's post-processing still runs off the CSV already on disk, and + the run must say plainly that there was nothing to resume. + """ + log = logging.getLogger("resume.r01") + run = tmp_path / "run" + run.mkdir() + + # Artifacts of a completed epochs=2 run: a checkpoint recording epoch 2 + # (written by the in-loop save) plus its CSV rows. + t1 = _stub_trainer( + run, + train_from_scratch=False, + log=log, + epochs=2, + checkpoint_interval=1, + ) + t1.checkpoint_manager.save_checkpoint( + epoch=2, + val_loss_avg=0.5, + extras={ + "train_mask_values": None, + "global_step": 8, + "total_optimizer_steps": 8, + }, + ) + csv = run / "train_stats.csv" + csv.write_text(_HEADER + "\n") + _write_rows(csv, [1, 2]) + + ckpt = t1.checkpoint_manager.last_ckpt_path + before = ckpt.read_bytes() + + # The user reruns the generated restart command against the same dir. + t2 = _stub_trainer( + run, + train_from_scratch=False, + log=log, + epochs=2, + checkpoint_interval=1, + ) + t2.config.restart = True + t2.cleanup_or_resume() + assert t2.start_epoch == 3 # past the last epoch: nothing left to train + + with caplog.at_level(logging.WARNING): + t2.train() # must not raise + + # Nothing new was written: the existing checkpoint is byte-identical (a + # fresh save would serialize this trainer's own randomly-initialized model) + # and the CSV still holds exactly the original run's epochs. + assert ckpt.read_bytes() == before + epochs = [ln.split(",")[0] for ln in csv.read_text().splitlines()[1:]] + assert epochs == ["1", "2"] + assert "nothing to resume" in caplog.text.lower() + + def test_total_steps_resume_predates_dedicated_key(tmp_path): """A checkpoint recording only global_step still resumes the step total. From 8ba619109c047f6574a09a79036d634c5eea14d7 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:23:08 -0700 Subject: [PATCH 03/62] Persist the achieved validation dice so a converged run resumes converged The training loop runs while the validation dice is below target_dice but that score was never checkpointed, so every restart of a converged epochs:-1 run re-trained one full epoch before rediscovering it had converged, inflating the epoch count and the FOM's total train time. The score now rides along in the checkpoint extras and seeds the loop variable on resume. Round-2 review: R06. --- ScaFFold/utils/trainer.py | 24 ++++++++++++++- tests/test_resume.py | 61 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 84 insertions(+), 1 deletion(-) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 29e7222..cfd27ee 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -117,6 +117,10 @@ def __init__(self, model, config, device, log): self.global_step = 0 self.total_optimizer_steps = 0 self.start_epoch = -1 + # Validation dice already achieved by the epoch we resume from (0 for a + # fresh run). The training loop's exit condition is a threshold on this + # score, so it has to survive a restart: see cleanup_or_resume. + self.start_val_dice = 0.0 self.ps = getattr(self.config, "_parallel_strategy", None) self.spatial_mesh = None # Spatial mesh for use w/ DistConv self.data_num_replicas = self.world_size @@ -387,6 +391,7 @@ def cleanup_or_resume(self): pass self.start_epoch = 1 + self.start_val_dice = 0.0 else: # Load checkpoint via manager. An explicit restart must find a # checkpoint; a plain non-scratch launch may simply start fresh. @@ -399,6 +404,16 @@ def cleanup_or_resume(self): if "train_mask_values" in restored: self.train_set.mask_values = restored["train_mask_values"] + # Resume the convergence state, not just the weights. The loop runs + # while the validation dice is below target_dice, so a run that had + # already converged must re-enter the loop with the score it + # achieved -- starting from 0 would unconditionally re-train (and + # re-log, and re-checkpoint) one full epoch before the condition is + # re-tested, inflating the epoch count and the FOM's total train + # time. Checkpoints written before this key existed simply fall + # back to 0, i.e. the old behaviour. + self.start_val_dice = restored.get("val_dice", 0.0) + # Continue the optimizer-step counts from where the checkpoint # left off; otherwise a resumed run restarts them at 0 and # undercounts all pre-resume work in the reported step totals. @@ -766,7 +781,10 @@ def train(self, profiler=None): """ epoch = self.start_epoch - dice_score_train = 0 + # Seeded from the resumed checkpoint (0 for a fresh run) so an already + # converged run exits the loop immediately instead of training an + # extra epoch to rediscover that it converged. + dice_score_train = self.start_val_dice epoch_minibatch_times_s = [] # Track the last epoch checkpointed inside the loop so the final-save # decision below is identical on every rank. The in-loop checkpoint @@ -1027,6 +1045,9 @@ def train(self, profiler=None): "train_mask_values": self.train_set.mask_values, "global_step": self.global_step, "total_optimizer_steps": self.total_optimizer_steps, + # The convergence state: what a resume must restore to + # know this epoch already met (or missed) target_dice. + "val_dice": val_score, } self.checkpoint_manager.save_checkpoint(epoch, val_loss_avg, extras) last_checkpoint_epoch = epoch @@ -1073,6 +1094,7 @@ def train(self, profiler=None): "train_mask_values": self.train_set.mask_values, "global_step": self.global_step, "total_optimizer_steps": self.total_optimizer_steps, + "val_dice": val_score, } self.checkpoint_manager.save_checkpoint( completed_epochs, val_loss_avg, extras diff --git a/tests/test_resume.py b/tests/test_resume.py index 93acdaf..3e8aa4d 100644 --- a/tests/test_resume.py +++ b/tests/test_resume.py @@ -133,6 +133,7 @@ def cfg(): from types import SimpleNamespace # noqa: E402 +import ScaFFold.utils.trainer as trainer_mod # noqa: E402 from ScaFFold.utils.checkpointing import CheckpointManager # noqa: E402 from ScaFFold.utils.trainer import PyTorchTrainer # noqa: E402 @@ -418,6 +419,66 @@ def test_restart_of_completed_run_trains_and_saves_nothing(tmp_path, caplog): assert "nothing to resume" in caplog.text.lower() +def test_converged_resume_does_not_retrain(tiny_trainer, monkeypatch): + """A restart of an already-converged run must not train another epoch. + + The validation dice the checkpointed epoch achieved is persisted with the + checkpoint and restored on resume, so the ``while dice_score_train < + target_dice`` loop is never entered. Without it a converged ``epochs: -1`` + run re-trains, re-logs and re-checkpoints one full epoch on every restart, + inflating ``sum(epoch_duration)`` -- the FOM denominator -- and the + reported epoch count relative to the same run left un-restarted. + """ + overrides = { + "checkpoint_interval": 1, + "epochs": -1, + "target_dice": 0.95, + "train_from_scratch": 0, + } + + def converged_evaluate(*args, **kwargs): + # Two validation samples at hard dice 0.96, i.e. above target. + return (0.96 * 2, 0.1 * 2, 0.1, 2, 2) + + monkeypatch.setattr(trainer_mod, "evaluate", converged_evaluate) + + def stub_batches(trainer): + """Replace the DistConv forward path and count the batches it runs.""" + calls = {"n": 0} + + def _step(batch, **kwargs): + calls["n"] += 1 + return 1, torch.tensor(0.1), torch.tensor(0.96) + + monkeypatch.setattr(trainer, "_run_training_batch", _step) + return calls + + # The original run: converges in epoch 1 and checkpoints it. + first = tiny_trainer(config_overrides=overrides) + first_calls = stub_batches(first) + first.cleanup_or_resume() + first.train() + assert first_calls["n"] > 0 # it really did train + + ckpt_path = first.checkpoint_manager.last_ckpt_path + saved = torch.load(ckpt_path, map_location="cpu", weights_only=False) + assert saved["epoch"] == 1 + csv = Path(first.outfile_path) + assert [ln.split(",")[0] for ln in csv.read_text().splitlines()[1:]] == ["1"] + + # The user reruns the restart command: nothing is left to train. + second = tiny_trainer(config_overrides=overrides) + second_calls = stub_batches(second) + second.cleanup_or_resume() + assert second.start_epoch == 2 + second.train() + + assert second_calls["n"] == 0, "a converged run re-trained an extra epoch" + assert [ln.split(",")[0] for ln in csv.read_text().splitlines()[1:]] == ["1"] + reloaded = torch.load(ckpt_path, map_location="cpu", weights_only=False) + assert reloaded["epoch"] == 1 + + def test_total_steps_resume_predates_dedicated_key(tmp_path): """A checkpoint recording only global_step still resumes the step total. From 7bb34e9be5be44a1908aad8e03f49a0f03bc5599 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:24:27 -0700 Subject: [PATCH 04/62] Reset the cached best-loss state when cleaning up for a fresh run cleanup(train_from_scratch=True) deleted the checkpoint files but kept best_val_loss and last_saved_epoch, so the fresh run's is_best decisions stayed gated by the deleted run's best and it wrote no best checkpoint until it beat a score nothing backed any more. Round-2 review: R02. --- ScaFFold/utils/checkpointing.py | 9 ++++++++ tests/test_checkpointing.py | 39 +++++++++++++++++++++++++++++++++ 2 files changed, 48 insertions(+) diff --git a/ScaFFold/utils/checkpointing.py b/ScaFFold/utils/checkpointing.py index 823f154..932f6d1 100644 --- a/ScaFFold/utils/checkpointing.py +++ b/ScaFFold/utils/checkpointing.py @@ -107,6 +107,15 @@ def cleanup(self, train_from_scratch: bool) -> None: self._barrier() return + # Drop the cached state that described the run being deleted. Both + # fields are seeded from disk (or a previous save), so keeping them + # would let a deleted run's best gate this run's is_best decisions -- + # the fresh run would then never write a best checkpoint until it beat + # a score no file backs any more, leaving it with no best-checkpoint + # fallback. + self.best_val_loss = math.inf + self.last_saved_epoch = None + if self.world_rank == 0: for p in (self.last_ckpt_path, self.best_ckpt_path): if p.exists(): diff --git a/tests/test_checkpointing.py b/tests/test_checkpointing.py index 4e1737d..d2d7183 100644 --- a/tests/test_checkpointing.py +++ b/tests/test_checkpointing.py @@ -26,6 +26,7 @@ from __future__ import annotations +import math import re import time from pathlib import Path @@ -244,6 +245,44 @@ def spy_load(path, *args, **kwargs): assert final_best["val_loss_avg"] == pytest.approx(0.2) +# --------------------------------------------------------------------------- +# R02 -- a from-scratch cleanup drops the deleted run's best, not just its files +# --------------------------------------------------------------------------- + + +def test_cleanup_from_scratch_resets_best(tmp_path): + """``cleanup(train_from_scratch=True)`` resets the cached best-loss state. + + The manager seeds ``best_val_loss`` from ``checkpoint_best.pth`` at + construction so a resumed run does not call its first epoch "best". When + the same directory is then wiped for a fresh run, that cached score + outlives the file it came from: every ``is_best`` decision of the new run + is gated by a deleted run's score, so no ``checkpoint_best.pth`` is written + until the retrain beats it -- leaving the run with no best-checkpoint + fallback at all. + """ + mgr, _ = _make_manager(tmp_path) + mgr.save_checkpoint(epoch=1, val_loss_avg=0.01) + assert mgr.best_ckpt_path.exists() + + # A driver reusing the run directory: the new manager seeds from disk. + mgr2, _ = _make_manager(tmp_path) + assert mgr2.best_val_loss == pytest.approx(0.01) + mgr2.save_checkpoint(epoch=2, val_loss_avg=0.9) + assert mgr2.last_saved_epoch == 2 + + mgr2.cleanup(train_from_scratch=True) + + assert not mgr2.last_ckpt_path.exists() + assert not mgr2.best_ckpt_path.exists() + assert mgr2.best_val_loss == math.inf + assert mgr2.last_saved_epoch is None + + # The fresh run's first epoch is its best, and a best checkpoint exists. + assert mgr2.save_checkpoint(epoch=1, val_loss_avg=0.5) is True + assert mgr2.best_ckpt_path.exists() + + # --------------------------------------------------------------------------- # F71 -- CPU tensors are cloned into the snapshot # --------------------------------------------------------------------------- From f9e5a8f1bc1cb065f550480a8a7e266377296ab4 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:25:57 -0700 Subject: [PATCH 05/62] Abort on non-finite epoch losses instead of checkpointing a diverged model The existing NaN check tests the hard-argmax dice, which stays finite even for an all-NaN model, so it could never fire; a diverged run kept looping below target while NaN weights overwrote checkpoint_last.pth. The reduced train and validation losses are now checked right after the data-parallel reductions -- identical on every rank, so all ranks raise together. Round-2 review: R03. --- ScaFFold/utils/trainer.py | 19 +++++++++++++++ tests/test_checkpointing.py | 48 +++++++++++++++++++++++++++++++++++++ 2 files changed, 67 insertions(+) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index cfd27ee..07fdd2c 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -970,6 +970,25 @@ def train(self, profiler=None): # Reduced sample-weighted total and per-sample mean val loss. val_loss_epoch = val_info[1].item() val_loss_avg = val_loss_epoch / global_val_samples + + # Divergence check. The dice score below is computed from a + # hard argmax, so it stays finite even for an all-NaN model + # (argmax of NaN logits is 0 and the one-hots are finite) -- + # the loss is the only value that actually goes non-finite. Bail + # out before the CSV row and the checkpoint: continuing would + # keep overwriting checkpoint_last.pth with NaN weights (which + # then poison the next restart) while the loop's dice threshold + # can never be met. Both values come out of the data-parallel + # reductions above, so they are identical on every rank and + # every rank raises here together, leaving no unmatched + # collective behind. + if not (math.isfinite(overall_loss) and math.isfinite(val_loss_avg)): + raise ValueError( + f"Non-finite loss at epoch {epoch} " + f"(train_loss={overall_loss}, val_loss={val_loss_avg}): " + "training diverged, aborting before checkpointing." + ) + if not self.config.disable_scheduler: self.scheduler.step() else: diff --git a/tests/test_checkpointing.py b/tests/test_checkpointing.py index d2d7183..40151f6 100644 --- a/tests/test_checkpointing.py +++ b/tests/test_checkpointing.py @@ -346,6 +346,54 @@ def fake_evaluate(*args, **kwargs): assert saved["epoch"] == 2 +# --------------------------------------------------------------------------- +# R03 -- a diverged epoch aborts instead of checkpointing NaN weights +# --------------------------------------------------------------------------- + + +def test_divergence_aborts_before_poisoning_checkpoint(tiny_trainer, monkeypatch): + """Non-finite epoch losses abort the run before any checkpoint is written. + + The dice score is computed from a hard argmax, so an all-NaN model still + produces a *finite* dice (argmax of NaN logits is 0 and the one-hots are + finite): the dice check can never fire on divergence. Left unguarded, a + diverged ``epochs: -1`` run keeps looping on a finite plateau below target + while every checkpoint interval overwrites ``checkpoint_last.pth`` with NaN + weights, poisoning the next ``--restart``. The reduced losses are the + values that actually go non-finite, and being reductions they are identical + on every rank, so the check fires on all ranks together. + """ + trainer = tiny_trainer( + config_overrides={ + "checkpoint_interval": 1, + "epochs": 3, + "target_dice": 0.95, + } + ) + + # A diverged step: NaN loss, and a dice that stays finite. + monkeypatch.setattr( + trainer, + "_run_training_batch", + lambda batch, **kw: (1, torch.tensor(float("nan")), torch.tensor(0.0)), + ) + + def diverged_evaluate(*args, **kwargs): + # Exactly what evaluate() returns for an all-NaN model: a tiny but + # finite hard-argmax dice sum alongside a NaN validation loss. + return (7.4e-10, float("nan"), float("nan"), 2, 2) + + monkeypatch.setattr(trainer_mod, "evaluate", diverged_evaluate) + + trainer.cleanup_or_resume() + with pytest.raises(ValueError, match="[Nn]on-finite"): + trainer.train() + + # The run died before the diverged epoch could be checkpointed. + assert not trainer.checkpoint_manager.last_ckpt_path.exists() + assert not trainer.checkpoint_manager.best_ckpt_path.exists() + + # --------------------------------------------------------------------------- # F49 -- GradScaler-skipped steps do not advance the optimizer-step counter # --------------------------------------------------------------------------- From c2083a06427dd7cfbc16cf9b83898da06190f00a Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:31:42 -0700 Subject: [PATCH 06/62] Broadcast the checkpoint write outcome so all ranks fail together wait_for_save now re-raises the writer's exception instead of logging and dropping it, train() consumes the run's final save so a failed async write can no longer exit 0 with no checkpoint, and save_checkpoint broadcasts rank 0's outcome (success or error sentinel, including a deferred async failure) before anyone raises -- previously rank 0 raised ahead of its own broadcast and left the peers in an unmatched collective reporting a transport error instead of the disk error. Round-2 review: R04, R05. --- ScaFFold/utils/checkpointing.py | 230 ++++++++++++++++++++++++-------- ScaFFold/utils/trainer.py | 8 ++ tests/test_checkpointing.py | 181 +++++++++++++++++++++++++ 3 files changed, 364 insertions(+), 55 deletions(-) diff --git a/ScaFFold/utils/checkpointing.py b/ScaFFold/utils/checkpointing.py index 932f6d1..5fedc8c 100644 --- a/ScaFFold/utils/checkpointing.py +++ b/ScaFFold/utils/checkpointing.py @@ -25,6 +25,16 @@ import torch.distributed as dist +class CheckpointSaveError(RuntimeError): + """A checkpoint write failed. + + Raised identically on every rank. Only rank 0 writes, but its outcome is + broadcast, so the peers report the real disk error instead of the + unmatched-collective symptom (an opaque gloo transport error, or an NCCL + watchdog timeout minutes later) that a rank-0-only raise produces. + """ + + class CheckpointManager: """ Checkpoint Manager for DDP/Single-Process. @@ -90,6 +100,10 @@ def __init__( # Async handling self.executor = None self.future = None + # The exception behind the most recently reported save failure, kept + # only so rank 0 can chain it (and its traceback) onto the + # CheckpointSaveError every rank raises. + self._save_error_exc: Optional[BaseException] = None if self.async_save and self.world_rank == 0: # We only need 1 worker for serializing writes self.executor = ThreadPoolExecutor(max_workers=1) @@ -99,44 +113,106 @@ def __init__( self.base_dir.mkdir(parents=True, exist_ok=True) def cleanup(self, train_from_scratch: bool) -> None: - """Clear existing checkpoints if training from scratch.""" - # Ensure any pending async saves are finished before deleting - self.wait_for_save() + """Clear existing checkpoints if training from scratch. + + Rank-symmetric, like every other collective point here: any pending + async write is drained and its outcome broadcast, so a failure raises + on all ranks together (see ``save_checkpoint``). + """ + # Ensure any pending async save is finished before deleting. + error = self._drain_pending_save() + + if train_from_scratch: + # Drop the cached state that described the run being deleted. Both + # fields are seeded from disk (or a previous save), so keeping them + # would let a deleted run's best gate this run's is_best decisions + # -- the fresh run would then never write a best checkpoint until + # it beat a score no file backs any more, leaving it with no + # best-checkpoint fallback. + self.best_val_loss = math.inf + self.last_saved_epoch = None + + if self.world_rank == 0: + self._remove_checkpoint_files() + + error = self._broadcast_obj(error) + self._barrier() + if error is not None: + self._raise_save_error(error) + + def _remove_checkpoint_files(self) -> None: + """Delete this run's checkpoint files (rank 0 only).""" + for p in (self.last_ckpt_path, self.best_ckpt_path): + if p.exists(): + try: + p.unlink() + self._log(f"Removed existing checkpoint: {p}") + except Exception as e: + self._log(f"Failed to remove {p}: {e}") - if not train_from_scratch: - self._barrier() + def wait_for_save(self): + """Block until the background save (if any) is complete. + + The writer's exception is re-raised rather than logged and dropped: a + save that could not complete has to surface, otherwise the manager + state (and the process exit code) claims a checkpoint that is not on + disk. The future is consumed exactly once, so the failure is reported + at exactly one point. + + Callers that are inside a collective region must not let this + propagate directly -- use ``_drain_pending_save`` instead, per the + collective invariant documented on ``save_checkpoint``. + """ + if self.future is None: return + # Clear the handle *before* blocking on it so a failed write is + # reported once and does not re-raise at some later, arbitrary point. + future, self.future = self.future, None + if not future.done(): + self._log("Waiting for background checkpoint save to complete...") + future.result() # Blocks and re-raises whatever the writer raised + + def _drain_pending_save(self) -> Optional[str]: + """Consume the in-flight async save, reporting failure without raising. + + Returns a description of the writer's failure, or ``None``. Only rank 0 + ever has a pending write, so raising here directly would leave the + peers blocked in the next collective; callers broadcast this result and + raise on every rank together. + """ + try: + self.wait_for_save() + except Exception as e: + self._save_error_exc = e + return f"{type(e).__name__}: {e}" + return None - # Drop the cached state that described the run being deleted. Both - # fields are seeded from disk (or a previous save), so keeping them - # would let a deleted run's best gate this run's is_best decisions -- - # the fresh run would then never write a best checkpoint until it beat - # a score no file backs any more, leaving it with no best-checkpoint - # fallback. - self.best_val_loss = math.inf - self.last_saved_epoch = None + def _raise_save_error(self, description: str) -> None: + """Raise a broadcast save failure on this rank. - if self.world_rank == 0: - for p in (self.last_ckpt_path, self.best_ckpt_path): - if p.exists(): - try: - p.unlink() - self._log(f"Removed existing checkpoint: {p}") - except Exception as e: - self._log(f"Failed to remove {p}: {e}") + Rank 0 chains the original exception so its traceback survives; the + peers never saw it and raise the same message on its own. + """ + cause, self._save_error_exc = self._save_error_exc, None + raise CheckpointSaveError( + f"Checkpoint save failed on rank 0: {description}" + ) from cause + + def finalize_saves(self) -> None: + """Consume the outcome of the run's last save before the run ends. + + Nothing touches the manager after the training loop, so an + asynchronous write that failed there would never be observed: the + process would exit successfully having written no checkpoint (or left + a stale one), the benchmark would report success, and a later + ``--restart`` would resume from the wrong epoch or fail its pre-check. + Rank-symmetric, like ``save_checkpoint``. + """ + error = self._drain_pending_save() + error = self._broadcast_obj(error) self._barrier() - - def wait_for_save(self): - """Blocks until the background save (if any) is complete.""" - if self.future is not None: - # check if running - if not self.future.done(): - self._log("Waiting for background checkpoint save to complete...") - try: - self.future.result() # Blocks and raises exceptions if any occurred - except Exception as e: - self._log(f"Background save failed with error: {e}") - self.future = None + if error is not None: + self._raise_save_error(error) def snapshot_training_state(self) -> Dict[str, Any]: """Capture mutable in-memory training state without writing a checkpoint.""" @@ -184,7 +260,10 @@ def load_from_checkpoint(self, require_checkpoint: bool = False) -> int: With ``require_checkpoint`` (an explicit ``--restart``), a missing checkpoint raises instead of silently starting over. """ - self.wait_for_save() # Safety: don't load while writing + # Safety: don't load while writing. A failure from that write is + # folded into the decision broadcast below rather than raised here, so + # rank 0 never abandons its peers inside the broadcast. + error = self._drain_pending_save() # 1. Rank 0 is the sole reader: it selects the newest readable # checkpoint and deserializes it once, then broadcasts the loaded @@ -193,10 +272,16 @@ def load_from_checkpoint(self, require_checkpoint: bool = False) -> int: # read of one (multi-GB) file from the shared filesystem on restart -- # a restart I/O storm that serializes on the parallel FS. Peer ranks # therefore never open the checkpoint files at all. - result = self._select_and_load() if self.world_rank == 0 else None + result = None + if self.world_rank == 0: + result = ( + ("save_failed", error) if error is not None else self._select_and_load() + ) status, payload = self._broadcast_obj(result) # 2. Every rank acts on the same decision rank 0 reached. + if status == "save_failed": + self._raise_save_error(payload) if status == "empty": if require_checkpoint: # An explicit restart must resume real state; silently @@ -310,22 +395,55 @@ def save_checkpoint( """ Save checkpoint. If async_save is True, this returns immediately after CPU transfer. + + Collective invariant: every rank posts exactly one ``_broadcast_obj`` + followed by exactly one ``_barrier`` on every path through this method, + failures included. Rank 0 is the sole writer, but it never raises + before those collectives: what gets broadcast is the write's *outcome* + -- ``is_best`` on success, or an error description on failure, + including a *previous* asynchronous write whose failure is surfaced + here, at the next collective point. Every rank then raises the same + ``CheckpointSaveError`` together. Raising on rank 0 before the + broadcast would strand the peers in an unmatched collective, where a + plain disk error resurfaces as a gloo transport error or an NCCL + watchdog timeout that hides the real cause. """ - is_best = False + # Non-zero ranks contribute nothing; their placeholder is overwritten + # by rank 0's outcome in the broadcast below (and is a harmless no-op + # in the degenerate non-distributed case). + outcome = ("ok", False) if self.world_rank == 0: + outcome = self._rank0_save(epoch, val_loss_avg, extras) + + status, payload = self._broadcast_obj(outcome) + + # Barrier: ensure Rank 0 has finished the "Snapshot" phase before anyone continues. + # Even in async mode, we must wait for the CPU transfer to finish. + self._barrier() + + if status == "error": + self._raise_save_error(payload) + return payload + + def _rank0_save(self, epoch, val_loss_avg, extras): + """Perform rank 0's write and REPORT its outcome; never raises. + + Returns ``("ok", is_best)`` or ``("error", description)``. The caller + broadcasts that outcome so every rank fails together -- see the + collective invariant on ``save_checkpoint``. + """ + try: + # 1. Wait for previous async save to prevent OOM or race. If that + # write failed, this is where it surfaces. + if self.async_save: + self.wait_for_save() + # Decide is_best from the cached best loss (single source of truth), # not by re-reading checkpoint_best.pth from disk. The cache is # seeded once at construction and updated below, so the decision # never races the background writer that may still be replacing the # best checkpoint in async mode. - if val_loss_avg < self.best_val_loss: - is_best = True - self.best_val_loss = val_loss_avg - - if self.world_rank == 0: - # 1. Wait for previous async save to prevent OOM or race - if self.async_save: - self.wait_for_save() + is_best = val_loss_avg < self.best_val_loss model_to_save = ( self.model.module if hasattr(self.model, "module") else self.model @@ -350,10 +468,6 @@ def save_checkpoint( if extras: state_dict.update(extras) - # Record the epoch being written so callers can tell whether the - # last completed epoch has already been checkpointed. - self.last_saved_epoch = epoch - # 2. Save Trigger if self.async_save: # We must clone tensors to CPU now, because training will resume @@ -380,13 +494,19 @@ def save_checkpoint( self.log, ) - # Broadcast result (for logging elsewhere) - is_best = self._broadcast_obj(is_best) - - # Barrier: ensure Rank 0 has finished the "Snapshot" phase before anyone continues. - # Even in async mode, we must wait for the CPU transfer to finish. - self._barrier() - return is_best + # Only now claim the save: the bytes are on disk (sync) or handed + # to the writer (async). Recording the epoch lets callers tell + # whether the last completed epoch has already been checkpointed. + # An async write that fails later aborts the run at the next + # collective point, so this optimistic state is never observed by + # a run that keeps going. + if is_best: + self.best_val_loss = val_loss_avg + self.last_saved_epoch = epoch + return ("ok", is_best) + except Exception as e: + self._save_error_exc = e + return ("error", f"{type(e).__name__}: {e}") @staticmethod def _atomic_save(state_dict, path): diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 07fdd2c..da33c12 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -1119,6 +1119,14 @@ def train(self, profiler=None): completed_epochs, val_loss_avg, extras ) + # Nothing downstream of the training loop touches the checkpoint + # manager, so this is the last chance to observe the outcome of the + # run's final (possibly asynchronous) write. Without it a failed final + # save would let the process exit successfully with no checkpoint at + # all, and the next --restart would resume from a stale epoch or fail + # its pre-check. + self.checkpoint_manager.finalize_saves() + if epoch_minibatch_times_s: minibatch_time_s = statistics.median(epoch_minibatch_times_s) adiak_value("minibatch_time_s", minibatch_time_s) diff --git a/tests/test_checkpointing.py b/tests/test_checkpointing.py index 40151f6..f102a6d 100644 --- a/tests/test_checkpointing.py +++ b/tests/test_checkpointing.py @@ -28,7 +28,9 @@ import math import re +import textwrap import time +from concurrent.futures import ThreadPoolExecutor from pathlib import Path import pytest @@ -187,6 +189,185 @@ def always_raise(obj, f, *args, **kwargs): mgr.save_checkpoint(epoch=1, val_loss_avg=0.5) +# --------------------------------------------------------------------------- +# R04 -- async save failures are reported, and the run's LAST save is consumed +# --------------------------------------------------------------------------- + + +def _failing_torch_save(obj, f, *args, **kwargs): + raise RuntimeError("writer boom") + + +def test_async_save_failure_surfaces_at_next_save(tmp_path, monkeypatch): + """A background write that failed is reported at the next save, not dropped. + + In async mode ``save_checkpoint`` returns as soon as the CPU snapshot is + handed to the writer thread, so the failure can only be observed later. + Swallowing it leaves ``last_saved_epoch``/``best_val_loss`` claiming a + checkpoint that does not exist on disk. + """ + mgr, _ = _make_manager(tmp_path, async_save=True) + monkeypatch.setattr(torch, "save", _failing_torch_save) + + mgr.save_checkpoint(epoch=1, val_loss_avg=0.5) # write fails in background + + with pytest.raises(RuntimeError, match="writer boom"): + mgr.save_checkpoint(epoch=2, val_loss_avg=0.4) + + assert not mgr.last_ckpt_path.exists() + + +def test_async_save_failure_surfaces_at_finalize(tmp_path, monkeypatch): + """The run's final save has its outcome consumed before the run ends.""" + mgr, _ = _make_manager(tmp_path, async_save=True) + monkeypatch.setattr(torch, "save", _failing_torch_save) + + mgr.save_checkpoint(epoch=1, val_loss_avg=0.5) + + with pytest.raises(RuntimeError, match="writer boom"): + mgr.finalize_saves() + + assert not mgr.last_ckpt_path.exists() + + +def test_final_async_save_failure_fails_the_run(tiny_trainer, monkeypatch): + """``train()`` must not return successfully after a failed final save. + + Nothing downstream of the training loop touches the checkpoint manager, so + without an explicit wait the child process exits 0 with no checkpoint at + all: the benchmark reports success and a later ``--restart`` resumes from a + stale epoch or fails its pre-check. + """ + trainer = tiny_trainer( + config_overrides={ + "checkpoint_interval": 1, + "epochs": 1, + "target_dice": 0.95, + } + ) + # The config has no async knob at this scale; drive the manager directly. + mgr = trainer.checkpoint_manager + mgr.async_save = True + mgr.executor = ThreadPoolExecutor(max_workers=1) + + monkeypatch.setattr( + trainer, + "_run_training_batch", + lambda batch, **kw: (1, torch.tensor(0.3), torch.tensor(0.5)), + ) + monkeypatch.setattr( + trainer_mod, "evaluate", lambda *a, **k: (0.5 * 2, 0.4 * 2, 0.4, 2, 2) + ) + monkeypatch.setattr(torch, "save", _failing_torch_save) + + trainer.cleanup_or_resume() + try: + with pytest.raises(RuntimeError, match="writer boom"): + trainer.train() + finally: + mgr.executor.shutdown(wait=True) + + assert not mgr.last_ckpt_path.exists() + + +# --------------------------------------------------------------------------- +# R05 -- a rank-0 write failure fails every rank with the same error +# --------------------------------------------------------------------------- + +# Two-rank script: rank 0's torch.save fails. Both ranks must come out of +# save_checkpoint with the SAME real error rather than rank 0 raising the disk +# error while its peers die (gloo) or stall (NCCL) in an unmatched collective. +# Kept inline rather than in tests/helpers/rank_scripts/ because it is only +# meaningful together with the assertions below. +SAVE_FAIL_RANK_SCRIPT = textwrap.dedent( + '''\ + """Two-rank save-failure rank script (gloo, CPU).""" + + import os + import sys + + import torch + import torch.distributed as dist + + from ScaFFold.utils.checkpointing import CheckpointManager + + rank = int(os.environ["RANK"]) + dist.init_process_group(backend="gloo") + + torch.manual_seed(0) + mgr = CheckpointManager( + model=torch.nn.Linear(64, 64), + base_dir=os.environ["CKPT_DIR"], + world_rank=rank, + dist_enabled=True, + ) + + if rank == 0: + def _boom(obj, f, *args, **kwargs): + raise RuntimeError("injected disk failure on rank 0") + + torch.save = _boom + + # Markers are delimited (trailing '.', '<<...>>') because two ranks writing + # the same pipe can interleave without a newline between them. + try: + mgr.save_checkpoint(epoch=1, val_loss_avg=0.5) + except BaseException as exc: # noqa: BLE001 - the point is what we caught + print(f"RANK {rank} RAISED {type(exc).__name__}.", flush=True) + message = str(exc).replace(chr(10), " ") + print(f"RANK {rank} MESSAGE <<{message}>>", flush=True) + else: + print(f"RANK {rank} NO_RAISE.", flush=True) + + print(f"RANK {rank} DONE.", flush=True) + sys.stdout.flush() + try: + dist.destroy_process_group() + except Exception: + pass + ''' +) + + +@_requires_gloo +def test_rank0_save_failure_fails_all_ranks(tmp_path): + """Under 2 gloo ranks, a rank-0 disk error reaches the peer as itself. + + Rank 0 is the only writer, so raising its error before the broadcast the + peers are already waiting in leaves them in an unmatched collective: gloo + reports an opaque "Connection closed by peer" and NCCL stalls until the + watchdog timeout, in both cases hiding the disk error that actually + happened. The write's OUTCOME must be broadcast instead, so both ranks + raise the same error together. + """ + ckpt_dir = tmp_path / "checkpoints" + ckpt_dir.mkdir(parents=True, exist_ok=True) + script = tmp_path / "save_fail_2rank.py" + script.write_text(SAVE_FAIL_RANK_SCRIPT) + + rc, out, err = mpi_runner.torchrun_gloo( + str(script), n=2, timeout=90, env={"CKPT_DIR": str(ckpt_dir)} + ) + + done = set(re.findall(r"RANK (\d+) DONE\.", out)) + assert rc == 0 and {"0", "1"} <= done, ( + f"expected both ranks to fail cleanly and finish, rc={rc}\n" + f"stdout:\n{out}\nstderr:\n{err[-3000:]}" + ) + + raised = dict(re.findall(r"RANK (\d+) RAISED (\w+)\.", out)) + messages = dict(re.findall(r"RANK (\d+) MESSAGE <<(.*?)>>", out)) + assert set(raised) == {"0", "1"}, f"both ranks must raise\nstdout:\n{out}" + assert raised["0"] == raised["1"], ( + f"ranks raised different error types: {raised}\nstdout:\n{out}" + ) + for rank in ("0", "1"): + assert "injected disk failure on rank 0" in messages.get(rank, ""), ( + f"rank {rank} did not see the real disk error: " + f"{messages.get(rank)!r}\nstdout:\n{out}" + ) + + # --------------------------------------------------------------------------- # F41 -- race-free best decision (cached best loss, no per-save probe) # --------------------------------------------------------------------------- From 6af26c21ed96ad8f03b5e2fb408266c65dc52526 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:34:51 -0700 Subject: [PATCH 07/62] Clean up checkpoint temp and quarantine debris A killed write strands a full-checkpoint-sized checkpoint_*.pth.tmp. and a quarantined checkpoint keeps its .corrupt copy forever; nothing ever removed either. The from-scratch cleanup now deletes both, and manager construction sweeps orphaned temp files (skipping this pid's) since the kill/restart cycle that creates them always takes the resume path. Round-2 review: R07. --- ScaFFold/utils/checkpointing.py | 47 ++++++++++++++++++++++++++-- tests/test_checkpointing.py | 54 +++++++++++++++++++++++++++++++++ 2 files changed, 99 insertions(+), 2 deletions(-) diff --git a/ScaFFold/utils/checkpointing.py b/ScaFFold/utils/checkpointing.py index 5fedc8c..24f3fe6 100644 --- a/ScaFFold/utils/checkpointing.py +++ b/ScaFFold/utils/checkpointing.py @@ -111,6 +111,7 @@ def __init__( # Ensure base directory exists (Rank 0 only) if self.world_rank == 0: self.base_dir.mkdir(parents=True, exist_ok=True) + self._sweep_orphaned_tmp_files() def cleanup(self, train_from_scratch: bool) -> None: """Clear existing checkpoints if training from scratch. @@ -141,8 +142,20 @@ def cleanup(self, train_from_scratch: bool) -> None: self._raise_save_error(error) def _remove_checkpoint_files(self) -> None: - """Delete this run's checkpoint files (rank 0 only).""" - for p in (self.last_ckpt_path, self.best_ckpt_path): + """Delete this run's checkpoint files and debris (rank 0 only). + + Besides the two canonical files, a run directory can hold + ``checkpoint_*.pth.tmp.`` (an interrupted write whose Python-level + cleanup never ran) and ``checkpoint_*.pth.corrupt`` (a checkpoint + quarantined on resume). Both are full-checkpoint-sized and nothing else + removes them, so a "from scratch" cleanup that left them would claim to + have cleared the checkpoints while keeping their bytes on disk. + """ + debris = sorted( + set(self.base_dir.glob("checkpoint_*.pth.tmp.*")) + | set(self.base_dir.glob("checkpoint_*.pth.corrupt")) + ) + for p in (self.last_ckpt_path, self.best_ckpt_path, *debris): if p.exists(): try: p.unlink() @@ -150,6 +163,36 @@ def _remove_checkpoint_files(self) -> None: except Exception as e: self._log(f"Failed to remove {p}: {e}") + def _sweep_orphaned_tmp_files(self) -> None: + """Delete temp files stranded by checkpoint writes that were killed. + + ``_atomic_save`` unlinks its ``.tmp.`` file when the write + raises, but a SIGKILL (walltime, node failure) skips that Python-level + cleanup and strands a full-checkpoint-sized file. These accumulate one + per killed pid: the kill/restart cycle that produces them always takes + the *resume* path, so ``cleanup(train_from_scratch=True)`` never gets a + chance to clear them. + + Sweeping at construction is safe because run directories are per-run + and not shared between concurrently running jobs (F55), so any temp + file here belongs to a dead process -- except one from this pid, which + another manager in this process could still be writing. + + ``*.corrupt`` files are deliberately left alone here: + ``_quarantine_corrupt`` renames onto a fixed name, so at most two can + ever exist (they cannot accumulate) and they are the only evidence of + what a resume discarded. The from-scratch cleanup removes them. + """ + own_suffix = f".tmp.{os.getpid()}" + for path in sorted(self.base_dir.glob("checkpoint_*.pth.tmp.*")): + if path.name.endswith(own_suffix): + continue + try: + path.unlink() + self._log(f"Removed orphaned checkpoint temp file: {path}") + except OSError as e: + self._log(f"Failed to remove {path}: {e}") + def wait_for_save(self): """Block until the background save (if any) is complete. diff --git a/tests/test_checkpointing.py b/tests/test_checkpointing.py index f102a6d..2ef9cad 100644 --- a/tests/test_checkpointing.py +++ b/tests/test_checkpointing.py @@ -27,6 +27,7 @@ from __future__ import annotations import math +import os import re import textwrap import time @@ -464,6 +465,59 @@ def test_cleanup_from_scratch_resets_best(tmp_path): assert mgr2.best_ckpt_path.exists() +# --------------------------------------------------------------------------- +# R07 -- checkpoint debris (.tmp., .corrupt) does not accumulate +# --------------------------------------------------------------------------- + + +def test_cleanup_from_scratch_removes_stale_debris(tmp_path): + """A from-scratch cleanup clears checkpoint debris, not just the two files. + + ``_atomic_save``'s temp file survives a process kill (its unlink only runs + on a Python-level exception) and ``_quarantine_corrupt``'s ``.corrupt`` + rename is never undone. Both are full-checkpoint-sized, so a cleanup that + claims to have cleared the checkpoints while leaving them behind keeps + multi-GB files on the shared filesystem. + """ + mgr, _ = _make_manager(tmp_path) + mgr.save_checkpoint(epoch=1, val_loss_avg=0.5) + + stale_tmp = tmp_path / "checkpoint_last.pth.tmp.999999" + stale_tmp.write_bytes(b"partial checkpoint") + quarantined = tmp_path / "checkpoint_best.pth.corrupt" + quarantined.write_bytes(b"truncated checkpoint") + + mgr.cleanup(train_from_scratch=True) + + assert not stale_tmp.exists() + assert not quarantined.exists() + assert list(tmp_path.iterdir()) == [] + + +def test_init_sweeps_orphaned_tmp_files(tmp_path): + """Constructing a manager sweeps temp files left by killed writes. + + Repeated walltime kills of a long run take the *resume* path, never the + from-scratch cleanup, so without this sweep one stranded temp file per + killed pid piles up in the run directory. A temp file belonging to this + process is left alone (another manager here may still be writing it), and + the quarantined ``.corrupt`` file is kept: at most two can ever exist and + they are the only evidence of what a resume discarded. + """ + orphan = tmp_path / "checkpoint_last.pth.tmp.999999" + orphan.write_bytes(b"partial checkpoint") + own = tmp_path / f"checkpoint_last.pth.tmp.{os.getpid()}" + own.write_bytes(b"possibly in flight") + quarantined = tmp_path / "checkpoint_last.pth.corrupt" + quarantined.write_bytes(b"truncated checkpoint") + + _make_manager(tmp_path) + + assert not orphan.exists() + assert own.exists() + assert quarantined.exists() + + # --------------------------------------------------------------------------- # F71 -- CPU tensors are cloned into the snapshot # --------------------------------------------------------------------------- From 3c4f5b4612e1bd5adf8c4286f70369e0f70982b6 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:35:13 -0700 Subject: [PATCH 08/62] Delete the unused checkpoint_validators module The module had no import sites anywhere in the package, tests or scripts, its usage comments reference a train.py flow that no longer exists, and compare_state_dicts2 crashes on any real optimizer state dict (it truth-tests an elementwise tensor comparison). Round-2 review: R10. --- ScaFFold/utils/checkpoint_validators.py | 143 ------------------------ 1 file changed, 143 deletions(-) delete mode 100644 ScaFFold/utils/checkpoint_validators.py diff --git a/ScaFFold/utils/checkpoint_validators.py b/ScaFFold/utils/checkpoint_validators.py deleted file mode 100644 index c2212b6..0000000 --- a/ScaFFold/utils/checkpoint_validators.py +++ /dev/null @@ -1,143 +0,0 @@ -# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. -# Produced at the Lawrence Livermore National Laboratory. -# Written by the LBANN Research Team (B. Van Essen, et al.) listed in -# the CONTRIBUTORS file. See the top-level LICENSE file for details. -# -# LLNL-CODE-697807. -# All rights reserved. -# -# This file is part of LBANN: Livermore Big Artificial Neural Network -# Toolkit. For details, see http://software.llnl.gov/LBANN or -# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. -# -# SPDX-License-Identifier: (Apache-2.0) - -import os - -# import pdb - -if hasattr(os, "sched_getaffinity"): - _orig_affinity = os.sched_getaffinity(0) -else: - _orig_affinity = None - -import torch - - -def compare_state_dicts(quantity: str, dict1, dict2): - def _compare(dict1, dict2, prefix=""): - equal = True - if dict1.keys() != dict2.keys(): - missing_in_dict2 = dict1.keys() - dict2.keys() - missing_in_dict1 = dict2.keys() - dict1.keys() - if missing_in_dict2: - print( - f"train.py: {quantity} missing in dict2: {', '.join(missing_in_dict2)} at {prefix}" - ) - equal = False - if missing_in_dict1: - print( - f"train.py: {quantity} missing in dict1: {', '.join(missing_in_dict1)} at {prefix}" - ) - equal = False - - for key in dict1.keys() & dict2.keys(): - full_key = f"{prefix}.{key}" if prefix else key - if isinstance(dict1[key], torch.Tensor) and isinstance( - dict2[key], torch.Tensor - ): - if not torch.equal(dict1[key], dict2[key]): - print( - f"train.py: {quantity} tensor discrepancy at {full_key}: dict1[{key}] != dict2[{key}]" - ) - equal = False - elif isinstance(dict1[key], dict) and isinstance(dict2[key], dict): - if not _compare(dict1[key], dict2[key], prefix=full_key): - equal = False - else: - if dict1[key] != dict2[key]: - print( - f"train.py: {quantity} value discrepancy at {full_key}: dict1[{key}]={dict1[key]}, dict2[{key}]={dict2[key]}" - ) - equal = False - return equal - - return _compare(dict1, dict2) - - -def compare_state_dicts2(*dicts): - keys = dicts[0].keys() - for key in keys: - values = [d[key] for d in dicts] - tensor_comparisons = [ - ( - torch.equal(values[i], values[i + 1]) - if torch.is_tensor(values[i]) - else values[i] == values[i + 1] - ) - for i in range(len(values) - 1) - ] - if not all(tensor_comparisons): - return False - return True - - -def compare_tensors(tensor1, tensor2): - return torch.all(torch.eq(tensor1, tensor2)) - - -def compare_items(item1, item2): - if isinstance(item1, torch.Tensor) and isinstance(item2, torch.Tensor): - return compare_tensors(item1, item2) - elif isinstance(item1, dict) and isinstance(item2, dict): - return compare_dicts3(item1, item2) - else: - return item1 == item2 - - -def compare_dicts3(dict1, dict2): - if dict1.keys() != dict2.keys(): - return False - for key in dict1.keys(): - if not compare_items(dict1[key], dict2[key]): - return False - return True - - -# -# Usage in `train.py` below: -# - -# For debugging, write the saved optimizer state to file to compare to loaded state on restart -# timestamp = datetime.now().strftime('%Y%m%d_%H%M%S') -# optim_saved_path = f"{dir_checkpoint}/restarts/optim_saved_epoch{epoch}_{timestamp}.txt" -# with open(optim_saved_path, "w") as optim_f: -# optim_f.write(str(state_dict['optimizer_state_dict'])) - -# -# For debugging purposes, load that checkpoint into new model, optimizer, etc and compare to active -# -# prev_checkpoint = torch.load(checkpoint_path) -# newmodel = UNet(n_channels=3, n_classes=n_classes, trilinear=False, layers=unet_layers) -# newmodel = newmodel.to(memory_format=torch.channels_last_3d) -# newmodel.to(device=device) -# newmodel = torch.nn.parallel.DistributedDataParallel(newmodel, device_ids=[get_cuda_device()], output_device=get_cuda_device()) -# newmodel.module.load_state_dict(prev_checkpoint['model_state_dict']) -# newoptimizer = optim.RMSprop(newmodel.parameters(), -# lr=learning_rate, weight_decay=weight_decay, momentum=momentum, foreach=True) -# if optimizer_name == "ADAM": -# print(f"train.py(w{rank}|l{local_rank}): using ADAM optimizer .........") -# newoptimizer = optim.Adam(newmodel.parameters(), lr=learning_rate) -# elif optimizer_name == "SGD": -# print(f"train.py(w{rank}|l{local_rank}): using SGD optimizer .........") -# newoptimizer = optim.SGD(newmodel.parameters(), lr=learning_rate, momentum=0.9) -# newoptimizer.load_state_dict(prev_checkpoint['optimizer_state_dict']) -# newscheduler = optim.lr_scheduler.ReduceLROnPlateau(newoptimizer, 'max', patience=25) -# newscheduler.load_state_dict(prev_checkpoint['scheduler_state_dict']) - -# # Compare model state dicts -# model_compare = compare_state_dicts("model", model.state_dict(), newmodel.state_dict()) -# optimizer_compare = compare_state_dicts("optimizer", optimizer.state_dict(), newoptimizer.state_dict()) -# scheduler_compare = compare_state_dicts("scheduler", scheduler.state_dict(), newscheduler.state_dict()) -# print(f"train.py: model_compare={model_compare}, optimizer_compare={optimizer_compare}, scheduler_compare={scheduler_compare}") -# print(f"train.py: all equal? {all(compare_dicts3(sd, optimizer.state_dict()) for sd in [state_dict['optimizer_state_dict'], prev_checkpoint['optimizer_state_dict'], newoptimizer.state_dict()])}") From 638aac6b78c86c76d0f84480846fda5e70782f1d Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:39:10 -0700 Subject: [PATCH 09/62] State why no epoch was trained instead of guessing the remedy The nothing-to-resume warning suggested lowering target_dice, which is exactly backwards for the converged case, and read oddly for a fresh run that never entered the loop. Report the actual inputs (start epoch, epochs, starting val dice vs target) instead. Round-2 review: R01, R06 follow-up. --- ScaFFold/utils/trainer.py | 25 ++++++++++++++----------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index da33c12..81dd5f2 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -1086,18 +1086,21 @@ def train(self, profiler=None): completed_epochs = epoch - 1 if not completed_new_epoch: - # The loop exited without running a single new epoch: the resumed - # checkpoint already covers every epoch this run was asked for. - # There is nothing new to save (the existing checkpoint already - # records epoch `completed_epochs`) and none of the per-epoch - # metrics the final save would write were ever computed, so skip - # it and return normally -- the caller's post-processing still has - # the CSV the original run left behind. + # The loop exited without running a single epoch: the state we + # resumed either already covers every epoch this run was asked for, + # or already met target_dice. There is nothing new to save (the + # checkpoint on disk already records epoch `completed_epochs`) and + # none of the per-epoch metrics the final save would write were + # ever computed, so skip it and return normally -- the caller's + # post-processing still has the CSV the original run left behind. self.log.warning( - "Nothing to resume: the loaded checkpoint already covers epoch " - "%s, so no new epoch was trained and no checkpoint was written. " - "Increase 'epochs' (or lower 'target_dice') to train further.", - completed_epochs, + "No new epoch was trained (start epoch %s, 'epochs' %s, " + "starting val dice %s vs target_dice %s): there was nothing to " + "resume, and no checkpoint was written.", + self.start_epoch, + self.config.epochs, + self.start_val_dice, + self.config.target_dice, ) # Save a final checkpoint when the run exits (convergence or max epochs) # at an epoch that was not a checkpoint interval, so the converged From 9df9a1bd0e1083a8576591866987f2e099df2cf3 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:21:18 -0700 Subject: [PATCH 10/62] Guard rank-0 windows in the get_dataset consensus Any failure in the rank-0 reuse/generate decision or the final meta-write and rename is now broadcast as an error sentinel, so peers raise the same error instead of hanging in bcast/Barrier. The reuse scan also skips .tmp_* staging dirs and tolerates unreadable metadata, and staging dir names carry pid+uuid so same-second jobs cannot collide. R27 --- ScaFFold/datagen/get_dataset.py | 108 ++++++++++++--- tests/datagen/test_mpi_consensus.py | 198 ++++++++++++++++++++++++++++ 2 files changed, 286 insertions(+), 20 deletions(-) diff --git a/ScaFFold/datagen/get_dataset.py b/ScaFFold/datagen/get_dataset.py index c536b14..bda9572 100644 --- a/ScaFFold/datagen/get_dataset.py +++ b/ScaFFold/datagen/get_dataset.py @@ -16,9 +16,11 @@ import hashlib import json +import os import shutil import subprocess import time +import uuid from argparse import Namespace from pathlib import Path from typing import Any, Dict @@ -30,6 +32,11 @@ from ScaFFold.utils.utils import setup_mpi_logger META_FILENAME = "meta.yaml" +# Datasets are generated into a staging directory carrying this prefix and only +# renamed into their final ``__`` name once complete, so a +# reader never observes a half-written dataset. The prefix is also what the +# reuse scan skips and what the orphan cleanup collects. +TMP_PREFIX = ".tmp_" # Bumped from 2 to 3 when instance point clouds moved from float64 to float32: # the storage layout is unchanged, but float32 voxel binning shifts a handful of # boundary voxels, so a float64-era dataset must not be reused as if it were @@ -123,15 +130,44 @@ def _decide_reuse_or_generate( staging and final paths for a new generation. Making this decision in one place and broadcasting it prevents ranks from diverging when their views of the shared filesystem differ. + + The scan is deliberately forgiving: a candidate whose metadata is missing, + unreadable, or malformed is warned about and skipped rather than allowed to + raise. This function runs inside a window where every peer is already + waiting in the decision broadcast, so a crash here is a job-wide hang; a + poison directory (exactly what a killed job leaves behind) must never be + able to cause one. """ candidates = sorted( (p for p in base.iterdir() if p.is_dir()), key=lambda p: p.name, reverse=True ) for dataset_path in candidates: + # Staging dirs are not datasets: a job killed between the meta write and + # the rename leaves a complete meta.yaml inside one, and reusing it hands + # back a partially generated (and cleanup-eligible) directory. + if dataset_path.name.startswith(TMP_PREFIX): + continue meta_path = dataset_path / META_FILENAME if not meta_path.exists(): continue - meta = yaml.safe_load(meta_path.read_text()) + try: + meta = yaml.safe_load(meta_path.read_text()) + except Exception as exc: + log.warning( + "Skipping dataset candidate %s: unreadable %s (%s: %s)", + dataset_path, + META_FILENAME, + type(exc).__name__, + exc, + ) + continue + if not isinstance(meta, dict): + log.warning( + "Skipping dataset candidate %s: %s is empty or malformed", + dataset_path, + META_FILENAME, + ) + continue if meta.get("config_id") != config_id: continue if meta.get("dataset_format_version", 1) != DATASET_FORMAT_VERSION: @@ -145,7 +181,11 @@ def _decide_reuse_or_generate( log.info("No valid existing dataset found at %s. Generating new dataset.", base) ts = time.strftime("%Y%m%d-%H%M%S") dest = base / f"{ts}__{commit}" - tmp = base / f".tmp_{ts}" + # The staging name must be unique per job: a bare 1-second-granularity + # timestamp let two same-config jobs starting in the same second collide on + # ``mkdir(exist_ok=False)``, killing one of them mid-consensus. Adding the + # pid and a random suffix makes the name unique even across nodes. + tmp = base / f"{TMP_PREFIX}{ts}_{os.getpid()}_{uuid.uuid4().hex[:8]}" tmp.mkdir(parents=True, exist_ok=False) return ("generate", str(tmp), str(dest)) @@ -188,14 +228,28 @@ def get_dataset( # same branch. Scanning the shared filesystem independently per rank lets # divergent views (stale metadata caches, a racing job's rename) strand some # ranks in the generation collectives while others return early. + # Everything rank 0 does here happens while the peers are already blocked in + # the broadcast below, so a rank-0 exception would strand the whole job. + # Any failure is therefore turned into an error sentinel that travels + # through the same broadcast and makes every rank raise the same error. if rank == 0: - decision = _decide_reuse_or_generate( - base, config_id, commit, require_commit, log - ) + try: + decision = _decide_reuse_or_generate( + base, config_id, commit, require_commit, log + ) + except (Exception, SystemExit) as e: + decision = ( + "error", + f"rank 0 failed to select a dataset under {base}: " + f"{type(e).__name__}: {e}", + ) else: decision = None decision = comm.bcast(decision, root=0) + if decision[0] == "error": + raise RuntimeError(f"dataset selection failed: {decision[1]}") + if decision[0] == "reuse": return Path(decision[1]) @@ -230,21 +284,35 @@ def get_dataset( raise RuntimeError(f"dataset generation failed: {msgs or 'unknown error'}") # rank 0 writes metadata into the staging dir, then renames it into place so - # readers never observe a half-written dataset. + # readers never observe a half-written dataset. This is another rank-0-only + # window inside a collective sequence: the rename can fail (a racing job + # already published this name, quota, ...), so the outcome is broadcast + # rather than allowed to kill rank 0 while the peers wait for it. + finalize_err = "" if rank == 0: - meta = { - "config_id": config_id, - "dataset_format_version": DATASET_FORMAT_VERSION, - "config_subset": volume_config, - "include_keys": INCLUDE_KEYS, - "code_commit": commit, - "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), - } - (tmp / META_FILENAME).write_text( - yaml.safe_dump(meta, sort_keys=True, default_flow_style=False) - ) - tmp.rename(dest) + try: + meta = { + "config_id": config_id, + "dataset_format_version": DATASET_FORMAT_VERSION, + "config_subset": volume_config, + "include_keys": INCLUDE_KEYS, + "code_commit": commit, + "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + } + (tmp / META_FILENAME).write_text( + yaml.safe_dump(meta, sort_keys=True, default_flow_style=False) + ) + tmp.rename(dest) + except (Exception, SystemExit) as e: + finalize_err = ( + f"rank 0 failed to finalize dataset at {dest}: {type(e).__name__}: {e}" + ) + + # This broadcast doubles as the synchronization the old Barrier provided: no + # rank returns before rank 0 has published the rename (or reported that it + # could not), so nobody observes the staging path or a missing dataset. + finalize_err = comm.bcast(finalize_err, root=0) + if finalize_err: + raise RuntimeError(f"dataset generation failed: {finalize_err}") - # ensure the rename is visible everywhere before returning - comm.Barrier() return dest diff --git a/tests/datagen/test_mpi_consensus.py b/tests/datagen/test_mpi_consensus.py index 6a67e89..08bd0c9 100644 --- a/tests/datagen/test_mpi_consensus.py +++ b/tests/datagen/test_mpi_consensus.py @@ -38,6 +38,7 @@ from __future__ import annotations +import logging import re from argparse import Namespace from math import ceil @@ -340,6 +341,203 @@ def test_generation_success_finalizes_and_returns(tmp_path, monkeypatch): assert leftover == [] +# --------------------------------------------------------------------------- +# R27: rank 0 must never die between the collectives its peers have entered. +# Every rank-0-only step of the consensus (the reuse/generate decision and the +# final meta-write + rename) is wrapped so a failure travels to the peers as a +# broadcast sentinel instead of stranding them in ``bcast``/``Barrier``. +# --------------------------------------------------------------------------- + + +def _base_dir_for(config: Namespace) -> Path: + """The ``/`` directory ``get_dataset`` scans.""" + config_dict = vars(config).copy() + config_dict["dataset_format_version"] = gd.DATASET_FORMAT_VERSION + volume_config = gd._get_required_keys_dict(config_dict, gd.INCLUDE_KEYS) + return Path(config.dataset_dir) / gd._hash_volume_config(volume_config) + + +def test_decision_failure_is_broadcast_not_raised_before_bcast(tmp_path, monkeypatch): + """A rank-0 decision failure reaches peers through the broadcast. + + Any exception inside the rank-0-only decision (an unreadable base dir, a + staging ``mkdir`` hitting ENOSPC, ...) must be converted into an error + sentinel that is broadcast, so peers already waiting in ``bcast`` learn + about it and raise the same error. Before the fix rank 0 raised *before* + reaching the broadcast, leaving every peer blocked forever. + """ + config = _reuse_config(tmp_path / "datasets") + comm = FakeComm(rank=0, size=2) + monkeypatch.setattr(gd, "MPI", FakeMPI(comm)) + monkeypatch.setattr(gd, "_git_commit_short", lambda log: "abc123") + + def explode(*_args, **_kwargs): + raise OSError("No space left on device") + + monkeypatch.setattr(gd, "_decide_reuse_or_generate", explode) + + with pytest.raises(RuntimeError) as excinfo: + gd.get_dataset(config) + + # Rank 0 reached the broadcast before raising, and the payload is the + # error sentinel every peer will see. + assert comm.calls == ["bcast"] + assert comm.bcast_payloads[0][0] == "error" + assert "No space left on device" in str(excinfo.value) + + +def test_non_root_raises_on_broadcast_decision_error(tmp_path, monkeypatch): + """A peer receiving the error sentinel raises instead of generating.""" + config = _reuse_config(tmp_path / "datasets") + sentinel = ("error", "rank 0: OSError: No space left on device") + comm = FakeComm(rank=1, size=2, bcast_returns=[sentinel]) + monkeypatch.setattr(gd, "MPI", FakeMPI(comm)) + monkeypatch.setattr(gd, "_git_commit_short", lambda log: "abc123") + + with pytest.raises(RuntimeError) as excinfo: + gd.get_dataset(config) + + assert "No space left on device" in str(excinfo.value) + # The peer stopped at the decision broadcast: no generation collectives. + assert comm.calls == ["bcast"] + + +def test_reuse_scan_skips_staging_dirs(tmp_path, monkeypatch): + """A complete ``meta.yaml`` stranded in a ``.tmp_*`` dir is never reused. + + A job killed between the meta write and the rename leaves a fully valid + meta inside its staging dir. Treating that as a publishable dataset hands + back a half-generated directory (and one that cleanup may delete). + """ + config = _reuse_config(tmp_path / "datasets") + base = _base_dir_for(config) + config_id = base.name + stranded = base / ".tmp_20260101-000000_1234" + stranded.mkdir(parents=True) + (stranded / gd.META_FILENAME).write_text( + yaml.safe_dump( + { + "config_id": config_id, + "dataset_format_version": gd.DATASET_FORMAT_VERSION, + } + ) + ) + + comm = FakeComm(rank=0, size=1) + monkeypatch.setattr(gd, "MPI", FakeMPI(comm)) + monkeypatch.setattr(gd, "_git_commit_short", lambda log: "abc123") + log = logging.getLogger("test_reuse_scan_skips_staging_dirs") + + decision = gd._decide_reuse_or_generate(base, config_id, "abc123", False, log) + + assert decision[0] == "generate", f"staging dir was reused: {decision}" + + +def test_reuse_scan_tolerates_corrupt_meta(tmp_path, monkeypatch): + """A corrupt/unreadable candidate meta is skipped, not fatal. + + A 0-byte ``meta.yaml`` (``yaml.safe_load`` -> ``None``) or an unparseable + one used to raise inside the rank-0-only scan. The scan must warn, skip the + directory, and keep looking -- here finding the good dataset next to it. + """ + config = _reuse_config(tmp_path / "datasets") + base = _base_dir_for(config) + config_id = base.name + base.mkdir(parents=True) + + # Sorted-descending scan order visits these two poison dirs first. + (base / "20260301-000000__zzz").mkdir() + (base / "20260301-000000__zzz" / gd.META_FILENAME).write_text("") + (base / "20260201-000000__yyy").mkdir() + (base / "20260201-000000__yyy" / gd.META_FILENAME).write_text("{[not yaml") + + good = _write_reusable_dataset(base, config_id) + log = logging.getLogger("test_reuse_scan_tolerates_corrupt_meta") + + decision = gd._decide_reuse_or_generate(base, config_id, "abc123", False, log) + + assert decision[0] == "reuse" + assert Path(decision[1]) == good + + +def test_staging_dir_names_are_collision_proof(tmp_path, monkeypatch): + """Two decisions in the same second stage into different directories. + + The old name was ``.tmp_%Y%m%d-%H%M%S`` with ``mkdir(exist_ok=False)``, so + two same-config jobs starting in the same second raced to a + ``FileExistsError`` on one of them -- inside the unguarded rank-0 window. + """ + config = _reuse_config(tmp_path / "datasets") + base = _base_dir_for(config) + base.mkdir(parents=True) + log = logging.getLogger("test_staging_dir_names_are_collision_proof") + + # Pin the clock so both decisions share a timestamp: only a non-time + # component can keep the names apart. + monkeypatch.setattr(gd.time, "strftime", lambda *_args: "20260101-000000") + + first = gd._decide_reuse_or_generate(base, base.name, "abc123", False, log) + second = gd._decide_reuse_or_generate(base, base.name, "abc123", False, log) + + assert first[0] == "generate" and second[0] == "generate" + assert first[1] != second[1], "same-second staging dirs collided" + assert Path(first[1]).is_dir() and Path(second[1]).is_dir() + + +def test_finalize_failure_is_broadcast_not_left_to_barrier(tmp_path, monkeypatch): + """A rank-0 rename failure is broadcast; peers raise instead of hanging. + + The rename happens *after* the generation consensus, so a failure there + (e.g. a racing job already created the destination) used to kill rank 0 + while every peer sat in the final ``Barrier``. The fix carries the failure + through one more collective and raises everywhere. + """ + config = _reuse_config(tmp_path / "datasets") + comm = FakeComm(rank=0, size=2, allreduce_result=1) + monkeypatch.setattr(gd, "MPI", FakeMPI(comm)) + monkeypatch.setattr(gd, "_git_commit_short", lambda log: "abc123") + monkeypatch.setattr(volumegen, "main", lambda _config: None) + + # Pin the clock so the destination name is predictable, then have a + # "racing job" occupy it with a non-empty directory: the rename fails with + # ENOTEMPTY exactly as it did in the field. + monkeypatch.setattr(gd.time, "strftime", lambda *_args: "20260101-000000") + base = _base_dir_for(config) + dest = base / "20260101-000000__abc123" + dest.mkdir(parents=True) + (dest / "placeholder").write_text("created by a racing job") + + with pytest.raises(RuntimeError) as excinfo: + gd.get_dataset(config) + + message = str(excinfo.value) + assert "20260101-000000__abc123" in message + # The failure travelled through a collective *after* the generation + # consensus, so peers learn about it rather than waiting in the barrier. + assert "allgather" in comm.calls + assert comm.calls.index("allgather") < len(comm.calls) - 1 + + +def test_non_root_raises_on_broadcast_finalize_error(tmp_path, monkeypatch): + """A peer receiving the finalize error raises rather than returning dest.""" + config = _reuse_config(tmp_path / "datasets") + dest = tmp_path / "datasets" / "cid" / "20260101-000000__abc123" + comm, _tmp, _dest = _generate_decision_comm( + rank=1, size=2, dest=dest, allreduce_result=1 + ) + # Second bcast: root's finalize verdict (a failure message). + comm._bcast_returns.append("rank 0 failed to finalize: OSError: boom") + monkeypatch.setattr(gd, "MPI", FakeMPI(comm)) + monkeypatch.setattr(gd, "_git_commit_short", lambda log: "abc123") + monkeypatch.setattr(volumegen, "main", lambda _config: None) + + with pytest.raises(RuntimeError) as excinfo: + gd.get_dataset(config) + + assert "boom" in str(excinfo.value) + assert not dest.exists() + + # --------------------------------------------------------------------------- # A missing instance file raises FileNotFoundError (a catchable Exception) # rather than calling sys.exit(1) (a BaseException that bypasses consensus), and From 3d26b7b3f83952359ec17c8aebb517e2be45c621 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:23:36 -0700 Subject: [PATCH 11/62] Publish meta.yaml atomically and reject damaged metadata meta.yaml is now written to a temp file, fsynced, and renamed into place, so a killed job cannot leave a truncated document. The loader treats only a missing meta.yaml as legacy v1; a present-but-unreadable or version-less one raises instead of silently reinterpreting a modern dataset. R28 --- ScaFFold/datagen/get_dataset.py | 31 ++++++++++++++-- ScaFFold/utils/data_loading.py | 33 ++++++++++++++--- tests/datagen/test_mpi_consensus.py | 48 ++++++++++++++++++++++++ tests/test_data_loading.py | 57 +++++++++++++++++++++++++++++ 4 files changed, 160 insertions(+), 9 deletions(-) diff --git a/ScaFFold/datagen/get_dataset.py b/ScaFFold/datagen/get_dataset.py index bda9572..f200c61 100644 --- a/ScaFFold/datagen/get_dataset.py +++ b/ScaFFold/datagen/get_dataset.py @@ -115,6 +115,33 @@ def _git_commit_short(log) -> str: return "no-commit-id" +def _write_meta_atomic(meta_path: Path, meta: Dict[str, Any]) -> None: + """Write ``meta`` to ``meta_path`` atomically. + + ``meta.yaml`` is what the loader reads to decide how every sample in the + dataset is interpreted, so a partially written one is worse than none at + all: a truncated file parses as empty and silently reclassifies a modern + dataset as legacy v1. The document is therefore written to a temp file in + the same directory, flushed and fsynced, and only then ``os.replace``d onto + the final name -- an atomic rename within one filesystem. + """ + tmp_path = meta_path.parent / f".{meta_path.name}.tmp{os.getpid()}" + try: + with open(tmp_path, "w") as handle: + handle.write(yaml.safe_dump(meta, sort_keys=True, default_flow_style=False)) + handle.flush() + os.fsync(handle.fileno()) + os.replace(tmp_path, meta_path) + except BaseException: + # A failed write must not leave a temp file behind, and the final name + # must keep whatever complete document was already there. + try: + os.remove(tmp_path) + except OSError: + pass + raise + + def _decide_reuse_or_generate( base: Path, config_id: str, @@ -299,9 +326,7 @@ def get_dataset( "code_commit": commit, "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), } - (tmp / META_FILENAME).write_text( - yaml.safe_dump(meta, sort_keys=True, default_flow_style=False) - ) + _write_meta_atomic(tmp / META_FILENAME, meta) tmp.rename(dest) except (Exception, SystemExit) as e: finalize_err = ( diff --git a/ScaFFold/utils/data_loading.py b/ScaFFold/utils/data_loading.py index 564c380..dcc1a45 100644 --- a/ScaFFold/utils/data_loading.py +++ b/ScaFFold/utils/data_loading.py @@ -277,20 +277,41 @@ def _load_numpy_array(path, mmap_mode=None): return np.load(path, allow_pickle=False, mmap_mode=mmap_mode) def _load_dataset_format_version(self): + """Determine which on-disk layout this dataset uses. + + Only a *missing* ``meta.yaml`` means legacy v1: those datasets predate + the metadata file. A metadata file that exists but cannot be read or + does not carry a usable version is a damaged modern dataset, and + falling back to the legacy loader there silently transposes + channels-first volumes and remaps already-dense labels -- corrupt + training data with no error. Such a dataset is rejected instead, with a + message naming the file so it can be repaired or regenerated. + """ meta_path = self.dataset_root / META_FILENAME if not meta_path.exists(): return LEGACY_DATASET_FORMAT_VERSION try: with open(meta_path, "r") as meta_file: - meta = yaml.safe_load(meta_file) or {} + meta = yaml.safe_load(meta_file) except Exception as exc: - customlog( - f"Failed to read dataset metadata from {meta_path}: {exc}. Falling back to legacy loader." - ) - return LEGACY_DATASET_FORMAT_VERSION + raise ValueError( + f"Dataset metadata {meta_path} exists but could not be read " + f"({type(exc).__name__}: {exc}). A dataset carrying a " + f"{META_FILENAME} is not a legacy dataset; refusing to guess its " + "layout. Repair the file or regenerate the dataset." + ) from exc - return int(meta.get("dataset_format_version", LEGACY_DATASET_FORMAT_VERSION)) + version = meta.get("dataset_format_version") if isinstance(meta, dict) else None + try: + return int(version) + except (TypeError, ValueError): + raise ValueError( + f"Dataset metadata {meta_path} is missing a usable " + f"'dataset_format_version' (got {version!r}). A dataset carrying " + f"a {META_FILENAME} is not a legacy dataset; refusing to guess " + "its layout. Repair the file or regenerate the dataset." + ) from None @staticmethod def _prepare_legacy_image(img): diff --git a/tests/datagen/test_mpi_consensus.py b/tests/datagen/test_mpi_consensus.py index 08bd0c9..fa99a1b 100644 --- a/tests/datagen/test_mpi_consensus.py +++ b/tests/datagen/test_mpi_consensus.py @@ -538,6 +538,54 @@ def test_non_root_raises_on_broadcast_finalize_error(tmp_path, monkeypatch): assert not dest.exists() +# --------------------------------------------------------------------------- +# R28: meta.yaml is published atomically, so no reader ever sees a partial one. +# --------------------------------------------------------------------------- + + +def test_meta_write_is_atomic(tmp_path, monkeypatch): + """An interrupted meta write leaves the previous file intact and no temp. + + ``meta.yaml`` is the file that decides how every sample is interpreted (a + truncated one reclassifies the dataset as legacy v1), so it must appear at + its final name complete or not at all. + """ + target = tmp_path / gd.META_FILENAME + gd._write_meta_atomic(target, {"dataset_format_version": gd.DATASET_FORMAT_VERSION}) + good_bytes = target.read_bytes() + + # Interrupt the write after bytes have reached the temp file but before the + # rename -- the shape of a kill mid-write. + def boom(_fd): + raise OSError("simulated SIGKILL mid-write") + + monkeypatch.setattr(gd.os, "fsync", boom) + + with pytest.raises(OSError): + gd._write_meta_atomic(target, {"dataset_format_version": 99}) + + # The final name still holds the complete previous file, byte-for-byte, and + # no temp file is left behind for the reuse scan to trip over. + assert target.read_bytes() == good_bytes + assert [p.name for p in tmp_path.iterdir()] == [gd.META_FILENAME] + + +def test_published_dataset_has_no_partial_meta(tmp_path, monkeypatch): + """A successful generation publishes a parseable meta and no temp files.""" + config = _reuse_config(tmp_path / "datasets") + comm = FakeComm(rank=0, size=1, allreduce_result=1) + monkeypatch.setattr(gd, "MPI", FakeMPI(comm)) + monkeypatch.setattr(gd, "_git_commit_short", lambda log: "abc123") + monkeypatch.setattr(volumegen, "main", lambda _config: None) + + result = Path(gd.get_dataset(config)) + + meta = yaml.safe_load((result / gd.META_FILENAME).read_text()) + assert meta["dataset_format_version"] == gd.DATASET_FORMAT_VERSION + # Nothing hidden alongside it (a temp meta would be a dotted sibling). + assert [p.name for p in result.iterdir() if p.name.startswith(".")] == [] + + # --------------------------------------------------------------------------- # A missing instance file raises FileNotFoundError (a catchable Exception) # rather than calling sys.exit(1) (a BaseException that bypasses consensus), and diff --git a/tests/test_data_loading.py b/tests/test_data_loading.py index 5f3a623..b29d212 100644 --- a/tests/test_data_loading.py +++ b/tests/test_data_loading.py @@ -594,3 +594,60 @@ def counting_load(path, mmap_mode=None): assert loaded["image"] == 0 assert mask_only.dtype == torch.int16 assert torch.equal(mask_only, expected) + + +# --------------------------------------------------------------------------- +# R28: a *present but broken* meta.yaml must not be mistaken for a v1 dataset +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "broken_meta", + [ + pytest.param("", id="zero-byte"), + pytest.param("{[not: valid: yaml", id="unparseable"), + pytest.param("- just\n- a\n- list\n", id="not-a-mapping"), + pytest.param("config_id: abc123\n", id="version-key-missing"), + pytest.param("dataset_format_version: two\n", id="version-not-an-int"), + ], +) +def test_broken_meta_raises_instead_of_silent_legacy(tmp_path, broken_meta): + """A corrupt ``meta.yaml`` is an error, never a silent legacy downgrade. + + Treating a broken meta as "no meta" reclassifies a modern dataset as legacy + v1: the loader then transposes channels-first volumes (a (3,N,N,N) sample + comes back (N,3,N,N)) and remaps already-dense labels. Training proceeds on + silently corrupted data. The dataset directory itself is intact here -- only + the metadata is damaged -- so the failure must be loud and actionable. + """ + root = _build_v2_constant_dataset(tmp_path / "ds", n_volumes=2) + (root / "meta.yaml").write_text(broken_meta) + + with pytest.raises(ValueError) as excinfo: + FractalDataset( + root / "volumes" / "training", + root / "masks" / "training", + data_dir=root / "train_unique_mask_vals", + ) + + message = str(excinfo.value) + assert "meta.yaml" in message + # The message must point at the offending file so it can be repaired. + assert str(root) in message + + +def test_absent_meta_is_still_legacy_v1(tiny_v1_dataset): + """The genuine legacy case (no ``meta.yaml`` at all) is unchanged. + + Control for the test above: v1 datasets predate the metadata file, so a + *missing* meta must keep selecting the legacy loader rather than raising. + """ + root = tiny_v1_dataset(n_categories=2, n_train=2, n_val=1, n=8) + assert not (root / "meta.yaml").exists() + + ds = FractalDataset( + root / "volumes" / "training", + root / "masks" / "training", + data_dir=root / "train_unique_mask_vals", + ) + assert ds.dataset_format_version == 1 From 956b18930b255f058ae88b10479d3fdac6951155 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:25:33 -0700 Subject: [PATCH 12/62] Reclaim orphaned dataset staging directories Killed generations left their .tmp_* staging trees under the config_id base forever, so every retry stacked another copy. Rank 0 now removes staging dirs that have sat untouched past an age threshold, which keeps a concurrent job's (far younger) staging dir safe. R37 --- ScaFFold/datagen/get_dataset.py | 56 ++++++++++++++++++++++++ tests/datagen/test_mpi_consensus.py | 68 +++++++++++++++++++++++++++++ 2 files changed, 124 insertions(+) diff --git a/ScaFFold/datagen/get_dataset.py b/ScaFFold/datagen/get_dataset.py index f200c61..0b193f2 100644 --- a/ScaFFold/datagen/get_dataset.py +++ b/ScaFFold/datagen/get_dataset.py @@ -37,6 +37,10 @@ # reader never observes a half-written dataset. The prefix is also what the # reuse scan skips and what the orphan cleanup collects. TMP_PREFIX = ".tmp_" +# How long a staging directory must have sat untouched before it is treated as +# orphaned (left by a killed/OOM'd job) and reclaimed. See +# ``_cleanup_stale_staging_dirs`` for the safety argument behind the value. +STALE_STAGING_AGE_SECONDS = 24 * 60 * 60 # Bumped from 2 to 3 when instance point clouds moved from float64 to float32: # the storage layout is unchanged, but float32 voxel binning shifts a handful of # boundary voxels, so a float64-era dataset must not be reused as if it were @@ -115,6 +119,54 @@ def _git_commit_short(log) -> str: return "no-commit-id" +def _cleanup_stale_staging_dirs( + base: Path, log, max_age: float = STALE_STAGING_AGE_SECONDS +) -> None: + """Reclaim orphaned ``.tmp_*`` staging directories under one config base. + + A generation killed by a walltime limit, an OOM, or a node failure leaves + its whole staging tree behind, and nothing ever removed it: every retry + stacked another (potentially multi-terabyte) copy under the same config_id. + + Safety policy. Only directories that (a) live directly under *this* + config_id base, (b) carry the ``.tmp_`` prefix this module owns, and (c) + have been untouched for ``max_age`` are removed. The age gate is what keeps + a *concurrent* job's staging directory safe: unique staging names mean two + live jobs never share a directory, but they do share the base, so a live + peer's directory is visible here -- it is simply orders of magnitude younger + than the threshold (a day, against generations measured in minutes to + hours). Published datasets and anything outside ``base`` are never touched. + Failures are logged and ignored: cleanup is opportunistic and must never + break the decision it runs inside. + """ + now = time.time() + for path in base.iterdir(): + if not path.name.startswith(TMP_PREFIX) or not path.is_dir(): + continue + try: + # Newest mtime among the staging dir and its immediate children: a + # bounded, cheap probe (no recursive stat storm over a partially + # generated dataset) that still notices a job that has started + # laying down its split directories. + newest = path.stat().st_mtime + for child in path.iterdir(): + newest = max(newest, child.stat().st_mtime) + except OSError as exc: + log.warning("Could not stat staging dir %s: %s", path, exc) + continue + + age = now - newest + if age < max_age: + continue + + log.info( + "Removing orphaned dataset staging dir %s (untouched for %.1f hours)", + path, + age / 3600.0, + ) + shutil.rmtree(path, ignore_errors=True) + + def _write_meta_atomic(meta_path: Path, meta: Dict[str, Any]) -> None: """Write ``meta`` to ``meta_path`` atomically. @@ -165,6 +217,10 @@ def _decide_reuse_or_generate( poison directory (exactly what a killed job leaves behind) must never be able to cause one. """ + # Rank 0 is the only rank that touches this base, so this is also the one + # safe place to reclaim staging dirs orphaned by earlier killed jobs. + _cleanup_stale_staging_dirs(base, log) + candidates = sorted( (p for p in base.iterdir() if p.is_dir()), key=lambda p: p.name, reverse=True ) diff --git a/tests/datagen/test_mpi_consensus.py b/tests/datagen/test_mpi_consensus.py index fa99a1b..5086150 100644 --- a/tests/datagen/test_mpi_consensus.py +++ b/tests/datagen/test_mpi_consensus.py @@ -39,7 +39,9 @@ from __future__ import annotations import logging +import os import re +import time from argparse import Namespace from math import ceil from pathlib import Path @@ -538,6 +540,72 @@ def test_non_root_raises_on_broadcast_finalize_error(tmp_path, monkeypatch): assert not dest.exists() +# --------------------------------------------------------------------------- +# R37: orphaned staging dirs are reclaimed instead of accumulating forever. +# --------------------------------------------------------------------------- + + +def _age_tree(path: Path, seconds: float) -> None: + """Backdate ``path`` and everything under it by ``seconds``.""" + stamp = time.time() - seconds + for entry in sorted(path.rglob("*"), reverse=True): + os.utime(entry, (stamp, stamp)) + os.utime(path, (stamp, stamp)) + + +def test_stale_staging_dirs_are_cleaned(tmp_path, monkeypatch): + """A long-orphaned ``.tmp_*`` dir is reclaimed; a live one is not. + + Every killed generation leaves a full staging tree behind (potentially + terabytes) that nothing ever reclaims. Cleanup is age-gated so a *running* + job's staging dir -- which by construction is far younger than the + threshold -- is never deleted out from under it. + """ + config = _reuse_config(tmp_path / "datasets") + base = _base_dir_for(config) + base.mkdir(parents=True) + + orphan = base / f"{gd.TMP_PREFIX}20200101-000000_111_deadbeef" + (orphan / "volumes" / "training").mkdir(parents=True) + (orphan / "volumes" / "training" / "0.npy").write_bytes(b"stale payload") + _age_tree(orphan, 10 * gd.STALE_STAGING_AGE_SECONDS) + + live = base / f"{gd.TMP_PREFIX}20260101-000000_222_cafebabe" + (live / "volumes").mkdir(parents=True) + + log = logging.getLogger("test_stale_staging_dirs_are_cleaned") + gd._decide_reuse_or_generate(base, base.name, "abc123", False, log) + + assert not orphan.exists(), "orphaned staging dir was not reclaimed" + assert live.exists(), "a concurrent job's staging dir was deleted" + + +def test_cleanup_never_touches_published_datasets(tmp_path): + """Only ``.tmp_*`` dirs under this config_id base are ever removed. + + An old published dataset is precisely what reuse is for, and other configs' + (or other users') directories are none of this job's business. + """ + config = _reuse_config(tmp_path / "datasets") + base = _base_dir_for(config) + base.mkdir(parents=True) + + published = _write_reusable_dataset(base, "some-other-config") + _age_tree(published, 10 * gd.STALE_STAGING_AGE_SECONDS) + + # A staging dir belonging to a different config_id base entirely. + other_base = base.parent / "0123456789ab" + other_orphan = other_base / f"{gd.TMP_PREFIX}20200101-000000_333_f00d" + other_orphan.mkdir(parents=True) + _age_tree(other_orphan, 10 * gd.STALE_STAGING_AGE_SECONDS) + + log = logging.getLogger("test_cleanup_never_touches_published_datasets") + gd._decide_reuse_or_generate(base, base.name, "abc123", False, log) + + assert published.exists(), "an old published dataset was deleted" + assert other_orphan.exists(), "cleanup escaped this job's config_id base" + + # --------------------------------------------------------------------------- # R28: meta.yaml is published atomically, so no reader ever sees a partial one. # --------------------------------------------------------------------------- From 74b0fd490e2b72790148754ccea6bd43a94beedf Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:32:23 -0700 Subject: [PATCH 13/62] Key the fractal library directory by seed Categories and instances are derived from config.seed but resume is a pure file-existence check, so a run under a new seed silently adopted another seed's library and published a dataset stamped with the wrong seed. Library paths now carry seed, so cross-seed reuse is impossible; old-layout libraries are not found and are regenerated. R29 --- ScaFFold/datagen/category_search.py | 11 +- ScaFFold/datagen/instance.py | 13 +- ScaFFold/datagen/layout.py | 67 +++++ ScaFFold/datagen/volumegen.py | 11 +- tests/datagen/test_artifacts.py | 9 +- tests/datagen/test_library_layout.py | 248 ++++++++++++++++++ tests/datagen/test_mpi_consensus.py | 24 +- .../datagen_get_dataset_consensus.py | 14 +- .../datagen_instance_partition.py | 5 +- 9 files changed, 353 insertions(+), 49 deletions(-) create mode 100644 ScaFFold/datagen/layout.py create mode 100644 tests/datagen/test_library_layout.py diff --git a/ScaFFold/datagen/category_search.py b/ScaFFold/datagen/category_search.py index 246666f..ae3b8fb 100644 --- a/ScaFFold/datagen/category_search.py +++ b/ScaFFold/datagen/category_search.py @@ -26,6 +26,7 @@ import numpy as np from mpi4py import MPI +from ScaFFold.datagen import layout from ScaFFold.datagen.generate_fractal_points import generate_fractal_points from ScaFFold.datagen.rng import SEED_MASK, derive_seed, seed_numba from ScaFFold.utils.config_utils import Config @@ -403,11 +404,11 @@ def main(config: Config) -> None: log.info("MPI size = %s", size) - # Setup directories - fracts_sub_dir = f"var{config.variance_threshold}" - fracts_write_dir = os.path.join( - config.fract_base_dir, fracts_sub_dir, "3DIFS_param" - ) + # Setup directories. The library is keyed by seed (see + # ScaFFold.datagen.layout): categories are drawn from a seed-derived + # candidate stream, so a run under a different seed must never resume onto + # another seed's parameter files. + fracts_write_dir = layout.category_param_dir(config) if rank == 0: log.info("Writing fractals to %s", fracts_write_dir) if os.path.exists(fracts_write_dir) and config.datagen_from_scratch: diff --git a/ScaFFold/datagen/instance.py b/ScaFFold/datagen/instance.py index e28d104..de8eb0d 100644 --- a/ScaFFold/datagen/instance.py +++ b/ScaFFold/datagen/instance.py @@ -28,6 +28,7 @@ import numpy as np from mpi4py import MPI +from ScaFFold.datagen import layout from ScaFFold.datagen.generate_fractal_points import generate_fractal_points from ScaFFold.datagen.rng import derive_seed, seed_numba from ScaFFold.utils.config_utils import Config @@ -239,12 +240,12 @@ def main(config: Config): log.info("MPI size = %s", size) - # Setup directories - fracts_sub_dir = f"var{config.variance_threshold}" - fracts_read_dir = os.path.join(config.fract_base_dir, fracts_sub_dir, "3DIFS_param") - instance_write_dir = os.path.join( - config.fract_base_dir, fracts_sub_dir, "instances", f"np{config.point_num}" - ) + # Setup directories. The library is keyed by seed (see + # ScaFFold.datagen.layout): every instance is generated from + # (seed, category, instance), so resuming onto another seed's files would + # silently mix data from two different seeds into one dataset. + fracts_read_dir = layout.category_param_dir(config) + instance_write_dir = layout.instance_dir(config) if rank == 0: log.info( "Generating instances for num_points=%s, writing to %s", diff --git a/ScaFFold/datagen/layout.py b/ScaFFold/datagen/layout.py new file mode 100644 index 0000000..e558c6d --- /dev/null +++ b/ScaFFold/datagen/layout.py @@ -0,0 +1,67 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +"""On-disk layout of the fractal library. + +Every artifact in the library is a deterministic function of the seed: category +IFS parameters come from a ``(seed, rank, attempt)`` candidate stream, and each +instance point cloud is generated from ``(seed, category, instance)``. Resume, +by contrast, is a pure file-existence test -- an instance is "already done" if +its file is on disk. + +Those two facts only compose safely if the path itself carries the seed. +Without it, a run under a new seed found the previous seed's files, generated +nothing, and produced a dataset whose metadata advertised the new seed while +its contents came from the old one. Keying the directory by seed makes the +existence question seed-specific, so data from one seed can never be mistaken +for another's: + + /var/seed/3DIFS_param/ + /var/seed/instances/np/ + +Libraries written under the older, seed-agnostic layout are simply not found +and are regenerated in the new location. + +These helpers are the single definition of that layout; every producer and +consumer (``category_search``, ``instance``, ``volumegen``) goes through them +so the two sides cannot drift apart. +""" + +from __future__ import annotations + +import os + + +def library_root(config) -> str: + """Return the root of the fractal library for this config's seed.""" + return os.path.join( + str(config.fract_base_dir), + f"var{config.variance_threshold}", + f"seed{int(config.seed)}", + ) + + +def category_param_dir(config) -> str: + """Return the directory holding this seed's category IFS parameter CSVs.""" + return os.path.join(library_root(config), "3DIFS_param") + + +def instance_dir(config) -> str: + """Return the directory holding this seed's instance point clouds. + + Instances are additionally keyed by point count, which is a property of the + cloud rather than of the category, so several point counts can coexist for + one seed. + """ + return os.path.join(library_root(config), "instances", f"np{config.point_num}") diff --git a/ScaFFold/datagen/volumegen.py b/ScaFFold/datagen/volumegen.py index cfd6d1b..05567e0 100644 --- a/ScaFFold/datagen/volumegen.py +++ b/ScaFFold/datagen/volumegen.py @@ -22,6 +22,7 @@ import numpy as np from mpi4py import MPI +from ScaFFold.datagen import layout from ScaFFold.utils.config_utils import Config from ScaFFold.utils.data_types import MASK_DTYPE, VOLUME_DTYPE from ScaFFold.utils.utils import setup_mpi_logger @@ -240,7 +241,10 @@ def main(config: Dict): fractal_colors = np.random.rand(config.n_categories, 3) grid_size = resolve_grid_size(config) - fract_base_dir = str(config.fract_base_dir) + # The instance library is keyed by seed (see ScaFFold.datagen + # .layout), so a volume can only ever be built from point clouds + # this run's seed produced. Resolved once, outside the loop. + instances_dir = layout.instance_dir(config) # Generation loop start_time = time.time() @@ -269,12 +273,7 @@ def main(config: Dict): curr_instance = curr_vol[1 + 2 * curr_fract + 1] fractal_color = fractal_colors[curr_category] - instances_dir = ( - f"var{config.variance_threshold}/instances/np{config.point_num}" - ) - point_cloud_path = os.path.join( - fract_base_dir, instances_dir, f"{curr_category:06d}", f"{curr_category:06d}_{curr_instance:04d}.npy", diff --git a/tests/datagen/test_artifacts.py b/tests/datagen/test_artifacts.py index 224fefa..2f969a6 100644 --- a/tests/datagen/test_artifacts.py +++ b/tests/datagen/test_artifacts.py @@ -37,6 +37,7 @@ import pytest from ScaFFold.datagen import instance as inst +from ScaFFold.datagen import layout from ScaFFold.datagen import mask_detection as md from ScaFFold.datagen.volumegen import ( load_np_ptcloud, @@ -74,13 +75,15 @@ def _seed_category(fract_base: Path, *, point_num: int, keep: range) -> Path: Pre-seeding all but instance 0 means a ``main`` run only has to generate the single missing instance, keeping the test fast. Returns the instance dir. + The library lives under the seed-keyed layout, so the paths are derived from + the same config the run under test uses. """ - vt = 0.15 - param_dir = fract_base / f"var{vt}" / "3DIFS_param" + config = _make_config(fract_base, point_num=point_num) + param_dir = Path(layout.category_param_dir(config)) param_dir.mkdir(parents=True) np.savetxt(param_dir / "000000.csv", _contractive_params(), delimiter=",") - inst_dir = fract_base / f"var{vt}" / "instances" / f"np{point_num}" / "000000" + inst_dir = Path(layout.instance_dir(config)) / "000000" inst_dir.mkdir(parents=True) rng = np.random.default_rng(0) for i in keep: diff --git a/tests/datagen/test_library_layout.py b/tests/datagen/test_library_layout.py new file mode 100644 index 0000000..d09a524 --- /dev/null +++ b/tests/datagen/test_library_layout.py @@ -0,0 +1,248 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +"""The fractal library is keyed by the seed that produced it (R29). + +Categories and instances are *derived from* ``config.seed``: the IFS parameters +come from a seed-keyed candidate stream and every instance point cloud is +seeded from ``(seed, category, instance)``. Resume, however, is a pure +file-existence check, so a library laid out only by variance threshold and +point count let a run under one seed silently adopt another seed's data -- +and then publish a dataset whose metadata claimed the new seed. Two datasets +with identical provenance and different content. + +The fix puts the seed in the directory path, so the question "does this file +exist" is asked in a seed-specific place and can only ever be answered with +data that seed produced: + + /var/seed/3DIFS_param/ + /var/seed/instances/np/ + +Everything here runs single-process at tiny scale (one category, 60-point +clouds) so the real generators run in well under a second. +""" + +from __future__ import annotations + +import hashlib +from argparse import Namespace +from pathlib import Path + +import numpy as np +import pytest + +from ScaFFold.datagen import category_search as cs +from ScaFFold.datagen import instance as inst +from ScaFFold.datagen import layout, volumegen + +VT = 0.15 +POINT_NUM = 60 + + +def _contractive_params() -> np.ndarray: + """A 2-map IFS whose orbit stays bounded, so generation is fast and finite.""" + params = np.zeros((2, 13), dtype=np.float64) + params[:, 0] = params[:, 4] = params[:, 8] = 0.5 + params[1, 9] = params[1, 10] = params[1, 11] = 0.5 + params[0, 12] = 0.5 + return params + + +def _param_dir(fract_base: Path, seed: int) -> Path: + """The category directory the new layout mandates, spelled out literally.""" + return fract_base / f"var{VT}" / f"seed{seed}" / "3DIFS_param" + + +def _instance_dir(fract_base: Path, seed: int) -> Path: + """The instance directory the new layout mandates, spelled out literally.""" + return fract_base / f"var{VT}" / f"seed{seed}" / "instances" / f"np{POINT_NUM}" + + +def _seed_params(fract_base: Path, seed: int, n_categories: int = 1) -> Path: + param_dir = _param_dir(fract_base, seed) + param_dir.mkdir(parents=True, exist_ok=True) + for category in range(n_categories): + np.savetxt( + param_dir / f"{category:06d}.csv", _contractive_params(), delimiter="," + ) + return param_dir + + +def _inst_config(fract_base: Path, seed: int) -> Namespace: + return Namespace( + fract_base_dir=str(fract_base), + n_categories=1, + seed=seed, + variance_threshold=VT, + point_num=POINT_NUM, + datagen_from_scratch=False, + verbose=0, + ) + + +def _cs_config(fract_base: Path, seed: int) -> Namespace: + return Namespace( + fract_base_dir=str(fract_base), + n_categories=1, + seed=seed, + variance_threshold=VT, + point_num=POINT_NUM, + normalize=1, + datagen_from_scratch=False, + datagen_batch_size=4, + verbose=0, + ) + + +def _volumegen_config(dataset_dir: Path, fract_base: Path, seed: int) -> Namespace: + return Namespace( + dataset_dir=str(dataset_dir), + fract_base_dir=str(fract_base), + n_categories=1, + n_instances_used_per_fractal=1, + n_fracts_per_vol=1, + seed=seed, + variance_threshold=VT, + val_split=0, + vol_size=8, + point_num=POINT_NUM, + scale=1, + verbose=0, + ) + + +def _library_digest(instance_dir: Path) -> tuple[int, str]: + """(file count, content digest) for one instance directory.""" + digest = hashlib.sha256() + files = sorted(instance_dir.rglob("*.npy")) + for path in files: + digest.update(path.name.encode()) + digest.update(path.read_bytes()) + return len(files), digest.hexdigest() + + +# --------------------------------------------------------------------------- +# The layout helper is the single definition of the seed-keyed paths. +# --------------------------------------------------------------------------- + + +def test_layout_helpers_key_every_path_by_seed(tmp_path): + """Both library directories carry the seed, and differ across seeds.""" + fract_base = tmp_path / "fractals" + + for seed in (7, 999): + config = _inst_config(fract_base, seed) + assert Path(layout.category_param_dir(config)) == _param_dir(fract_base, seed) + assert Path(layout.instance_dir(config)) == _instance_dir(fract_base, seed) + + assert layout.category_param_dir(_inst_config(fract_base, 7)) != ( + layout.category_param_dir(_inst_config(fract_base, 999)) + ) + assert layout.instance_dir(_inst_config(fract_base, 7)) != ( + layout.instance_dir(_inst_config(fract_base, 999)) + ) + + +# --------------------------------------------------------------------------- +# Producers write under the seed; consumers read from under the seed. +# --------------------------------------------------------------------------- + + +def test_category_search_writes_under_the_seed_dir(tmp_path): + """A generated category CSV lands in this seed's parameter directory.""" + fract_base = tmp_path / "fractals" + + cs.main(_cs_config(fract_base, seed=42)) + + assert (_param_dir(fract_base, 42) / "000000.csv").exists() + # Nothing was written to a seed-agnostic location. + assert not (fract_base / f"var{VT}" / "3DIFS_param").exists() + + +def test_instances_are_written_under_the_seed_dir(tmp_path): + """Instance point clouds land under this seed's instance directory.""" + fract_base = tmp_path / "fractals" + _seed_params(fract_base, seed=7) + + inst.main(_inst_config(fract_base, seed=7)) + + count, _digest = _library_digest(_instance_dir(fract_base, 7)) + assert count == 145 + assert not (fract_base / f"var{VT}" / "instances").exists() + + +def test_same_seed_resume_generates_nothing_new(tmp_path): + """Re-running under the same seed reuses the library byte-for-byte. + + The resume path must stay cheap: the whole point of the library is that a + second run under the same configuration regenerates nothing. + """ + fract_base = tmp_path / "fractals" + _seed_params(fract_base, seed=7) + inst.main(_inst_config(fract_base, seed=7)) + + instance_dir = _instance_dir(fract_base, 7) + count_before, digest_before = _library_digest(instance_dir) + mtimes_before = {p: p.stat().st_mtime_ns for p in sorted(instance_dir.rglob("*"))} + + inst.main(_inst_config(fract_base, seed=7)) + + count_after, digest_after = _library_digest(instance_dir) + assert (count_after, digest_after) == (count_before, digest_before) + # No file was rewritten (0 new instances generated). + assert {p: p.stat().st_mtime_ns for p in sorted(instance_dir.rglob("*"))} == ( + mtimes_before + ) + + +def test_different_seed_cannot_reuse_another_seeds_instances(tmp_path): + """A second seed generates its own library instead of adopting the first. + + Before the fix this run reported "Generated 0 instances" and left the + first seed's bytes in place, so the dataset built on top of it carried the + wrong seed's data under the new seed's provenance. + """ + fract_base = tmp_path / "fractals" + _seed_params(fract_base, seed=7) + _seed_params(fract_base, seed=999) + + inst.main(_inst_config(fract_base, seed=7)) + count_7, digest_7 = _library_digest(_instance_dir(fract_base, 7)) + + inst.main(_inst_config(fract_base, seed=999)) + count_999, digest_999 = _library_digest(_instance_dir(fract_base, 999)) + + # Both libraries are complete and independent... + assert count_7 == count_999 == 145 + assert digest_999 != digest_7, "seed 999 reused seed 7's instances" + # ...and the first seed's data was left untouched. + assert _library_digest(_instance_dir(fract_base, 7)) == (count_7, digest_7) + + +def test_volumegen_reads_instances_for_its_own_seed(tmp_path): + """volumegen resolves point clouds under the seed it was configured with.""" + fract_base = tmp_path / "fractals" + _seed_params(fract_base, seed=7) + inst.main(_inst_config(fract_base, seed=7)) + + # Seed 7: the instances are where volumegen looks, so generation succeeds. + volumegen.main(_volumegen_config(tmp_path / "ds7", fract_base, seed=7)) + assert list((tmp_path / "ds7" / "volumes").rglob("*.npy")) + + # Seed 999: a different library entirely. Nothing has been generated for + # it, so volumegen must report the missing file rather than quietly + # rasterizing seed 7's clouds. + with pytest.raises(RuntimeError) as excinfo: + volumegen.main(_volumegen_config(tmp_path / "ds999", fract_base, seed=999)) + assert "seed999" in str(excinfo.value) diff --git a/tests/datagen/test_mpi_consensus.py b/tests/datagen/test_mpi_consensus.py index 5086150..c98fc5b 100644 --- a/tests/datagen/test_mpi_consensus.py +++ b/tests/datagen/test_mpi_consensus.py @@ -53,7 +53,7 @@ from ScaFFold.datagen import get_dataset as gd from ScaFFold.datagen import instance as inst -from ScaFFold.datagen import volumegen +from ScaFFold.datagen import layout, volumegen RANK_SCRIPTS = Path(__file__).resolve().parents[1] / "helpers" / "rank_scripts" @@ -684,15 +684,10 @@ def _seed_one_instance(fract_base: Path, config: Namespace, *, present: bool) -> volumegen selects instance indices with ``random.sample(range(145), ...)`` seeded by ``config.seed``; to be robust we populate every one of the 145 - instance slots for category 0 when ``present`` is True. + instance slots for category 0 when ``present`` is True. The library path is + seed-keyed, so it is derived from the same config the run under test uses. """ - inst_dir = ( - fract_base - / f"var{config.variance_threshold}" - / "instances" - / f"np{config.point_num}" - / "000000" - ) + inst_dir = Path(layout.instance_dir(config)) / "000000" inst_dir.mkdir(parents=True, exist_ok=True) if present: rng = np.random.default_rng(0) @@ -803,7 +798,7 @@ def _instance_config(fract_base: Path) -> Namespace: def _seed_ifs_params(fract_base: Path, config: Namespace, n_categories: int) -> None: """Write a contractive IFS param CSV per category so generation stays fast.""" - param_dir = fract_base / f"var{config.variance_threshold}" / "3DIFS_param" + param_dir = Path(layout.category_param_dir(config)) param_dir.mkdir(parents=True, exist_ok=True) params = np.zeros((2, 13), dtype=np.float64) params[:, 0] = params[:, 4] = params[:, 8] = 0.5 @@ -869,14 +864,7 @@ def no_scan(*_args, **_kwargs): assert rc == 0 # Rank 1 received the broadcast list and generated its share (pair [1, 0]). - generated = ( - fract_base - / f"var{config.variance_threshold}" - / "instances" - / f"np{config.point_num}" - / "000001" - / "000001_0000.npy" - ) + generated = Path(layout.instance_dir(config)) / "000001" / "000001_0000.npy" assert generated.exists() diff --git a/tests/helpers/rank_scripts/datagen_get_dataset_consensus.py b/tests/helpers/rank_scripts/datagen_get_dataset_consensus.py index 3799874..bc7e7a4 100644 --- a/tests/helpers/rank_scripts/datagen_get_dataset_consensus.py +++ b/tests/helpers/rank_scripts/datagen_get_dataset_consensus.py @@ -45,7 +45,7 @@ from mpi4py import MPI import ScaFFold.datagen.get_dataset as gd -from ScaFFold.datagen import volumegen +from ScaFFold.datagen import layout, volumegen VT = 0.15 PN = 64 @@ -71,14 +71,10 @@ def _config(dataset_dir: Path, fract_base: Path) -> Namespace: def _instance_path(fract_base: Path, cat: int, inst: int) -> Path: - return ( - fract_base - / f"var{VT}" - / "instances" - / f"np{PN}" - / f"{cat:06d}" - / f"{cat:06d}_{inst:04d}.npy" - ) + # The instance library is keyed by seed; derive the path from the same + # config the run under test uses. + inst_root = Path(layout.instance_dir(_config(Path("unused"), fract_base))) + return inst_root / f"{cat:06d}" / f"{cat:06d}_{inst:04d}.npy" def _seed_instances(fract_base: Path) -> None: diff --git a/tests/helpers/rank_scripts/datagen_instance_partition.py b/tests/helpers/rank_scripts/datagen_instance_partition.py index 10df10f..37d0c98 100644 --- a/tests/helpers/rank_scripts/datagen_instance_partition.py +++ b/tests/helpers/rank_scripts/datagen_instance_partition.py @@ -44,6 +44,7 @@ from mpi4py import MPI import ScaFFold.datagen.instance as inst +from ScaFFold.datagen import layout VT = 0.15 PN = 64 @@ -62,7 +63,7 @@ def _config(fract_base: Path) -> Namespace: def _seed_ifs_params(fract_base: Path) -> None: - param_dir = fract_base / f"var{VT}" / "3DIFS_param" + param_dir = Path(layout.category_param_dir(_config(fract_base))) param_dir.mkdir(parents=True, exist_ok=True) params = np.zeros((2, 13), dtype=np.float64) params[:, 0] = params[:, 4] = params[:, 8] = 0.5 @@ -82,7 +83,7 @@ def main() -> None: _seed_ifs_params(fract_base) comm.Barrier() - inst_root = fract_base / f"var{VT}" / "instances" / f"np{PN}" + inst_root = Path(layout.instance_dir(_config(fract_base))) # Give each rank a divergent view of pre-existing instances. Only rank 0's # view should matter after the fix (its list is broadcast); rank 1's phantom From aa088c176bf8c6e79ca1185c1e4ddaffa2e1cec8 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:33:57 -0700 Subject: [PATCH 14/62] Scan for existing categories once on rank 0 and broadcast category_search derived its loop-gating remaining count from a per-rank filesystem scan, so ranks with divergent views could post mismatched collectives (bcast against reduce) and hang. Rank 0 now scans and broadcasts the existing-index list, mirroring the instance.py work-list fix. R30 --- ScaFFold/datagen/category_search.py | 17 +++- tests/datagen/test_category_search.py | 139 +++++++++++++++++++++++++- 2 files changed, 152 insertions(+), 4 deletions(-) diff --git a/ScaFFold/datagen/category_search.py b/ScaFFold/datagen/category_search.py index ae3b8fb..98e2342 100644 --- a/ScaFFold/datagen/category_search.py +++ b/ScaFFold/datagen/category_search.py @@ -424,9 +424,20 @@ def main(config: Config) -> None: # the ones a fresh run produced. attempt_index = read_attempt_counter(fracts_write_dir, rank) - # Parse existing category files (rank 0 owns saving/dedup). Free indices are - # derived from these parsed names -- filling holes, never overwriting. - existing_indices = parse_category_indices(fracts_write_dir) + # Parse existing category files on rank 0 alone and broadcast the result. + # Free indices are derived from these parsed names -- filling holes, never + # overwriting -- and, critically, the count derived below gates a loop that + # contains collectives. Scanning the shared filesystem independently per + # rank lets divergent views (stale metadata caches, a concurrent job, a + # partially visible directory) put one rank inside the loop while another is + # past it, so the two post mismatched collectives on COMM_WORLD and the job + # hangs. One scan, one broadcast, one shared verdict. + if rank == 0: + existing_indices = parse_category_indices(fracts_write_dir) + else: + existing_indices = None + existing_indices = comm.bcast(existing_indices, root=0) + existing_params = [] if rank == 0: for idx in existing_indices: diff --git a/tests/datagen/test_category_search.py b/tests/datagen/test_category_search.py index 8608884..90236a8 100644 --- a/tests/datagen/test_category_search.py +++ b/tests/datagen/test_category_search.py @@ -12,8 +12,15 @@ # # SPDX-License-Identifier: (Apache-2.0) -"""Tests for category-search round sizing.""" +"""Tests for category-search round sizing and its work-scan consensus.""" +from argparse import Namespace +from pathlib import Path + +import numpy as np + +from ScaFFold.datagen import category_search as cs +from ScaFFold.datagen import layout from ScaFFold.datagen.category_search import compute_round_attempts @@ -53,3 +60,133 @@ def test_at_least_one_attempt_per_rank_when_work_remains(): # A positive remaining count always yields at least one attempt per rank so # the loop makes progress and can keep learning the acceptance rate. assert compute_round_attempts(1, 1024, 10000, 0.99) >= 1 + + +# --------------------------------------------------------------------------- +# R30: the initial work scan is made once on rank 0 and broadcast. +# +# ``categories_remaining`` gates a while loop that contains collectives, so it +# must be identical on every rank. Deriving it from a per-rank filesystem scan +# lets divergent views (a stale metadata cache, a racing job, a partially +# visible directory) put one rank inside the loop issuing ``bcast`` while +# another is past it issuing ``reduce`` -- mismatched collectives on +# COMM_WORLD, i.e. a hang. This mirrors the fix already applied in +# ``instance.py``: rank 0 scans, everyone else consumes the broadcast. +# --------------------------------------------------------------------------- + + +class CategorySearchComm: + """Single-process stand-in for COMM_WORLD recording the collective order.""" + + def __init__(self, rank=0, size=1, bcast_returns=None): + self.rank = rank + self.size = size + self.calls = [] + self.bcast_payloads = [] + self._bcast_returns = list(bcast_returns or []) + + def Get_rank(self): + return self.rank + + def Get_size(self): + return self.size + + def Barrier(self): + self.calls.append("Barrier") + + def bcast(self, obj, root=0): + self.calls.append("bcast") + self.bcast_payloads.append(obj) + if self.rank == root: + return obj + return self._bcast_returns.pop(0) + + def gather(self, obj, root=0): + self.calls.append("gather") + return [obj] if self.rank == root else None + + def reduce(self, value, op=None, root=0): + self.calls.append("reduce") + return value if self.rank == root else None + + +class FakeMPI: + """Namespace mimicking ``mpi4py.MPI`` for one ``CategorySearchComm``.""" + + def __init__(self, comm): + import mpi4py.MPI as real_mpi + + self.COMM_WORLD = comm + self.SUM = real_mpi.SUM + + +def _cs_config(fract_base: Path) -> Namespace: + return Namespace( + fract_base_dir=str(fract_base), + n_categories=1, + seed=42, + variance_threshold=0.15, + point_num=60, + normalize=1, + datagen_from_scratch=False, + datagen_batch_size=4, + verbose=0, + ) + + +def _seed_one_category(config: Namespace) -> None: + """Write the single category CSV this config asks for.""" + param_dir = Path(layout.category_param_dir(config)) + param_dir.mkdir(parents=True, exist_ok=True) + params = np.zeros((2, 13), dtype=np.float64) + params[:, 0] = params[:, 4] = params[:, 8] = 0.5 + params[1, 9] = params[1, 10] = params[1, 11] = 0.5 + params[0, 12] = 0.5 + np.savetxt(param_dir / "000000.csv", params, delimiter=",") + + +def test_work_scan_is_root_only_and_broadcast(tmp_path, monkeypatch): + """A non-root rank never scans; it consumes root's broadcast index list.""" + config = _cs_config(tmp_path / "fractals") + comm = CategorySearchComm(rank=1, size=2, bcast_returns=[[0]]) + monkeypatch.setattr(cs, "MPI", FakeMPI(comm)) + + def no_scan(*_args, **_kwargs): + raise AssertionError("non-root rank must not scan the filesystem") + + monkeypatch.setattr(cs, "parse_category_indices", no_scan) + + cs.main(config) + + # Root said category 0 already exists, so this rank has nothing to do and + # went straight to the post-loop reductions. + assert comm.calls == ["Barrier", "bcast", "reduce", "reduce", "reduce", "reduce"] + # It contributed nothing to the scan broadcast (it is not the scanner). + assert comm.bcast_payloads[0] is None + + +def test_divergent_fs_views_take_the_same_collective_path(tmp_path, monkeypatch): + """Ranks disagreeing about the directory still issue identical collectives. + + Rank 0 sees the finished category; rank 1's own view is empty. Before the + fix rank 1 entered the work loop (``bcast``) while rank 0 was already past + it (``reduce``). With the scan broadcast, rank 1's view is irrelevant. + """ + # Rank 0: the category is on disk, so its scan finds it. + root_config = _cs_config(tmp_path / "root_view") + _seed_one_category(root_config) + root_comm = CategorySearchComm(rank=0, size=2) + monkeypatch.setattr(cs, "MPI", FakeMPI(root_comm)) + cs.main(root_config) + + # Rank 1: an empty directory (a divergent view), but root broadcast [0]. + peer_config = _cs_config(tmp_path / "peer_view") + peer_comm = CategorySearchComm(rank=1, size=2, bcast_returns=[[0]]) + monkeypatch.setattr(cs, "MPI", FakeMPI(peer_comm)) + cs.main(peer_config) + + assert root_comm.bcast_payloads[0] == [0] + assert peer_comm.calls == root_comm.calls, ( + "ranks with divergent filesystem views issued different collectives: " + f"rank 0 {root_comm.calls} vs rank 1 {peer_comm.calls}" + ) From 26bd51239b632c0558259132bd9e89f84114d689 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:35:32 -0700 Subject: [PATCH 15/62] Write category parameter CSVs atomically A category CSV truncated by a killed job kept its six-digit name, so the resume scan counted it as done forever while instance generation and the search's own resume both died parsing it. Categories are now staged under a temp name that no scan matches, fsynced, and renamed into place. R31 --- ScaFFold/datagen/category_search.py | 31 ++++++++++++- tests/datagen/test_category_search.py | 64 +++++++++++++++++++++++++++ 2 files changed, 94 insertions(+), 1 deletion(-) diff --git a/ScaFFold/datagen/category_search.py b/ScaFFold/datagen/category_search.py index 98e2342..af53a9c 100644 --- a/ScaFFold/datagen/category_search.py +++ b/ScaFFold/datagen/category_search.py @@ -335,12 +335,41 @@ def save_valid_category( target = os.path.join(fracts_write_dir, "%06d.csv" % idx) if os.path.exists(target): raise FileExistsError(f"Refusing to overwrite existing category file: {target}") - np.savetxt(target, params, delimiter=",") + _savetxt_atomic(target, params) existing_indices.append(idx) existing_params.append(params) return idx +def _savetxt_atomic(target: str, params: np.array) -> None: + """Write one category's parameters to ``target`` atomically. + + A category CSV truncated by a killed job is poison: the six-digit name is + all the resume scan looks at, so the category counts as done forever, while + every consumer (instance generation, and the search's own resume) dies + parsing it. The file is therefore written to a temp name in the same + directory -- one that neither the resume glob (``NNNNNN.csv``) nor the + instance loader's ``*.csv`` filter can match -- flushed, fsynced, and only + then ``os.replace``d onto the final name. + """ + directory, name = os.path.split(target) + tmp_path = os.path.join(directory, f".{name}.tmp{os.getpid()}") + try: + with open(tmp_path, "w") as handle: + np.savetxt(handle, params, delimiter=",") + handle.flush() + os.fsync(handle.fileno()) + os.replace(tmp_path, target) + except BaseException: + # A failed write must leave nothing behind: no temp file, and no + # partial file under the name resume would accept. + try: + os.remove(tmp_path) + except OSError: + pass + raise + + def _attempt_state_path(fracts_write_dir: str, rank: int) -> str: return os.path.join(fracts_write_dir, f".rng_attempt_rank{rank}") diff --git a/tests/datagen/test_category_search.py b/tests/datagen/test_category_search.py index 90236a8..2cc43b1 100644 --- a/tests/datagen/test_category_search.py +++ b/tests/datagen/test_category_search.py @@ -18,6 +18,7 @@ from pathlib import Path import numpy as np +import pytest from ScaFFold.datagen import category_search as cs from ScaFFold.datagen import layout @@ -190,3 +191,66 @@ def test_divergent_fs_views_take_the_same_collective_path(tmp_path, monkeypatch) "ranks with divergent filesystem views issued different collectives: " f"rank 0 {root_comm.calls} vs rank 1 {peer_comm.calls}" ) + + +# --------------------------------------------------------------------------- +# R31: category CSVs appear complete or not at all. +# +# A category file truncated by a killed job is still counted as "done" by the +# resume scan, so nothing ever regenerates it: instance generation then dies +# parsing it, and the category search's own resume dies re-loading it. The +# pipeline cannot self-heal -- the file has to be deleted by hand. +# --------------------------------------------------------------------------- + + +def _params() -> np.ndarray: + params = np.zeros((2, 13), dtype=np.float64) + params[:, 0] = params[:, 4] = params[:, 8] = 0.5 + params[0, 12] = 0.5 + return params + + +def test_category_csv_write_is_atomic(tmp_path, monkeypatch): + """A killed mid-write leaves no category file under a name resume accepts.""" + param_dir = tmp_path / "3DIFS_param" + param_dir.mkdir() + + # One complete category, saved normally. + indices, saved = [], [] + first = _params() + assert cs.save_valid_category(str(param_dir), first, indices, saved) == 0 + assert np.loadtxt(param_dir / "000000.csv", delimiter=",").shape == (2, 13) + + # The next save is interrupted after some bytes have been written. + observed = {} + + def partial_then_raise(fname, arr, *args, **kwargs): + observed["listing"] = sorted(p.name for p in param_dir.iterdir()) + handle = fname if hasattr(fname, "write") else open(fname, "w") + handle.write("0.5,0.5,0.5\n") + handle.flush() + if handle is not fname: + handle.close() + raise OSError("simulated SIGKILL mid-write") + + monkeypatch.setattr(cs.np, "savetxt", partial_then_raise) + + second = _params() + second[0, 0] = 0.25 + with pytest.raises(OSError): + cs.save_valid_category(str(param_dir), second, indices, saved) + + # No truncated category is visible: the resume scan still sees exactly the + # one complete category, and every file it names parses. + assert cs.parse_category_indices(str(param_dir)) == [0] + assert not (param_dir / "000001.csv").exists() + for idx in cs.parse_category_indices(str(param_dir)): + assert np.loadtxt(param_dir / f"{idx:06d}.csv", delimiter=",").shape == (2, 13) + + # Mid-write, the partial data lived under a name neither the resume scan + # nor the instance loader (which takes every ``*.csv``) would pick up. + partial_names = [n for n in observed["listing"] if n != "000000.csv"] + assert all(not name.endswith(".csv") for name in partial_names), partial_names + + # And nothing was left behind afterwards. + assert sorted(p.name for p in param_dir.iterdir()) == ["000000.csv"] From 9b028fa856d8f27b85d67c8b448b83b636f6ac02 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:39:26 -0700 Subject: [PATCH 16/62] Correct the half-voxel shift in voxelization The centering offset subtracted an extra voxel, biasing every cloud toward the origin: the first half-voxel of each filled axis floored to -1 and was clipped onto plane 0 (1.5x the interior density, with the far plane at 0.5x). DATASET_FORMAT_VERSION is bumped so misregistered datasets are regenerated rather than reused. R33 --- ScaFFold/datagen/get_dataset.py | 14 +++-- ScaFFold/datagen/volumegen.py | 10 +++- tests/datagen/test_artifacts.py | 97 ++++++++++++++++++++++++++++++--- 3 files changed, 105 insertions(+), 16 deletions(-) diff --git a/ScaFFold/datagen/get_dataset.py b/ScaFFold/datagen/get_dataset.py index 0b193f2..f76f208 100644 --- a/ScaFFold/datagen/get_dataset.py +++ b/ScaFFold/datagen/get_dataset.py @@ -44,11 +44,15 @@ # Bumped from 2 to 3 when instance point clouds moved from float64 to float32: # the storage layout is unchanged, but float32 voxel binning shifts a handful of # boundary voxels, so a float64-era dataset must not be reused as if it were -# float32. This version stamps new datasets, gates reuse below, and feeds the -# config_id hash, so an older dataset is neither matched nor scanned. The loader -# in data_loading.py keeps its own (lower) minimum-layout version and still reads -# v3 through the modern dense path. -DATASET_FORMAT_VERSION = 3 +# float32. Bumped from 3 to 4 when the voxel centering offset was corrected from +# (grid_size - 1 - span)/2 to (grid_size - span)/2: every volume and mask +# generated before that was misregistered by half a voxel (with the first +# half-voxel of each axis clipped onto plane 0), so those datasets must be +# regenerated rather than reused. This version stamps new datasets, gates reuse +# below, and feeds the config_id hash, so an older dataset is neither matched nor +# scanned. The loader in data_loading.py keeps its own (lower) minimum-layout +# version and still reads v4 through the modern dense path. +DATASET_FORMAT_VERSION = 4 INCLUDE_KEYS = [ "dataset_format_version", "n_categories", diff --git a/ScaFFold/datagen/volumegen.py b/ScaFFold/datagen/volumegen.py index 05567e0..34e17f9 100644 --- a/ScaFFold/datagen/volumegen.py +++ b/ScaFFold/datagen/volumegen.py @@ -82,9 +82,15 @@ def points_to_voxel_indices( scaled = (points - mins) / voxel_size # 4) Center the occupied region: the largest axis fills the grid while the - # shorter axes are offset so their span sits in the middle. + # shorter axes are offset so their span sits in the middle. The free + # space to split between the two margins is (grid_size - span) voxels, + # measured in the same voxel units as ``scaled``; subtracting an extra 1 + # (as if the offset were an index rather than a length) shifted every + # cloud half a voxel toward the origin, floored the first half-voxel of + # each filled axis to -1, and let the clip below fold those points onto + # plane 0. span = scaled.max(axis=0) - offset = (grid_size - 1 - span) / 2.0 + offset = (grid_size - span) / 2.0 idx = np.floor(scaled + offset).astype(int) # 5) Clip to valid range (guards float rounding at the boundaries). diff --git a/tests/datagen/test_artifacts.py b/tests/datagen/test_artifacts.py index 2f969a6..c1c38e3 100644 --- a/tests/datagen/test_artifacts.py +++ b/tests/datagen/test_artifacts.py @@ -312,20 +312,22 @@ def test_scale_config_rejected(): # --------------------------------------------------------------------------- -def _dense_reference_indices(points: np.ndarray, grid_size: int, eps=1e-6): - """The pre-refactor index computation, kept verbatim as a reference. - - This is the exact math the old dense ``points_to_voxelgrid`` ran before - scattering ``True`` into a full ``grid_size**3`` boolean array; the scatter - API must reproduce the identical occupied-voxel set and painted values. +def _reference_indices(points: np.ndarray, grid_size: int, eps=1e-6, *, clip=True): + """The index computation of ``points_to_voxel_indices``, spelled out. + + Two tests need to see inside the function: the scatter-vs-dense equivalence + check (which needs the per-point indices the dense grid was built from) and + the centering check (which needs the indices *before* ``np.clip`` hides + out-of-range bins). Every test using this asserts the replica reproduces the + real function's output, so it cannot silently drift from it. """ mins = points.min(axis=0) maxs = points.max(axis=0) voxel_size = (float((maxs - mins).max()) + eps) / grid_size scaled = (points - mins) / voxel_size - offset = (grid_size - 1 - scaled.max(axis=0)) / 2.0 + offset = (grid_size - scaled.max(axis=0)) / 2.0 idx = np.floor(scaled + offset).astype(int) - return np.clip(idx, 0, grid_size - 1) + return np.clip(idx, 0, grid_size - 1) if clip else idx def test_voxel_indices_match_dense_grid(): @@ -337,7 +339,7 @@ def test_voxel_indices_match_dense_grid(): idx = points_to_voxel_indices(points, grid_size) # Reference dense grid built the old way, from the reference index math. - ref_idx = _dense_reference_indices(points, grid_size) + ref_idx = _reference_indices(points, grid_size) reference = np.zeros((grid_size,) * 3, dtype=bool) reference[ref_idx[:, 0], ref_idx[:, 1], ref_idx[:, 2]] = True @@ -421,3 +423,80 @@ def test_dataset_version_bumped_past_float64_era(): from ScaFFold.datagen import get_dataset as gd assert gd.DATASET_FORMAT_VERSION > 2 + + +# --------------------------------------------------------------------------- +# R33: voxel centering is a whole voxel, not half of one +# --------------------------------------------------------------------------- + + +def _assert_replica_tracks_real(grid_size: int = 16) -> None: + """Pin ``_reference_indices`` to the real function on a sparse cloud. + + A sparse cloud is essential here: a dense one occupies every voxel under + any offset, so the comparison would pass vacuously. + """ + sparse = np.random.default_rng(7).random((300, 3)).astype(np.float32) + assert np.array_equal( + np.unique(_reference_indices(sparse, grid_size), axis=0), + points_to_voxel_indices(sparse, grid_size), + ), "the replicated index arithmetic no longer matches points_to_voxel_indices" + + +def test_voxelization_never_bins_outside_the_grid(): + """No point lands outside ``[0, grid_size)`` before the safety clip. + + The centering offset positions a span of ``span`` voxels inside a grid of + ``grid_size`` voxels, so the free space to split between the two margins is + ``grid_size - span``. Using ``grid_size - 1 - span`` shifted every cloud + half a voxel toward the origin: points in the first half-voxel of each + filled axis floored to -1, and ``np.clip`` quietly folded them into bin 0. + """ + grid_size = 16 + _assert_replica_tracks_real(grid_size) + rng = np.random.default_rng(1234) + points = rng.random((200_000, 3)).astype(np.float32) + + pre_clip = _reference_indices(points, grid_size, clip=False) + assert pre_clip.min() >= 0, ( + f"{int((pre_clip < 0).any(axis=1).sum())} of {len(points)} points floored " + "below bin 0 and were clipped back in" + ) + assert pre_clip.max() <= grid_size - 1 + + +def test_voxelization_density_is_uniform_at_the_boundaries(): + """A uniform cloud fills the boundary planes like the interior ones. + + The half-voxel shift piled the clipped points onto plane 0 (1.5x the + interior density) and starved the far plane (0.5x), a systematic + misregistration in every generated volume and mask. + """ + grid_size = 16 + _assert_replica_tracks_real(grid_size) + rng = np.random.default_rng(1234) + points = rng.random((200_000, 3)).astype(np.float32) + + idx = _reference_indices(points, grid_size) + counts = np.bincount(idx[:, 0], minlength=grid_size) + interior = counts[2:-2].mean() + assert 0.9 <= counts[0] / interior <= 1.1, ( + f"boundary plane 0 holds {counts[0] / interior:.2f}x the interior density" + ) + assert 0.9 <= counts[-1] / interior <= 1.1, ( + f"boundary plane {grid_size - 1} holds {counts[-1] / interior:.2f}x the " + "interior density" + ) + + +def test_dataset_version_bumped_past_half_voxel_era(): + """The reuse marker advanced past 3: pre-fix datasets are misregistered. + + Correcting the offset changes the voxel contents of every generated volume + and mask, so a dataset built before the fix must not be handed to a run + after it. Bumping the version both stops the reuse scan from matching those + directories and changes the config_id they hash to. + """ + from ScaFFold.datagen import get_dataset as gd + + assert gd.DATASET_FORMAT_VERSION > 3 From d52dd07a9a976ba11cfd20d1610df712cb4376ee Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:41:00 -0700 Subject: [PATCH 17/62] Bound the int16 mask carrier by the largest class id The guard compared the class count against the int16 limit, but v2 masks ship raw category ids and a sparse split lists only the categories it contains, so a two-entry table holding id 40000 passed and then wrapped negative. The bound is now the largest id the carrier will hold, still the remapped count for legacy datasets. R34 --- ScaFFold/utils/data_loading.py | 32 ++++++++--- tests/test_data_loading.py | 97 ++++++++++++++++++++++++++++++++++ 2 files changed, 123 insertions(+), 6 deletions(-) diff --git a/ScaFFold/utils/data_loading.py b/ScaFFold/utils/data_loading.py index dcc1a45..8a75b1d 100644 --- a/ScaFFold/utils/data_loading.py +++ b/ScaFFold/utils/data_loading.py @@ -165,16 +165,36 @@ def __init__( customlog(f"Dataset format version: {self.dataset_format_version}") # Masks are handed off in a signed 16-bit carrier (widened to long on - # the compute device), so every class id must fit that range. Legacy - # masks are remapped to 0..len(mask_values)-1; optimized masks store - # dense ids that stay within the same bound. - max_class_id = len(self.mask_values) - 1 + # the compute device), so the largest class id the carrier will hold + # must fit that range. + max_class_id = self._max_class_id() if max_class_id > np.iinfo(np.int16).max: raise ValueError( - f"{len(self.mask_values)} classes exceed the int16 mask carrier " - f"limit ({np.iinfo(np.int16).max})" + f"Mask class id {max_class_id} (from {len(self.mask_values)} " + f"classes) exceeds the int16 mask carrier limit " + f"({np.iinfo(np.int16).max}); it would wrap negative" ) + def _max_class_id(self): + """Return the largest class id ``_to_mask_carrier`` will have to carry. + + The bound differs by format, and using the wrong one is unsafe in one + direction and needlessly strict in the other. v2+ masks ship *raw* + ``category + 1`` ids, and the per-split table lists only the categories + present in that split -- so a sparse split can declare two classes while + holding an id in the tens of thousands, which the class *count* check + happily waved through. Legacy masks, by contrast, are remapped to + ``0..len(mask_values)-1``, so the count is exactly right there and their + (arbitrarily large) raw values are irrelevant. + """ + if self.dataset_format_version < DATASET_FORMAT_VERSION: + return len(self.mask_values) - 1 + + ids = np.asarray(self.mask_values) + if ids.size == 0: + return 0 + return int(ids.max()) + def _load_mask_values(self, data_dir): """Return the label-remap table for this split. diff --git a/tests/test_data_loading.py b/tests/test_data_loading.py index b29d212..38746bb 100644 --- a/tests/test_data_loading.py +++ b/tests/test_data_loading.py @@ -636,6 +636,103 @@ def test_broken_meta_raises_instead_of_silent_legacy(tmp_path, broken_meta): assert str(root) in message +def _build_v2_sparse_label_dataset(root: Path, label: int) -> Path: + """A v2 dataset whose only foreground label is the (large) ``label``. + + v2 masks store raw ``category + 1`` ids and the per-split pickle lists only + the categories actually present in that split, so a sparse split can hold a + handful of very large ids. + """ + vol_dir = root / "volumes" / "training" + mask_dir = root / "masks" / "training" + vol_dir.mkdir(parents=True) + mask_dir.mkdir(parents=True) + np.save(vol_dir / "0.npy", np.zeros((3, 4, 4, 4), dtype=VOLUME_DTYPE)) + mask = np.zeros((4, 4, 4), dtype=MASK_DTYPE) + mask[0, 0, 0] = label + np.save(mask_dir / "0_mask.npy", mask) + with open(root / "train_unique_mask_vals", "wb") as handle: + pickle.dump({"mask_values": [0, label]}, handle) + (root / "meta.yaml").write_text("dataset_format_version: 2\n") + return root + + +# --------------------------------------------------------------------------- +# R34: the int16 carrier guard must bound the largest class *id*, not the count +# --------------------------------------------------------------------------- + + +def test_int16_guard_checks_the_largest_class_id(tmp_path): + """A v2 split holding an id above the int16 range is rejected. + + The guard compared ``len(mask_values) - 1`` -- the class *count* -- against + the carrier limit, but v2 masks ship raw ids. A split listing only + ``[0, 40000]`` passed a two-class check and then wrapped 40000 to -25536 in + the int16 carrier, which survives the downstream ``.long()`` cast as a + negative label. + """ + root = _build_v2_sparse_label_dataset(tmp_path / "sparse", label=40000) + + with pytest.raises(ValueError) as excinfo: + FractalDataset( + root / "volumes" / "training", + root / "masks" / "training", + data_dir=root / "train_unique_mask_vals", + ) + + message = str(excinfo.value) + # The message names the offending id and how many classes the split has. + assert "40000" in message + assert "2" in message + assert str(np.iinfo(np.int16).max) in message + + +def test_int16_guard_accepts_ids_inside_the_range(tmp_path): + """A large-but-representable id still loads, and does not wrap negative.""" + label = int(np.iinfo(np.int16).max) + root = _build_v2_sparse_label_dataset(tmp_path / "edge", label=label) + + ds = FractalDataset( + root / "volumes" / "training", + root / "masks" / "training", + data_dir=root / "train_unique_mask_vals", + ) + carrier = ds[0]["mask"] + assert carrier.dtype == torch.int16 + assert int(carrier.min()) >= 0 + assert int(carrier.max()) == label + + +def test_int16_guard_uses_remapped_ids_for_legacy_datasets(tmp_path): + """v1 raw values are remapped to 0..n-1, so huge raw values are fine. + + Guarding on ``max(mask_values)`` alone would reject a legacy dataset that + the loader handles perfectly well: its carrier only ever holds the remapped + index, not the raw voxel value. + """ + raw_mask = np.zeros((4, 4, 4), dtype=MASK_DTYPE) + raw_mask[0, 0, 0] = 40000 + volume = np.zeros((4, 4, 4, 3), dtype=VOLUME_DTYPE) + root = _build_v1_split_dataset( + tmp_path / "legacy", + raw_mask, + volume, + train_vals=[0, 40000], + val_vals=[0, 40000], + ) + + ds = FractalDataset( + root / "volumes" / "training", + root / "masks" / "training", + data_dir=root / "train_unique_mask_vals", + ) + assert ds.dataset_format_version == 1 + carrier = ds[0]["mask"] + # 40000 was remapped to class index 1; nothing wrapped. + assert int(carrier.max()) == 1 + assert int(carrier.min()) == 0 + + def test_absent_meta_is_still_legacy_v1(tiny_v1_dataset): """The genuine legacy case (no ``meta.yaml`` at all) is unchanged. From 07723e619d01cbcb8742bde74a0805f544494ed7 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:42:19 -0700 Subject: [PATCH 18/62] Read dataset provenance from the ScaFFold source tree _git_commit_short ran git in the process working directory, so meta.yaml, the published directory name, and the commit-based reuse gate carried whatever repo the job was launched from. It now runs git in the package directory; a non-checkout install still degrades to no-commit-id. R35 --- ScaFFold/datagen/get_dataset.py | 17 ++++- tests/datagen/test_provenance.py | 109 +++++++++++++++++++++++++++++++ 2 files changed, 125 insertions(+), 1 deletion(-) create mode 100644 tests/datagen/test_provenance.py diff --git a/ScaFFold/datagen/get_dataset.py b/ScaFFold/datagen/get_dataset.py index f76f208..09e4b33 100644 --- a/ScaFFold/datagen/get_dataset.py +++ b/ScaFFold/datagen/get_dataset.py @@ -99,11 +99,26 @@ def _hash_volume_config(volume_config: Dict[str, Any]) -> str: return hashlib.sha256(s).hexdigest()[:12] -def _git_commit_short(log) -> str: +def _git_commit_short(log, source_dir: Path | None = None) -> str: + """Return the short commit of the ScaFFold checkout, or ``"no-commit-id"``. + + The commit identifies *the code that generated a dataset*: it is stamped + into ``meta.yaml``, into the published directory name, and is what + ``dataset_reuse_enforce_commit_id`` compares against. It must therefore be + read from the ScaFFold source tree rather than from the process working + directory, which is wherever the job was launched (a site workflow repo, a + scratch directory, ...) and has nothing to do with this code. + + ``source_dir`` overrides the directory git runs in; it defaults to this + module's own location and exists so the non-checkout case can be tested. + """ + if source_dir is None: + source_dir = Path(__file__).resolve().parent try: return ( subprocess.check_output( ["git", "rev-parse", "--short", "HEAD"], + cwd=str(source_dir), stderr=subprocess.DEVNULL, # Don't show console output to user ) .decode() diff --git a/tests/datagen/test_provenance.py b/tests/datagen/test_provenance.py new file mode 100644 index 0000000..d53bd33 --- /dev/null +++ b/tests/datagen/test_provenance.py @@ -0,0 +1,109 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +"""Dataset provenance: the commit stamped on a dataset is ScaFFold's (R35). + +``meta.yaml``'s ``code_commit``, the published ``__`` +directory name, and the ``dataset_reuse_enforce_commit_id`` gate all key off +one string. Reading it from the *launch* directory made it a property of +wherever the job happened to start -- a site workflow repo, a scratch +directory -- instead of the code that generated the data. Reuse was then gated +on an unrelated repo's churn while real ScaFFold changes went undetected. +""" + +from __future__ import annotations + +import logging +import subprocess +from pathlib import Path + +from ScaFFold.datagen import get_dataset as gd + +LOG = logging.getLogger("test_provenance") + +# The ScaFFold source tree: the checkout whose commit must be stamped. +PACKAGE_DIR = Path(gd.__file__).resolve().parent + + +def _head_of(repo: Path) -> str: + return ( + subprocess.check_output(["git", "rev-parse", "--short", "HEAD"], cwd=repo) + .decode() + .strip() + ) + + +def _make_repo(path: Path) -> str: + """Create a throwaway git repo with one commit; return its short HEAD.""" + path.mkdir(parents=True, exist_ok=True) + subprocess.run(["git", "init", "-q"], cwd=path, check=True) + (path / "README").write_text("an unrelated project\n") + subprocess.run(["git", "add", "README"], cwd=path, check=True) + subprocess.run( + [ + "git", + "-c", + "user.name=test", + "-c", + "user.email=test@example.invalid", + "-c", + "commit.gpgsign=false", + "commit", + "-qm", + "init", + ], + cwd=path, + check=True, + ) + return _head_of(path) + + +def test_commit_is_read_from_the_scaffold_tree_not_the_cwd(tmp_path, monkeypatch): + """Running from an unrelated repo still stamps ScaFFold's commit.""" + expected = _head_of(PACKAGE_DIR) + + other = tmp_path / "workflow-repo" + other_head = _make_repo(other) + assert other_head != expected, "the throwaway repo must differ from ScaFFold" + + monkeypatch.chdir(other) + assert gd._git_commit_short(LOG) == expected + + +def test_commit_survives_a_non_repo_working_directory(tmp_path, monkeypatch): + """A scratch launch directory does not degrade provenance to no-commit-id.""" + expected = _head_of(PACKAGE_DIR) + + scratch = tmp_path / "scratch-cwd" + scratch.mkdir() + monkeypatch.chdir(scratch) + + assert gd._git_commit_short(LOG) == expected + + +def test_non_repo_install_reports_no_commit_id(tmp_path): + """An installed (non-git) ScaFFold still degrades gracefully. + + Provenance is best-effort: when the source tree is not a checkout there is + no commit to record, and reuse simply is not gated on one. + """ + not_a_repo = tmp_path / "site-packages" / "ScaFFold" / "datagen" + not_a_repo.mkdir(parents=True) + + assert gd._git_commit_short(LOG, source_dir=not_a_repo) == "no-commit-id" + + +def test_missing_source_dir_reports_no_commit_id(tmp_path): + """A source directory that does not exist is handled, not raised.""" + assert gd._git_commit_short(LOG, source_dir=tmp_path / "gone") == "no-commit-id" From 2f72060dd4c98c94a4ed2331ea7b68e3eb332a2e Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:43:13 -0700 Subject: [PATCH 19/62] Document the uneven spatial shard hazard with an xfail test Nothing validates vol_size against dc_num_shards, so 16 over 3 shards is accepted as 6/6/4 and per-shard pooling silently diverges from the global result. The fix belongs in DistConv, so this records the hazard as a strict xfail rather than working around it. R32 --- tests/test_data_loading.py | 66 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 66 insertions(+) diff --git a/tests/test_data_loading.py b/tests/test_data_loading.py index 38746bb..1f60055 100644 --- a/tests/test_data_loading.py +++ b/tests/test_data_loading.py @@ -636,6 +636,72 @@ def test_broken_meta_raises_instead_of_silent_legacy(tmp_path, broken_meta): assert str(root) in message +# --------------------------------------------------------------------------- +# R32: uneven spatial shards are accepted but computed wrong (known, unfixed) +# --------------------------------------------------------------------------- + + +@pytest.mark.xfail( + reason="uneven spatial shards mishandled; to be fixed upstream in DistConv", + strict=True, +) +def test_uneven_spatial_shards_pool_like_the_unsharded_volume(): + """Per-shard pooling must agree with pooling the whole volume. + + ``SpatialShardSpec`` slices with ``torch.chunk`` semantics and only rejects + an *empty* shard, so ``vol_size=16`` over 3 shards is accepted as 6/6/4 -- + nothing anywhere validates ``vol_size % num_shards`` or the per-U-Net-level + evenness that pooling needs. DistConv's ``DCTensor`` intercepts only + ``aten.convolution``, so ``max_pool3d`` runs independently on each local + shard; with unequal (or odd) shards the local pooling windows stop lining up + with the global ones, and by the second level whole planes are dropped and + values appear that are the max of no global window at all. The same happens + for equal-but-odd shards (``vol_size=10`` over 2 shards -> 5/5), and an odd + local shard entering a strided conv trips DistConv's own divisibility check + with a cryptic error deep in the first forward. + + A volume ramp is used so a pooled value names the plane it came from. + + XFAIL: the fix belongs in DistConv (its sharded ops must handle uneven + spatial decompositions), not in the loader, which is why this documents the + hazard instead of asserting a workaround. ``strict=True``: the comparison is + plain deterministic CPU pooling, so it cannot pass by chance -- if it ever + passes, DistConv/ScaFFold has changed and this test must be revisited. + """ + vol_size, num_shards = 16, 3 + ramp = torch.zeros(1, 1, vol_size, vol_size, vol_size) + for plane in range(vol_size): + ramp[0, 0, plane] = plane + + # What the dataset hands each rank: uneven shards, accepted without a word. + volume = np.arange(vol_size**3, dtype=np.float32).reshape((vol_size,) * 3) + shard_sizes = [ + dl.SpatialShardSpec( + shard_dims=(2,), num_shards=(num_shards,), shard_indices=(index,) + ) + .slice_array(volume, {2: 0, 3: 1, 4: 2}, "mask") + .shape[0] + for index in range(num_shards) + ] + assert shard_sizes == [6, 6, 4] + + # Two U-Net levels of pooling, globally versus per local shard. + pool = torch.nn.MaxPool3d(2) + global_pooled = pool(pool(ramp)) + shards = list(torch.split(ramp, shard_sizes, dim=2)) + shards = [pool(pool(shard)) for shard in shards] + sharded_pooled = torch.cat(shards, dim=2) + + assert sharded_pooled.shape == global_pooled.shape, ( + f"sharded pooling produced {list(sharded_pooled.shape[2:])} planes vs " + f"{list(global_pooled.shape[2:])} globally" + ) + assert torch.equal(sharded_pooled, global_pooled), ( + f"per-shard planes {sharded_pooled[0, 0, :, 0, 0].tolist()} vs global " + f"{global_pooled[0, 0, :, 0, 0].tolist()}" + ) + + def _build_v2_sparse_label_dataset(root: Path, label: int) -> Path: """A v2 dataset whose only foreground label is the (large) ``label``. From 2d77083a64d3f1ec8eaa1f203329ddb76b524ac9 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:56:43 -0700 Subject: [PATCH 20/62] Abort when the MPI world does not span the job cli.py decides run dirs and restart state on MPI rank 0 while training uses the launcher's rank environment; a plain torchrun makes mpi4py a singleton in every process, so each claims its own run dir and the job hangs. Cross-check the two world sizes at the CLI entry and abort with an explanation of the launcher wiring. R13 --- ScaFFold/cli.py | 44 +++++++++++ tests/test_cli.py | 187 ++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 231 insertions(+) create mode 100644 tests/test_cli.py diff --git a/ScaFFold/cli.py b/ScaFFold/cli.py index 858e6ba..01119c4 100644 --- a/ScaFFold/cli.py +++ b/ScaFFold/cli.py @@ -24,9 +24,49 @@ from ScaFFold.utils import config_utils from ScaFFold.utils.collect_scheduler_info import collect_scheduler_metadata from ScaFFold.utils.create_restart_script import create_restart_script +from ScaFFold.utils.distributed import get_world_size from ScaFFold.utils.utils import setup_mpi_logger +def check_launcher_world_size(mpi_world_size): + """Verify that the MPI world spans the whole job. + + ScaFFold uses two different sources of truth for the job shape: the CLI and + the benchmark driver make their job-wide decisions on ``MPI.COMM_WORLD`` + rank 0, while the training path takes its rank and world size from the + launcher's environment (``get_world_rank`` / ``get_world_size``). Those + agree only when the job was started by an MPI-aware launcher. + + Under a plain ``torchrun`` (or a bare ``python`` invocation of several + processes) mpi4py initializes as an independent singleton in every process, + so every process believes it is MPI rank 0: each one runs the rank-0 block, + atomically claims its *own* timestamped run directory, and the job then + diverges -- non-zero launcher ranks crash on a broadcast that never + happened while rank 0 blocks in the first collective until it times out. + + Fail loudly here, before any run directory is created, instead of leaving + that mess behind. ``get_world_size`` falls back to the MPI communicator + when the environment reports nothing, so an unlauncher-ed single process + trivially agrees with itself. + """ + env_world_size = get_world_size() + if env_world_size != mpi_world_size: + raise RuntimeError( + f"Launcher/MPI world size mismatch: MPI.COMM_WORLD reports " + f"{mpi_world_size} rank(s) but the launcher environment reports " + f"{env_world_size}. ScaFFold decides run directories and restart " + "state on MPI rank 0 and broadcasts them, so an MPI world that " + "does not span the job is unrecoverable: every process acts as " + "rank 0, each claims a separate run directory, and the job hangs " + "in the first collective. This is what a plain 'torchrun' (or " + "launching the processes directly) produces, because mpi4py then " + "initializes as a singleton in every process. Launch ScaFFold " + "with an MPI-aware launcher (torchrun-hpc, flux run, srun, " + "mpirun) so that MPI spans all ranks, or run a single process " + "with no launcher environment set." + ) + + def _make_fresh_run_dir(base_run_dir, timestamp): """Create a fresh timestamped run directory without clobbering an existing one. @@ -299,6 +339,10 @@ def main(): comm = MPI.COMM_WORLD rank = comm.Get_rank() + # Every rank runs this identically, before any run directory is created, + # so a mis-launched job aborts uniformly instead of leaving per-rank run + # dirs behind and hanging. + check_launcher_world_size(comm.Get_size()) # Parse the command-line arguments. args = parser.parse_args() log = setup_mpi_logger(__file__, args.verbose) diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..9691499 --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,187 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +"""Tests for the ``scaffold`` CLI entry point (``ScaFFold.cli.main``). + +``cli.main`` makes every job-wide decision on MPI rank 0 and broadcasts it, so +these tests drive it with a *fake* communicator: the real ``MPI.COMM_WORLD`` in +this environment is always a one-rank singleton, which cannot express the +multi-rank shapes the CLI must get right (rank-0-decides-and-broadcasts, and +the mismatch between the MPI world and the launcher's environment). + +``_FakeComm`` records what rank 0 broadcasts and, for a non-zero rank, replays +a scripted sequence of values as if rank 0 had sent them. That makes it +possible to assert that a non-zero rank *uses the broadcast decision* instead +of consulting its own view of the filesystem. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest +import yaml + +import ScaFFold.cli as cli + +REPO_ROOT = Path(__file__).resolve().parents[1] +DEFAULT_CONFIG = REPO_ROOT / "ScaFFold" / "configs" / "benchmark_default.yml" + + +# --------------------------------------------------------------------------- +# harness +# --------------------------------------------------------------------------- + + +class _FakeComm: + """A stand-in for ``MPI.COMM_WORLD`` with a settable rank and size. + + ``bcast`` returns the caller's object on rank 0 (recording it); on a + non-zero rank it pops the next scripted value from ``bcast_returns``, + falling back to the caller's object when the script is exhausted. + """ + + def __init__(self, rank: int = 0, size: int = 1, bcast_returns=None): + self._rank = rank + self._size = size + self._scripted = list(bcast_returns or []) + self.broadcast = [] + self.barriers = 0 + + def Get_rank(self) -> int: + return self._rank + + def Get_size(self) -> int: + return self._size + + def Barrier(self) -> None: + self.barriers += 1 + + def bcast(self, obj, root=0): + self.broadcast.append(obj) + if self._rank == root: + return obj + if self._scripted: + return self._scripted.pop(0) + return obj + + +class _FakeMPI: + """Minimal ``mpi4py.MPI`` stand-in exposing only ``COMM_WORLD``.""" + + def __init__(self, comm): + self.COMM_WORLD = comm + + +def write_config(tmp_path, updates=None, name="bench.yml"): + """Write a complete benchmark config into ``tmp_path`` and return its path.""" + config = yaml.safe_load(DEFAULT_CONFIG.read_text()) + config["base_run_dir"] = str(tmp_path / "runs") + config["dataset_dir"] = str(tmp_path / "datasets") + config["fract_base_dir"] = str(tmp_path / "fractals") + if updates: + config.update(updates) + path = tmp_path / name + path.write_text(yaml.dump(config)) + return path + + +def run_cli(monkeypatch, argv, *, comm=None): + """Run ``cli.main`` with a fake communicator and stubbed subcommand drivers. + + Returns ``(comm, calls)`` where ``calls`` maps the subcommand name to the + list of config dicts its driver was invoked with. + """ + import ScaFFold.benchmark as benchmark_mod + import ScaFFold.generate_fractals as generate_fractals_mod + + comm = comm if comm is not None else _FakeComm() + calls = {"benchmark": [], "generate_fractals": []} + + monkeypatch.setattr(sys, "argv", list(argv)) + monkeypatch.setattr(cli, "MPI", _FakeMPI(comm)) + monkeypatch.setattr( + benchmark_mod, + "main", + lambda kwargs_dict={}: calls["benchmark"].append(dict(kwargs_dict)), + ) + monkeypatch.setattr( + generate_fractals_mod, + "main", + lambda kwargs_dict={}: calls["generate_fractals"].append(dict(kwargs_dict)), + ) + cli.main() + return comm, calls + + +# --------------------------------------------------------------------------- +# R13: the MPI world must span the whole job +# --------------------------------------------------------------------------- + + +def test_mpi_singleton_under_multirank_launcher_aborts(monkeypatch, tmp_path): + """A 1-rank MPI world inside a 2-rank launcher job aborts with an explanation. + + This is the plain-``torchrun`` shape: every process is an mpi4py singleton + while the launcher says WORLD_SIZE=2. Left unchecked, every process runs + cli.py's rank-0 block, claims its own run directory, and the job then hangs + in the first real collective. + """ + monkeypatch.setenv("WORLD_SIZE", "2") + monkeypatch.setenv("RANK", "0") + cfg = write_config(tmp_path) + + with pytest.raises(RuntimeError) as excinfo: + run_cli( + monkeypatch, + ["scaffold", "benchmark", "-c", str(cfg)], + comm=_FakeComm(rank=0, size=1), + ) + + message = str(excinfo.value) + assert "1" in message and "2" in message + assert "torchrun" in message.lower() + # The abort happens before any run directory is claimed. + assert not (tmp_path / "runs").exists() + + +def test_matching_world_sizes_are_accepted(monkeypatch, tmp_path): + """An MPI world that matches the launcher environment runs normally.""" + monkeypatch.setenv("WORLD_SIZE", "4") + monkeypatch.setenv("RANK", "0") + cfg = write_config(tmp_path) + + _, calls = run_cli( + monkeypatch, + ["scaffold", "benchmark", "-c", str(cfg)], + comm=_FakeComm(rank=0, size=4), + ) + + assert len(calls["benchmark"]) == 1 + + +def test_no_launcher_env_is_not_a_mismatch(monkeypatch, tmp_path): + """With no launcher variables set, the MPI world alone defines the size.""" + for var in ("WORLD_SIZE", "RANK", "LOCAL_RANK", "SLURM_NTASKS", "FLUX_JOB_SIZE"): + monkeypatch.delenv(var, raising=False) + cfg = write_config(tmp_path) + + _, calls = run_cli( + monkeypatch, + ["scaffold", "benchmark", "-c", str(cfg)], + comm=_FakeComm(rank=0, size=1), + ) + + assert len(calls["benchmark"]) == 1 From c1e7a3dac80c3a4a9651deb64e1ea4194cd4550e Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:57:24 -0700 Subject: [PATCH 21/62] Substitute restart placeholders in combined flag tokens A run launched with --config=PATH emitted a literal --config=__CFG__ into restart.sh, because placeholder substitution only matched whole tokens; the restart then died with "Config file '__CFG__' not found". Substitute the value half of any --flag=PLACEHOLDER token as well. R14 --- ScaFFold/utils/create_restart_script.py | 24 ++++++-- tests/test_restart_script.py | 75 +++++++++++++++++++++++++ 2 files changed, 93 insertions(+), 6 deletions(-) diff --git a/ScaFFold/utils/create_restart_script.py b/ScaFFold/utils/create_restart_script.py index b43c0e3..18d6a9f 100644 --- a/ScaFFold/utils/create_restart_script.py +++ b/ScaFFold/utils/create_restart_script.py @@ -78,14 +78,26 @@ def _rewrite_config_and_add_restart(cli_args: List[str]) -> List[str]: return new_args +def _substitute_placeholder(tok: str, var_subs: dict[str, str]) -> str: + """Return ``tok`` with a placeholder replaced by its Bash expansion. + + A placeholder may be a whole token (``--config __CFG__``) or the value half + of a combined token (``--config=__CFG__``); argparse accepts both spellings + on the command line, so the rewriter can emit either. Anything else is + shell-quoted verbatim. + """ + if tok in var_subs: + return var_subs[tok] # e.g., "$RUN_DIR/config.yaml" + flag, sep, value = tok.partition("=") + if sep and value in var_subs: + # --config=__CFG__ -> --config="$RUN_DIR/config.yaml" + return shlex.quote(flag + sep) + var_subs[value] + return shlex.quote(tok) + + def _bash_array(var_name: str, argv: List[str], var_subs: dict[str, str]) -> str: """Render a Bash array declaration VAR=( ... ), safely quoted, with simple placeholder substitution.""" - parts = [] - for tok in argv: - if tok in var_subs: - parts.append(var_subs[tok]) # e.g., "$RUN_DIR/config.yaml" - else: - parts.append(shlex.quote(tok)) + parts = [_substitute_placeholder(tok, var_subs) for tok in argv] return f"{var_name}=( " + " ".join(parts) + " )" diff --git a/tests/test_restart_script.py b/tests/test_restart_script.py index 1059c7e..8e9f81c 100644 --- a/tests/test_restart_script.py +++ b/tests/test_restart_script.py @@ -182,3 +182,78 @@ def test_generated_script_is_valid_bash(monkeypatch, tmp_path): assert result.returncode == 0, ( f"bash -n failed for variant {i}:\n{result.stderr}" ) + + +# --------------------------------------------------------------------------- +# R14: combined ``--flag=value`` tokens +# --------------------------------------------------------------------------- + + +def test_config_equals_form_is_substituted(monkeypatch, tmp_path): + """``--config=PATH`` is repointed at the run dir, not left as a placeholder. + + The rewriter emits the placeholder as part of a combined token, so a + substitution that only matches whole tokens leaves ``--config=__CFG__`` in + the script and the restart dies with "Config file '__CFG__' not found". + """ + _isolate_env(monkeypatch) + argv = [ + "/usr/bin/scaffold", + "benchmark", + "--config=/some/where/config.yml", + "--epochs", + "10", + ] + + script = _generate(monkeypatch, tmp_path / "run", argv=argv) + + assert "__CFG__" not in script + assert '--config="$RUN_DIR/config.yaml"' in script + assert "/some/where/config.yml" not in script + + +@pytest.mark.parametrize( + "config_argv", + [ + ["-c", "/some/where/config.yml"], + ["--config", "/some/where/config.yml"], + ["--config=/some/where/config.yml"], + ], + ids=["short", "long-space", "long-equals"], +) +@pytest.mark.skipif(shutil.which("bash") is None, reason="bash not available") +def test_every_config_spelling_expands_to_the_run_dir_config( + monkeypatch, tmp_path, config_argv +): + """Bash expands every ``--config`` spelling to ``$RUN_DIR/config.yaml``. + + The generated PY array is sourced and expanded by a real shell so the test + asserts on the arguments the restarted CLI actually receives. + """ + _isolate_env(monkeypatch) + argv = ["/usr/bin/scaffold", "benchmark"] + config_argv + + script = _generate(monkeypatch, tmp_path / "run", argv=argv) + + py_decl = next(line for line in script.splitlines() if line.startswith("PY=(")) + probe = tmp_path / f"probe_{config_argv[0][-1]}.sh" + probe.write_text(f'RUN_DIR=/run/dir\n{py_decl}\nprintf "%s\\n" "${{PY[@]}}"\n') + result = subprocess.run( + ["bash", str(probe)], capture_output=True, text=True, check=True + ) + + tokens = result.stdout.split("\n") + assert "/run/dir/config.yaml" in tokens or ( + "--config=/run/dir/config.yaml" in tokens + ) + assert not any("__CFG__" in tok for tok in tokens) + + +def test_run_dir_placeholder_is_substituted(monkeypatch, tmp_path): + """The appended ``--run-dir`` placeholder still resolves (control).""" + _isolate_env(monkeypatch) + + script = _generate(monkeypatch, tmp_path / "run") + + assert "__RUN_DIR__" not in script + assert '--run-dir "$RUN_DIR"' in script From 8f9487cafe95ebf06465c2d701046cfbcb8d336c Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:58:53 -0700 Subject: [PATCH 22/62] Generate restart scripts at the true job scale The CLI holds MPI.COMM_WORLD but did not pass its size, and the generator's environment sniffing missed launcher variables the rank side honors (MV2, PALS, and bare SLURM/FLUX task counts), so e.g. a Cray PALS job got a single-process restart.sh. Pass the communicator size and share one world-size variable list with get_world_size. R17 --- ScaFFold/cli.py | 6 ++-- ScaFFold/utils/create_restart_script.py | 24 +++++++++++-- tests/test_cli.py | 33 +++++++++++++++++ tests/test_restart_script.py | 48 +++++++++++++++++++++++++ 4 files changed, 106 insertions(+), 5 deletions(-) diff --git a/ScaFFold/cli.py b/ScaFFold/cli.py index 01119c4..f3ea545 100644 --- a/ScaFFold/cli.py +++ b/ScaFFold/cli.py @@ -441,8 +441,10 @@ def main(): with open(benchmark_run_dir / "config.yaml", "w") as file: yaml.dump(combined_config, file) - # 4. Generate/Update the restart script in the directory - create_restart_script(benchmark_run_dir) + # 4. Generate/Update the restart script in the directory. The + # communicator size is ground truth for the job scale; environment + # sniffing is only the fallback for callers that lack it. + create_restart_script(benchmark_run_dir, world_size=comm.Get_size()) comm.Barrier() combined_config = comm.bcast(combined_config, root=0) diff --git a/ScaFFold/utils/create_restart_script.py b/ScaFFold/utils/create_restart_script.py index 18d6a9f..3f93ade 100644 --- a/ScaFFold/utils/create_restart_script.py +++ b/ScaFFold/utils/create_restart_script.py @@ -27,6 +27,21 @@ # were active in the generating run. Names mirror ScaFFold.utils.perf_measure. _PROFILING_ENV_VARS = ("PROFILE_TORCH", "CALI_CONFIG") +# Launcher variables carrying the total rank count, in the same priority order +# as ScaFFold.utils.distributed.get_world_size. The rank side and the restart +# generator must recognize the same set, or a job launched under a launcher +# only one of them knows about (e.g. Cray PALS) gets a restart script for the +# wrong number of ranks. +_WORLD_SIZE_ENV_VARS = ( + "WORLD_SIZE", + "MV2_COMM_WORLD_SIZE", + "OMPI_COMM_WORLD_SIZE", + "PMI_SIZE", + "PALS_NRANKS", + "SLURM_NTASKS", + "FLUX_JOB_SIZE", +) + def _rewrite_config_and_add_restart(cli_args: List[str]) -> List[str]: """ @@ -243,8 +258,11 @@ def _sniff_launch_shape(env: Mapping[str, str]) -> tuple[int | None, int, int]: Reads, in priority order: 1. Flux (FLUX_JOB_SIZE is total tasks, FLUX_JOB_NNODES is node count), 2. Slurm (SLURM_NTASKS / SLURM_NPROCS total tasks, SLURM_*NODES nodes), - 3. generic launcher hints for total rank count: torchrun's WORLD_SIZE, - Open MPI's OMPI_COMM_WORLD_SIZE, and PMI's PMI_SIZE. + 3. generic launcher hints for the total rank count, in the same order + and covering the same variables as + ``ScaFFold.utils.distributed.get_world_size``: keeping the two in + sync is what stops a restart script from relaunching the job at the + wrong scale. ``nodes`` is None when the environment does not report a node count. ``world_size`` is the best available total-rank estimate (>= 1). @@ -260,7 +278,7 @@ def _sniff_launch_shape(env: Mapping[str, str]) -> tuple[int | None, int, int]: total_tasks = int(env.get("SLURM_NTASKS") or env.get("SLURM_NPROCS") or 1) else: # No scheduler: fall back to generic launcher hints for the rank count. - for key in ("WORLD_SIZE", "OMPI_COMM_WORLD_SIZE", "PMI_SIZE"): + for key in _WORLD_SIZE_ENV_VARS: val = env.get(key) if val: total_tasks = int(val) diff --git a/tests/test_cli.py b/tests/test_cli.py index 9691499..ca47cdf 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -185,3 +185,36 @@ def test_no_launcher_env_is_not_a_mismatch(monkeypatch, tmp_path): ) assert len(calls["benchmark"]) == 1 + + +# --------------------------------------------------------------------------- +# R17: the restart script is generated at the true job scale +# --------------------------------------------------------------------------- + + +def test_restart_script_gets_the_mpi_world_size(monkeypatch, tmp_path): + """The CLI passes its communicator size to the restart-script generator. + + Without it the generator falls back to sniffing the environment, which + misses launcher variables the rank side honors (e.g. PALS_NRANKS) and + silently emits a single-process restart script for a multi-rank job. + """ + recorded = {} + + def _recorder(run_dir, world_size=None): + recorded["run_dir"] = run_dir + recorded["world_size"] = world_size + return Path(run_dir) / "restart.sh" + + monkeypatch.setattr(cli, "create_restart_script", _recorder) + monkeypatch.setenv("WORLD_SIZE", "4") + monkeypatch.setenv("RANK", "0") + cfg = write_config(tmp_path) + + run_cli( + monkeypatch, + ["scaffold", "benchmark", "-c", str(cfg)], + comm=_FakeComm(rank=0, size=4), + ) + + assert recorded["world_size"] == 4 diff --git a/tests/test_restart_script.py b/tests/test_restart_script.py index 8e9f81c..a8f3e08 100644 --- a/tests/test_restart_script.py +++ b/tests/test_restart_script.py @@ -25,6 +25,7 @@ from __future__ import annotations +import os import shutil import subprocess import sys @@ -45,8 +46,10 @@ "SLURM_JOB_NUM_NODES", "SLURM_NNODES", "WORLD_SIZE", + "MV2_COMM_WORLD_SIZE", "OMPI_COMM_WORLD_SIZE", "PMI_SIZE", + "PALS_NRANKS", ) # Profiling variables re-exported only when set in the generating run. @@ -257,3 +260,48 @@ def test_run_dir_placeholder_is_substituted(monkeypatch, tmp_path): assert "__RUN_DIR__" not in script assert '--run-dir "$RUN_DIR"' in script + + +# --------------------------------------------------------------------------- +# R17: launch-shape sniffing must match the rank side +# --------------------------------------------------------------------------- + +# Every variable ``ScaFFold.utils.distributed.get_world_size`` honors. The +# restart generator must derive the same world size from each of them, or a +# restart script silently relaunches the job at the wrong scale. +_WORLD_SIZE_ENV_VARS = ( + "WORLD_SIZE", + "MV2_COMM_WORLD_SIZE", + "OMPI_COMM_WORLD_SIZE", + "PMI_SIZE", + "PALS_NRANKS", + "SLURM_NTASKS", + "FLUX_JOB_SIZE", +) + + +@pytest.mark.parametrize("var", _WORLD_SIZE_ENV_VARS) +def test_sniffed_world_size_matches_rank_side(monkeypatch, var): + """The generator and ``get_world_size`` agree on every launcher variable.""" + from ScaFFold.utils.distributed import get_world_size + + _isolate_env(monkeypatch) + for other in _WORLD_SIZE_ENV_VARS: + monkeypatch.delenv(other, raising=False) + monkeypatch.setenv(var, "8") + + _, _, sniffed = crs._sniff_launch_shape(os.environ) + + assert sniffed == 8, f"{var} not recognized by the restart generator" + assert sniffed == get_world_size() + + +def test_pals_job_gets_a_multirank_restart_script(monkeypatch, tmp_path): + """A Cray PALS launch (PALS_NRANKS) emits the multi-rank template.""" + _isolate_env(monkeypatch) + monkeypatch.setenv("PALS_NRANKS", "8") + + script = _generate(monkeypatch, tmp_path / "run") + + assert "torchrun-hpc" in script + assert 'exec "${PY[@]}"' not in script From 7563905e2156bed06fed770bc9b8d27393202e60 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 16:00:57 -0700 Subject: [PATCH 23/62] Decide the restart pre-check on rank 0 and broadcast it Every rank stat-ed the shared filesystem for a checkpoint and raised on its own verdict, so a divergent view (stale attribute cache) either strands the peers in benchmark.py's timeout-less barrier or lets them run on after rank 0 aborted. Make it a rank-0 decision broadcast to all ranks, matching the other CLI decisions. R18 --- ScaFFold/cli.py | 48 ++++++++++++++++-------- tests/test_cli.py | 96 ++++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 128 insertions(+), 16 deletions(-) diff --git a/ScaFFold/cli.py b/ScaFFold/cli.py index f3ea545..4e354c4 100644 --- a/ScaFFold/cli.py +++ b/ScaFFold/cli.py @@ -90,6 +90,26 @@ def _make_fresh_run_dir(base_run_dir, timestamp): candidate = base_run_dir / f"{timestamp}-{suffix}" +def missing_checkpoint_error(combined_config): + """Return the "nothing to resume from" message, or None if a restart can run. + + Reports the first problem found rather than raising, so the caller can make + this a rank-0 decision and broadcast the verdict instead of letting every + rank stat the shared filesystem and possibly disagree. + """ + checkpoint_dir = Path(combined_config["run_dir"]) / combined_config.get( + "checkpoint_dir", "checkpoints" + ) + expected_checkpoints = ( + checkpoint_dir / "checkpoint_last.pth", + checkpoint_dir / "checkpoint_best.pth", + ) + if any(path.exists() for path in expected_checkpoints): + return None + expected = " or ".join(str(path) for path in expected_checkpoints) + return f"Restart requested but no checkpoint was found. Expected {expected}." + + def resolve_run_dir(args_dict, combined_config): """Decide the benchmark run directory and whether this launch resumes a run. @@ -449,23 +469,21 @@ def main(): comm.Barrier() combined_config = comm.bcast(combined_config, root=0) + # Restart pre-check. Like every other decision here it is made once, on + # rank 0, and broadcast: the check reads the filesystem, and ranks can see + # different views of a shared filesystem (stale NFS/Lustre attribute + # caches). A rank that decided for itself would either abort alone -- + # stranding its peers in benchmark.py's timeout-less barrier -- or keep + # running after rank 0 had already aborted. + restart_precheck_error = None if combined_config.get("restart", False): - run_dir = combined_config.get("run_dir") - if not run_dir: + if not combined_config.get("run_dir"): raise ValueError("--restart requires --run-dir") - - checkpoint_dir = Path(run_dir) / combined_config.get( - "checkpoint_dir", "checkpoints" - ) - expected_checkpoints = ( - checkpoint_dir / "checkpoint_last.pth", - checkpoint_dir / "checkpoint_best.pth", - ) - if not any(path.exists() for path in expected_checkpoints): - expected = " or ".join(str(path) for path in expected_checkpoints) - raise FileNotFoundError( - f"Restart requested but no checkpoint was found. Expected {expected}." - ) + if rank == 0: + restart_precheck_error = missing_checkpoint_error(combined_config) + restart_precheck_error = comm.bcast(restart_precheck_error, root=0) + if restart_precheck_error is not None: + raise FileNotFoundError(restart_precheck_error) if rank == 0: log.debug("combined_config = %s", combined_config) diff --git a/tests/test_cli.py b/tests/test_cli.py index ca47cdf..f5ddd56 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -98,9 +98,14 @@ def write_config(tmp_path, updates=None, name="bench.yml"): return path -def run_cli(monkeypatch, argv, *, comm=None): +def run_cli(monkeypatch, argv, *, comm=None, sync_env=True): """Run ``cli.main`` with a fake communicator and stubbed subcommand drivers. + ``sync_env`` makes the launcher environment agree with the fake + communicator, which is what a correctly launched job looks like; tests of + the mismatch check itself pass ``sync_env=False`` and set the environment + themselves. + Returns ``(comm, calls)`` where ``calls`` maps the subcommand name to the list of config dicts its driver was invoked with. """ @@ -110,6 +115,9 @@ def run_cli(monkeypatch, argv, *, comm=None): comm = comm if comm is not None else _FakeComm() calls = {"benchmark": [], "generate_fractals": []} + if sync_env: + monkeypatch.setenv("WORLD_SIZE", str(comm.Get_size())) + monkeypatch.setenv("RANK", str(comm.Get_rank())) monkeypatch.setattr(sys, "argv", list(argv)) monkeypatch.setattr(cli, "MPI", _FakeMPI(comm)) monkeypatch.setattr( @@ -148,6 +156,7 @@ def test_mpi_singleton_under_multirank_launcher_aborts(monkeypatch, tmp_path): monkeypatch, ["scaffold", "benchmark", "-c", str(cfg)], comm=_FakeComm(rank=0, size=1), + sync_env=False, ) message = str(excinfo.value) @@ -182,6 +191,7 @@ def test_no_launcher_env_is_not_a_mismatch(monkeypatch, tmp_path): monkeypatch, ["scaffold", "benchmark", "-c", str(cfg)], comm=_FakeComm(rank=0, size=1), + sync_env=False, ) assert len(calls["benchmark"]) == 1 @@ -218,3 +228,87 @@ def _recorder(run_dir, world_size=None): ) assert recorded["world_size"] == 4 + + +# --------------------------------------------------------------------------- +# R18: the restart pre-check is a rank-0 decision, broadcast to everyone +# --------------------------------------------------------------------------- + + +def _restart_argv(cfg, run_dir): + return [ + "scaffold", + "benchmark", + "-c", + str(cfg), + "--restart", + "--run-dir", + str(run_dir), + ] + + +def _make_checkpoint(run_dir): + ckpt_dir = run_dir / "checkpoints" + ckpt_dir.mkdir(parents=True, exist_ok=True) + (ckpt_dir / "checkpoint_last.pth").write_bytes(b"") + return ckpt_dir + + +def test_restart_without_checkpoint_is_rejected_on_rank0(monkeypatch, tmp_path): + """Rank 0 still rejects a restart with no checkpoint, naming the paths.""" + cfg = write_config(tmp_path) + run_dir = tmp_path / "prior" + run_dir.mkdir() + + with pytest.raises(FileNotFoundError) as excinfo: + run_cli(monkeypatch, _restart_argv(cfg, run_dir), comm=_FakeComm(rank=0)) + + assert "checkpoint_last.pth" in str(excinfo.value) + + +def test_restart_precheck_follows_the_broadcast_decision(monkeypatch, tmp_path): + """A non-zero rank trusts rank 0's verdict instead of stat-ing the FS itself. + + Simulates a stale attribute cache: rank 0 saw the checkpoint and broadcast + "go", while this rank's view of the shared filesystem shows nothing. A rank + that re-decides locally raises alone and strands its peers in the next + barrier. + """ + cfg = write_config(tmp_path) + run_dir = tmp_path / "prior" + run_dir.mkdir() # deliberately empty: this rank sees no checkpoint + rank0_config = { + "restart": True, + "run_dir": str(run_dir), + "checkpoint_dir": "checkpoints", + "verbose": 0, + } + + comm = _FakeComm(rank=1, size=2, bcast_returns=[rank0_config, None]) + _, calls = run_cli(monkeypatch, _restart_argv(cfg, run_dir), comm=comm) + + assert len(calls["benchmark"]) == 1 + + +def test_restart_precheck_failure_raises_on_every_rank(monkeypatch, tmp_path): + """Rank 0's rejection is broadcast, so non-zero ranks raise too. + + Here the local filesystem view *does* show a checkpoint; the rank must + still fail, because rank 0 -- the only rank whose verdict counts -- did not + find one. Otherwise the job splits: rank 0 aborts and the rest run on. + """ + cfg = write_config(tmp_path) + run_dir = tmp_path / "prior" + run_dir.mkdir() + _make_checkpoint(run_dir) + rank0_config = { + "restart": True, + "run_dir": str(run_dir), + "checkpoint_dir": "checkpoints", + "verbose": 0, + } + rank0_error = "Restart requested but no checkpoint was found. Expected /nope." + + comm = _FakeComm(rank=1, size=2, bcast_returns=[rank0_config, rank0_error]) + with pytest.raises(FileNotFoundError, match="no checkpoint"): + run_cli(monkeypatch, _restart_argv(cfg, run_dir), comm=comm) From e4aded529b1d0b378789c14a3528b56c1764da8d Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 16:01:59 -0700 Subject: [PATCH 24/62] Scope the run directory to the benchmark subcommand generate_fractals ran the CLI's rank-0 block too, littering base_run_dir with a timestamped benchmark directory whose restart.sh replayed generate_fractals with --restart/--run-dir, flags that subparser rejects (exit 2). Create the run dir, config dumps and restart script only for benchmark. R19 --- ScaFFold/cli.py | 54 ++++++++++++++++++++++++++--------------------- tests/test_cli.py | 51 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 81 insertions(+), 24 deletions(-) diff --git a/ScaFFold/cli.py b/ScaFFold/cli.py index 4e354c4..3ed0a20 100644 --- a/ScaFFold/cli.py +++ b/ScaFFold/cli.py @@ -441,30 +441,36 @@ def main(): combined_config["vol_size"] = pow(2, combined_config["problem_scale"]) combined_config["point_num"] = int(combined_config["vol_size"] ** 3 / 256) - # Resolve the run directory and whether this launch resumes a run. - # This sets combined_config["benchmark_run_dir"] on every path and, - # when resuming, forces train_from_scratch off / restart on. - benchmark_run_dir, restarting = resolve_run_dir(vars(args), combined_config) - if restarting: - log.info("Resuming in existing directory: %s", benchmark_run_dir) - - # Add scheduler metadata and machine name to config.yaml - combined_config["scheduler_metadata"] = collect_scheduler_metadata() - combined_config["machine_name"] = socket.gethostname() - - # Dump configs (Overwrite is okay/desired on restart to capture new job IDs) - overrides = { - k: v for k, v in cli_args.items() if v is not None and k != "command" - } - with open(benchmark_run_dir / "overrides.yaml", "w") as file: - yaml.dump(overrides, file) - with open(benchmark_run_dir / "config.yaml", "w") as file: - yaml.dump(combined_config, file) - - # 4. Generate/Update the restart script in the directory. The - # communicator size is ground truth for the job scale; environment - # sniffing is only the fallback for callers that lack it. - create_restart_script(benchmark_run_dir, world_size=comm.Get_size()) + # The run directory, its config dumps and its restart script belong to + # the benchmark subcommand alone. Fractal generation writes nothing + # there, and the restart script it used to get replayed + # `generate_fractals --restart --run-dir ...` -- flags that subparser + # rejects, so the script could only ever exit 2. + if args.command == "benchmark": + # Resolve the run directory and whether this launch resumes a run. + # This sets combined_config["benchmark_run_dir"] on every path and, + # when resuming, forces train_from_scratch off / restart on. + benchmark_run_dir, restarting = resolve_run_dir(vars(args), combined_config) + if restarting: + log.info("Resuming in existing directory: %s", benchmark_run_dir) + + # Add scheduler metadata and machine name to config.yaml + combined_config["scheduler_metadata"] = collect_scheduler_metadata() + combined_config["machine_name"] = socket.gethostname() + + # Dump configs (Overwrite is okay/desired on restart to capture new job IDs) + overrides = { + k: v for k, v in cli_args.items() if v is not None and k != "command" + } + with open(benchmark_run_dir / "overrides.yaml", "w") as file: + yaml.dump(overrides, file) + with open(benchmark_run_dir / "config.yaml", "w") as file: + yaml.dump(combined_config, file) + + # 4. Generate/Update the restart script in the directory. The + # communicator size is ground truth for the job scale; environment + # sniffing is only the fallback for callers that lack it. + create_restart_script(benchmark_run_dir, world_size=comm.Get_size()) comm.Barrier() combined_config = comm.bcast(combined_config, root=0) diff --git a/tests/test_cli.py b/tests/test_cli.py index f5ddd56..f379aaa 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -312,3 +312,54 @@ def test_restart_precheck_failure_raises_on_every_rank(monkeypatch, tmp_path): comm = _FakeComm(rank=1, size=2, bcast_returns=[rank0_config, rank0_error]) with pytest.raises(FileNotFoundError, match="no checkpoint"): run_cli(monkeypatch, _restart_argv(cfg, run_dir), comm=comm) + + +# --------------------------------------------------------------------------- +# R19: generate_fractals is not a benchmark run +# --------------------------------------------------------------------------- + + +def test_generate_fractals_creates_no_benchmark_run_dir(monkeypatch, tmp_path): + """Fractal generation leaves no benchmark run dir and no restart script. + + The rank-0 block used to run for every subcommand, so a generation job + littered base_run_dir with a timestamped benchmark directory holding a + restart.sh that replays ``generate_fractals --restart --run-dir ...`` -- + flags the generate_fractals subparser rejects, so the script exits 2. + """ + cfg = write_config(tmp_path) + + _, calls = run_cli(monkeypatch, ["scaffold", "generate_fractals", "-c", str(cfg)]) + + assert len(calls["generate_fractals"]) == 1 + assert not (tmp_path / "runs").exists(), "generation created a benchmark run dir" + assert list(tmp_path.rglob("restart.sh")) == [] + assert list(tmp_path.rglob("overrides.yaml")) == [] + + +def test_generate_fractals_config_reaches_the_driver(monkeypatch, tmp_path): + """The merged config still reaches the generation driver (control).""" + cfg = write_config(tmp_path) + + _, calls = run_cli( + monkeypatch, + ["scaffold", "generate_fractals", "-c", str(cfg), "--n-categories", "3"], + ) + + (config,) = calls["generate_fractals"] + assert config["n_categories"] == 3 + assert config["fract_base_dir"] == str(tmp_path / "fractals") + + +def test_benchmark_still_creates_its_run_dir(monkeypatch, tmp_path): + """The benchmark subcommand keeps its run dir, dumps and restart script.""" + cfg = write_config(tmp_path) + + _, calls = run_cli(monkeypatch, ["scaffold", "benchmark", "-c", str(cfg)]) + + (config,) = calls["benchmark"] + run_dir = Path(config["benchmark_run_dir"]) + assert run_dir.is_dir() + assert (run_dir / "config.yaml").exists() + assert (run_dir / "overrides.yaml").exists() + assert (run_dir / "restart.sh").exists() From 13acd12398c3e7cdc80131ccb3fece82977d3d40 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 16:03:43 -0700 Subject: [PATCH 25/62] Honor auxiliary config keys set in YAML datagen_batch_size and verbose are accepted config keys, but Config never stores them, so the merge silently replaced them with argparse defaults. Carry the file's values into the merged config and let only options actually given on the command line override them: CLI flag > config file > argparse default. R20 --- ScaFFold/cli.py | 47 ++++++++++++++++++++++++++++++-- tests/test_cli.py | 68 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 113 insertions(+), 2 deletions(-) diff --git a/ScaFFold/cli.py b/ScaFFold/cli.py index 3ed0a20..c951d02 100644 --- a/ScaFFold/cli.py +++ b/ScaFFold/cli.py @@ -90,6 +90,31 @@ def _make_fresh_run_dir(base_run_dir, timestamp): candidate = base_run_dir / f"{timestamp}-{suffix}" +def explicit_cli_keys(args, parsers): + """Return the names of the options actually given on the command line. + + argparse does not record which options were supplied, so a value counts as + explicit when it differs from the default of the first parser in + ``parsers`` that defines one (subcommand parser first, then the top-level + parser). Only these may outrank a config-file setting; everything else in + the namespace is an argparse default, which is the weakest source. + + The one ambiguity is a flag passed with exactly its default value: it looks + absent, so a config-file entry wins over it. Both spellings then agree on + the default, which is the only value the flag could have contributed. + """ + explicit = set() + for name, value in vars(args).items(): + default = None + for parser in parsers: + default = parser.get_default(name) + if default is not None: + break + if value != default: + explicit.add(name) + return explicit + + def missing_checkpoint_error(combined_config): """Return the "nothing to resume from" message, or None if a restart can run. @@ -365,6 +390,11 @@ def main(): check_launcher_world_size(comm.Get_size()) # Parse the command-line arguments. args = parser.parse_args() + subcommand_parsers = { + "benchmark": benchmark_parser, + "generate_fractals": generate_fractals_parser, + } + active_parser = subcommand_parsers[args.command] log = setup_mpi_logger(__file__, args.verbose) combined_config = None @@ -398,12 +428,23 @@ def main(): # into the run dir); keep the base config there. cli_args["config"] = config_paths[0] - # Combine configs: CLI args override config file values + # Combine configs, in increasing order of precedence: + # argparse default < config file < explicit command-line flag. combined_config = bench_config_dict.copy() + # Config only keeps the keys it consumes; the auxiliary keys it accepts + # (verbose, datagen_batch_size, ...) never become attributes, so put + # the file's values back first. Without this they are absent below and + # the argparse default overwrites what the user wrote in the config. + for key, value in merged_dict.items(): + combined_config.setdefault(key, value) + + explicit_cli = explicit_cli_keys(args, (active_parser, parser)) for key, value in cli_args.items(): + if key == "command": + continue if key not in combined_config: combined_config[key] = value - elif value is not None and key != "command": + elif key in explicit_cli and value is not None: log.info( "Overriding '%s=%s' with '%s=%s'", key, @@ -412,6 +453,8 @@ def main(): value, ) combined_config[key] = value + # The subcommand is always owned by the command line. + combined_config["command"] = cli_args["command"] # Recalculate unet_layers to capture any CLI overrides combined_config["unet_layers"] = ( diff --git a/tests/test_cli.py b/tests/test_cli.py index f379aaa..88782f5 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -363,3 +363,71 @@ def test_benchmark_still_creates_its_run_dir(monkeypatch, tmp_path): assert (run_dir / "config.yaml").exists() assert (run_dir / "overrides.yaml").exists() assert (run_dir / "restart.sh").exists() + + +# --------------------------------------------------------------------------- +# R20: auxiliary keys set in YAML must survive; CLI > YAML > argparse default +# --------------------------------------------------------------------------- + + +def test_yaml_aux_keys_reach_the_driver(monkeypatch, tmp_path): + """Auxiliary keys set in the config file are not replaced by defaults. + + ``Config`` accepts ``datagen_batch_size``/``verbose`` but does not store + them, so they used to vanish from the merged config and the argparse + default was installed instead -- making them settable only on the command + line despite being documented, validated config keys. + """ + cfg = write_config(tmp_path, {"datagen_batch_size": 500, "verbose": 1}) + + _, calls = run_cli(monkeypatch, ["scaffold", "generate_fractals", "-c", str(cfg)]) + + (config,) = calls["generate_fractals"] + assert config["datagen_batch_size"] == 500 + assert config["verbose"] == 1 + + +def test_cli_flag_outranks_yaml_aux_key(monkeypatch, tmp_path): + """An explicit command-line flag still wins over the config file.""" + cfg = write_config(tmp_path, {"datagen_batch_size": 500, "verbose": 0}) + + _, calls = run_cli( + monkeypatch, + [ + "scaffold", + "-v", + "generate_fractals", + "-c", + str(cfg), + "--datagen-batch-size", + "250", + ], + ) + + (config,) = calls["generate_fractals"] + assert config["datagen_batch_size"] == 250 + assert config["verbose"] == 1 + + +def test_argparse_default_used_when_yaml_is_silent(monkeypatch, tmp_path): + """With neither a flag nor a config entry, the argparse default applies.""" + cfg = write_config(tmp_path) + + _, calls = run_cli(monkeypatch, ["scaffold", "generate_fractals", "-c", str(cfg)]) + + (config,) = calls["generate_fractals"] + assert config["datagen_batch_size"] == 10000 + assert config["verbose"] == 0 + + +def test_run_config_records_the_effective_aux_values(monkeypatch, tmp_path): + """The run dir's config.yaml records what the run actually used.""" + cfg = write_config(tmp_path, {"verbose": 1}) + + _, calls = run_cli(monkeypatch, ["scaffold", "benchmark", "-c", str(cfg)]) + + (config,) = calls["benchmark"] + dumped = yaml.safe_load( + (Path(config["benchmark_run_dir"]) / "config.yaml").read_text() + ) + assert dumped["verbose"] == 1 From f494ec8b7bf71bab3e78989c974984062092f157 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 16:05:05 -0700 Subject: [PATCH 26/62] Keep memory diagnostics from crashing CPU-only runs mem_stats called torch.cuda.current_device() unguarded, so gather_and_print_mem -- which BaseTrainer.__init__ invokes unconditionally -- killed any CPU/gloo run launched with -v. Report the missing device and log a fallback instead; the early return is uniform across ranks, so no collective is skipped on one rank only. R21 --- ScaFFold/utils/utils.py | 23 +++++++++++++++++- tests/conftest.py | 8 +++--- tests/test_infra.py | 54 +++++++++++++++++++++++++++++++++++++++++ 3 files changed, 81 insertions(+), 4 deletions(-) diff --git a/ScaFFold/utils/utils.py b/ScaFFold/utils/utils.py index 5d9e334..12c196b 100644 --- a/ScaFFold/utils/utils.py +++ b/ScaFFold/utils/utils.py @@ -113,12 +113,24 @@ def setup_mpi_logger( def mem_stats(): + """Return this rank's GPU memory counters. + + On a host with no visible GPU (a CPU/gloo run, or a job launched with the + devices masked off) there are no counters to read: report that instead of + raising, so a diagnostic call cannot take down a run that is otherwise + perfectly able to proceed. + """ + rank = dist.get_rank() if dist.is_initialized() else 0 + if not torch.cuda.is_available(): + return {"rank": rank, "device": "cpu", "cuda_available": False} + dev = torch.cuda.current_device() free, total = torch.cuda.mem_get_info() # device-level (driver) view stats = torch.cuda.memory_stats(dev) # allocator internals return { - "rank": dist.get_rank() if dist.is_initialized() else 0, + "rank": rank, "device": dev, + "cuda_available": True, "allocated": torch.cuda.memory_allocated( dev ), # bytes currently used by tensors @@ -135,6 +147,15 @@ def mem_stats(): def gather_and_print_mem(log, tag=""): if log.getEffectiveLevel() > 10: # 10 -> DEBUG return + if not torch.cuda.is_available(): + # Uniform across ranks, so returning here skips the all_gather on every + # rank rather than desynchronizing them. + log.debug( + "=== %s === no CUDA device visible on this rank; " + "GPU memory statistics unavailable", + tag, + ) + return stats = mem_stats() if dist.is_initialized(): world = dist.get_world_size() diff --git a/tests/conftest.py b/tests/conftest.py index 376afbc..e35f739 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -467,6 +467,7 @@ def make( n_train: int = 4, n_val: int = 2, n: int = 16, + log_level: int = logging.INFO, config_overrides: Optional[dict] = None, ) -> "PyTorchTrainer": dataset_root = tiny_dataset( @@ -495,9 +496,10 @@ def make( device = torch.device("cpu") log = logging.getLogger(f"tiny_trainer.{id(config)}") - # INFO (20) > DEBUG (10) => gather_and_print_mem short-circuits and - # never touches CUDA / torch.distributed. - log.setLevel(logging.INFO) + # At the INFO default (20 > DEBUG's 10) gather_and_print_mem + # short-circuits and never touches CUDA / torch.distributed; pass + # log_level=logging.DEBUG to exercise the memory diagnostics. + log.setLevel(log_level) return PyTorchTrainer(model, config, device, log) diff --git a/tests/test_infra.py b/tests/test_infra.py index bf0fbeb..c279f34 100644 --- a/tests/test_infra.py +++ b/tests/test_infra.py @@ -22,6 +22,7 @@ from __future__ import annotations +import logging import os import numpy as np @@ -29,6 +30,7 @@ from ScaFFold.utils.data_loading import FractalDataset from ScaFFold.utils.data_types import MASK_DTYPE, VOLUME_DTYPE +from ScaFFold.utils.utils import gather_and_print_mem, mem_stats from tests.helpers import mpi_runner # --------------------------------------------------------------------------- @@ -230,3 +232,55 @@ def test_torchrun_gloo_two_ranks(tmp_path): # Both ranks reported; all_reduce of ranks {0,1} sums to 1. assert "RANK 0/2 sum=1.0" in out assert "RANK 1/2 sum=1.0" in out + + +# --------------------------------------------------------------------------- +# R21: memory diagnostics on a CPU-only run +# --------------------------------------------------------------------------- + + +def _debug_logger(name): + log = logging.getLogger(name) + log.setLevel(logging.DEBUG) + return log + + +def test_mem_stats_without_cuda(caplog): + """``mem_stats`` reports "no GPU" instead of raising on a CPU-only host.""" + if torch.cuda.is_available(): + import pytest + + pytest.skip("test covers the CPU-only path") + + stats = mem_stats() + + assert stats["cuda_available"] is False + assert "rank" in stats + + +def test_gather_and_print_mem_without_cuda(caplog): + """A DEBUG-level CPU run logs a fallback instead of crashing. + + ``BaseTrainer.__init__`` calls this unconditionally, so a CPU/gloo run with + ``-v`` used to die in trainer construction with "No CUDA GPUs are + available". + """ + if torch.cuda.is_available(): + import pytest + + pytest.skip("test covers the CPU-only path") + + log = _debug_logger("test_gather_and_print_mem_without_cuda") + with caplog.at_level(logging.DEBUG, logger=log.name): + gather_and_print_mem(log, "after_trainer_setup") + + messages = " ".join(record.getMessage() for record in caplog.records) + assert "after_trainer_setup" in messages + assert "cuda" in messages.lower() or "gpu" in messages.lower() + + +def test_trainer_constructs_with_debug_logging(tiny_trainer): + """The real call site survives: a trainer builds with a DEBUG logger.""" + trainer = tiny_trainer(log_level=logging.DEBUG) + + assert trainer.log.getEffectiveLevel() == logging.DEBUG From ffe7c0b49b6027d41a3fcff77316ffe2f4325b39 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 16:06:03 -0700 Subject: [PATCH 27/62] Preserve the base config under its own name in the run dir Copying the base config into the run dir kept its original filename, so a base config named config.yaml overwrote the merged config.yaml the CLI had just written -- and restart.sh points -c at that file. Always copy it to base_config.yaml. R22 --- README.md | 2 +- ScaFFold/benchmark.py | 9 +++++-- tests/test_config.py | 57 +++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 65 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 3206b02..86acb7b 100644 --- a/README.md +++ b/README.md @@ -61,7 +61,7 @@ The model is trained from a random initialization until convergence, which is de ScaFFold benchmark training always uses PyTorch distributed execution with DistConv spatial parallelism. For a singleton run, launch one distributed rank rather than disabling distributed execution. -Each `benchmark` invocation performs exactly one benchmark run, in a run folder created under `base_run_dir` set in the config file. Every run parameter must be single-valued; a list (e.g. `problem_scale: [6, 7]`) is rejected by name, since parameter sweeps are not supported. To compare parameter settings, launch one benchmark run per setting. For reproducibility, the run folder holds a copy of the benchmark config yml plus the fully merged `config.yaml` for that run. +Each `benchmark` invocation performs exactly one benchmark run, in a run folder created under `base_run_dir` set in the config file. Every run parameter must be single-valued; a list (e.g. `problem_scale: [6, 7]`) is rejected by name, since parameter sweeps are not supported. To compare parameter settings, launch one benchmark run per setting. For reproducibility, the run folder holds a copy of the benchmark config yml as `base_config.yaml` plus the fully merged `config.yaml` for that run. After the run completes, statistics from the run are stored in `train_stats.csv`. Additionally, users can inspect plots of the training and validation losses over time in ` Date: Fri, 31 Jul 2026 16:07:18 -0700 Subject: [PATCH 28/62] Validate the U-Net bottleneck against the problem scale Config accepted any unet_bottleneck_dim, so a value outside 0..problem_scale-1 built a U-Net with zero or too many pooling levels and failed much later with an opaque max_pool3d size error naming no config key. Reject it at config time, in Config and again after the CLI applies overrides. R25 --- ScaFFold/cli.py | 6 +++++- ScaFFold/utils/config_utils.py | 38 ++++++++++++++++++++++++++++++++- tests/test_cli.py | 33 ++++++++++++++++++++++++++++ tests/test_config.py | 39 ++++++++++++++++++++++++++++++++++ 4 files changed, 114 insertions(+), 2 deletions(-) diff --git a/ScaFFold/cli.py b/ScaFFold/cli.py index c951d02..623f868 100644 --- a/ScaFFold/cli.py +++ b/ScaFFold/cli.py @@ -456,7 +456,11 @@ def main(): # The subcommand is always owned by the command line. combined_config["command"] = cli_args["command"] - # Recalculate unet_layers to capture any CLI overrides + # Recalculate unet_layers to capture any CLI overrides. The overridden + # pair has to be re-validated: Config only saw the config-file values. + config_utils.validate_unet_dims( + combined_config["problem_scale"], combined_config["unet_bottleneck_dim"] + ) combined_config["unet_layers"] = ( combined_config["problem_scale"] - combined_config["unet_bottleneck_dim"] ) diff --git a/ScaFFold/utils/config_utils.py b/ScaFFold/utils/config_utils.py index 3289b42..92af20c 100644 --- a/ScaFFold/utils/config_utils.py +++ b/ScaFFold/utils/config_utils.py @@ -27,6 +27,40 @@ def require_positive_int(name: str, value: int) -> int: return value +def validate_unet_dims(problem_scale, unet_bottleneck_dim) -> int: + """Check that ``problem_scale``/``unet_bottleneck_dim`` describe a real U-Net. + + The U-Net has ``unet_layers = problem_scale - unet_bottleneck_dim`` + down/up levels over a ``2**problem_scale`` volume, so the bottleneck + exponent must satisfy ``0 <= unet_bottleneck_dim <= problem_scale - 1``: + a larger value asks for a bottleneck no smaller than the input (zero or + negative layers) and a negative one asks for more pooling levels than the + volume has. Both are only discovered later as an opaque + ``max_pool3d`` size error -- in production, after the whole dataset has + been generated -- so reject them here, at config time, naming the two keys + that have to change. + + Returns the validated bottleneck dimension. + """ + if isinstance(unet_bottleneck_dim, bool) or not isinstance( + unet_bottleneck_dim, int + ): + raise ValueError( + f"unet_bottleneck_dim must be an integer; got {unet_bottleneck_dim!r}" + ) + unet_layers = problem_scale - unet_bottleneck_dim + if unet_bottleneck_dim < 0 or unet_layers < 1: + raise ValueError( + f"unet_bottleneck_dim={unet_bottleneck_dim} is out of range for " + f"problem_scale={problem_scale}: it must satisfy " + f"0 <= unet_bottleneck_dim <= problem_scale - 1 " + f"(i.e. <= {problem_scale - 1}) so that the U-Net has at least one " + f"layer, but unet_layers = problem_scale - unet_bottleneck_dim = " + f"{unet_layers}. Raise problem_scale or lower unet_bottleneck_dim." + ) + return unet_bottleneck_dim + + class Config: """ A class for storing configuration settings for a specific run. @@ -182,7 +216,9 @@ def __init__(self, config_dict, strict=True): "WARNING: problem_scale found to be non-integer. Truncating to nearest int." ) self.problem_scale = math.floor(self.problem_scale) - self.unet_bottleneck_dim = config_dict["unet_bottleneck_dim"] + self.unet_bottleneck_dim = validate_unet_dims( + self.problem_scale, config_dict["unet_bottleneck_dim"] + ) self.unet_layers = self.problem_scale - self.unet_bottleneck_dim self.n_fracts_per_vol = config_dict["n_fracts_per_vol"] self.n_instances_used_per_fractal = config_dict["n_instances_used_per_fractal"] diff --git a/tests/test_cli.py b/tests/test_cli.py index 88782f5..b215e44 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -431,3 +431,36 @@ def test_run_config_records_the_effective_aux_values(monkeypatch, tmp_path): (Path(config["benchmark_run_dir"]) / "config.yaml").read_text() ) assert dumped["verbose"] == 1 + + +# --------------------------------------------------------------------------- +# R25: an out-of-range bottleneck is rejected before any work starts +# --------------------------------------------------------------------------- + + +def test_cli_override_bottleneck_out_of_range_rejected(monkeypatch, tmp_path): + """A command-line override that empties the U-Net is caught at config time. + + The CLI recomputes unet_layers after applying overrides, so the check has + to run there too -- not only inside Config. + """ + cfg = write_config(tmp_path) + + with pytest.raises(ValueError) as excinfo: + run_cli( + monkeypatch, + [ + "scaffold", + "benchmark", + "-c", + str(cfg), + "--problem-scale", + "4", + "--unet-bottleneck-dim", + "4", + ], + ) + + message = str(excinfo.value) + assert "unet_bottleneck_dim" in message + assert "problem_scale" in message diff --git a/tests/test_config.py b/tests/test_config.py index 0c8976b..28cc15b 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -329,3 +329,42 @@ def test_base_config_copy_never_clobbers_merged_config( preserved = yaml.safe_load((run_dir / "base_config.yaml").read_text()) assert preserved["local_batch_size"] == BASE["local_batch_size"] assert "machine_name" not in preserved + + +# --------------------------------------------------------------------------- +# unet_bottleneck_dim range (R25) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "problem_scale, bottleneck", + [(5, -1), (5, 5), (5, 6)], + ids=["negative", "zero-layers", "negative-layers"], +) +def test_bottleneck_out_of_range_rejected(problem_scale, bottleneck): + """An out-of-range bottleneck fails at config time, naming both keys. + + Left unvalidated it produced a U-Net with more pooling levels than the + volume has, and the run died hours later inside max_pool3d with "Given + input size: (2048x1x1x1)" -- naming no config key at all. + """ + bad = {**BASE, "problem_scale": problem_scale, "unet_bottleneck_dim": bottleneck} + + with pytest.raises(ValueError) as excinfo: + config_utils.Config(bad) + + message = str(excinfo.value) + assert "unet_bottleneck_dim" in message + assert "problem_scale" in message + assert str(bottleneck) in message + assert str(problem_scale) in message + + +@pytest.mark.parametrize("bottleneck", [0, 3, 4]) +def test_bottleneck_in_range_accepted(bottleneck): + """The full valid range (at least one U-Net layer) is accepted.""" + cfg = config_utils.Config( + {**BASE, "problem_scale": 5, "unet_bottleneck_dim": bottleneck} + ) + assert cfg.unet_layers == 5 - bottleneck + assert cfg.unet_layers >= 1 From 2ac96437f1d0b6e6289635ff7b56d018b9bdcf48 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 16:08:57 -0700 Subject: [PATCH 29/62] Never let a trace export strand the other ranks export_chrome_trace runs before the barrier that precedes rank-0 post-processing and raises when the run had zero profiled steps (or the write fails), killing the profiling rank and blocking every other rank in that barrier until timeout. Move it into a helper that logs the failure and returns. R15 --- ScaFFold/worker.py | 40 +++++++++++++++++++--- tests/test_reporting.py | 76 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 112 insertions(+), 4 deletions(-) diff --git a/ScaFFold/worker.py b/ScaFFold/worker.py index 459a8cf..8eafa51 100644 --- a/ScaFFold/worker.py +++ b/ScaFFold/worker.py @@ -96,6 +96,41 @@ def wrap_model_ddp(model, device, ps): ) +def export_profiler_trace(prof, config, log, rank, world_size, ranks_per_node): + """Write this rank's chrome trace, reporting failures instead of raising. + + A profiling rank reaches this while the others are already heading for the + barrier that precedes rank-0 post-processing, so an exception here does not + just lose a trace: it kills this rank and leaves every other rank blocked + in that barrier until the collective times out. Failures are real (a run + with zero training batches never starts the profiler, and export_chrome_trace + then raises; the trace can also fill the filesystem), so log them and let + the job finish. + + Returns the path written, or None if the trace could not be written. + """ + tracename = ( + f"torch-{socket.gethostname()}-r{rank}" + f"-N{world_size // ranks_per_node}-n{world_size}" + f"-ps{config.problem_scale}-e{config.epochs}" + f"-nipf{config.n_instances_used_per_fractal}-{int(time.time())}.json" + ) + try: + prof.export_chrome_trace(tracename) + except Exception as e: + log.error( + "Could not write PyTorch trace '%s': %s: %s. Continuing so the " + "run can finish; a run with zero profiled steps never starts the " + "profiler and has no trace to export.", + tracename, + type(e).__name__, + e, + ) + return None + log.info("Wrote PyTorch trace '%s'", tracename) + return tracename + + @annotate() def main(kwargs_dict: dict = {}): # @@ -261,10 +296,7 @@ def main(kwargs_dict: dict = {}): trainer.train(profiler=prof if TORCH_PERF_LOCAL else None) end_code_region("train") if TORCH_PERF_LOCAL: - hostname = socket.gethostname() - tracename = f"torch-{hostname}-r{rank}-N{world_size // ranks_per_node}-n{world_size}-ps{config.problem_scale}-e{config.epochs}-nipf{config.n_instances_used_per_fractal}-{int(time.time())}.json" - prof.export_chrome_trace(tracename) - log.info("Wrote PyTorch trace '%s'", tracename) + export_profiler_trace(prof, config, log, rank, world_size, ranks_per_node) # Results are final here; synchronize before rank-0 post-processing so a # post-processing failure on rank 0 cannot strand the other ranks in a diff --git a/tests/test_reporting.py b/tests/test_reporting.py index a6f0106..f77fb19 100644 --- a/tests/test_reporting.py +++ b/tests/test_reporting.py @@ -270,3 +270,79 @@ def test_torch_profiler_independent_of_caliper(self, monkeypatch): ) finally: importlib.reload(perf_measure) + + +class TestProfilerTraceExport: + """R15: a failed trace export must not strand the other ranks.""" + + @staticmethod + def _unstepped_profiler(): + """A profiler whose window never opened (a run with zero batches).""" + from torch.profiler import ProfilerActivity, profile, schedule + + prof = profile( + activities=[ProfilerActivity.CPU], + schedule=schedule(wait=1, warmup=1, active=3, repeat=1), + ) + with prof: + pass # no prof.step(): the schedule never leaves its wait phase + return prof + + @staticmethod + def _config(run_dir): + return SimpleNamespace( + problem_scale=4, + epochs=1, + n_instances_used_per_fractal=2, + run_dir=str(run_dir), + ) + + def test_zero_step_export_is_reported_not_raised(self, tmp_path, caplog): + """Exporting an unstepped profiler logs an error instead of raising. + + The export runs before the ``dist.barrier()`` that precedes rank-0 + post-processing, so a raise here kills the profiling rank and leaves + every other rank blocked in that barrier until the collective timeout. + """ + import logging + + import ScaFFold.worker as worker + + log = logging.getLogger("test_zero_step_export") + with caplog.at_level(logging.DEBUG, logger=log.name): + result = worker.export_profiler_trace( + self._unstepped_profiler(), + self._config(tmp_path), + log, + rank=0, + world_size=1, + ranks_per_node=1, + ) + + assert result is None + messages = " ".join(record.getMessage() for record in caplog.records) + assert "trace" in messages.lower() + + def test_successful_export_writes_a_trace(self, tmp_path, caplog): + """A profiler with a completed window still writes its trace (control).""" + import logging + + from torch.profiler import ProfilerActivity, profile, schedule + + import ScaFFold.worker as worker + + prof = profile( + activities=[ProfilerActivity.CPU], + schedule=schedule(wait=1, warmup=1, active=1, repeat=1), + ) + with prof: + for _ in range(4): + prof.step() + + log = logging.getLogger("test_successful_export") + path = worker.export_profiler_trace( + prof, self._config(tmp_path), log, rank=0, world_size=1, ranks_per_node=1 + ) + + assert path is not None + assert Path(path).exists() From f722a9afdc7cc6e1fd97caf4b7af9d8c336c7592 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 16:11:11 -0700 Subject: [PATCH 30/62] Detect the local rank count and place traces in the run dir get_local_size ignored LOCAL_WORLD_SIZE/PMI_LOCAL_SIZE/PALS_LOCAL_SIZE that get_local_rank honors, so it returned 1 and the per-node profiler gate selected every rank while the trace name claimed one node per rank. Add the missing variables, round the node count up, and write the trace into config.run_dir instead of the working directory. R23 --- ScaFFold/utils/distributed.py | 13 +++++++- ScaFFold/worker.py | 19 +++++++++--- tests/test_reporting.py | 58 +++++++++++++++++++++++++++++++++++ tests/test_worker_dist.py | 50 ++++++++++++++++++++++++++++++ 4 files changed, 134 insertions(+), 6 deletions(-) diff --git a/ScaFFold/utils/distributed.py b/ScaFFold/utils/distributed.py index c46a5bd..afad5ab 100644 --- a/ScaFFold/utils/distributed.py +++ b/ScaFFold/utils/distributed.py @@ -62,11 +62,22 @@ def get_local_rank(required: bool = False) -> int: def get_local_size(required: bool = False) -> int: - """Return the number of local MPI ranks.""" + """Return the number of local MPI ranks. + + Recognizes the same launchers as ``get_local_rank``: a variable honored + there but not here silently yields 1, which makes per-node logic (e.g. the + profiler's one-rank-per-node gate) treat every rank as node-local. + """ + if "LOCAL_WORLD_SIZE" in os.environ: + return int(os.environ["LOCAL_WORLD_SIZE"]) if "MV2_COMM_WORLD_LOCAL_SIZE" in os.environ: return int(os.environ["MV2_COMM_WORLD_LOCAL_SIZE"]) if "OMPI_COMM_WORLD_LOCAL_SIZE" in os.environ: return int(os.environ["OMPI_COMM_WORLD_LOCAL_SIZE"]) + if "PMI_LOCAL_SIZE" in os.environ: + return int(os.environ["PMI_LOCAL_SIZE"]) + if "PALS_LOCAL_SIZE" in os.environ: + return int(os.environ["PALS_LOCAL_SIZE"]) if "SLURM_NNODES" in os.environ and "SLURM_NTASKS" in os.environ: return int(os.environ["SLURM_NTASKS"]) // int(os.environ["SLURM_NNODES"]) # Flux does not have an env variable for this, so we assume an diff --git a/ScaFFold/worker.py b/ScaFFold/worker.py index 8eafa51..2f36b9a 100644 --- a/ScaFFold/worker.py +++ b/ScaFFold/worker.py @@ -17,6 +17,7 @@ import socket import time from argparse import Namespace +from pathlib import Path import numpy as np import psutil @@ -107,28 +108,36 @@ def export_profiler_trace(prof, config, log, rank, world_size, ranks_per_node): then raises; the trace can also fill the filesystem), so log them and let the job finish. + The trace is written into ``config.run_dir`` so it lands with the rest of + the run's artifacts instead of wherever the job happened to be launched + from. + Returns the path written, or None if the trace could not be written. """ + # Round up: with a partly-filled last node, flooring would report one node + # too few (and a ranks_per_node larger than the job would report none). + nodes = max(1, math.ceil(world_size / max(1, ranks_per_node))) tracename = ( f"torch-{socket.gethostname()}-r{rank}" - f"-N{world_size // ranks_per_node}-n{world_size}" + f"-N{nodes}-n{world_size}" f"-ps{config.problem_scale}-e{config.epochs}" f"-nipf{config.n_instances_used_per_fractal}-{int(time.time())}.json" ) + tracepath = Path(getattr(config, "run_dir", None) or os.getcwd()) / tracename try: - prof.export_chrome_trace(tracename) + prof.export_chrome_trace(str(tracepath)) except Exception as e: log.error( "Could not write PyTorch trace '%s': %s: %s. Continuing so the " "run can finish; a run with zero profiled steps never starts the " "profiler and has no trace to export.", - tracename, + tracepath, type(e).__name__, e, ) return None - log.info("Wrote PyTorch trace '%s'", tracename) - return tracename + log.info("Wrote PyTorch trace '%s'", tracepath) + return tracepath @annotate() diff --git a/tests/test_reporting.py b/tests/test_reporting.py index f77fb19..a4948c4 100644 --- a/tests/test_reporting.py +++ b/tests/test_reporting.py @@ -18,6 +18,7 @@ import matplotlib import numpy as np +import pytest matplotlib.use("Agg") @@ -288,6 +289,20 @@ def _unstepped_profiler(): pass # no prof.step(): the schedule never leaves its wait phase return prof + @staticmethod + def _stepped_profiler(): + """A profiler with a completed capture window.""" + from torch.profiler import ProfilerActivity, profile, schedule + + prof = profile( + activities=[ProfilerActivity.CPU], + schedule=schedule(wait=1, warmup=1, active=1, repeat=1), + ) + with prof: + for _ in range(4): + prof.step() + return prof + @staticmethod def _config(run_dir): return SimpleNamespace( @@ -346,3 +361,46 @@ def test_successful_export_writes_a_trace(self, tmp_path, caplog): assert path is not None assert Path(path).exists() + + def test_trace_lands_in_the_run_dir(self, tmp_path, caplog): + """R23: the trace goes to the run dir, not whatever CWD happens to be.""" + import logging + + import ScaFFold.worker as worker + + prof = self._stepped_profiler() + log = logging.getLogger("test_trace_lands_in_the_run_dir") + + path = worker.export_profiler_trace( + prof, self._config(tmp_path), log, rank=0, world_size=1, ranks_per_node=1 + ) + + assert Path(path).parent == tmp_path + assert list(tmp_path.glob("torch-*.json")) == [Path(path)] + + @pytest.mark.parametrize( + "world_size, ranks_per_node, expected", + [(8, 4, "-N2-n8-"), (6, 4, "-N2-n6-"), (1, 1, "-N1-n1-")], + ids=["even", "ragged-last-node", "singleton"], + ) + def test_trace_name_counts_nodes_not_ranks( + self, tmp_path, world_size, ranks_per_node, expected + ): + """R23: the N field is a node count, and never rounds a node away.""" + import logging + + import ScaFFold.worker as worker + + prof = self._stepped_profiler() + log = logging.getLogger("test_trace_name_counts_nodes") + + path = worker.export_profiler_trace( + prof, + self._config(tmp_path), + log, + rank=0, + world_size=world_size, + ranks_per_node=ranks_per_node, + ) + + assert expected in Path(path).name diff --git a/tests/test_worker_dist.py b/tests/test_worker_dist.py index 40be05b..b1ae524 100644 --- a/tests/test_worker_dist.py +++ b/tests/test_worker_dist.py @@ -226,3 +226,53 @@ def fake_train(self, profiler=None): assert trainer.config.global_batch_size == trainer.config.local_batch_size # The worker destroyed the process group before rank-0 post-processing. assert not torch.distributed.is_initialized() + + +# --------------------------------------------------------------------------- +# Local size detection (R23) +# --------------------------------------------------------------------------- + +_LOCAL_SIZE_CASES = [ + # torchrun exports LOCAL_WORLD_SIZE alongside LOCAL_RANK. + ({"LOCAL_WORLD_SIZE": "4"}, 4), + ({"MV2_COMM_WORLD_LOCAL_SIZE": "4"}, 4), + ({"OMPI_COMM_WORLD_LOCAL_SIZE": "4"}, 4), + ({"PMI_LOCAL_SIZE": "4"}, 4), + ({"PALS_LOCAL_SIZE": "4"}, 4), + ({"SLURM_NTASKS": "8", "SLURM_NNODES": "2"}, 4), + ({"FLUX_JOB_SIZE": "8", "FLUX_JOB_NNODES": "2"}, 4), +] + +_LOCAL_SIZE_VARS = [ + "LOCAL_WORLD_SIZE", + "MV2_COMM_WORLD_LOCAL_SIZE", + "OMPI_COMM_WORLD_LOCAL_SIZE", + "PMI_LOCAL_SIZE", + "PALS_LOCAL_SIZE", + "SLURM_NTASKS", + "SLURM_NNODES", + "FLUX_JOB_SIZE", + "FLUX_JOB_NNODES", +] + + +def test_local_size_detection_matrix(monkeypatch): + """Every launcher that reports a local rank has its local size honored too. + + An unrecognized variable silently yields 1, which makes the per-node + profiler gate ``rank % ranks_per_node == 0`` select *every* rank and + mislabels the trace's node count. + """ + for env, want_local_size in _LOCAL_SIZE_CASES: + for var in _LOCAL_SIZE_VARS: + monkeypatch.delenv(var, raising=False) + for key, value in env.items(): + monkeypatch.setenv(key, value) + assert distributed_mod.get_local_size() == want_local_size, env + + +def test_local_size_defaults_to_one(monkeypatch): + """With nothing to go on, one rank per node is still the assumption.""" + for var in _LOCAL_SIZE_VARS: + monkeypatch.delenv(var, raising=False) + assert distributed_mod.get_local_size() == 1 From 9e87457850562379a1309df9481c750f50ad6b8c Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 16:12:21 -0700 Subject: [PATCH 31/62] Parse PROFILE_TORCH like the other profiler flags The gate enabled profiling for any value except the literal "off", so PROFILE_TORCH=0, =false, =no and ="" all turned the profiler ON -- the opposite of what they say, and inconsistent with the sub-option flags in the same module. BEHAVIOR CHANGE: profiling is now enabled only by 1/true/on/yes (any case); every other value, including previously-enabling ones such as "enabled", leaves it off. R24 --- README.md | 2 +- ScaFFold/utils/perf_measure.py | 20 ++++++----- tests/test_reporting.py | 65 ++++++++++++++++++++++++++++++++++ 3 files changed, 78 insertions(+), 9 deletions(-) diff --git a/README.md b/README.md index 86acb7b..2442494 100644 --- a/README.md +++ b/README.md @@ -172,7 +172,7 @@ For n  in n_volumes: #### 1. Profiling with the PyTorch Profiler -Set `PROFILE_TORCH=ON` to generate a PyTorch profiling trace that can be read into [Perfetto](https://ui.perfetto.dev/). +Set `PROFILE_TORCH=ON` to generate a PyTorch profiling trace that can be read into [Perfetto](https://ui.perfetto.dev/). The trace is written into the run directory. `1`, `true`, `on` and `yes` (any case) enable profiling; every other value, including `0`, `false`, `no` and `off`, leaves it disabled. #### 2. Profiling with Caliper & Adiak diff --git a/ScaFFold/utils/perf_measure.py b/ScaFFold/utils/perf_measure.py index c69b4b6..2f17cc9 100644 --- a/ScaFFold/utils/perf_measure.py +++ b/ScaFFold/utils/perf_measure.py @@ -18,6 +18,17 @@ CALI_PERF_ENV_VAR = "CALI_CONFIG" TORCH_PERF_ENV_VAR = "PROFILE_TORCH" + +def _profiler_env_flag(name): + """Return True only for an affirmative value of the environment variable. + + Every profiler toggle -- the master switch and its sub-options alike -- + goes through this, so "0"/"false"/"no"/"off"/"" all mean off and there is + no spelling that means the opposite of what it says. + """ + return os.environ.get(name, "").lower() in ("1", "true", "on", "yes") + + _CALI_PERF_ENABLED = False TORCH_PERF_ENABLED = False if CALI_PERF_ENV_VAR in os.environ: @@ -33,10 +44,7 @@ # The torch profiler is gated purely on its own environment variable: Caliper # and the torch profiler may both be enabled at once. -if ( - TORCH_PERF_ENV_VAR in os.environ - and os.environ.get(TORCH_PERF_ENV_VAR).lower() != "off" -): +if _profiler_env_flag(TORCH_PERF_ENV_VAR): try: from torch.profiler import ProfilerActivity from torch.profiler import profile as torchprofile @@ -100,10 +108,6 @@ def _profiler_env_int(name, default): return default -def _profiler_env_flag(name): - return os.environ.get(name, "").lower() in ("1", "true", "on", "yes") - - def get_torch_context(ranks_per_node, rank): if TORCH_PERF_ENABLED: TORCH_PERF_LOCAL = TORCH_PERF_ENABLED and (rank % ranks_per_node == 0) diff --git a/tests/test_reporting.py b/tests/test_reporting.py index a4948c4..e9cead1 100644 --- a/tests/test_reporting.py +++ b/tests/test_reporting.py @@ -404,3 +404,68 @@ def test_trace_name_counts_nodes_not_ranks( ) assert expected in Path(path).name + + +class TestProfileTorchGate: + """R24: PROFILE_TORCH is parsed like every other profiler flag.""" + + @staticmethod + def _reload_with(monkeypatch_context, value): + import importlib + + import ScaFFold.utils.perf_measure as perf_measure + + if value is None: + monkeypatch_context.delenv("PROFILE_TORCH", raising=False) + else: + monkeypatch_context.setenv("PROFILE_TORCH", value) + monkeypatch_context.delenv("CALI_CONFIG", raising=False) + importlib.reload(perf_measure) + return perf_measure + + @pytest.mark.parametrize("value", [None, "", "0", "false", "no", "off", "OFF"]) + def test_disabled_values(self, monkeypatch, value): + """Anything that is not an affirmative value leaves profiling off. + + ``PROFILE_TORCH=0`` used to *enable* the profiler: the gate only + rejected the literal "off", so every conventional way of saying "no" + silently turned profiling on. + """ + import importlib + + import ScaFFold.utils.perf_measure as perf_measure + + try: + with monkeypatch.context() as m: + assert not self._reload_with(m, value).TORCH_PERF_ENABLED + finally: + importlib.reload(perf_measure) + + @pytest.mark.parametrize("value", ["1", "true", "on", "ON", "yes", "TRUE"]) + def test_enabled_values(self, monkeypatch, value): + """The affirmative spellings still enable the profiler.""" + import importlib + + import ScaFFold.utils.perf_measure as perf_measure + + try: + with monkeypatch.context() as m: + assert self._reload_with(m, value).TORCH_PERF_ENABLED + finally: + importlib.reload(perf_measure) + + def test_gate_matches_the_sub_option_parser(self, monkeypatch): + """The master switch and the sub-option flags agree on every spelling.""" + import importlib + + import ScaFFold.utils.perf_measure as perf_measure + + try: + for value in ("1", "true", "on", "yes", "0", "false", "no", "off", ""): + with monkeypatch.context() as m: + module = self._reload_with(m, value) + assert module.TORCH_PERF_ENABLED == module._profiler_env_flag( + "PROFILE_TORCH" + ), value + finally: + importlib.reload(perf_measure) From b9a42aa2529559e818b521c61c57c057a06e6521 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 16:13:20 -0700 Subject: [PATCH 32/62] Require the profiler window to skip at least one step The profiler context wraps checkpoint cleanup and every warmup batch while prof.step() advances only per training batch, so PROFILE_TORCH_WAIT=0 buffered that entire prologue in host memory as one unbounded step. Clamp wait to 1 and say so, rather than reordering the context or stepping it from warmup. R26 --- ScaFFold/utils/perf_measure.py | 15 +++++++ tests/test_reporting.py | 80 ++++++++++++++++++++++++++++++++++ 2 files changed, 95 insertions(+) diff --git a/ScaFFold/utils/perf_measure.py b/ScaFFold/utils/perf_measure.py index 2f17cc9..7236e69 100644 --- a/ScaFFold/utils/perf_measure.py +++ b/ScaFFold/utils/perf_measure.py @@ -122,7 +122,22 @@ def get_torch_context(ranks_per_node, rank): # trace. The window (skip `wait`, prime `warmup`, capture `active`, once) # is tunable via the environment. Callers must drive it with # ``prof.step()`` once per training step for the schedule to advance. + # The context is entered around checkpoint cleanup and the warmup + # batches, but prof.step() only advances once per *training* batch, so + # everything before the first training batch lands in step 0. The + # window must therefore skip at least one step: with wait=0 that whole + # prologue -- warmup_batches forward+backward passes per rank -- is + # buffered in host memory as a single unbounded step, which is the very + # thing the bounded window exists to prevent. wait = _profiler_env_int("PROFILE_TORCH_WAIT", 1) + if wait < 1: + print( + "PROFILE_TORCH_WAIT must be at least 1: the profiler window " + "opens before the warmup batches, whose work would otherwise " + "accumulate in host memory as one unbounded step. Using " + "PROFILE_TORCH_WAIT=1." + ) + wait = 1 warmup = _profiler_env_int("PROFILE_TORCH_WARMUP", 1) active = _profiler_env_int("PROFILE_TORCH_ACTIVE", 3) or 1 diff --git a/tests/test_reporting.py b/tests/test_reporting.py index e9cead1..3b658fd 100644 --- a/tests/test_reporting.py +++ b/tests/test_reporting.py @@ -469,3 +469,83 @@ def test_gate_matches_the_sub_option_parser(self, monkeypatch): ), value finally: importlib.reload(perf_measure) + + +class TestProfilerSchedule: + """R26: the schedule must not record everything before the first step.""" + + @staticmethod + def _context_with(monkeypatch_context, env): + """Reload perf_measure with ``env`` applied and build a profiler context.""" + import importlib + + import ScaFFold.utils.perf_measure as perf_measure + + monkeypatch_context.setenv("PROFILE_TORCH", "1") + monkeypatch_context.delenv("CALI_CONFIG", raising=False) + for name in ("PROFILE_TORCH_WAIT", "PROFILE_TORCH_WARMUP"): + monkeypatch_context.delenv(name, raising=False) + for key, value in env.items(): + monkeypatch_context.setenv(key, value) + importlib.reload(perf_measure) + assert perf_measure.TORCH_PERF_ENABLED + ctx, is_local = perf_measure.get_torch_context(1, 0) + assert is_local + return ctx + + def test_wait_zero_does_not_record_step_zero(self, monkeypatch, capsys): + """PROFILE_TORCH_WAIT=0 is clamped so step 0 records nothing. + + worker.main enters the profiler context around checkpoint cleanup and + every warmup batch, and ``prof.step()`` only advances once per training + batch -- so a schedule that is already active at step 0 buffers all of + that as a single unbounded step, which is exactly what the bounded + window exists to prevent. + """ + import importlib + + from torch.profiler import ProfilerAction + + import ScaFFold.utils.perf_measure as perf_measure + + try: + with monkeypatch.context() as m: + ctx = self._context_with(m, {"PROFILE_TORCH_WAIT": "0"}) + assert ctx.schedule(0) == ProfilerAction.NONE + output = capsys.readouterr().out + assert "PROFILE_TORCH_WAIT" in output + finally: + importlib.reload(perf_measure) + + def test_default_schedule_skips_step_zero(self, monkeypatch): + """The default window already skips step 0 (control).""" + import importlib + + from torch.profiler import ProfilerAction + + import ScaFFold.utils.perf_measure as perf_measure + + try: + with monkeypatch.context() as m: + ctx = self._context_with(m, {}) + assert ctx.schedule(0) == ProfilerAction.NONE + finally: + importlib.reload(perf_measure) + + def test_larger_wait_is_preserved(self, monkeypatch): + """A wait longer than the minimum is left alone.""" + import importlib + + from torch.profiler import ProfilerAction + + import ScaFFold.utils.perf_measure as perf_measure + + try: + with monkeypatch.context() as m: + ctx = self._context_with( + m, {"PROFILE_TORCH_WAIT": "3", "PROFILE_TORCH_WARMUP": "1"} + ) + assert ctx.schedule(2) == ProfilerAction.NONE + assert ctx.schedule(3) == ProfilerAction.WARMUP + finally: + importlib.reload(perf_measure) From 6a0ae3a9c46f9efa896f5f5fd83fc7b405abe6de Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 16:16:44 -0700 Subject: [PATCH 33/62] Cover the config round-trip across a restart Regression guard for the merge order: a restart driven by the run dir's config.yaml (what restart.sh emits) must reproduce the first run's CLI overrides and auxiliary config keys and still take the resume path. R20 R22 --- tests/test_cli.py | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/tests/test_cli.py b/tests/test_cli.py index b215e44..381744f 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -464,3 +464,36 @@ def test_cli_override_bottleneck_out_of_range_rejected(monkeypatch, tmp_path): message = str(excinfo.value) assert "unet_bottleneck_dim" in message assert "problem_scale" in message + + +# --------------------------------------------------------------------------- +# The whole config path survives a restart (R20/R22 together) +# --------------------------------------------------------------------------- + + +def test_config_round_trips_through_a_restart(monkeypatch, tmp_path): + """A restart driven by the run dir's config.yaml reproduces the run. + + This is exactly what the generated restart.sh does: ``-c + $RUN_DIR/config.yaml --restart --run-dir $RUN_DIR``. It exercises the merge + in both directions -- CLI overrides and auxiliary config keys have to come + back out of the dumped config, and the resume flags have to win. + """ + cfg = write_config(tmp_path, {"verbose": 1, "datagen_batch_size": 500}) + + _, calls = run_cli( + monkeypatch, + ["scaffold", "benchmark", "-c", str(cfg), "--local-batch-size", "2"], + ) + run_dir = Path(calls["benchmark"][0]["benchmark_run_dir"]) + _make_checkpoint(run_dir) + + _, resumed = run_cli(monkeypatch, _restart_argv(run_dir / "config.yaml", run_dir)) + + (config,) = resumed["benchmark"] + assert config["local_batch_size"] == 2 # CLI override from the first run + assert config["verbose"] == 1 # auxiliary key set in YAML + assert config["datagen_batch_size"] == 500 + assert config["restart"] is True + assert config["train_from_scratch"] is False + assert config["benchmark_run_dir"] == str(run_dir) From d6c7836eff30b62ebf45041ef4e824a0ff7a7d72 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 16:47:25 -0700 Subject: [PATCH 34/62] Fence the remaining rank-0 windows in checkpointing Rank 0 is the only rank that touches the run directory, and it does so while its peers are already committed to the next collective, so every one of those windows has to report failure *through* that collective. Three were left open. cleanup() ran _remove_checkpoint_files between the drain and the broadcast: individual unlinks were tolerated, but the glob and stat around them can raise on a shared filesystem (ESTALE, EACCES), re-creating the R05 hazard in a narrow window. load_from_checkpoint() folded only the drain error into its decision broadcast, leaving _select_and_load -- which stats, deserializes and renames -- able to kill rank 0 before the broadcast the peers were waiting in. Both now travel as decisions; cleanup's payload carries the phase so the peers name the operation that actually failed rather than reporting a "save". __init__'s orphan sweep is rank-0-only outside any collective, so a directory it cannot list would abort rank 0 alone and strand the peers at cleanup's first collective. It only reclaims space, so it is now best-effort with a warning. VA-1, VA-2, VA-3 --- ScaFFold/utils/checkpointing.py | 95 +++++++++++++++++++++++++------- tests/test_checkpointing.py | 98 ++++++++++++++++++++++++++++++++- 2 files changed, 171 insertions(+), 22 deletions(-) diff --git a/ScaFFold/utils/checkpointing.py b/ScaFFold/utils/checkpointing.py index 24f3fe6..22c6051 100644 --- a/ScaFFold/utils/checkpointing.py +++ b/ScaFFold/utils/checkpointing.py @@ -26,12 +26,13 @@ class CheckpointSaveError(RuntimeError): - """A checkpoint write failed. + """A rank-0 checkpoint operation failed (a write, a cleanup, a load). - Raised identically on every rank. Only rank 0 writes, but its outcome is - broadcast, so the peers report the real disk error instead of the - unmatched-collective symptom (an opaque gloo transport error, or an NCCL - watchdog timeout minutes later) that a rank-0-only raise produces. + Raised identically on every rank. Only rank 0 touches the run directory, + but its outcome is broadcast, so the peers report the real disk error + instead of the unmatched-collective symptom (an opaque gloo transport + error, or an NCCL watchdog timeout minutes later) that a rank-0-only raise + produces. """ @@ -111,17 +112,33 @@ def __init__( # Ensure base directory exists (Rank 0 only) if self.world_rank == 0: self.base_dir.mkdir(parents=True, exist_ok=True) - self._sweep_orphaned_tmp_files() + try: + self._sweep_orphaned_tmp_files() + except Exception as e: + # Construction is rank-0-only work outside any collective, so a + # raise here aborts rank 0 alone and leaves the peers waiting in + # the manager's first collective (``cleanup``). The sweep only + # reclaims space, so a directory that cannot be listed (a stale + # NFS/Lustre handle, a permissions oddity) degrades to a warning + # rather than taking the job down asymmetrically. + self._log( + f"Could not sweep orphaned checkpoint temp files in " + f"{self.base_dir}: {type(e).__name__}: {e}" + ) def cleanup(self, train_from_scratch: bool) -> None: """Clear existing checkpoints if training from scratch. Rank-symmetric, like every other collective point here: any pending - async write is drained and its outcome broadcast, so a failure raises - on all ranks together (see ``save_checkpoint``). + async write is drained, the rank-0 deletion is fenced, and whichever + failed is broadcast, so a failure raises on all ranks together (see + ``save_checkpoint``). The broadcast payload is ``(phase, description)`` + so the peers -- which saw neither the write nor the deletion -- report + the same operation rank 0 did. """ # Ensure any pending async save is finished before deleting. error = self._drain_pending_save() + failure = None if error is None else ("save", error) if train_from_scratch: # Drop the cached state that described the run being deleted. Both @@ -134,12 +151,29 @@ def cleanup(self, train_from_scratch: bool) -> None: self.last_saved_epoch = None if self.world_rank == 0: - self._remove_checkpoint_files() - - error = self._broadcast_obj(error) + # Rank 0 alone touches the filesystem here, while every peer is + # already committed to the broadcast below. Individual unlinks + # are tolerated inside, but the glob/stat around them can still + # raise on a shared filesystem (ESTALE, EACCES), and raising in + # this window strands the peers in an unmatched collective -- + # the R05 hazard. Report the failure through the broadcast, like + # a failed write. + try: + self._remove_checkpoint_files() + except Exception as e: + if failure is None: + self._save_error_exc = e + failure = ("cleanup", f"{type(e).__name__}: {e}") + else: + # A drained write already failed; that outcome is the + # one being reported, so this is only logged. + self._log(f"Clearing existing checkpoints also failed: {e}") + + failure = self._broadcast_obj(failure) self._barrier() - if error is not None: - self._raise_save_error(error) + if failure is not None: + phase, description = failure + self._raise_save_error(description, phase=phase) def _remove_checkpoint_files(self) -> None: """Delete this run's checkpoint files and debris (rank 0 only). @@ -230,15 +264,17 @@ def _drain_pending_save(self) -> Optional[str]: return f"{type(e).__name__}: {e}" return None - def _raise_save_error(self, description: str) -> None: - """Raise a broadcast save failure on this rank. + def _raise_save_error(self, description: str, phase: str = "save") -> None: + """Raise a broadcast rank-0 failure on this rank. - Rank 0 chains the original exception so its traceback survives; the - peers never saw it and raise the same message on its own. + ``phase`` names the operation that failed (``save``, ``cleanup``, + ``load``) so the message describes what actually went wrong. Rank 0 + chains the original exception so its traceback survives; the peers never + saw it and raise the same message on its own. """ cause, self._save_error_exc = self._save_error_exc, None raise CheckpointSaveError( - f"Checkpoint save failed on rank 0: {description}" + f"Checkpoint {phase} failed on rank 0: {description}" ) from cause def finalize_saves(self) -> None: @@ -317,14 +353,26 @@ def load_from_checkpoint(self, require_checkpoint: bool = False) -> int: # therefore never open the checkpoint files at all. result = None if self.world_rank == 0: - result = ( - ("save_failed", error) if error is not None else self._select_and_load() - ) + if error is not None: + result = ("save_failed", error) + else: + # The selection itself stats, deserializes and renames files + # while the peers are already blocked in the broadcast below, so + # anything it raises (a stale handle on ``exists``, a rename + # denied, an unpickling MemoryError) has to become a decision + # rather than a rank-0-only death. + try: + result = self._select_and_load() + except Exception as e: + self._save_error_exc = e + result = ("load_failed", f"{type(e).__name__}: {e}") status, payload = self._broadcast_obj(result) # 2. Every rank acts on the same decision rank 0 reached. if status == "save_failed": self._raise_save_error(payload) + if status == "load_failed": + self._raise_save_error(payload, phase="load") if status == "empty": if require_checkpoint: # An explicit restart must resume real state; silently @@ -410,6 +458,11 @@ def _select_and_load(self): * ``("empty", None)`` -- no checkpoint files exist; * ``("ok", checkpoint_dict)`` -- a candidate deserialized cleanly; * ``("unreadable", [paths])`` -- every candidate was corrupt. + + The filesystem calls around those decisions (``exists``, the quarantine + rename) can still fail outright; the caller runs this inside a guard + that turns such a failure into a broadcast ``("load_failed", ...)`` + decision, because raising here would strand the peers. """ candidates = [] if self.last_ckpt_path.exists(): diff --git a/tests/test_checkpointing.py b/tests/test_checkpointing.py index 2ef9cad..acca106 100644 --- a/tests/test_checkpointing.py +++ b/tests/test_checkpointing.py @@ -39,7 +39,7 @@ import torch.distributed as dist import ScaFFold.utils.trainer as trainer_mod -from ScaFFold.utils.checkpointing import CheckpointManager +from ScaFFold.utils.checkpointing import CheckpointManager, CheckpointSaveError from ScaFFold.utils.trainer import PyTorchTrainer from tests.helpers import mpi_runner @@ -518,6 +518,102 @@ def test_init_sweeps_orphaned_tmp_files(tmp_path): assert quarantined.exists() +# --------------------------------------------------------------------------- +# VA-1/VA-2/VA-3 -- the remaining rank-0 filesystem windows are fenced +# +# Rank 0 is the only rank that touches the run directory, and it does so while +# its peers are already committed to the next collective. Every such window must +# therefore report its failure *through* that collective; raising inside it +# leaves the peers in an unmatched collective, where a plain disk error +# resurfaces as an opaque gloo transport error or an NCCL watchdog timeout. +# --------------------------------------------------------------------------- + + +def _record_collectives(monkeypatch): + """Stub the process-group collectives, recording what a rank posts.""" + posted = {"broadcasts": [], "barriers": 0} + + def fake_broadcast(objs, src=0): + posted["broadcasts"].append(objs[0]) + + def fake_barrier(*args, **kwargs): + posted["barriers"] += 1 + + monkeypatch.setattr(dist, "broadcast_object_list", fake_broadcast) + monkeypatch.setattr(dist, "barrier", fake_barrier) + return posted + + +def _raise_stale(*args, **kwargs): + raise OSError("[Errno 116] Stale file handle") + + +def test_cleanup_rank0_fs_error_travels_through_the_broadcast(tmp_path, monkeypatch): + """A failure while clearing checkpoints fails every rank, not just rank 0. + + ``_remove_checkpoint_files`` globs and stats the run directory; on a shared + filesystem those can raise (ESTALE, EACCES) even though each individual + unlink is already tolerated. That happens between the drain and the + broadcast, so an unfenced raise re-creates exactly the hazard R05 closed. + """ + mgr, _ = _make_manager(tmp_path) + mgr.dist_enabled = True + posted = _record_collectives(monkeypatch) + monkeypatch.setattr(CheckpointManager, "_remove_checkpoint_files", _raise_stale) + + with pytest.raises(CheckpointSaveError) as excinfo: + mgr.cleanup(train_from_scratch=True) + + assert "Stale file handle" in str(excinfo.value) + # The peers' collectives were posted before the raise, so they fail with the + # same error instead of hanging. + assert len(posted["broadcasts"]) == 1 + assert "Stale file handle" in str(posted["broadcasts"][0]) + assert posted["barriers"] == 1 + + +def test_load_rank0_selection_error_travels_through_the_broadcast( + tmp_path, monkeypatch +): + """A failure inside ``_select_and_load`` is a broadcast decision too. + + The load path folded only the *drain* error into its decision broadcast, so + rank 0 stat-ing or renaming a checkpoint candidate could still die before + the broadcast the peers were already waiting in. + """ + mgr, _ = _make_manager(tmp_path) + mgr.dist_enabled = True + posted = _record_collectives(monkeypatch) + monkeypatch.setattr(CheckpointManager, "_select_and_load", _raise_stale) + + with pytest.raises(CheckpointSaveError) as excinfo: + mgr.load_from_checkpoint() + + assert "Stale file handle" in str(excinfo.value) + assert len(posted["broadcasts"]) == 1 + assert posted["broadcasts"][0][0] == "load_failed" + + +def test_init_survives_an_unlistable_run_dir(tmp_path, monkeypatch, capsys): + """A run directory that cannot be listed warns instead of killing __init__. + + The orphan sweep runs on rank 0 only and outside any collective, so a raise + here aborts rank 0 alone and strands the peers at the manager's first + collective (``cleanup``). ``pathlib`` already swallows PermissionError, but + a stale NFS/Lustre handle propagates; the sweep is opportunistic, so it must + degrade to a warning. + """ + base = tmp_path / "checkpoints" + base.mkdir() + (base / "checkpoint_last.pth.tmp.999999").write_bytes(b"partial checkpoint") + monkeypatch.setattr(Path, "glob", _raise_stale) + + mgr, _ = _make_manager(base) + + assert mgr.base_dir == base + assert "sweep" in capsys.readouterr().out.lower() + + # --------------------------------------------------------------------------- # F71 -- CPU tensors are cloned into the snapshot # --------------------------------------------------------------------------- From 2fb1bf53661670bc1cfa2e734799ecf6ae02c3f6 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 16:51:15 -0700 Subject: [PATCH 35/62] Fence the remaining rank-0 windows in datagen The category search broadcasts its scan of the existing category files, but rank 0 then loaded those files' parameters -- for its duplicate guard -- just *after* that broadcast, unfenced. A CSV that will not parse (ragged, or hand-edited) killed rank 0 while every peer had already taken the broadcast and entered the work loop, stranding them in its next collective: the same window class the broadcast was introduced to close. Scan and load are now one guarded decision reported through one broadcast, as get_dataset reports its selection. The datagen consensus guards caught (Exception, SystemExit), which is not the same as "everything": a KeyboardInterrupt delivered to rank 0 alone (Ctrl-C on the launching terminal, a watchdog SIGINT) unwound straight past the broadcast and hung the peers. They now catch BaseException. SystemExit keeps being converted -- it must never escape get_dataset, since a peer would read the silent unwind as success -- while the rank that was interrupted re-raises the interrupt after posting the sentinel, so it keeps the operator's exit status and its peers still learn to stop. volumegen's guard, the same pattern feeding the same allreduce, is fixed with it. VB-2, VB-3 --- ScaFFold/datagen/category_search.py | 45 +++++++++++++++++------ ScaFFold/datagen/get_dataset.py | 42 ++++++++++++++++++--- ScaFFold/datagen/volumegen.py | 12 +++++- tests/datagen/test_category_search.py | 53 +++++++++++++++++++++++++-- tests/datagen/test_mpi_consensus.py | 51 ++++++++++++++++++++++++++ 5 files changed, 182 insertions(+), 21 deletions(-) diff --git a/ScaFFold/datagen/category_search.py b/ScaFFold/datagen/category_search.py index af53a9c..f3b3496 100644 --- a/ScaFFold/datagen/category_search.py +++ b/ScaFFold/datagen/category_search.py @@ -461,21 +461,44 @@ def main(config: Config) -> None: # partially visible directory) put one rank inside the loop while another is # past it, so the two post mismatched collectives on COMM_WORLD and the job # hangs. One scan, one broadcast, one shared verdict. - if rank == 0: - existing_indices = parse_category_indices(fracts_write_dir) - else: - existing_indices = None - existing_indices = comm.bcast(existing_indices, root=0) - + # + # Rank 0 also loads the parameters of the categories already on disk (only + # it writes, so only it needs them for the duplicate guard). That read is + # part of the same rank-0-only window: a category CSV that will not parse -- + # ragged, or hand-edited -- would otherwise kill rank 0 *after* the peers + # had already taken the broadcast and moved on to the next collective. Scan + # and load are therefore one guarded decision, reported through one + # broadcast, exactly as ``get_dataset`` reports its selection. existing_params = [] + interrupt = None if rank == 0: - for idx in existing_indices: - existing_params.append( - np.loadtxt( - os.path.join(fracts_write_dir, "%06d.csv" % idx), - delimiter=",", + try: + existing_indices = parse_category_indices(fracts_write_dir) + for idx in existing_indices: + existing_params.append( + np.loadtxt( + os.path.join(fracts_write_dir, "%06d.csv" % idx), + delimiter=",", + ) ) + scan = ("ok", existing_indices) + except BaseException as e: + existing_params = [] + scan = ( + "error", + f"rank 0 failed to scan existing categories in " + f"{fracts_write_dir}: {type(e).__name__}: {e}", ) + interrupt = e if isinstance(e, KeyboardInterrupt) else None + else: + scan = None + status, payload = comm.bcast(scan, root=0) + if status == "error": + # Rank 0 keeps an operator's interrupt; every rank aborts either way. + if interrupt is not None: + raise interrupt + raise RuntimeError(f"category search failed: {payload}") + existing_indices = payload # Calculate number of remaining fractal categories to generate existing_categories = len(existing_indices) diff --git a/ScaFFold/datagen/get_dataset.py b/ScaFFold/datagen/get_dataset.py index 09e4b33..e54b29f 100644 --- a/ScaFFold/datagen/get_dataset.py +++ b/ScaFFold/datagen/get_dataset.py @@ -292,6 +292,21 @@ def _decide_reuse_or_generate( return ("generate", str(tmp), str(dest)) +def _reraisable(exc: BaseException) -> BaseException | None: + """Return ``exc`` when the rank that caught it should re-raise it verbatim. + + The consensus guards below turn any failure into a sentinel so every rank + aborts together, and every rank then raises a ``RuntimeError`` carrying the + collected messages. That is the right report for a genuine error -- and for + ``SystemExit``, which must never escape ``get_dataset`` (a peer would treat + the silent unwind as success). A ``KeyboardInterrupt`` is different: it is + not an error but an operator abort, so the rank that received it re-raises + it after posting the sentinel, keeping the interrupt's own semantics (and + exit status) while its peers still learn to stop. + """ + return exc if isinstance(exc, KeyboardInterrupt) else None + + def get_dataset( config: Namespace, require_commit: bool = False, # default: ignore commit mismatches for reuse @@ -334,22 +349,32 @@ def get_dataset( # the broadcast below, so a rank-0 exception would strand the whole job. # Any failure is therefore turned into an error sentinel that travels # through the same broadcast and makes every rank raise the same error. + interrupt = None if rank == 0: try: decision = _decide_reuse_or_generate( base, config_id, commit, require_commit, log ) - except (Exception, SystemExit) as e: + except BaseException as e: + # BaseException, not (Exception, SystemExit): a KeyboardInterrupt + # delivered to rank 0 alone (Ctrl-C on the launching terminal, a + # site watchdog SIGINT) would otherwise skip the broadcast and hang + # every peer -- the exact failure this guard exists to prevent. decision = ( "error", f"rank 0 failed to select a dataset under {base}: " f"{type(e).__name__}: {e}", ) + interrupt = _reraisable(e) else: decision = None decision = comm.bcast(decision, root=0) if decision[0] == "error": + # Rank 0 keeps the abort signal it was actually given; the peers, which + # only ever saw the sentinel, report it as a generation failure. + if interrupt is not None: + raise interrupt raise RuntimeError(f"dataset selection failed: {decision[1]}") if decision[0] == "reuse": @@ -364,13 +389,15 @@ def get_dataset( err = "" # A worker failure must not skip any collective below: catch everything - # (including SystemExit, which is a BaseException and would otherwise bypass - # the consensus) so every rank always reaches the allreduce and gather. + # (BaseException, so neither SystemExit nor a KeyboardInterrupt delivered to + # one rank can bypass the consensus) so every rank always reaches the + # allreduce and gather. try: volumegen.main(config) - except (Exception, SystemExit) as e: + except BaseException as e: ok = False err = f"volumegen attempt failed: rank {rank}: {type(e).__name__}: {e}" + interrupt = _reraisable(e) # Reach a global verdict, then have every rank participate in the error # gather so no rank is left in a mismatched collective on the failure path. @@ -382,6 +409,8 @@ def get_dataset( shutil.rmtree(tmp, ignore_errors=True) # Every rank raises with the collected messages, so a non-root rank # never returns an unfinalized dataset path. + if interrupt is not None: + raise interrupt msgs = "; ".join(e for e in errs if e) raise RuntimeError(f"dataset generation failed: {msgs or 'unknown error'}") @@ -403,16 +432,19 @@ def get_dataset( } _write_meta_atomic(tmp / META_FILENAME, meta) tmp.rename(dest) - except (Exception, SystemExit) as e: + except BaseException as e: finalize_err = ( f"rank 0 failed to finalize dataset at {dest}: {type(e).__name__}: {e}" ) + interrupt = _reraisable(e) # This broadcast doubles as the synchronization the old Barrier provided: no # rank returns before rank 0 has published the rename (or reported that it # could not), so nobody observes the staging path or a missing dataset. finalize_err = comm.bcast(finalize_err, root=0) if finalize_err: + if interrupt is not None: + raise interrupt raise RuntimeError(f"dataset generation failed: {finalize_err}") return dest diff --git a/ScaFFold/datagen/volumegen.py b/ScaFFold/datagen/volumegen.py index 34e17f9..abb05fe 100644 --- a/ScaFFold/datagen/volumegen.py +++ b/ScaFFold/datagen/volumegen.py @@ -229,6 +229,7 @@ def main(config: Dict): # the failure is then propagated to all ranks via an allreduce. ok = True err = "" + interrupt = None try: if start_idx >= end_idx: @@ -337,11 +338,14 @@ def main(config: Dict): total_time, len(volumes_contents_subset) / total_time, ) - except (Exception, SystemExit) as e: + except BaseException as e: # Capture the failure locally instead of letting it unwind past the - # collective below, which would desynchronize the ranks. + # collective below, which would desynchronize the ranks. BaseException, + # not (Exception, SystemExit): a KeyboardInterrupt delivered to one rank + # would otherwise skip the consensus and hang the others. ok = False err = f"rank {rank}: {type(e).__name__}: {e}" + interrupt = e if isinstance(e, KeyboardInterrupt) else None # Consensus on the generation status. This replaces a bare Barrier: every # rank always executes exactly this collective (regardless of success or @@ -350,6 +354,10 @@ def main(config: Dict): all_ok = comm.allreduce(1 if ok else 0, op=MPI.MIN) == 1 errs = comm.allgather(err) if not all_ok: + # The interrupted rank re-raises the operator's abort verbatim; the + # others report the gathered failure. + if interrupt is not None: + raise interrupt msgs = "; ".join(e for e in errs if e) raise RuntimeError(f"volume generation failed: {msgs or 'unknown error'}") diff --git a/tests/datagen/test_category_search.py b/tests/datagen/test_category_search.py index 2cc43b1..1b861c3 100644 --- a/tests/datagen/test_category_search.py +++ b/tests/datagen/test_category_search.py @@ -149,7 +149,7 @@ def _seed_one_category(config: Namespace) -> None: def test_work_scan_is_root_only_and_broadcast(tmp_path, monkeypatch): """A non-root rank never scans; it consumes root's broadcast index list.""" config = _cs_config(tmp_path / "fractals") - comm = CategorySearchComm(rank=1, size=2, bcast_returns=[[0]]) + comm = CategorySearchComm(rank=1, size=2, bcast_returns=[("ok", [0])]) monkeypatch.setattr(cs, "MPI", FakeMPI(comm)) def no_scan(*_args, **_kwargs): @@ -182,17 +182,64 @@ def test_divergent_fs_views_take_the_same_collective_path(tmp_path, monkeypatch) # Rank 1: an empty directory (a divergent view), but root broadcast [0]. peer_config = _cs_config(tmp_path / "peer_view") - peer_comm = CategorySearchComm(rank=1, size=2, bcast_returns=[[0]]) + peer_comm = CategorySearchComm(rank=1, size=2, bcast_returns=[("ok", [0])]) monkeypatch.setattr(cs, "MPI", FakeMPI(peer_comm)) cs.main(peer_config) - assert root_comm.bcast_payloads[0] == [0] + assert root_comm.bcast_payloads[0] == ("ok", [0]) assert peer_comm.calls == root_comm.calls, ( "ranks with divergent filesystem views issued different collectives: " f"rank 0 {root_comm.calls} vs rank 1 {peer_comm.calls}" ) +# --------------------------------------------------------------------------- +# VB-2: the whole rank-0 scan window is fenced, not just the index parse. +# +# Rank 0 also reads the parameters of every category already on disk, right +# after the scan broadcast. A CSV that will not parse therefore killed rank 0 +# while its peers had already consumed the broadcast and moved on -- the same +# stranding the scan broadcast was introduced to prevent. +# --------------------------------------------------------------------------- + + +def test_unparseable_existing_category_is_broadcast_not_raised_on_root( + tmp_path, monkeypatch +): + """A ragged category CSV becomes a broadcast error, not a rank-0-only death.""" + config = _cs_config(tmp_path / "fractals") + param_dir = Path(layout.category_param_dir(config)) + param_dir.mkdir(parents=True, exist_ok=True) + # Six-digit name, so the scan accepts it; ragged rows, so loadtxt raises. + (param_dir / "000000.csv").write_text("0.5,0.5,0.5\n0.5,0.5\n") + + comm = CategorySearchComm(rank=0, size=2) + monkeypatch.setattr(cs, "MPI", FakeMPI(comm)) + + with pytest.raises(RuntimeError) as excinfo: + cs.main(config) + + # Rank 0 stopped at the scan broadcast, and what it broadcast is the error + # sentinel its peers need in order to abort with it. + assert comm.calls == ["Barrier", "bcast"] + assert comm.bcast_payloads[-1][0] == "error" + assert "000000.csv" in str(excinfo.value) or "3DIFS_param" in str(excinfo.value) + + +def test_peer_raises_on_broadcast_scan_error(tmp_path, monkeypatch): + """A peer receiving the scan sentinel raises instead of entering the loop.""" + config = _cs_config(tmp_path / "fractals") + sentinel = ("error", "rank 0 failed to scan existing categories: ValueError: boom") + comm = CategorySearchComm(rank=1, size=2, bcast_returns=[sentinel]) + monkeypatch.setattr(cs, "MPI", FakeMPI(comm)) + + with pytest.raises(RuntimeError, match="boom"): + cs.main(config) + + # It never reached the work loop's collectives. + assert comm.calls == ["Barrier", "bcast"] + + # --------------------------------------------------------------------------- # R31: category CSVs appear complete or not at all. # diff --git a/tests/datagen/test_mpi_consensus.py b/tests/datagen/test_mpi_consensus.py index c98fc5b..a4c954a 100644 --- a/tests/datagen/test_mpi_consensus.py +++ b/tests/datagen/test_mpi_consensus.py @@ -540,6 +540,57 @@ def test_non_root_raises_on_broadcast_finalize_error(tmp_path, monkeypatch): assert not dest.exists() +# --------------------------------------------------------------------------- +# VB-3: the consensus guards catch BaseException, not (Exception, SystemExit). +# +# ``KeyboardInterrupt`` is neither, so an interrupt delivered to rank 0 alone +# (Ctrl-C on the launching terminal, a site watchdog's SIGINT) unwound straight +# past the broadcast and hung every peer -- the failure mode the guard exists to +# prevent, arriving through the one exception class it did not cover. +# --------------------------------------------------------------------------- + + +def test_rank0_interrupt_still_posts_the_decision_sentinel(tmp_path, monkeypatch): + """A KeyboardInterrupt on rank 0 reaches the peers as an error sentinel.""" + config = _reuse_config(tmp_path / "datasets") + comm = FakeComm(rank=0, size=2) + monkeypatch.setattr(gd, "MPI", FakeMPI(comm)) + monkeypatch.setattr(gd, "_git_commit_short", lambda log: "abc123") + + def interrupted(*_args, **_kwargs): + raise KeyboardInterrupt + + monkeypatch.setattr(gd, "_decide_reuse_or_generate", interrupted) + + # Rank 0 keeps the operator's abort ... + with pytest.raises(KeyboardInterrupt): + gd.get_dataset(config) + + # ... but only after telling the peers to stop. + assert comm.calls == ["bcast"] + assert comm.bcast_payloads[0][0] == "error" + assert "KeyboardInterrupt" in comm.bcast_payloads[0][1] + + +def test_interrupt_during_generation_reaches_the_consensus(tmp_path, monkeypatch): + """An interrupt inside volumegen still drives the allreduce and allgather.""" + config = _reuse_config(tmp_path / "datasets") + comm = FakeComm(rank=0, size=2, allreduce_result=0, allgather_peers=[""]) + monkeypatch.setattr(gd, "MPI", FakeMPI(comm)) + monkeypatch.setattr(gd, "_git_commit_short", lambda log: "abc123") + + def interrupted(_config): + raise KeyboardInterrupt + + monkeypatch.setattr(volumegen, "main", interrupted) + + with pytest.raises(KeyboardInterrupt): + gd.get_dataset(config) + + assert "allreduce" in comm.calls and "allgather" in comm.calls + assert "KeyboardInterrupt" in comm.allgather_payloads[0] + + # --------------------------------------------------------------------------- # R37: orphaned staging dirs are reclaimed instead of accumulating forever. # --------------------------------------------------------------------------- From 1cce9bb10c5101c039f64d8cb9b3634d810a834d Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 16:55:36 -0700 Subject: [PATCH 36/62] Age staging dirs by their deepest recent write The orphan cleanup judged a staging directory by the newest mtime among the directory and its immediate children. Volume writing lands at depth >= 2 (volumes//N.npy), and generation is not bounded by the staleness threshold -- at the larger scales it runs for days -- so a perfectly healthy generation stopped touching anything the probe could see and read as "untouched for 48 hours". A concurrent same-config start then rmtree'd it out from under its peers, which died on FileNotFoundError. Liveness is now reported rather than inferred: every rank writing volumes refreshes /.heartbeat every five minutes, and cleanup keeps any directory whose heartbeat is younger than the threshold. The marker is removed before the staging dir is published, so a dataset does not carry it. Backing the heartbeat up (for staging dirs written before it existed, or killed before the first beat) the mtime probe now walks a bounded number of directory levels instead of one. It stats directories, whose mtimes change when entries are created in them, so it notices a writer without a stat storm over the volume files, and it stops at the first recent entry -- the live case, the one that must not be misjudged, is the cheap one. A tree too wide to walk within the bound is called live: failing to reclaim disk is recoverable, deleting a running job's dataset is not. VB-1 --- ScaFFold/datagen/get_dataset.py | 119 ++++++++++++++++++++++------ ScaFFold/datagen/volumegen.py | 58 +++++++++++++- tests/datagen/test_mpi_consensus.py | 114 ++++++++++++++++++++++++++ 3 files changed, 266 insertions(+), 25 deletions(-) diff --git a/ScaFFold/datagen/get_dataset.py b/ScaFFold/datagen/get_dataset.py index e54b29f..168f002 100644 --- a/ScaFFold/datagen/get_dataset.py +++ b/ScaFFold/datagen/get_dataset.py @@ -37,10 +37,17 @@ # reader never observes a half-written dataset. The prefix is also what the # reuse scan skips and what the orphan cleanup collects. TMP_PREFIX = ".tmp_" -# How long a staging directory must have sat untouched before it is treated as -# orphaned (left by a killed/OOM'd job) and reclaimed. See +# How long a staging directory must have shown no sign of life before it is +# treated as orphaned (left by a killed/OOM'd job) and reclaimed. See # ``_cleanup_stale_staging_dirs`` for the safety argument behind the value. STALE_STAGING_AGE_SECONDS = 24 * 60 * 60 +# Bounds on the liveness probe that backs up the heartbeat: how many directory +# levels below a staging dir it looks at, and how many directories it is willing +# to visit before it gives up and calls the tree live. Directory mtimes change +# when entries are created in them, so a couple of levels is enough to notice a +# writer without stat-ing every volume file. +_STALE_PROBE_MAX_DEPTH = 3 +_STALE_PROBE_MAX_DIRS = 10000 # Bumped from 2 to 3 when instance point clouds moved from float64 to float32: # the storage layout is unchanged, but float32 voxel binning shifts a handful of # boundary voxels, so a float64-era dataset must not be reused as if it were @@ -138,6 +145,67 @@ def _git_commit_short(log, source_dir: Path | None = None) -> str: return "no-commit-id" +def _has_recent_write(path: Path, cutoff: float) -> bool: + """Return True if the top levels of ``path`` were written after ``cutoff``. + + A directory's mtime changes whenever an entry is created in it, so the + directory mtimes near the top of a staging tree are a cheap proxy for "a + writer is active down there" -- no stat of the (possibly hundreds of + thousands of) volume files is needed. The walk therefore stats the staging + directory, its immediate children, and directories down to + ``_STALE_PROBE_MAX_DEPTH``, and stops at the first recent entry, which makes + the live case (the one that must not be misjudged) the cheap one. + + A tree wide enough to exceed ``_STALE_PROBE_MAX_DIRS`` is reported as recent + rather than walked further: failing to reclaim disk is recoverable, deleting + a running job's dataset is not. + """ + stack = [(path, 0)] + dirs_seen = 0 + while stack: + current, depth = stack.pop() + if current.stat().st_mtime > cutoff: + return True + if depth >= _STALE_PROBE_MAX_DEPTH: + continue + with os.scandir(current) as entries: + for entry in entries: + if entry.is_dir(follow_symlinks=False): + dirs_seen += 1 + if dirs_seen > _STALE_PROBE_MAX_DIRS: + return True + stack.append((Path(entry.path), depth + 1)) + elif depth == 0 and entry.stat().st_mtime > cutoff: + # Files directly in the staging dir (the heartbeat, + # volumes_contents.csv, meta.yaml) are few and cheap. + return True + return False + + +def _staging_dir_is_live(path: Path, cutoff: float) -> bool: + """Return True if ``path`` shows any sign of a generation still running. + + Two signals, in order of authority: + + 1. ``/.heartbeat``, refreshed by every writing rank every few + minutes for as long as volumes are being written (see + ``volumegen.StagingHeartbeat``). This is the reliable one, because it + does not depend on where in the tree the writers currently are. + 2. a bounded-depth mtime probe, which covers staging directories written + before the heartbeat existed, or killed before the first beat. + + Raises ``OSError`` if the directory cannot be examined; the caller treats + that as "cannot tell" and leaves the directory alone. + """ + heartbeat = path / volumegen.STAGING_HEARTBEAT_NAME + try: + if heartbeat.stat().st_mtime > cutoff: + return True + except OSError: + pass # No heartbeat: fall back to the mtime probe. + return _has_recent_write(path, cutoff) + + def _cleanup_stale_staging_dirs( base: Path, log, max_age: float = STALE_STAGING_AGE_SECONDS ) -> None: @@ -149,39 +217,39 @@ def _cleanup_stale_staging_dirs( Safety policy. Only directories that (a) live directly under *this* config_id base, (b) carry the ``.tmp_`` prefix this module owns, and (c) - have been untouched for ``max_age`` are removed. The age gate is what keeps - a *concurrent* job's staging directory safe: unique staging names mean two - live jobs never share a directory, but they do share the base, so a live - peer's directory is visible here -- it is simply orders of magnitude younger - than the threshold (a day, against generations measured in minutes to - hours). Published datasets and anything outside ``base`` are never touched. - Failures are logged and ignored: cleanup is opportunistic and must never - break the decision it runs inside. + show no sign of life for ``max_age`` are removed. Published datasets and + anything outside ``base`` are never touched. Failures are logged and + ignored: cleanup is opportunistic and must never break the decision it runs + inside. + + "No sign of life" is the delicate part, because the staging directory of a + *running* job is visible here (unique staging names mean two live jobs never + share a directory, but they do share the base). Age alone is not enough: + generation is not bounded by a day -- at the larger scales it is measured in + days -- and after the first minutes it writes only at depth >= 2, so the top + of the tree stops changing while the job is perfectly healthy. Judging by + the top-level mtimes alone therefore let a concurrent same-config start + rmtree a live generation out from under its peers. ``_staging_dir_is_live`` + is the answer: an explicit heartbeat maintained by the writers, backed by a + bounded-depth mtime probe for directories that predate it. """ now = time.time() + cutoff = now - max_age for path in base.iterdir(): if not path.name.startswith(TMP_PREFIX) or not path.is_dir(): continue try: - # Newest mtime among the staging dir and its immediate children: a - # bounded, cheap probe (no recursive stat storm over a partially - # generated dataset) that still notices a job that has started - # laying down its split directories. - newest = path.stat().st_mtime - for child in path.iterdir(): - newest = max(newest, child.stat().st_mtime) + if _staging_dir_is_live(path, cutoff): + continue except OSError as exc: - log.warning("Could not stat staging dir %s: %s", path, exc) - continue - - age = now - newest - if age < max_age: + log.warning("Could not examine staging dir %s: %s", path, exc) continue log.info( - "Removing orphaned dataset staging dir %s (untouched for %.1f hours)", + "Removing orphaned dataset staging dir %s (no write in the last " + "%.1f hours, and no live generation heartbeat)", path, - age / 3600.0, + max_age / 3600.0, ) shutil.rmtree(path, ignore_errors=True) @@ -430,6 +498,9 @@ def get_dataset( "code_commit": commit, "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), } + # The liveness marker described this directory while it was being + # written; it has no meaning in a published dataset. + (tmp / volumegen.STAGING_HEARTBEAT_NAME).unlink(missing_ok=True) _write_meta_atomic(tmp / META_FILENAME, meta) tmp.rename(dest) except BaseException as e: diff --git a/ScaFFold/datagen/volumegen.py b/ScaFFold/datagen/volumegen.py index abb05fe..997a73f 100644 --- a/ScaFFold/datagen/volumegen.py +++ b/ScaFFold/datagen/volumegen.py @@ -27,6 +27,55 @@ from ScaFFold.utils.data_types import MASK_DTYPE, VOLUME_DTYPE from ScaFFold.utils.utils import setup_mpi_logger +# Liveness marker for the directory being generated into. Volume writing is the +# long phase of a generation and it happens deep inside the tree +# (``volumes//N.npy``), so the top of the staging directory can look +# untouched for many hours while the job is perfectly healthy. Every writing +# rank therefore refreshes this file periodically, and +# ``get_dataset._staging_dir_is_live`` reads it instead of trying to infer +# liveness from mtimes it cannot cheaply see. The name is owned here, next to +# the writer; ``get_dataset`` (which already imports this module) reads it from +# here so the two sides cannot drift. +STAGING_HEARTBEAT_NAME = ".heartbeat" +# Refresh interval. Small enough to be negligible against the staleness +# threshold (a day), large enough that it is one utime per rank per few minutes +# no matter how fast volumes are written. +STAGING_HEARTBEAT_INTERVAL_SECONDS = 5 * 60 + + +class StagingHeartbeat: + """Periodically touch a staging directory's heartbeat file. + + ``beat()`` is called from the volume loop and is a no-op until the interval + has elapsed, so it costs one comparison per volume. Every rank writing into + the directory beats the same file: the marker means "somebody is still + working here", and a last-writer-wins utime is exactly the semantics wanted. + Failures are swallowed -- a heartbeat that cannot be written must never take + down a generation that is otherwise fine (the cleanup's second signal, the + bounded mtime probe, still applies). + """ + + def __init__( + self, staging_dir, interval: float = STAGING_HEARTBEAT_INTERVAL_SECONDS + ) -> None: + self.path = os.path.join(str(staging_dir), STAGING_HEARTBEAT_NAME) + self.interval = interval + self._last_beat = float("-inf") + + def beat(self, now: float | None = None) -> bool: + """Touch the marker if the interval has elapsed; return whether it did.""" + now = time.time() if now is None else now + if now - self._last_beat < self.interval: + return False + self._last_beat = now + try: + with open(self.path, "a"): + pass + os.utime(self.path, (now, now)) + except OSError: + return False + return True + def load_np_ptcloud(path: str) -> np.ndarray: """ @@ -253,9 +302,16 @@ def main(config: Dict): # this run's seed produced. Resolved once, outside the loop. instances_dir = layout.instance_dir(config) - # Generation loop + # Generation loop. Every rank reports that this staging directory + # is still being written to, so a concurrent job's orphan cleanup + # can tell a live multi-hour generation from one killed a day ago + # (the volumes themselves land two levels down, where a cheap + # top-level mtime probe cannot see them). + heartbeat = StagingHeartbeat(dataset_dir) + heartbeat.beat() start_time = time.time() for i, curr_vol in enumerate(volumes_contents_subset): + heartbeat.beat() if i % 10 == 0: log.debug("Rank %s processing local volume %s", rank, i) diff --git a/tests/datagen/test_mpi_consensus.py b/tests/datagen/test_mpi_consensus.py index a4c954a..5185eb3 100644 --- a/tests/datagen/test_mpi_consensus.py +++ b/tests/datagen/test_mpi_consensus.py @@ -657,6 +657,120 @@ def test_cleanup_never_touches_published_datasets(tmp_path): assert other_orphan.exists(), "cleanup escaped this job's config_id base" +# --------------------------------------------------------------------------- +# VB-1: a staging dir is aged by its DEEPEST recent write, not its top level. +# +# Generation is not bounded by the staleness threshold -- at the larger scales +# it runs for days -- and after the first minutes it writes only at depth >= 2 +# (``volumes//N.npy``). Judging liveness from the staging dir and its +# immediate children alone therefore reported a healthy multi-day generation as +# "untouched for 48 hours", and a concurrent same-config start rmtree'd it out +# from under its peers (which then died on FileNotFoundError). +# --------------------------------------------------------------------------- + + +def _cleanup(base: Path, name: str) -> None: + gd._cleanup_stale_staging_dirs(base, logging.getLogger(name)) + + +def test_live_generation_survives_a_deep_write(tmp_path): + """A >24h-old staging dir with a fresh deep write is NOT reclaimed.""" + base = tmp_path / "cid" + base.mkdir(parents=True) + + live = base / f"{gd.TMP_PREFIX}20260101-000000_222_cafebabe" + split = live / "volumes" / "training" + split.mkdir(parents=True) + _age_tree(live, 2 * gd.STALE_STAGING_AGE_SECONDS) + + # The job is alive and still laying down volumes: the write lands two levels + # down, so only that directory's mtime is current. + (split / "0.npy").write_bytes(b"payload from a running job") + + _cleanup(base, "test_live_generation_survives_a_deep_write") + + assert live.exists(), "a live generation's staging dir was reclaimed" + + +def test_live_generation_survives_on_its_heartbeat_alone(tmp_path): + """A fresh heartbeat keeps a staging dir whose whole tree looks ancient. + + The mtime probe is bounded, so a generation writing deeper than it looks + (or on a filesystem with coarse directory mtimes) still has to be safe. The + heartbeat is the signal that does not depend on the shape of the tree. + """ + base = tmp_path / "cid" + base.mkdir(parents=True) + + live = base / f"{gd.TMP_PREFIX}20260101-000000_333_f00d" + (live / "volumes" / "training" / "deep" / "deeper").mkdir(parents=True) + heartbeat = live / volumegen.STAGING_HEARTBEAT_NAME + heartbeat.write_text("") + _age_tree(live, 2 * gd.STALE_STAGING_AGE_SECONDS) + now = time.time() + os.utime(heartbeat, (now, now)) + + _cleanup(base, "test_live_generation_survives_on_its_heartbeat_alone") + + assert live.exists(), "a heartbeating generation's staging dir was reclaimed" + + +def test_dead_staging_dir_with_stale_heartbeat_is_reclaimed(tmp_path): + """The heartbeat must not turn cleanup into a no-op (R37 still holds).""" + base = tmp_path / "cid" + base.mkdir(parents=True) + + orphan = base / f"{gd.TMP_PREFIX}20200101-000000_111_deadbeef" + split = orphan / "volumes" / "training" + split.mkdir(parents=True) + (split / "0.npy").write_bytes(b"stale payload") + (orphan / volumegen.STAGING_HEARTBEAT_NAME).write_text("") + _age_tree(orphan, 2 * gd.STALE_STAGING_AGE_SECONDS) + + _cleanup(base, "test_dead_staging_dir_with_stale_heartbeat_is_reclaimed") + + assert not orphan.exists(), "an orphaned staging dir was not reclaimed" + + +def test_generation_writes_and_then_drops_the_heartbeat(tmp_path, monkeypatch): + """volumegen marks the staging dir live; publishing removes the marker.""" + config = _reuse_config(tmp_path / "datasets") + comm = FakeComm(rank=0, size=1, allreduce_result=1) + monkeypatch.setattr(gd, "MPI", FakeMPI(comm)) + monkeypatch.setattr(gd, "_git_commit_short", lambda log: "abc123") + + beating = {} + + def fake_volumegen(cfg): + # Stand in for the write loop: report the staging dir as live. + volumegen.StagingHeartbeat(cfg.dataset_dir).beat() + beating["path"] = Path(cfg.dataset_dir) / volumegen.STAGING_HEARTBEAT_NAME + beating["existed_during_generation"] = beating["path"].exists() + + monkeypatch.setattr(volumegen, "main", fake_volumegen) + + published = Path(gd.get_dataset(config)) + + assert beating["existed_during_generation"], "no heartbeat during generation" + assert not (published / volumegen.STAGING_HEARTBEAT_NAME).exists(), ( + "the staging heartbeat was published with the dataset" + ) + + +def test_heartbeat_respects_its_interval(tmp_path): + """``beat`` is a no-op until the interval elapses, then refreshes.""" + staging = tmp_path / "staging" + staging.mkdir() + heartbeat = volumegen.StagingHeartbeat(staging, interval=60) + + assert heartbeat.beat(now=1000.0) is True + first = Path(heartbeat.path).stat().st_mtime + assert heartbeat.beat(now=1030.0) is False # inside the interval + assert Path(heartbeat.path).stat().st_mtime == first + assert heartbeat.beat(now=1090.0) is True + assert Path(heartbeat.path).stat().st_mtime > first + + # --------------------------------------------------------------------------- # R28: meta.yaml is published atomically, so no reader ever sees a partial one. # --------------------------------------------------------------------------- From 5cfa82f977257b7f433843707521626f28755edb Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 16:57:19 -0700 Subject: [PATCH 37/62] Sweep category-search temp files and warn on old-layout libraries Both atomic writers in the category search name their temp file after the writing pid and unlink it when the write raises -- but a SIGKILL (walltime, OOM, node failure) skips that Python-level cleanup. Nothing ever looked at the strays again, so .NNNNNN.csv.tmp and .rng_attempt_rank*.tmp piled up one per killed process. Rank 0 now sweeps them at startup, best-effort, the way instance.py already sweeps its equivalents. The seed-keyed relayout is silent about libraries in the old location: the existence check simply does not find them. From outside that looks like a library that was there yesterday being regenerated for no reason -- hours of work at scale. One warning naming the old directory and the new one explains it. VB-4, VB-6 --- ScaFFold/datagen/category_search.py | 34 +++++++++++++ ScaFFold/datagen/layout.py | 35 +++++++++++++ tests/datagen/test_category_search.py | 72 +++++++++++++++++++++++++++ 3 files changed, 141 insertions(+) diff --git a/ScaFFold/datagen/category_search.py b/ScaFFold/datagen/category_search.py index f3b3496..3cc1f81 100644 --- a/ScaFFold/datagen/category_search.py +++ b/ScaFFold/datagen/category_search.py @@ -403,6 +403,35 @@ def write_attempt_counter(fracts_write_dir: str, rank: int, attempt_index: int) os.replace(tmp, path) +def _sweep_stale_temp_files(fracts_write_dir: str, log) -> None: + """Remove temp files stranded by killed writes in the category directory. + + Both atomic writers here (``_savetxt_atomic`` and ``write_attempt_counter``) + unlink their temp file when the write raises, but a SIGKILL -- walltime, an + OOM, a node failure -- skips that Python-level cleanup and strands it. The + names carry the writer's pid, so they accumulate one per killed process and + nothing else ever removes them; ``instance.py`` sweeps its equivalents for + exactly this reason. + + Called on rank 0 before any rank has written anything this run, and + best-effort: this is housekeeping, and it runs just before a Barrier the + peers are heading into, so it must not raise. + """ + patterns = ( + # .NNNNNN.csv.tmp -- a partially written category CSV. + f"{fracts_write_dir}/.*.csv.tmp*", + # .rng_attempt_rank.tmp -- a partially written attempt counter. + f"{fracts_write_dir}/.rng_attempt_rank*.tmp*", + ) + for pattern in patterns: + for stale in glob.glob(pattern): + try: + os.remove(stale) + log.info("Removed stale category-search temp file %s", stale) + except OSError as exc: + log.warning("Could not remove stale temp file %s: %s", stale, exc) + + def main(config: Config) -> None: """ Generate fractal categories. @@ -440,10 +469,15 @@ def main(config: Config) -> None: fracts_write_dir = layout.category_param_dir(config) if rank == 0: log.info("Writing fractals to %s", fracts_write_dir) + # A library in the pre-seed layout is invisible to everything below, so + # say why it is being ignored rather than appearing to regenerate work + # that is plainly still on disk. + layout.warn_if_legacy_library(config, log) if os.path.exists(fracts_write_dir) and config.datagen_from_scratch: log.info("Removing existing fractals directory") shutil.rmtree(fracts_write_dir) os.makedirs(fracts_write_dir, exist_ok=True) + _sweep_stale_temp_files(fracts_write_dir, log) # Wait until dir setup completes comm.Barrier() diff --git a/ScaFFold/datagen/layout.py b/ScaFFold/datagen/layout.py index e558c6d..eff02f8 100644 --- a/ScaFFold/datagen/layout.py +++ b/ScaFFold/datagen/layout.py @@ -57,6 +57,41 @@ def category_param_dir(config) -> str: return os.path.join(library_root(config), "3DIFS_param") +def legacy_category_param_dir(config) -> str: + """Return where the category CSVs lived before the layout was seed-keyed.""" + return os.path.join( + str(config.fract_base_dir), + f"var{config.variance_threshold}", + "3DIFS_param", + ) + + +def warn_if_legacy_library(config, log) -> bool: + """Warn when a library in the old, seed-agnostic layout is being ignored. + + The relayout is deliberately silent about old data -- an existence check in + the seed-keyed location simply does not find it -- which from the outside + looks like a library that was there yesterday being regenerated for no + reason (at large scales, hours of work). One line naming both directories + turns that into an explained, expected event. Returns whether the old + layout was present, so callers can test the condition directly. + """ + legacy = legacy_category_param_dir(config) + if not os.path.isdir(legacy): + return False + log.warning( + "Found a fractal library in the old, seed-agnostic layout at %s. " + "Libraries are now keyed by seed, so this one cannot be reused (a run " + "under a different seed would silently adopt another seed's data) and " + "the categories for seed %s will be generated at %s. Delete the old " + "directory once you no longer need it.", + legacy, + int(config.seed), + category_param_dir(config), + ) + return True + + def instance_dir(config) -> str: """Return the directory holding this seed's instance point clouds. diff --git a/tests/datagen/test_category_search.py b/tests/datagen/test_category_search.py index 1b861c3..3ff19a2 100644 --- a/tests/datagen/test_category_search.py +++ b/tests/datagen/test_category_search.py @@ -193,6 +193,78 @@ def test_divergent_fs_views_take_the_same_collective_path(tmp_path, monkeypatch) ) +# --------------------------------------------------------------------------- +# VB-4: temp files stranded by killed writes are swept, as in instance.py. +# +# Both atomic writers name their temp file after the writing pid, so a job +# killed mid-write leaves one behind per killed process, forever: nothing in +# this module ever looked at them again. +# --------------------------------------------------------------------------- + + +def test_stale_temp_files_are_swept(tmp_path, monkeypatch): + """Category and attempt-counter temps from dead pids are removed.""" + config = _cs_config(tmp_path / "fractals") + param_dir = Path(layout.category_param_dir(config)) + param_dir.mkdir(parents=True, exist_ok=True) + _seed_one_category(config) # n_categories=1, so the search has no work + + stale_csv = param_dir / ".000001.csv.tmp999999" + stale_csv.write_text("0.5,0.5\n") + stale_counter = param_dir / ".rng_attempt_rank3.tmp999999" + stale_counter.write_text("17") + + comm = CategorySearchComm(rank=0, size=1) + monkeypatch.setattr(cs, "MPI", FakeMPI(comm)) + + cs.main(config) + + assert not stale_csv.exists(), "a stranded category temp file was kept" + assert not stale_counter.exists(), "a stranded attempt-counter temp was kept" + # The real artifact is untouched: only the temp names are swept. + assert (param_dir / "000000.csv").exists() + + +# --------------------------------------------------------------------------- +# VB-6: a library in the old, seed-agnostic layout is explained, not ignored. +# --------------------------------------------------------------------------- + + +def test_old_layout_library_is_reported(tmp_path, monkeypatch, caplog): + """A pre-relayout library produces one warning naming both directories.""" + config = _cs_config(tmp_path / "fractals") + legacy = Path(layout.legacy_category_param_dir(config)) + legacy.mkdir(parents=True) + (legacy / "000000.csv").write_text("") + _seed_one_category(config) # nothing to generate under the new layout + + comm = CategorySearchComm(rank=0, size=1) + monkeypatch.setattr(cs, "MPI", FakeMPI(comm)) + + with caplog.at_level("WARNING"): + cs.main(config) + + messages = " ".join(record.getMessage() for record in caplog.records) + assert str(legacy) in messages + assert layout.category_param_dir(config) in messages + + +def test_no_warning_without_an_old_layout(tmp_path, monkeypatch, caplog): + """The warning does not fire for a fresh library (control).""" + config = _cs_config(tmp_path / "fractals") + _seed_one_category(config) + + comm = CategorySearchComm(rank=0, size=1) + monkeypatch.setattr(cs, "MPI", FakeMPI(comm)) + + with caplog.at_level("WARNING"): + cs.main(config) + + assert "seed-agnostic" not in " ".join( + record.getMessage() for record in caplog.records + ) + + # --------------------------------------------------------------------------- # VB-2: the whole rank-0 scan window is fenced, not just the index parse. # From fcdcee40ccfa72d48e47064830a6baa832425fc1 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 17:00:43 -0700 Subject: [PATCH 38/62] Resolve restart state before the pre-check restart and run_dir can come from the config file as well as the command line -- the generated restart.sh replays a dumped config.yaml, which carries both -- and an absent --restart cannot outrank a file that sets it, since an unset store_true flag is indistinguishable from its default. resolve_run_dir read only the command line while the pre-check read the merged config, so the two disagreed in both directions. A config-file `restart: true` with no run directory therefore created a fresh timestamped directory, wrote its config dumps and restart script into it, and only then died for want of a run dir; it now fails before anything is claimed. And a run_dir inherited from a reused config.yaml stayed in the config while a fresh directory was created to train in, so the pre-check passed on another run's checkpoints and the job died hours later with its dataset generated. resolve_run_dir now reads the merged values, and writes its answer back to benchmark_run_dir, restart and run_dir (cleared for a fresh run), so nothing downstream can re-derive a different one; missing_checkpoint_error keys off the resolved benchmark run dir. The pre-check is also gated on the benchmark subcommand, which is the only one that has run directories at all. VC-1 --- ScaFFold/cli.py | 73 +++++++++++++++++++++++++++++++++------------ tests/test_cli.py | 75 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 129 insertions(+), 19 deletions(-) diff --git a/ScaFFold/cli.py b/ScaFFold/cli.py index 623f868..6a3923b 100644 --- a/ScaFFold/cli.py +++ b/ScaFFold/cli.py @@ -121,8 +121,22 @@ def missing_checkpoint_error(combined_config): Reports the first problem found rather than raising, so the caller can make this a rank-0 decision and broadcast the verdict instead of letting every rank stat the shared filesystem and possibly disagree. + + The directory checked is the *resolved* benchmark run dir -- the one this + launch will actually train in -- and not the raw ``run_dir`` key. The two + can differ: a config.yaml dumped by a restarted run carries that run's + ``run_dir``, so reusing it as a base config had this check stat another + run's checkpoints, pass, and let the job die hours later with its dataset + already generated. ``resolve_run_dir`` keeps the two in agreement; this + prefers the resolved value so they cannot drift apart again. """ - checkpoint_dir = Path(combined_config["run_dir"]) / combined_config.get( + run_dir = combined_config.get("benchmark_run_dir") or combined_config.get("run_dir") + if not run_dir: + return ( + "Restart requested but no run directory was resolved. Pass " + "'--run-dir ' (or set run_dir in the config file)." + ) + checkpoint_dir = Path(run_dir) / combined_config.get( "checkpoint_dir", "checkpoints" ) expected_checkpoints = ( @@ -140,20 +154,34 @@ def resolve_run_dir(args_dict, combined_config): The semantics are fixed and unambiguous: - * ``--run-dir DIR`` (with or without ``--restart``): resume in that exact - directory. ``train_from_scratch`` is forced off and ``restart`` on so the - downstream benchmark driver takes its restart path. - * ``--restart`` without ``--run-dir``: rejected with a clear error. The run - directory to resume must be named explicitly; the most recent directory - is never guessed. - * neither flag: create a fresh timestamped directory under ``base_run_dir``, + * a run directory (``--run-dir DIR``, or ``run_dir`` in the config file), + with or without a restart flag: resume in that exact directory. + ``train_from_scratch`` is forced off and ``restart`` on so the downstream + benchmark driver takes its restart path. + * a restart requested with no run directory anywhere: rejected with a clear + error, before any directory is created. The directory to resume must be + named explicitly; the most recent one is never guessed. + * neither: create a fresh timestamped directory under ``base_run_dir``, retrying with a numeric suffix on a same-second name collision. - ``combined_config['benchmark_run_dir']`` is set in every path so the driver - can always read it. Returns ``(benchmark_run_dir: Path, restarting: bool)``. + Both keys are read from the *merged* config, not from the command line + alone. A config file is a first-class source for them -- the generated + ``restart.sh`` replays a dumped ``config.yaml``, which carries both -- and + an absent ``--restart`` cannot outrank a file that sets it, because an + unset ``store_true`` flag is indistinguishable from its default. Reading + only the command line here let the two disagree: a config-file restart + created a *fresh* run directory and only then failed for want of a run dir, + and a stale ``run_dir`` inherited from a reused ``config.yaml`` was left in + the config for the restart pre-check to stat while training happened + somewhere else entirely. + + The resolved answer is written back to ``benchmark_run_dir``, ``restart`` + and ``run_dir``, so every later reader -- the pre-check, the dumped + ``config.yaml``, the benchmark driver -- sees exactly what was decided here. + Returns ``(benchmark_run_dir: Path, restarting: bool)``. """ - restart_flag = bool(args_dict.get("restart")) - run_dir_arg = args_dict.get("run_dir") + restart_flag = bool(combined_config.get("restart") or args_dict.get("restart")) + run_dir_arg = combined_config.get("run_dir") or args_dict.get("run_dir") if run_dir_arg is not None: benchmark_run_dir = Path(run_dir_arg) @@ -163,9 +191,10 @@ def resolve_run_dir(args_dict, combined_config): restarting = True elif restart_flag: raise ValueError( - "--restart requires --run-dir: pass the directory of the run to " - "resume (e.g. '--restart --run-dir '). The most recent run " - "directory is not resolved automatically." + "A restart was requested (--restart, or restart: true in the " + "config file) but no run directory was given: pass the directory " + "of the run to resume (e.g. '--restart --run-dir '). The " + "most recent run directory is not resolved automatically." ) else: base_run_dir = Path(combined_config["base_run_dir"]) @@ -179,10 +208,15 @@ def resolve_run_dir(args_dict, combined_config): ) restarting = False + # Write the resolution back, so nothing downstream can re-derive a + # different answer from the raw inputs. ``run_dir`` is cleared for a fresh + # run: left set, a value inherited from a reused config.yaml names a + # directory this run has nothing to do with. combined_config["benchmark_run_dir"] = str(benchmark_run_dir) + combined_config["restart"] = restarting + combined_config["run_dir"] = str(benchmark_run_dir) if restarting else None if restarting: combined_config["train_from_scratch"] = False - combined_config["restart"] = True return benchmark_run_dir, restarting @@ -528,10 +562,11 @@ def main(): # caches). A rank that decided for itself would either abort alone -- # stranding its peers in benchmark.py's timeout-less barrier -- or keep # running after rank 0 had already aborted. + # Only the benchmark subcommand has run directories or checkpoints, so it + # is the only one this applies to: fractal generation reading a benchmark's + # config.yaml must not be judged on that run's restart state. restart_precheck_error = None - if combined_config.get("restart", False): - if not combined_config.get("run_dir"): - raise ValueError("--restart requires --run-dir") + if args.command == "benchmark" and combined_config.get("restart", False): if rank == 0: restart_precheck_error = missing_checkpoint_error(combined_config) restart_precheck_error = comm.bcast(restart_precheck_error, root=0) diff --git a/tests/test_cli.py b/tests/test_cli.py index 381744f..1db35af 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -314,6 +314,81 @@ def test_restart_precheck_failure_raises_on_every_rank(monkeypatch, tmp_path): run_cli(monkeypatch, _restart_argv(cfg, run_dir), comm=comm) +# --------------------------------------------------------------------------- +# VC-1: the restart state is resolved once, before anything acts on it. +# +# ``restart`` and ``run_dir`` can come from the config file as well as the +# command line (the generated restart.sh replays a dumped config.yaml, which +# carries both), and an absent ``--restart`` cannot outrank a file that sets +# it. Resolving the run directory from the command line alone while the +# pre-check read the merged config made the two disagree. +# --------------------------------------------------------------------------- + + +def test_yaml_restart_without_run_dir_fails_before_creating_a_run_dir( + monkeypatch, tmp_path +): + """``restart: true`` with no run dir aborts without claiming a directory. + + The run directory was resolved from the command line, which said nothing + about a restart, so a fresh timestamped directory was created and populated + -- and only then did the merged config's ``restart`` trip the pre-check. + """ + cfg = write_config(tmp_path, {"restart": True}) + + with pytest.raises(ValueError, match="run directory"): + run_cli(monkeypatch, ["scaffold", "benchmark", "-c", str(cfg)]) + + assert not (tmp_path / "runs").exists(), "a run dir was created before the abort" + + +def test_reused_restart_config_resolves_to_one_run_dir(monkeypatch, tmp_path): + """A config.yaml from a restarted run cannot split the run across two dirs. + + Reusing such a file as a base config left ``run_dir`` pointing at the run + it was dumped by while a *fresh* directory was created to train in. The + pre-check then passed on the old run's checkpoints, and the job died hours + later with its dataset already generated. Whatever the file resolves to, + the directory the pre-check judges and the directory the run uses must be + the same one. + """ + cfg = write_config(tmp_path) + _, first = run_cli(monkeypatch, ["scaffold", "benchmark", "-c", str(cfg)]) + first_run = Path(first["benchmark"][0]["benchmark_run_dir"]) + _make_checkpoint(first_run) + + # Restart it once, so its config.yaml records the restart state. + run_cli(monkeypatch, _restart_argv(first_run / "config.yaml", first_run)) + runs_before = sorted(p.name for p in (tmp_path / "runs").iterdir()) + + # Now reuse that config.yaml as a plain base config, with no flags at all. + _, reused = run_cli( + monkeypatch, ["scaffold", "benchmark", "-c", str(first_run / "config.yaml")] + ) + + (config,) = reused["benchmark"] + assert config["run_dir"] == config["benchmark_run_dir"], ( + "the pre-check's run_dir and the run's directory disagree" + ) + assert Path(config["benchmark_run_dir"]) == first_run + assert sorted(p.name for p in (tmp_path / "runs").iterdir()) == runs_before + + +def test_fresh_run_records_no_restart_state(monkeypatch, tmp_path): + """A fresh run's config.yaml carries no restart state to inherit.""" + cfg = write_config(tmp_path) + + _, calls = run_cli(monkeypatch, ["scaffold", "benchmark", "-c", str(cfg)]) + + (config,) = calls["benchmark"] + assert config["restart"] is False + assert config["run_dir"] is None + dumped = yaml.safe_load( + (Path(config["benchmark_run_dir"]) / "config.yaml").read_text() + ) + assert dumped["restart"] is False and dumped["run_dir"] is None + + # --------------------------------------------------------------------------- # R19: generate_fractals is not a benchmark run # --------------------------------------------------------------------------- From e338e6ae617fd93c2f112f30eecfbfe2c5ff74e7 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 17:03:28 -0700 Subject: [PATCH 39/62] Broadcast config validation outcomes Rank 0 builds the whole job's config while every peer waits in a barrier, so anything raising in there stranded them: the recomputed bottleneck check and the n_categories check added with R25, the Config() validation that has always been there, and now the run-dir resolution. The user saw a hang where an error message belonged. The rank-0 block moves into build_run_config and runs inside a guard whose outcome is broadcast, exactly like the restart pre-check below it. Rank 0 re-raises the original exception, keeping its traceback; the peers rebuild it from the type name and message that crossed the wire, so a builtin type comes back as itself and anything else degrades to a RuntimeError naming the original. The exception object itself is deliberately not pickled across: a failure to unpickle on the receiving side would turn the error being reported back into the hang it was reported to avoid. VC-2 --- ScaFFold/cli.py | 262 ++++++++++++++++++++++++++++------------------ tests/test_cli.py | 66 +++++++++++- 2 files changed, 222 insertions(+), 106 deletions(-) diff --git a/ScaFFold/cli.py b/ScaFFold/cli.py index 6a3923b..88bf1f9 100644 --- a/ScaFFold/cli.py +++ b/ScaFFold/cli.py @@ -13,6 +13,7 @@ # SPDX-License-Identifier: (Apache-2.0) import argparse +import builtins import socket import sys from datetime import datetime @@ -220,6 +221,142 @@ def resolve_run_dir(args_dict, combined_config): return benchmark_run_dir, restarting +def rebuild_error(type_name, message): + """Rebuild rank 0's configuration error on a peer that never saw it. + + Only the type name and message cross the wire: an arbitrary exception + object may not survive a pickle round trip, and a failure to unpickle on + the receiving side would turn the error being reported back into the hang + it was reported to avoid. Builtin exception types are rebuilt as + themselves, so a caller's ``except ValueError`` still catches what rank 0 + raised; anything else degrades to ``RuntimeError`` naming the original + type. + """ + cls = getattr(builtins, type_name, None) + if isinstance(cls, type) and issubclass(cls, Exception): + return cls(message) + return RuntimeError(f"{type_name}: {message}") + + +def build_run_config(args, parsers, log, world_size): + """Build the job-wide config, and lay down the benchmark's run directory. + + Rank-0-only work: it reads and merges the config files, applies the + command-line overrides, validates the result, and -- for the benchmark + subcommand -- resolves the run directory, dumps the configs into it and + writes its restart script. Returns the merged config the caller broadcasts. + + Raises whatever the validation or the filesystem raises; the caller runs + this inside a guard and broadcasts the outcome, since every peer is already + waiting in the barrier that follows. + """ + log.debug("args = %s", args) + + # --config may be a single path (generate_fractals) or a list of + # paths (benchmark, action="append"): base config plus overrides. + config_paths = args.config if isinstance(args.config, list) else [args.config] + merged_dict = config_utils.load_config_files(config_paths) + # Validate the merged result and derive dependent settings. Every run + # parameter must be single-valued; a list is rejected here by name. + bench_config = config_utils.Config(merged_dict) + bench_config_dict = vars(bench_config) + cli_args = vars(args) + # Downstream consumers expect a single config path (e.g. to copy it + # into the run dir); keep the base config there. + cli_args["config"] = config_paths[0] + + # Combine configs, in increasing order of precedence: + # argparse default < config file < explicit command-line flag. + combined_config = bench_config_dict.copy() + # Config only keeps the keys it consumes; the auxiliary keys it accepts + # (verbose, datagen_batch_size, ...) never become attributes, so put + # the file's values back first. Without this they are absent below and + # the argparse default overwrites what the user wrote in the config. + for key, value in merged_dict.items(): + combined_config.setdefault(key, value) + + explicit_cli = explicit_cli_keys(args, parsers) + for key, value in cli_args.items(): + if key == "command": + continue + if key not in combined_config: + combined_config[key] = value + elif key in explicit_cli and value is not None: + log.info( + "Overriding '%s=%s' with '%s=%s'", + key, + combined_config[key], + key, + value, + ) + combined_config[key] = value + # The subcommand is always owned by the command line. + combined_config["command"] = cli_args["command"] + + # Recalculate unet_layers to capture any CLI overrides. The overridden + # pair has to be re-validated: Config only saw the config-file values. + config_utils.validate_unet_dims( + combined_config["problem_scale"], combined_config["unet_bottleneck_dim"] + ) + combined_config["unet_layers"] = ( + combined_config["problem_scale"] - combined_config["unet_bottleneck_dim"] + ) + config_utils.require_positive_int("n_categories", combined_config["n_categories"]) + + # Resolve paths to absolute, matching Config() behavior + if "base_run_dir" in combined_config and combined_config["base_run_dir"]: + combined_config["base_run_dir"] = str( + Path(combined_config["base_run_dir"]).resolve() + ) + + if "dataset_dir" in combined_config and combined_config["dataset_dir"]: + combined_config["dataset_dir"] = str( + Path(combined_config["dataset_dir"]).resolve() + ) + + if "fract_base_dir" in combined_config and combined_config["fract_base_dir"]: + combined_config["fract_base_dir"] = str( + Path(combined_config["fract_base_dir"]).resolve() + ) + + # Calculate these variables after override + combined_config["vol_size"] = pow(2, combined_config["problem_scale"]) + combined_config["point_num"] = int(combined_config["vol_size"] ** 3 / 256) + + # The run directory, its config dumps and its restart script belong to + # the benchmark subcommand alone. Fractal generation writes nothing + # there, and the restart script it used to get replayed + # `generate_fractals --restart --run-dir ...` -- flags that subparser + # rejects, so the script could only ever exit 2. + if args.command == "benchmark": + # Resolve the run directory and whether this launch resumes a run. + # This sets combined_config["benchmark_run_dir"] on every path, and + # writes the resolved restart/run_dir back into the config. + benchmark_run_dir, restarting = resolve_run_dir(cli_args, combined_config) + if restarting: + log.info("Resuming in existing directory: %s", benchmark_run_dir) + + # Add scheduler metadata and machine name to config.yaml + combined_config["scheduler_metadata"] = collect_scheduler_metadata() + combined_config["machine_name"] = socket.gethostname() + + # Dump configs (Overwrite is okay/desired on restart to capture new job IDs) + overrides = { + k: v for k, v in cli_args.items() if v is not None and k != "command" + } + with open(benchmark_run_dir / "overrides.yaml", "w") as file: + yaml.dump(overrides, file) + with open(benchmark_run_dir / "config.yaml", "w") as file: + yaml.dump(combined_config, file) + + # 4. Generate/Update the restart script in the directory. The + # communicator size is ground truth for the job scale; environment + # sniffing is only the fallback for callers that lack it. + create_restart_script(benchmark_run_dir, world_size=world_size) + + return combined_config + + def main(): """ Command line interface for ScaFFold. @@ -446,114 +583,31 @@ def main(): "resume (e.g. '--restart --run-dir ')." ) + # Rank 0 builds the config for the whole job while every other rank waits + # in the barrier below. Everything in there can fail on user input (an + # unknown config key, an out-of-range bottleneck, a restart with no run + # dir) or on the filesystem, and a rank-0-only raise leaves the peers + # blocked in that barrier -- a hang instead of the error message the user + # needs. The outcome is therefore broadcast, exactly like the restart + # pre-check below, and every rank raises the same error together. + config_error = None + rank0_error = None if rank == 0: - log.debug("args = %s", args) - - # --config may be a single path (generate_fractals) or a list of - # paths (benchmark, action="append"): base config plus overrides. - config_paths = args.config if isinstance(args.config, list) else [args.config] - merged_dict = config_utils.load_config_files(config_paths) - # Validate the merged result and derive dependent settings. Every run - # parameter must be single-valued; a list is rejected here by name. - bench_config = config_utils.Config(merged_dict) - bench_config_dict = vars(bench_config) - cli_args = vars(args) - # Downstream consumers expect a single config path (e.g. to copy it - # into the run dir); keep the base config there. - cli_args["config"] = config_paths[0] - - # Combine configs, in increasing order of precedence: - # argparse default < config file < explicit command-line flag. - combined_config = bench_config_dict.copy() - # Config only keeps the keys it consumes; the auxiliary keys it accepts - # (verbose, datagen_batch_size, ...) never become attributes, so put - # the file's values back first. Without this they are absent below and - # the argparse default overwrites what the user wrote in the config. - for key, value in merged_dict.items(): - combined_config.setdefault(key, value) - - explicit_cli = explicit_cli_keys(args, (active_parser, parser)) - for key, value in cli_args.items(): - if key == "command": - continue - if key not in combined_config: - combined_config[key] = value - elif key in explicit_cli and value is not None: - log.info( - "Overriding '%s=%s' with '%s=%s'", - key, - combined_config[key], - key, - value, - ) - combined_config[key] = value - # The subcommand is always owned by the command line. - combined_config["command"] = cli_args["command"] - - # Recalculate unet_layers to capture any CLI overrides. The overridden - # pair has to be re-validated: Config only saw the config-file values. - config_utils.validate_unet_dims( - combined_config["problem_scale"], combined_config["unet_bottleneck_dim"] - ) - combined_config["unet_layers"] = ( - combined_config["problem_scale"] - combined_config["unet_bottleneck_dim"] - ) - config_utils.require_positive_int( - "n_categories", combined_config["n_categories"] - ) - - # Resolve paths to absolute, matching Config() behavior - if "base_run_dir" in combined_config and combined_config["base_run_dir"]: - combined_config["base_run_dir"] = str( - Path(combined_config["base_run_dir"]).resolve() - ) - - if "dataset_dir" in combined_config and combined_config["dataset_dir"]: - combined_config["dataset_dir"] = str( - Path(combined_config["dataset_dir"]).resolve() - ) - - if "fract_base_dir" in combined_config and combined_config["fract_base_dir"]: - combined_config["fract_base_dir"] = str( - Path(combined_config["fract_base_dir"]).resolve() + try: + combined_config = build_run_config( + args, (active_parser, parser), log, comm.Get_size() ) - - # Calculate these variables after override - combined_config["vol_size"] = pow(2, combined_config["problem_scale"]) - combined_config["point_num"] = int(combined_config["vol_size"] ** 3 / 256) - - # The run directory, its config dumps and its restart script belong to - # the benchmark subcommand alone. Fractal generation writes nothing - # there, and the restart script it used to get replayed - # `generate_fractals --restart --run-dir ...` -- flags that subparser - # rejects, so the script could only ever exit 2. - if args.command == "benchmark": - # Resolve the run directory and whether this launch resumes a run. - # This sets combined_config["benchmark_run_dir"] on every path and, - # when resuming, forces train_from_scratch off / restart on. - benchmark_run_dir, restarting = resolve_run_dir(vars(args), combined_config) - if restarting: - log.info("Resuming in existing directory: %s", benchmark_run_dir) - - # Add scheduler metadata and machine name to config.yaml - combined_config["scheduler_metadata"] = collect_scheduler_metadata() - combined_config["machine_name"] = socket.gethostname() - - # Dump configs (Overwrite is okay/desired on restart to capture new job IDs) - overrides = { - k: v for k, v in cli_args.items() if v is not None and k != "command" - } - with open(benchmark_run_dir / "overrides.yaml", "w") as file: - yaml.dump(overrides, file) - with open(benchmark_run_dir / "config.yaml", "w") as file: - yaml.dump(combined_config, file) - - # 4. Generate/Update the restart script in the directory. The - # communicator size is ground truth for the job scale; environment - # sniffing is only the fallback for callers that lack it. - create_restart_script(benchmark_run_dir, world_size=comm.Get_size()) + except Exception as e: + combined_config = None + rank0_error = e + config_error = (type(e).__name__, str(e)) comm.Barrier() + config_error = comm.bcast(config_error, root=0) + if config_error is not None: + # Rank 0 re-raises the original (keeping its traceback); the peers + # rebuild it from what crossed the wire. + raise rank0_error if rank0_error is not None else rebuild_error(*config_error) combined_config = comm.bcast(combined_config, root=0) # Restart pre-check. Like every other decision here it is made once, on diff --git a/tests/test_cli.py b/tests/test_cli.py index 1db35af..e91955a 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -284,7 +284,9 @@ def test_restart_precheck_follows_the_broadcast_decision(monkeypatch, tmp_path): "verbose": 0, } - comm = _FakeComm(rank=1, size=2, bcast_returns=[rank0_config, None]) + # Rank 0 broadcasts, in order: no config error, the config, no pre-check + # error. + comm = _FakeComm(rank=1, size=2, bcast_returns=[None, rank0_config, None]) _, calls = run_cli(monkeypatch, _restart_argv(cfg, run_dir), comm=comm) assert len(calls["benchmark"]) == 1 @@ -309,7 +311,7 @@ def test_restart_precheck_failure_raises_on_every_rank(monkeypatch, tmp_path): } rank0_error = "Restart requested but no checkpoint was found. Expected /nope." - comm = _FakeComm(rank=1, size=2, bcast_returns=[rank0_config, rank0_error]) + comm = _FakeComm(rank=1, size=2, bcast_returns=[None, rank0_config, rank0_error]) with pytest.raises(FileNotFoundError, match="no checkpoint"): run_cli(monkeypatch, _restart_argv(cfg, run_dir), comm=comm) @@ -541,6 +543,66 @@ def test_cli_override_bottleneck_out_of_range_rejected(monkeypatch, tmp_path): assert "problem_scale" in message +# --------------------------------------------------------------------------- +# VC-2: a rank-0 config failure is broadcast, not left to the barrier. +# +# Rank 0 builds the whole job's config while every peer waits in a barrier, so +# anything that raises in there -- the config-file validation, the recomputed +# bottleneck check, the run-dir resolution -- has to travel to the peers as a +# decision. Otherwise the user gets a hang instead of the error message. +# --------------------------------------------------------------------------- + + +def test_config_failure_is_broadcast_before_the_barrier(monkeypatch, tmp_path): + """Rank 0 posts its collectives and broadcasts the error it hit.""" + cfg = write_config(tmp_path) + comm = _FakeComm(rank=0, size=2) + + with pytest.raises(ValueError): + run_cli( + monkeypatch, + [ + "scaffold", + "benchmark", + "-c", + str(cfg), + "--problem-scale", + "4", + "--unet-bottleneck-dim", + "4", + ], + comm=comm, + ) + + # The peers' barrier was matched, and what they receive names the failure. + assert comm.barriers == 1 + errors = [ + payload + for payload in comm.broadcast + if isinstance(payload, tuple) and payload[0] == "ValueError" + ] + assert errors, f"no error sentinel was broadcast: {comm.broadcast}" + assert "unet_bottleneck_dim" in errors[0][1] + + +def test_peer_raises_the_broadcast_config_error(monkeypatch, tmp_path): + """A peer rebuilds rank 0's error instead of running with no config.""" + cfg = write_config(tmp_path) + rank0_error = ("ValueError", "unet_bottleneck_dim (4) must be < problem_scale (4)") + + comm = _FakeComm(rank=1, size=2, bcast_returns=[rank0_error]) + with pytest.raises(ValueError, match="unet_bottleneck_dim"): + run_cli(monkeypatch, ["scaffold", "benchmark", "-c", str(cfg)], comm=comm) + + +def test_unknown_error_types_degrade_to_runtime_error(): + """A non-builtin exception type is still reported, as a RuntimeError.""" + assert type(cli.rebuild_error("FileNotFoundError", "gone")) is FileNotFoundError + rebuilt = cli.rebuild_error("SomeSiteSpecificError", "boom") + assert isinstance(rebuilt, RuntimeError) + assert "SomeSiteSpecificError" in str(rebuilt) and "boom" in str(rebuilt) + + # --------------------------------------------------------------------------- # The whole config path survives a restart (R20/R22 together) # --------------------------------------------------------------------------- From f27aa2d23aae6911c038445b22baeef4fecbb55a Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 17:06:20 -0700 Subject: [PATCH 40/62] Harden launcher env parsing and CLI entry ordering The launcher-environment helpers called int() on whatever a variable held. A site wrapper that exports WORLD_SIZE= (an unset shell variable) or a placeholder like "auto" therefore killed every scaffold invocation with a bare ValueError naming neither the variable nor a remedy -- and it fired from the top of the entry point, so even `scaffold --help` died. Each lookup now goes through a helper that treats an unusable value as absent and consults the next source in the priority order, warning when the value looked deliberate (as _sniff_launch_shape already did with `if val:`). The Slurm and Flux tasks-per-node divisions no longer trust the node count either: zero is as unusable as a word. The world-size cross-check also moved after parse_args. Asking what the flags are is not a job launch, and answering it with a launcher mismatch -- exactly the environment someone debugging one is sitting in -- helps nobody. It still runs before any run directory is created, which is the property R13 needs. VC-3, VC-4 --- ScaFFold/cli.py | 12 ++- ScaFFold/utils/distributed.py | 165 ++++++++++++++++++++++------------ tests/test_cli.py | 22 +++++ tests/test_worker_dist.py | 53 +++++++++++ 4 files changed, 190 insertions(+), 62 deletions(-) diff --git a/ScaFFold/cli.py b/ScaFFold/cli.py index 88bf1f9..650a126 100644 --- a/ScaFFold/cli.py +++ b/ScaFFold/cli.py @@ -555,12 +555,16 @@ def main(): comm = MPI.COMM_WORLD rank = comm.Get_rank() - # Every rank runs this identically, before any run directory is created, - # so a mis-launched job aborts uniformly instead of leaving per-rank run - # dirs behind and hanging. - check_launcher_world_size(comm.Get_size()) # Parse the command-line arguments. args = parser.parse_args() + # Every rank runs this identically, before any run directory is created, so + # a mis-launched job aborts uniformly instead of leaving per-rank run dirs + # behind and hanging. It runs *after* parsing so that the arguments argparse + # handles by itself -- ``--help``, a usage error -- still behave: asking + # what the flags are is not a job launch, and answering it with a launcher + # mismatch (which is exactly the environment someone debugging one is + # sitting in) helps nobody. + check_launcher_world_size(comm.Get_size()) subcommand_parsers = { "benchmark": benchmark_parser, "generate_fractals": generate_fractals_parser, diff --git a/ScaFFold/utils/distributed.py b/ScaFFold/utils/distributed.py index afad5ab..5cb0700 100644 --- a/ScaFFold/utils/distributed.py +++ b/ScaFFold/utils/distributed.py @@ -12,6 +12,7 @@ # # SPDX-License-Identifier: (Apache-2.0) +import logging import os import os.path import socket @@ -21,6 +22,52 @@ import torch import torch.distributed +logger = logging.getLogger(__name__) + + +def _env_int(name: str) -> Optional[int]: + """Return the launcher variable ``name`` as an int, or None if unusable. + + Launcher variables are not always what they claim to be. A site wrapper + that exports ``WORLD_SIZE=`` (empty, e.g. from an unset shell variable) or + a placeholder like ``auto`` is common enough, and a bare ``int()`` turned + it into a ``ValueError`` raised from the first of these helpers anything + called -- killing every ``scaffold`` invocation, ``--help`` included, with + a traceback that named neither the variable nor a remedy. + + An unusable value is treated as absent so the next source in the priority + order is consulted (ultimately the MPI communicator, or the documented + default). A non-empty value that is not an integer is warned about, because + unlike an empty one it looks deliberate and the fallback may not be what + its author intended. ``create_restart_script._sniff_launch_shape`` skips + empty values for the same reason. + """ + value = os.environ.get(name) + if value is None: + return None + value = value.strip() + if not value: + return None + try: + return int(value) + except ValueError: + logger.warning( + "Ignoring launcher variable %s=%r: not an integer. Falling back to " + "the next source for the job shape.", + name, + value, + ) + return None + + +def _first_env_int(names) -> Optional[int]: + """Return the first usable integer among ``names``, in priority order.""" + for name in names: + value = _env_int(name) + if value is not None: + return value + return None + def get_num_gpus() -> int: """Return the number of GPUs on this node.""" @@ -40,22 +87,50 @@ def _mpi_comm_world(): return None +_LOCAL_RANK_VARS = ( + "LOCAL_RANK", + "MV2_COMM_WORLD_LOCAL_RANK", + "OMPI_COMM_WORLD_LOCAL_RANK", + "PMI_LOCAL_RANK", + "PALS_LOCAL_RANKID", + "SLURM_LOCALID", + "FLUX_TASK_LOCAL_ID", +) + +_LOCAL_SIZE_VARS = ( + "LOCAL_WORLD_SIZE", + "MV2_COMM_WORLD_LOCAL_SIZE", + "OMPI_COMM_WORLD_LOCAL_SIZE", + "PMI_LOCAL_SIZE", + "PALS_LOCAL_SIZE", +) + +_WORLD_RANK_VARS = ( + "RANK", + "MV2_COMM_WORLD_RANK", + "OMPI_COMM_WORLD_RANK", + "PMI_RANK", + "PALS_RANKID", + "SLURM_PROCID", + "FLUX_TASK_RANK", +) + +_WORLD_SIZE_VARS = ( + "WORLD_SIZE", + "MV2_COMM_WORLD_SIZE", + "OMPI_COMM_WORLD_SIZE", + "PMI_SIZE", + "PALS_NRANKS", + "SLURM_NTASKS", + "FLUX_JOB_SIZE", +) + + def get_local_rank(required: bool = False) -> int: """Return the local MPI rank.""" - if "LOCAL_RANK" in os.environ: - return int(os.environ["LOCAL_RANK"]) - if "MV2_COMM_WORLD_LOCAL_RANK" in os.environ: - return int(os.environ["MV2_COMM_WORLD_LOCAL_RANK"]) - if "OMPI_COMM_WORLD_LOCAL_RANK" in os.environ: - return int(os.environ["OMPI_COMM_WORLD_LOCAL_RANK"]) - if "PMI_LOCAL_RANK" in os.environ: - return int(os.environ["PMI_LOCAL_RANK"]) - if "PALS_LOCAL_RANKID" in os.environ: - return int(os.environ["PALS_LOCAL_RANKID"]) - if "SLURM_LOCALID" in os.environ: - return int(os.environ["SLURM_LOCALID"]) - if "FLUX_TASK_LOCAL_ID" in os.environ: - return int(os.environ["FLUX_TASK_LOCAL_ID"]) + value = _first_env_int(_LOCAL_RANK_VARS) + if value is not None: + return value if required: raise RuntimeError("Could not get local rank") return 0 @@ -68,22 +143,18 @@ def get_local_size(required: bool = False) -> int: there but not here silently yields 1, which makes per-node logic (e.g. the profiler's one-rank-per-node gate) treat every rank as node-local. """ - if "LOCAL_WORLD_SIZE" in os.environ: - return int(os.environ["LOCAL_WORLD_SIZE"]) - if "MV2_COMM_WORLD_LOCAL_SIZE" in os.environ: - return int(os.environ["MV2_COMM_WORLD_LOCAL_SIZE"]) - if "OMPI_COMM_WORLD_LOCAL_SIZE" in os.environ: - return int(os.environ["OMPI_COMM_WORLD_LOCAL_SIZE"]) - if "PMI_LOCAL_SIZE" in os.environ: - return int(os.environ["PMI_LOCAL_SIZE"]) - if "PALS_LOCAL_SIZE" in os.environ: - return int(os.environ["PALS_LOCAL_SIZE"]) - if "SLURM_NNODES" in os.environ and "SLURM_NTASKS" in os.environ: - return int(os.environ["SLURM_NTASKS"]) // int(os.environ["SLURM_NNODES"]) - # Flux does not have an env variable for this, so we assume an - # even distribution. - if "FLUX_JOB_SIZE" in os.environ and "FLUX_JOB_NNODES" in os.environ: - return int(os.environ["FLUX_JOB_SIZE"]) // int(os.environ["FLUX_JOB_NNODES"]) + value = _first_env_int(_LOCAL_SIZE_VARS) + if value is not None: + return value + # Slurm and Flux report only totals; assume an even distribution. A zero + # node count is as unusable as a non-numeric one, so it falls through + # rather than raising ZeroDivisionError. + ntasks, nnodes = _env_int("SLURM_NTASKS"), _env_int("SLURM_NNODES") + if ntasks is not None and nnodes: + return ntasks // nnodes + job_size, job_nodes = _env_int("FLUX_JOB_SIZE"), _env_int("FLUX_JOB_NNODES") + if job_size is not None and job_nodes: + return job_size // job_nodes if required: raise RuntimeError("Could not get local size") return 1 @@ -91,20 +162,9 @@ def get_local_size(required: bool = False) -> int: def get_world_rank(required: bool = False) -> int: """Return the global MPI rank.""" - if "RANK" in os.environ: - return int(os.environ["RANK"]) - if "MV2_COMM_WORLD_RANK" in os.environ: - return int(os.environ["MV2_COMM_WORLD_RANK"]) - if "OMPI_COMM_WORLD_RANK" in os.environ: - return int(os.environ["OMPI_COMM_WORLD_RANK"]) - if "PMI_RANK" in os.environ: - return int(os.environ["PMI_RANK"]) - if "PALS_RANKID" in os.environ: - return int(os.environ["PALS_RANKID"]) - if "SLURM_PROCID" in os.environ: - return int(os.environ["SLURM_PROCID"]) - if "FLUX_TASK_RANK" in os.environ: - return int(os.environ["FLUX_TASK_RANK"]) + value = _first_env_int(_WORLD_RANK_VARS) + if value is not None: + return value comm = _mpi_comm_world() if comm is not None: return comm.Get_rank() @@ -115,20 +175,9 @@ def get_world_rank(required: bool = False) -> int: def get_world_size(required: bool = False) -> int: """Return the number of MPI ranks.""" - if "WORLD_SIZE" in os.environ: - return int(os.environ["WORLD_SIZE"]) - if "MV2_COMM_WORLD_SIZE" in os.environ: - return int(os.environ["MV2_COMM_WORLD_SIZE"]) - if "OMPI_COMM_WORLD_SIZE" in os.environ: - return int(os.environ["OMPI_COMM_WORLD_SIZE"]) - if "PMI_SIZE" in os.environ: - return int(os.environ["PMI_SIZE"]) - if "PALS_NRANKS" in os.environ: - return int(os.environ["PALS_NRANKS"]) - if "SLURM_NTASKS" in os.environ: - return int(os.environ["SLURM_NTASKS"]) - if "FLUX_JOB_SIZE" in os.environ: - return int(os.environ["FLUX_JOB_SIZE"]) + value = _first_env_int(_WORLD_SIZE_VARS) + if value is not None: + return value comm = _mpi_comm_world() if comm is not None: return comm.Get_size() diff --git a/tests/test_cli.py b/tests/test_cli.py index e91955a..f7d83ca 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -181,6 +181,28 @@ def test_matching_world_sizes_are_accepted(monkeypatch, tmp_path): assert len(calls["benchmark"]) == 1 +def test_help_works_under_a_mismatched_launcher_env(monkeypatch, tmp_path, capsys): + """``--help`` is answered even when the launcher environment disagrees. + + The cross-check ran before ``parse_args``, so asking what the flags are + raised the launcher-mismatch error -- in exactly the environment (a + half-configured shell) where someone is most likely to be asking. + """ + monkeypatch.setenv("WORLD_SIZE", "2") + monkeypatch.setenv("RANK", "0") + + with pytest.raises(SystemExit) as excinfo: + run_cli( + monkeypatch, + ["scaffold", "--help"], + comm=_FakeComm(rank=0, size=1), + sync_env=False, + ) + + assert excinfo.value.code == 0 + assert "usage" in capsys.readouterr().out.lower() + + def test_no_launcher_env_is_not_a_mismatch(monkeypatch, tmp_path): """With no launcher variables set, the MPI world alone defines the size.""" for var in ("WORLD_SIZE", "RANK", "LOCAL_RANK", "SLURM_NTASKS", "FLUX_JOB_SIZE"): diff --git a/tests/test_worker_dist.py b/tests/test_worker_dist.py index b1ae524..e3441ea 100644 --- a/tests/test_worker_dist.py +++ b/tests/test_worker_dist.py @@ -21,6 +21,7 @@ import logging +import pytest import torch import ScaFFold.utils.distributed as distributed_mod @@ -276,3 +277,55 @@ def test_local_size_defaults_to_one(monkeypatch): for var in _LOCAL_SIZE_VARS: monkeypatch.delenv(var, raising=False) assert distributed_mod.get_local_size() == 1 + + +# --------------------------------------------------------------------------- +# VC-3: an unusable launcher variable is ignored, not fatal +# +# These helpers run at the top of every entry point, so a bare int() on a +# variable a site wrapper exported empty ("WORLD_SIZE=") or as a placeholder +# ("auto") killed the invocation -- ``scaffold --help`` included -- with a +# ValueError naming neither the variable nor a remedy. +# --------------------------------------------------------------------------- + + +def _clear_launcher_env(monkeypatch): + for var in set(_LAUNCHER_VARS) | set(_LOCAL_SIZE_VARS): + monkeypatch.delenv(var, raising=False) + + +@pytest.mark.parametrize("value", ["", " ", "auto"]) +def test_unusable_launcher_values_fall_through(monkeypatch, value): + """Empty and non-numeric values are treated as absent, never raise.""" + _clear_launcher_env(monkeypatch) + for var in ("WORLD_SIZE", "RANK", "LOCAL_RANK", "LOCAL_WORLD_SIZE"): + monkeypatch.setenv(var, value) + + # Falls through to the next source -- here the (singleton) communicator and + # the documented defaults. + assert distributed_mod.get_world_size() == 1 + assert distributed_mod.get_world_rank() == 0 + assert distributed_mod.get_local_rank() == 0 + assert distributed_mod.get_local_size() == 1 + + +def test_unusable_value_defers_to_the_next_launcher_variable(monkeypatch, caplog): + """A garbage value does not mask a usable variable further down the order.""" + _clear_launcher_env(monkeypatch) + monkeypatch.setenv("WORLD_SIZE", "auto") + monkeypatch.setenv("PALS_NRANKS", "8") + + with caplog.at_level(logging.WARNING, logger=distributed_mod.logger.name): + assert distributed_mod.get_world_size() == 8 + + messages = " ".join(record.getMessage() for record in caplog.records) + assert "WORLD_SIZE" in messages, "the ignored value was not reported" + + +def test_zero_node_count_does_not_divide_by_zero(monkeypatch): + """A nonsense node count falls through instead of raising.""" + _clear_launcher_env(monkeypatch) + monkeypatch.setenv("SLURM_NTASKS", "8") + monkeypatch.setenv("SLURM_NNODES", "0") + + assert distributed_mod.get_local_size() == 1 From 9492d9ad636d5891e147194cc21114b5f1ab7de0 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 17:08:01 -0700 Subject: [PATCH 41/62] Derive restart node shape from local size Flux and Slurm state their node count, so the restart script reproduced their shape correctly. Everything else -- Cray PALS, plain torchrun -- fell back to NODES=1, which relaunched an 8-rank job spread over 2 nodes as NODES=1 TASKS_PER_NODE=8: an oversubscribed node, or a job the scheduler rejects. get_local_size already reads those launchers' per-node variables (PALS_LOCAL_ SIZE, LOCAL_WORLD_SIZE, ...), so the node count is derived from it, ceil-ing the division as the rank side does. required=True keeps "one rank per node" distinct from "nothing reported a per-node count"; only the latter falls back to the historical single-node assumption. VC-5 --- ScaFFold/utils/create_restart_script.py | 28 +++++++++++++-- tests/test_restart_script.py | 48 +++++++++++++++++++++++++ 2 files changed, 74 insertions(+), 2 deletions(-) diff --git a/ScaFFold/utils/create_restart_script.py b/ScaFFold/utils/create_restart_script.py index 3f93ade..1d5caf1 100644 --- a/ScaFFold/utils/create_restart_script.py +++ b/ScaFFold/utils/create_restart_script.py @@ -15,6 +15,7 @@ # restart_script.py from __future__ import annotations +import math import os import shlex import stat @@ -23,6 +24,8 @@ from pathlib import Path from typing import List, Union +from ScaFFold.utils.distributed import get_local_size + # Profiling toggles that must be reproduced on restart -- but only when they # were active in the generating run. Names mirror ScaFFold.utils.perf_measure. _PROFILING_ENV_VARS = ("PROFILE_TORCH", "CALI_CONFIG") @@ -299,6 +302,11 @@ def create_restart_script(run_dir: str | Path, world_size: int | None = None) -> torchrun, Open MPI, PMI). The multi-rank torchrun-hpc template is emitted whenever the resulting world size is greater than one; the local single-process template is used only for a world size of one. + + The node count comes from the scheduler when it reports one (Flux, Slurm); + otherwise it is derived from the per-node rank count + ``ScaFFold.utils.distributed.get_local_size`` reads, so a PALS or torchrun + job is not relaunched with every rank crammed onto one node. """ run_dir = Path(run_dir) run_dir.mkdir(parents=True, exist_ok=True) @@ -339,8 +347,24 @@ def create_restart_script(run_dir: str | Path, world_size: int | None = None) -> if use_torchrun: # Calculate tasks per node for torchrun (-n arg). if nodes is None: - nodes = 1 - tasks_per_node = max(1, total_tasks // nodes) + # No scheduler reported a node count: this is a PALS or plain + # torchrun launch. Assuming one node put the job's whole rank count + # on a single node (an 8-rank job across 2 nodes came back as + # NODES=1 TASKS_PER_NODE=8), which either oversubscribes one node or + # is rejected outright. The rank side's ``get_local_size`` reads the + # same launchers' per-node variables, so ask it how many ranks share + # this node and derive the node count from that. ``required=True`` + # distinguishes "one rank per node" from "nothing reported a + # per-node count", where the historical single-node assumption is + # still the best guess available. + try: + local_size = get_local_size(required=True) + except RuntimeError: + local_size = total_tasks + tasks_per_node = max(1, min(local_size, total_tasks)) + nodes = math.ceil(total_tasks / tasks_per_node) + else: + tasks_per_node = max(1, total_tasks // nodes) script = _render_torchrun_hpc_restart( py_array_decl, nodes, tasks_per_node, env_setup diff --git a/tests/test_restart_script.py b/tests/test_restart_script.py index a8f3e08..291c57a 100644 --- a/tests/test_restart_script.py +++ b/tests/test_restart_script.py @@ -305,3 +305,51 @@ def test_pals_job_gets_a_multirank_restart_script(monkeypatch, tmp_path): assert "torchrun-hpc" in script assert 'exec "${PY[@]}"' not in script + + +# --------------------------------------------------------------------------- +# VC-5: the node shape comes from the local rank count when no scheduler +# reports one. Flux and Slurm state their node count; PALS and plain torchrun +# do not, and assuming one node relaunched an 8-rank/2-node job as +# NODES=1 TASKS_PER_NODE=8 -- an oversubscribed node, or a rejected job. +# --------------------------------------------------------------------------- + + +def test_pals_multinode_shape_uses_the_local_rank_count(monkeypatch, tmp_path): + """8 PALS ranks, 4 per node -> NODES=2, TASKS_PER_NODE=4.""" + _isolate_env(monkeypatch) + monkeypatch.delenv("PALS_LOCAL_SIZE", raising=False) + monkeypatch.delenv("LOCAL_WORLD_SIZE", raising=False) + monkeypatch.setenv("PALS_NRANKS", "8") + monkeypatch.setenv("PALS_LOCAL_SIZE", "4") + + script = _generate(monkeypatch, tmp_path / "run") + + assert 'NODES="2"' in script + assert 'TASKS_PER_NODE="4"' in script + + +def test_single_node_torchrun_shape_is_unchanged(monkeypatch, tmp_path): + """8 torchrun ranks all on one node stay NODES=1, TASKS_PER_NODE=8.""" + _isolate_env(monkeypatch) + monkeypatch.delenv("PALS_LOCAL_SIZE", raising=False) + monkeypatch.setenv("WORLD_SIZE", "8") + monkeypatch.setenv("LOCAL_WORLD_SIZE", "8") + + script = _generate(monkeypatch, tmp_path / "run") + + assert 'NODES="1"' in script + assert 'TASKS_PER_NODE="8"' in script + + +def test_unknown_local_size_keeps_the_single_node_assumption(monkeypatch, tmp_path): + """With nothing reporting a per-node count, the old assumption stands.""" + _isolate_env(monkeypatch) + for var in ("PALS_LOCAL_SIZE", "LOCAL_WORLD_SIZE", "PMI_LOCAL_SIZE"): + monkeypatch.delenv(var, raising=False) + monkeypatch.setenv("PALS_NRANKS", "8") + + script = _generate(monkeypatch, tmp_path / "run") + + assert 'NODES="1"' in script + assert 'TASKS_PER_NODE="8"' in script From 7cdef0302b702b77e8063689666d6a80516956d6 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 17:12:10 -0700 Subject: [PATCH 42/62] Wording and cosmetics The "no new epoch was trained" warning covers two ways of entering train() with nothing to do, and described one of them as the other: a fresh epochs:0 run was told "there was nothing to resume", sending its user to look for a checkpoint that was never part of the story. It now says the run had no epoch left to run, which is true of a completed resume and a zero-epoch fresh start alike. (The deeper problem in that scenario -- worker.py's rank-0-only genfromtxt raising after destroy_process_group, so rank 0 exits non-zero while its peers exit 0 -- is pre-existing structure, unchanged here.) explicit_cli_keys' docstring claimed that a flag passed with exactly its default value costs nothing because "both spellings then agree on the default". They do not: --datagen-batch-size 10000 next to datagen_batch_size: 500 in the config yields 500. The rationale is corrected to name the real trade-off rather than deny it; the defaults stay where they are, which is what the R20 tests pin. perf_measure now has a logger instead of printing. Its messages are all "your profiling request was not honored as written", which belongs on a channel a caller can filter or capture, not in the middle of the run's stdout. The R26 test follows it to caplog. And the R21 tests' function-body `import pytest` moves to module level as a skipif marker. The config-error re-raise added a commit ago is spelled as a statement rather than a conditional expression. VA-4, VC-6, VC-cosmetics --- ScaFFold/cli.py | 17 +++++++++++++---- ScaFFold/utils/perf_measure.py | 22 ++++++++++++++++----- ScaFFold/utils/trainer.py | 20 ++++++++++--------- tests/test_infra.py | 20 +++++++++---------- tests/test_reporting.py | 22 ++++++++------------- tests/test_resume.py | 35 ++++++++++++++++++++++++++++++++-- 6 files changed, 92 insertions(+), 44 deletions(-) diff --git a/ScaFFold/cli.py b/ScaFFold/cli.py index 650a126..105e87a 100644 --- a/ScaFFold/cli.py +++ b/ScaFFold/cli.py @@ -100,9 +100,16 @@ def explicit_cli_keys(args, parsers): parser). Only these may outrank a config-file setting; everything else in the namespace is an argparse default, which is the weakest source. - The one ambiguity is a flag passed with exactly its default value: it looks - absent, so a config-file entry wins over it. Both spellings then agree on - the default, which is the only value the flag could have contributed. + The one ambiguity is a flag passed with exactly its default value: it is + indistinguishable from an absent flag, so a config-file entry outranks it. + Where the flag has no default (``None``) that is harmless -- passing a + value always makes it explicit -- but the two flags that do have one, + ``--datagen-batch-size`` (10000) and ``-v`` (0), lose the argument in that + one case: ``--datagen-batch-size 10000`` next to ``datagen_batch_size: 500`` + in the config file yields 500. The alternative is to give every flag a + ``None`` default and re-derive the real defaults elsewhere, which buys a + narrow correctness win by scattering the defaults; the ambiguity is + documented instead. """ explicit = set() for name, value in vars(args).items(): @@ -611,7 +618,9 @@ def main(): if config_error is not None: # Rank 0 re-raises the original (keeping its traceback); the peers # rebuild it from what crossed the wire. - raise rank0_error if rank0_error is not None else rebuild_error(*config_error) + if rank0_error is not None: + raise rank0_error + raise rebuild_error(*config_error) combined_config = comm.bcast(combined_config, root=0) # Restart pre-check. Like every other decision here it is made once, on diff --git a/ScaFFold/utils/perf_measure.py b/ScaFFold/utils/perf_measure.py index 7236e69..9d87a2b 100644 --- a/ScaFFold/utils/perf_measure.py +++ b/ScaFFold/utils/perf_measure.py @@ -12,12 +12,20 @@ # # SPDX-License-Identifier: (Apache-2.0) +import logging import os from contextlib import nullcontext CALI_PERF_ENV_VAR = "CALI_CONFIG" TORCH_PERF_ENV_VAR = "PROFILE_TORCH" +# This module is imported before (and independently of) the run's MPI logger, +# so it keeps its own. Everything it has to say is about the user's profiling +# request not being honored as written, which belongs on a diagnostic channel +# that a caller can filter or capture -- not on stdout, where it lands in the +# middle of whatever the run is printing. +logger = logging.getLogger(__name__) + def _profiler_env_flag(name): """Return True only for an affirmative value of the environment variable. @@ -39,8 +47,11 @@ def _profiler_env_flag(name): _CALI_PERF_ENABLED = True except Exception as e: - print("User requested Caliper annotations, but could not import Caliper") - print(f"Exception: {e}") + logger.warning( + "User requested Caliper annotations, but could not import Caliper: %s: %s", + type(e).__name__, + e, + ) # The torch profiler is gated purely on its own environment variable: Caliper # and the torch profiler may both be enabled at once. @@ -51,8 +62,9 @@ def _profiler_env_flag(name): TORCH_PERF_ENABLED = True except Exception: - print( - "User requested PyTorch profiling, but could not import the PyTorch profiler" + logger.warning( + "User requested PyTorch profiling, but could not import the " + "PyTorch profiler" ) @@ -131,7 +143,7 @@ def get_torch_context(ranks_per_node, rank): # thing the bounded window exists to prevent. wait = _profiler_env_int("PROFILE_TORCH_WAIT", 1) if wait < 1: - print( + logger.warning( "PROFILE_TORCH_WAIT must be at least 1: the profiler window " "opens before the warmup batches, whose work would otherwise " "accumulate in host memory as one unbounded step. Using " diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 81dd5f2..3498cc1 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -1086,17 +1086,19 @@ def train(self, profiler=None): completed_epochs = epoch - 1 if not completed_new_epoch: - # The loop exited without running a single epoch: the state we - # resumed either already covers every epoch this run was asked for, - # or already met target_dice. There is nothing new to save (the - # checkpoint on disk already records epoch `completed_epochs`) and - # none of the per-epoch metrics the final save would write were - # ever computed, so skip it and return normally -- the caller's - # post-processing still has the CSV the original run left behind. + # The loop exited without running a single epoch: the epoch budget + # was already exhausted (a resume whose checkpoint covers every + # epoch asked for, or a fresh run configured with none) or the + # starting state already met target_dice. There is nothing new to + # save -- any checkpoint on disk already records epoch + # `completed_epochs`, and none of the per-epoch metrics the final + # save would write were ever computed -- so skip it and return + # normally; the caller's post-processing still has whatever CSV is + # there. self.log.warning( "No new epoch was trained (start epoch %s, 'epochs' %s, " - "starting val dice %s vs target_dice %s): there was nothing to " - "resume, and no checkpoint was written.", + "starting val dice %s vs target_dice %s): this run had no epoch " + "left to run, and no checkpoint was written.", self.start_epoch, self.config.epochs, self.start_val_dice, diff --git a/tests/test_infra.py b/tests/test_infra.py index c279f34..aaf3a73 100644 --- a/tests/test_infra.py +++ b/tests/test_infra.py @@ -26,6 +26,7 @@ import os import numpy as np +import pytest import torch from ScaFFold.utils.data_loading import FractalDataset @@ -245,19 +246,23 @@ def _debug_logger(name): return log -def test_mem_stats_without_cuda(caplog): - """``mem_stats`` reports "no GPU" instead of raising on a CPU-only host.""" - if torch.cuda.is_available(): - import pytest +# The bug was a CUDA-free host taking the CUDA path, so these only mean +# something where CUDA is genuinely unavailable. +_requires_no_cuda = pytest.mark.skipif( + torch.cuda.is_available(), reason="covers the CPU-only path" +) - pytest.skip("test covers the CPU-only path") +@_requires_no_cuda +def test_mem_stats_without_cuda(caplog): + """``mem_stats`` reports "no GPU" instead of raising on a CPU-only host.""" stats = mem_stats() assert stats["cuda_available"] is False assert "rank" in stats +@_requires_no_cuda def test_gather_and_print_mem_without_cuda(caplog): """A DEBUG-level CPU run logs a fallback instead of crashing. @@ -265,11 +270,6 @@ def test_gather_and_print_mem_without_cuda(caplog): ``-v`` used to die in trainer construction with "No CUDA GPUs are available". """ - if torch.cuda.is_available(): - import pytest - - pytest.skip("test covers the CPU-only path") - log = _debug_logger("test_gather_and_print_mem_without_cuda") with caplog.at_level(logging.DEBUG, logger=log.name): gather_and_print_mem(log, "after_trainer_setup") diff --git a/tests/test_reporting.py b/tests/test_reporting.py index 3b658fd..4a462ca 100644 --- a/tests/test_reporting.py +++ b/tests/test_reporting.py @@ -13,6 +13,7 @@ # SPDX-License-Identifier: (Apache-2.0) import csv +import logging from pathlib import Path from types import SimpleNamespace @@ -319,8 +320,6 @@ def test_zero_step_export_is_reported_not_raised(self, tmp_path, caplog): post-processing, so a raise here kills the profiling rank and leaves every other rank blocked in that barrier until the collective timeout. """ - import logging - import ScaFFold.worker as worker log = logging.getLogger("test_zero_step_export") @@ -340,8 +339,6 @@ def test_zero_step_export_is_reported_not_raised(self, tmp_path, caplog): def test_successful_export_writes_a_trace(self, tmp_path, caplog): """A profiler with a completed window still writes its trace (control).""" - import logging - from torch.profiler import ProfilerActivity, profile, schedule import ScaFFold.worker as worker @@ -364,8 +361,6 @@ def test_successful_export_writes_a_trace(self, tmp_path, caplog): def test_trace_lands_in_the_run_dir(self, tmp_path, caplog): """R23: the trace goes to the run dir, not whatever CWD happens to be.""" - import logging - import ScaFFold.worker as worker prof = self._stepped_profiler() @@ -387,8 +382,6 @@ def test_trace_name_counts_nodes_not_ranks( self, tmp_path, world_size, ranks_per_node, expected ): """R23: the N field is a node count, and never rounds a node away.""" - import logging - import ScaFFold.worker as worker prof = self._stepped_profiler() @@ -493,7 +486,7 @@ def _context_with(monkeypatch_context, env): assert is_local return ctx - def test_wait_zero_does_not_record_step_zero(self, monkeypatch, capsys): + def test_wait_zero_does_not_record_step_zero(self, monkeypatch, caplog): """PROFILE_TORCH_WAIT=0 is clamped so step 0 records nothing. worker.main enters the profiler context around checkpoint cleanup and @@ -509,11 +502,12 @@ def test_wait_zero_does_not_record_step_zero(self, monkeypatch, capsys): import ScaFFold.utils.perf_measure as perf_measure try: - with monkeypatch.context() as m: - ctx = self._context_with(m, {"PROFILE_TORCH_WAIT": "0"}) - assert ctx.schedule(0) == ProfilerAction.NONE - output = capsys.readouterr().out - assert "PROFILE_TORCH_WAIT" in output + with caplog.at_level(logging.WARNING, logger=perf_measure.logger.name): + with monkeypatch.context() as m: + ctx = self._context_with(m, {"PROFILE_TORCH_WAIT": "0"}) + assert ctx.schedule(0) == ProfilerAction.NONE + messages = " ".join(record.getMessage() for record in caplog.records) + assert "PROFILE_TORCH_WAIT" in messages finally: importlib.reload(perf_measure) diff --git a/tests/test_resume.py b/tests/test_resume.py index 3e8aa4d..dfe7b2c 100644 --- a/tests/test_resume.py +++ b/tests/test_resume.py @@ -364,7 +364,7 @@ def test_restart_of_completed_run_trains_and_saves_nothing(tmp_path, caplog): (the last checkpoint already covers the completed epochs) rather than saving with an unbound ``val_loss_avg``. ``train()`` has to return normally so the worker's post-processing still runs off the CSV already on disk, and - the run must say plainly that there was nothing to resume. + the run must say plainly that it had no epoch left to run. """ log = logging.getLogger("resume.r01") run = tmp_path / "run" @@ -416,7 +416,38 @@ def test_restart_of_completed_run_trains_and_saves_nothing(tmp_path, caplog): assert ckpt.read_bytes() == before epochs = [ln.split(",")[0] for ln in csv.read_text().splitlines()[1:]] assert epochs == ["1", "2"] - assert "nothing to resume" in caplog.text.lower() + assert "no epoch left to run" in caplog.text.lower() + + +def test_fresh_run_with_no_epochs_does_not_claim_a_resume(tmp_path, caplog): + """A fresh ``epochs: 0`` run reports the truth: nothing was resumed. + + The same message covers both ways of entering ``train()`` with no epoch to + run, so it must not describe one of them as the other. This run has no + checkpoint and never asked for one -- telling its user "there was nothing + to resume" sends them looking for a checkpoint that was never part of the + story. + """ + log = logging.getLogger("resume.va4") + run = tmp_path / "run" + run.mkdir() + + trainer = _stub_trainer( + run, + train_from_scratch=True, + log=log, + epochs=0, + checkpoint_interval=1, + ) + trainer.cleanup_or_resume() + assert trainer.start_epoch == 1 # a fresh run: nothing was resumed + + with caplog.at_level(logging.WARNING): + trainer.train() + + assert "no epoch left to run" in caplog.text.lower() + assert "nothing to resume" not in caplog.text.lower() + assert not trainer.checkpoint_manager.last_ckpt_path.exists() def test_converged_resume_does_not_retrain(tiny_trainer, monkeypatch): From bd380271d5c7ebe0880a93b98e362ef4c5589668 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 17:18:59 -0700 Subject: [PATCH 43/62] Pin the non-repo provenance test to its own directory test_non_repo_install_reports_no_commit_id builds a "not a checkout" directory under tmp_path, but git walks upwards until it finds a repository: run with --basetemp inside a ScaFFold checkout, that directory inherits the checkout's HEAD and the test fails on where it was run rather than on what it tests. GIT_CEILING_DIRECTORIES stops the walk at tmp_path. VB-7 --- tests/datagen/test_provenance.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/datagen/test_provenance.py b/tests/datagen/test_provenance.py index d53bd33..d907c39 100644 --- a/tests/datagen/test_provenance.py +++ b/tests/datagen/test_provenance.py @@ -92,12 +92,19 @@ def test_commit_survives_a_non_repo_working_directory(tmp_path, monkeypatch): assert gd._git_commit_short(LOG) == expected -def test_non_repo_install_reports_no_commit_id(tmp_path): +def test_non_repo_install_reports_no_commit_id(tmp_path, monkeypatch): """An installed (non-git) ScaFFold still degrades gracefully. Provenance is best-effort: when the source tree is not a checkout there is no commit to record, and reuse simply is not gated on one. + + "Not a checkout" has to be made true of the *whole path*, not just the leaf: + git walks upwards until it finds a repository, so with ``--basetemp`` inside + a ScaFFold checkout this directory inherits that checkout's HEAD and the + test fails on where it was run rather than on what it tests. The ceiling + stops the walk at ``tmp_path``. """ + monkeypatch.setenv("GIT_CEILING_DIRECTORIES", str(tmp_path)) not_a_repo = tmp_path / "site-packages" / "ScaFFold" / "datagen" not_a_repo.mkdir(parents=True) From b11569d1d6c378708f0c541a91b8ff94925dbd0c Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 17:26:24 -0700 Subject: [PATCH 44/62] Copy the committed checkpoint instead of re-serializing it for best An improving epoch pickled and fsynced the identical state dict twice, once for checkpoint_last.pth and again for checkpoint_best.pth. The best file is now copied (tmp + fsync + os.replace) from the last file the same writer just committed, halving the serialization CPU and checkpoint bytes. R08 --- ScaFFold/utils/checkpointing.py | 43 ++++++++++++++++++-- tests/test_checkpointing.py | 71 +++++++++++++++++++++++++++++++++ 2 files changed, 111 insertions(+), 3 deletions(-) diff --git a/ScaFFold/utils/checkpointing.py b/ScaFFold/utils/checkpointing.py index 22c6051..84ed96a 100644 --- a/ScaFFold/utils/checkpointing.py +++ b/ScaFFold/utils/checkpointing.py @@ -15,6 +15,7 @@ import math import os import random +import shutil import traceback from concurrent.futures import ThreadPoolExecutor from pathlib import Path @@ -630,6 +631,32 @@ def _atomic_save(state_dict, path): pass raise + @staticmethod + def _atomic_copy(src, dst): + """Copy an already-committed checkpoint onto another name atomically. + + Same discipline as ``_atomic_save`` -- copy into a temp file in the + same directory, fsync it, then ``os.replace`` -- so ``dst`` is never + observed half-written and a crash mid-copy cannot damage the previous + good file there. + """ + src = Path(src) + dst = Path(dst) + tmp_path = dst.with_name(f"{dst.name}.tmp.{os.getpid()}") + try: + with open(src, "rb") as fsrc, open(tmp_path, "wb") as fdst: + shutil.copyfileobj(fsrc, fdst) + fdst.flush() + os.fsync(fdst.fileno()) + os.replace(tmp_path, dst) + except Exception: + try: + if tmp_path.exists(): + tmp_path.unlink() + except OSError: + pass + raise + @classmethod def _write_to_disk(cls, state_dict, last_path, best_path, is_best, log): """Worker function to perform actual disk I/O. @@ -642,10 +669,20 @@ def _write_to_disk(cls, state_dict, last_path, best_path, is_best, log): try: # Save 'last' atomically. cls._atomic_save(state_dict, last_path) - # Save 'best' atomically (re-serialize rather than copy a file that - # a concurrent writer might still be replacing). + # 'best' is byte-identical to the 'last' just committed, so copy + # that file instead of pickling and fsyncing the same state a + # second time (double the serialization CPU and double the bytes + # pushed at the shared filesystem on every improving epoch). + # + # There is no concurrent writer to race: checkpoint writes are + # serialized through a single writer -- the caller's thread in sync + # mode, or the one-worker ThreadPoolExecutor in async mode, whose + # previous write ``_rank0_save`` drains before submitting the next + # -- and only rank 0 ever writes. So this very thread performed the + # ``os.replace`` onto ``last_path`` a moment ago and nothing else + # can be replacing it now. if is_best: - cls._atomic_save(state_dict, best_path) + cls._atomic_copy(last_path, best_path) except Exception: if log is not None: log.error("Saving checkpoint failed:\n%s", traceback.format_exc()) diff --git a/tests/test_checkpointing.py b/tests/test_checkpointing.py index acca106..82ab353 100644 --- a/tests/test_checkpointing.py +++ b/tests/test_checkpointing.py @@ -518,6 +518,77 @@ def test_init_sweeps_orphaned_tmp_files(tmp_path): assert quarantined.exists() +# --------------------------------------------------------------------------- +# R08 -- an improving epoch serializes the state dict once, not twice +# --------------------------------------------------------------------------- + + +def _count_torch_saves(monkeypatch): + """Count ``torch.save`` calls made by the checkpoint writer.""" + calls = [] + real_save = torch.save + + def counting_save(obj, f, *args, **kwargs): + calls.append(str(getattr(f, "name", f))) + return real_save(obj, f, *args, **kwargs) + + monkeypatch.setattr(torch, "save", counting_save) + return calls + + +def test_best_checkpoint_copied_not_reserialized(tmp_path, monkeypatch): + """The best checkpoint reuses the bytes just written to 'last'. + + An improving epoch used to pickle *and fsync* the identical state dict a + second time, doubling both the serialization CPU and the checkpoint bytes + pushed at the shared filesystem. Copying the file that was just committed + costs one read (usually from page cache) and one write instead. + """ + mgr, model = _make_manager(tmp_path) + saves = _count_torch_saves(monkeypatch) + + assert mgr.save_checkpoint(epoch=1, val_loss_avg=0.5) is True + assert len(saves) == 1 + assert saves == [str(mgr.last_ckpt_path) + f".tmp.{os.getpid()}"] + + # A second improving epoch: still exactly one serialization. + saves.clear() + assert mgr.save_checkpoint(epoch=2, val_loss_avg=0.25) is True + assert len(saves) == 1 + + # A non-improving epoch writes 'last' only and leaves 'best' alone. + best_bytes = mgr.best_ckpt_path.read_bytes() + saves.clear() + assert mgr.save_checkpoint(epoch=3, val_loss_avg=0.9) is False + assert len(saves) == 1 + assert mgr.best_ckpt_path.read_bytes() == best_bytes + + # And the best checkpoint is still a real, loadable checkpoint holding the + # epoch-2 state -- byte-identical to the 'last' file it was copied from. + monkeypatch.undo() + best = torch.load(mgr.best_ckpt_path, map_location="cpu", weights_only=False) + assert best["epoch"] == 2 + assert best["val_loss_avg"] == pytest.approx(0.25) + for name, tensor in model.state_dict().items(): + assert torch.equal(best["model_state_dict"][name], tensor) + + +def test_best_copy_failure_is_not_silent(tmp_path, monkeypatch): + """A failed best-checkpoint copy surfaces and leaves no partial file.""" + mgr, _ = _make_manager(tmp_path) + + def boom(*args, **kwargs): + raise OSError("[Errno 28] No space left on device") + + monkeypatch.setattr("ScaFFold.utils.checkpointing.shutil.copyfileobj", boom) + + with pytest.raises(CheckpointSaveError, match="No space left on device"): + mgr.save_checkpoint(epoch=1, val_loss_avg=0.5) + + assert not mgr.best_ckpt_path.exists() + assert list(tmp_path.glob("checkpoint_*.tmp.*")) == [] + + # --------------------------------------------------------------------------- # VA-1/VA-2/VA-3 -- the remaining rank-0 filesystem windows are fenced # From 913c92c57b518348d0986c358a12c92ad0be863a Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 17:29:52 -0700 Subject: [PATCH 45/62] Take the warmup training-state snapshot on the host snapshot_training_state cloned model and optimizer state device-to-device, holding ~3x parameter bytes of accelerator memory across the whole warmup phase (and warmup's own memory peak). It now copies to CPU; load_state_dict puts the state back on each parameter's device on restore. R09 --- ScaFFold/utils/checkpointing.py | 47 ++++++++------ tests/test_checkpointing.py | 109 ++++++++++++++++++++++++++++++++ 2 files changed, 135 insertions(+), 21 deletions(-) diff --git a/ScaFFold/utils/checkpointing.py b/ScaFFold/utils/checkpointing.py index 84ed96a..5a04ecd 100644 --- a/ScaFFold/utils/checkpointing.py +++ b/ScaFFold/utils/checkpointing.py @@ -295,17 +295,29 @@ def finalize_saves(self) -> None: self._raise_save_error(error) def snapshot_training_state(self) -> Dict[str, Any]: - """Capture mutable in-memory training state without writing a checkpoint.""" + """Capture mutable in-memory training state without writing a checkpoint. + + The copy is taken on the *host*. Cloning device-to-device instead would + hold roughly 3x parameter bytes of accelerator memory (the weights plus + Adam's two moment buffers) for as long as the snapshot lives -- which is + the whole of warmup, precisely where the run establishes its peak + memory. ``restore_training_state`` puts the state back on the model's + own device, so the detour is invisible to callers. + """ model_ref = self.model.module if hasattr(self.model, "module") else self.model return { - "model_state_dict": self._clone_state_dict(model_ref.state_dict()), - "optimizer_state_dict": self._clone_state_dict(self.optimizer.state_dict()) + "model_state_dict": self._transfer_dict_to_cpu(model_ref.state_dict()), + "optimizer_state_dict": self._transfer_dict_to_cpu( + self.optimizer.state_dict() + ) if self.optimizer else None, - "scheduler_state_dict": self._clone_state_dict(self.scheduler.state_dict()) + "scheduler_state_dict": self._transfer_dict_to_cpu( + self.scheduler.state_dict() + ) if self.scheduler else None, - "grad_scaler_state_dict": self._clone_state_dict( + "grad_scaler_state_dict": self._transfer_dict_to_cpu( self.grad_scaler.state_dict() ) if self.grad_scaler @@ -315,7 +327,13 @@ def snapshot_training_state(self) -> Dict[str, Any]: } def restore_training_state(self, snapshot: Dict[str, Any]) -> None: - """Restore an in-memory training snapshot.""" + """Restore an in-memory training snapshot. + + The snapshot is host-resident (see ``snapshot_training_state``); the + ``load_state_dict`` calls below copy into the live parameters and move + optimizer state onto each parameter's device, so the restored state + ends up exactly where it started. + """ model_ref = self.model.module if hasattr(self.model, "module") else self.model model_ref.load_state_dict(snapshot["model_state_dict"]) @@ -696,8 +714,8 @@ def _transfer_dict_to_cpu(self, obj): ``Tensor.cpu()`` is a no-op for tensors already on CPU (it returns the same object), so CPU-resident state must be cloned explicitly; - otherwise the async writer would serialize tensors the training loop - keeps mutating in place. + otherwise the async writer -- or the warmup snapshot -- would keep an + alias of tensors the training loop mutates in place. """ if torch.is_tensor(obj): t = obj.detach() @@ -711,19 +729,6 @@ def _transfer_dict_to_cpu(self, obj): else: return obj - def _clone_state_dict(self, obj): - """Recursively clone tensors so in-memory snapshots are isolated.""" - if torch.is_tensor(obj): - return obj.detach().clone() - elif isinstance(obj, dict): - return {k: self._clone_state_dict(v) for k, v in obj.items()} - elif isinstance(obj, list): - return [self._clone_state_dict(v) for v in obj] - elif isinstance(obj, tuple): - return tuple(self._clone_state_dict(v) for v in obj) - else: - return obj - def _quarantine_corrupt(self, path): """Rename an unreadable checkpoint aside so it is not retried on the next restart (a persistent corrupt 'last' would otherwise break every diff --git a/tests/test_checkpointing.py b/tests/test_checkpointing.py index 82ab353..5d7afa7 100644 --- a/tests/test_checkpointing.py +++ b/tests/test_checkpointing.py @@ -705,6 +705,115 @@ def test_cpu_tensors_cloned(tmp_path): assert torch.equal(snapshot["nested"][0], torch.ones(3)) +# --------------------------------------------------------------------------- +# R09 -- the warmup snapshot is host-resident, not a device-side copy +# --------------------------------------------------------------------------- + + +def _all_tensors(obj): + """Yield every tensor reachable from a snapshot payload.""" + if torch.is_tensor(obj): + yield obj + elif isinstance(obj, dict): + for value in obj.values(): + yield from _all_tensors(value) + elif isinstance(obj, (list, tuple)): + for value in obj: + yield from _all_tensors(value) + + +class _StubOptimizer: + """Stand-in exposing only the state_dict the snapshot reads.""" + + def __init__(self, state): + self._state = state + + def state_dict(self): + return self._state + + +def test_snapshot_holds_no_device_resident_tensors(tmp_path): + """``snapshot_training_state`` copies to the host, not device-to-device. + + The snapshot used to clone model and optimizer state on whatever device + they lived on, pinning ~3x parameter bytes of accelerator memory (model + clone + Adam's two moment buffers) from before the first warmup batch + until the restore -- i.e. straight across warmup's own peak, which is + where a memory-marginal configuration OOMs. Device tensors are simulated + with FakeTensorMode so this runs without a GPU. + """ + from torch._subclasses.fake_tensor import FakeTensorMode + + mgr, _ = _make_manager(tmp_path) + + with FakeTensorMode(): + mgr.model = torch.nn.Linear(8, 4, device="cuda") + mgr.optimizer = _StubOptimizer( + { + "state": { + 0: { + "step": torch.zeros(1, device="cuda"), + "exp_avg": torch.zeros(4, 8, device="cuda"), + "exp_avg_sq": torch.zeros(4, 8, device="cuda"), + } + }, + "param_groups": [{"params": [0], "lr": 0.1}], + } + ) + snapshot = mgr.snapshot_training_state() + + devices = {t.device.type for t in _all_tensors(snapshot)} + + assert devices == {"cpu"}, f"snapshot kept device-resident tensors: {devices}" + + +def test_snapshot_restore_round_trips_values_and_devices(tmp_path): + """A host-resident snapshot still restores exact state on its own device. + + ``load_state_dict`` copies into the live parameters and moves optimizer + state back to each parameter's device, so nothing downstream has to know + the snapshot took a detour through the host. + """ + mgr, model = _make_manager(tmp_path) + optimizer = mgr.optimizer + + # One applied step so the optimizer carries real per-parameter state. + model(torch.randn(2, 8)).sum().backward() + optimizer.step() + optimizer.zero_grad(set_to_none=True) + + reference_params = {k: v.detach().clone() for k, v in model.state_dict().items()} + reference_state = { + pid: {k: v.detach().clone() for k, v in state.items() if torch.is_tensor(v)} + for pid, state in optimizer.state_dict()["state"].items() + } + param_devices = {k: v.device for k, v in model.state_dict().items()} + + snapshot = mgr.snapshot_training_state() + + # Warmup-shaped mutation: more steps, then roll back. + for _ in range(3): + model(torch.randn(2, 8)).sum().backward() + optimizer.step() + optimizer.zero_grad(set_to_none=True) + assert not torch.equal(model.state_dict()["weight"], reference_params["weight"]) + + mgr.restore_training_state(snapshot) + + for name, tensor in model.state_dict().items(): + assert torch.equal(tensor, reference_params[name]), name + assert tensor.device == param_devices[name], name + restored_state = optimizer.state_dict()["state"] + for pid, state in reference_state.items(): + for key, value in state.items(): + assert torch.equal(restored_state[pid][key], value), (pid, key) + for group, param in zip(optimizer.param_groups, model.parameters()): + del group + for value in optimizer.state[param].values(): + if torch.is_tensor(value) and value.dim() > 0: + assert value.device == param.device + + # --------------------------------------------------------------------------- # F50 -- a final checkpoint is written when the run exits between intervals # --------------------------------------------------------------------------- From a4e9e047222490e8c00f233a996857b71d792b0b Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 17:31:33 -0700 Subject: [PATCH 46/62] Cache the parsed category IFS parameters across a category's instances instance.main re-ran np.genfromtxt on the same category CSV for every (category, instance) work item -- 145 identical shared-filesystem reads per category. Work items for a category are contiguous in the block partition, so a one-entry cache gives one parse per category per rank. R36 --- ScaFFold/datagen/instance.py | 24 ++++++++++---- tests/datagen/test_artifacts.py | 57 +++++++++++++++++++++++++++++++++ 2 files changed, 75 insertions(+), 6 deletions(-) diff --git a/ScaFFold/datagen/instance.py b/ScaFFold/datagen/instance.py index de8eb0d..ea84a91 100644 --- a/ScaFFold/datagen/instance.py +++ b/ScaFFold/datagen/instance.py @@ -346,14 +346,26 @@ def main(config: Config): start_time = time.time() + # One-entry cache of the most recently parsed category CSV. The work list + # is built category-major and block-sliced, so every rank's items for a + # given category are contiguous: a single entry is enough to turn the + # per-item re-parse (145 identical reads of the same small file off the + # shared filesystem, per category) into one parse per category per rank. + # ``generate_instance_points`` copies before scaling, so sharing the parsed + # array across instances cannot leak weights from one item into the next. + cached_category = None + params = None + for i, category_instance_pair in enumerate(instances_to_generate_for_this_rank): category, instance = category_instance_pair - category_IFS_params = IFS_param_csv_names[category] - params = np.genfromtxt( - f"{fracts_read_dir}/{category_IFS_params}", - dtype=DEFAULT_NP_DTYPE, - delimiter=",", - ) + if category != cached_category: + category_IFS_params = IFS_param_csv_names[category] + params = np.genfromtxt( + f"{fracts_read_dir}/{category_IFS_params}", + dtype=DEFAULT_NP_DTYPE, + delimiter=",", + ) + cached_category = category weights = weights_all[instance] # Generate a validated, weighted point cloud. Weighting can turn a diff --git a/tests/datagen/test_artifacts.py b/tests/datagen/test_artifacts.py index c1c38e3..f07bd3a 100644 --- a/tests/datagen/test_artifacts.py +++ b/tests/datagen/test_artifacts.py @@ -234,6 +234,63 @@ def test_resume_rejects_truncated(tmp_path): assert np.isfinite(np.load(victim)).all() +# --------------------------------------------------------------------------- +# R36: a category's IFS parameters are parsed once, not once per instance +# --------------------------------------------------------------------------- + + +def test_category_params_parsed_once_per_category(tmp_path, monkeypatch): + """``main`` parses each category CSV once per rank, not once per work item. + + The parse used to sit inside the per-item loop, so a full generation read + and re-parsed the same small CSV 145 times per category off the shared + filesystem. Work items for a category are contiguous in the block + partition, so a one-entry cache collapses that to one parse per category. + """ + fract_base = tmp_path / "fractals" + point_num = 60 + n_categories = 2 + missing_per_category = 3 + + config = _make_config(fract_base, point_num=point_num) + config.n_categories = n_categories + + param_dir = Path(layout.category_param_dir(config)) + param_dir.mkdir(parents=True) + instance_root = Path(layout.instance_dir(config)) + rng = np.random.default_rng(0) + for category in range(n_categories): + np.savetxt( + param_dir / f"{category:06d}.csv", _contractive_params(), delimiter="," + ) + # Pre-seed all but a few instances so the run stays fast; the ones left + # missing are what the loop (and the parse) actually iterates over. + inst_dir = instance_root / f"{category:06d}" + inst_dir.mkdir(parents=True) + for i in range(missing_per_category, 145): + np.save(inst_dir / f"{category:06d}_{i:04d}.npy", rng.random((10, 3))) + + parses = [] + real_genfromtxt = np.genfromtxt + + def counting_genfromtxt(fname, *args, **kwargs): + parses.append(Path(str(fname)).name) + return real_genfromtxt(fname, *args, **kwargs) + + monkeypatch.setattr(inst.np, "genfromtxt", counting_genfromtxt) + inst.main(config) + + category_parses = [name for name in parses if name[0].isdigit()] + # Every missing instance was generated ... + for category in range(n_categories): + for i in range(missing_per_category): + assert ( + instance_root / f"{category:06d}" / f"{category:06d}_{i:04d}.npy" + ).exists() + # ... from n_categories parses, not one per (category, instance) item. + assert sorted(category_parses) == ["000000.csv", "000001.csv"] + + # --------------------------------------------------------------------------- # F62: mask scanner requires exactly one file per id # --------------------------------------------------------------------------- From be4a01b891351bf031830893b5f728b8efe62674 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 17:43:40 -0700 Subject: [PATCH 47/62] Warm the ragged final batch shapes during warmup Warmup ran only leading, full-size batches, so with local_batch_size>1 and an indivisible shard the epoch's partial batch met cuDNN/MIOpen for the first time inside the first timed epoch -- a measured 95 s stall in epoch_duration, the FOM denominator. Warmup now runs one extra iteration per distinct ragged size (agreed across ranks, since validation shards are unpadded) inside the existing snapshot/restore envelope. R41 --- ScaFFold/utils/trainer.py | 69 ++++++++++++++++++++ tests/test_perf_hotpath.py | 127 +++++++++++++++++++++++++++++++++++++ 2 files changed, 196 insertions(+) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 3498cc1..9e53ecc 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -713,6 +713,71 @@ def _sync_gather_minibatch_timer(self, minibatch_events): minibatch_time_s = statistics.median(minibatch_times.cpu().tolist()) return minibatch_time_s + def _warmup_ragged_batches(self, batch): + """Warm the narrower final batch each loader ends its epoch with. + + Warmup runs only the *leading* batches of the train loader, which are + all ``local_batch_size`` wide, and neither loader drops its last batch. + So when a rank's sample count is not a multiple of the batch size, the + final batch of every epoch presents a set of convolution problems + nothing has ever run: with ``cudnn.benchmark`` on, its algorithm search + (MIOpen find) then happens inside the *first timed epoch* -- a one-off + stall measured at 95 s that lands in ``epoch_duration``, i.e. straight + in the FOM denominator, while staying invisible to the per-minibatch + timer (which excludes partial batches). + + One extra iteration per distinct ragged size fixes that, cut from a + batch warmup already fetched so no additional I/O is needed. The + validation shapes are covered by the same (training) step: the forward + convolutions are what validation shares, and this runs inside warmup's + snapshot/restore envelope, so the extra step cannot affect training. + ``local_batch_size = 1`` never has a ragged batch and does no extra work. + + The set of sizes is agreed across ranks first. The training shards are + padded to equal length, so their remainder is already identical + everywhere, but validation is sharded *unpadded* on purpose (F-series: + an unbiased SUM-reduced metric), so per-rank counts -- and their + remainders -- differ. A rank that decided on its own would run a step + its peers did not, and the collectives inside that step (the gradient + all-reduce, the sharded loss reductions) would deadlock. + """ + if batch is None: + return + local_batch_size = self.config.local_batch_size + available = batch["image"].shape[0] + + ragged_sizes = {len(self.train_sampler) % local_batch_size} + local_val_ragged = torch.tensor( + [len(self.val_sampler) % local_batch_size], device=self.device + ) + gathered_val_ragged = [ + torch.empty_like(local_val_ragged) for _ in range(self.world_size) + ] + torch.distributed.all_gather(gathered_val_ragged, local_val_ragged) + ragged_sizes.update(int(size.item()) for size in gathered_val_ragged) + + # The batches already run are all ``available`` wide, and ``available`` + # is itself rank-invariant (the padded training shards give every rank + # the same leading batch size), so this loop is identical on all ranks. + warmed = {0, available} + for ragged in sorted(ragged_sizes): + if ragged in warmed: + continue + if ragged > available: + self.log.debug( + f" warmup: cannot build a {ragged}-sample batch from a " + f"{available}-sample batch; skipping that shape" + ) + continue + warmed.add(ragged) + self.log.debug( + f" warmup: running the ragged batch shape ({ragged} samples)" + ) + self._run_training_batch( + {key: value[:ragged] for key, value in batch.items()}, + log_prefix=f"warmup ragged ({ragged}): ", + ) + def warmup(self): """Run warmup iterations before the main training loop.""" warmup_batches = self.config.warmup_batches @@ -735,10 +800,12 @@ def warmup(self): self.optimizer.zero_grad(set_to_none=True) try: + last_batch = None for batch_idx, batch in enumerate(self.train_loader): if batch_idx >= max_batches: break + last_batch = batch self._run_training_batch( batch, log_prefix="warmup: ", @@ -749,6 +816,8 @@ def warmup(self): f" warmup: batch {batch_idx} completed in {batch_t_end - start_warmup} seconds" ) + self._warmup_ragged_batches(last_batch) + self.val_loader.sampler.set_epoch(0) if max_val_batches > 0: diff --git a/tests/test_perf_hotpath.py b/tests/test_perf_hotpath.py index 840bc3d..0eb35d0 100644 --- a/tests/test_perf_hotpath.py +++ b/tests/test_perf_hotpath.py @@ -188,3 +188,130 @@ def test_training_batch_gathers_mem_only_first_batch(): # first-batch predicate. assert "gather_mem_stats=True" not in src assert "first_batch" in src + + +# --------------------------------------------------------------------------- +# R41: warmup covers the ragged final batch +# +# Warmup only ever runs the *leading* batches of the train loader, which are +# all local_batch_size wide, and neither loader drops its last batch. When a +# rank's sample count is not a multiple of the batch size, the narrower final +# batch is therefore a set of convolution shapes nothing has warmed, and with +# cudnn.benchmark on the algorithm search for it runs inside the first *timed* +# epoch (measured: 95 s) -- straight into epoch_duration, the FOM denominator. +# --------------------------------------------------------------------------- + + +def _stub_warmup_steps(trainer, monkeypatch, *, mutate=False): + """Record the batch sizes warmup runs; optionally mutate model state. + + The real training step is hardwired through DistConv and cannot run on the + CPU ``ps=None`` fixture, so the step itself is stubbed; what is under test + is which batches warmup feeds it. + """ + import ScaFFold.utils.trainer as tr + + sizes = [] + + def fake_training_batch(batch, **kwargs): + sizes.append(int(batch["image"].shape[0])) + if mutate: + with torch.no_grad(): + for param in trainer.model.parameters(): + param.add_(1.0) + return int(batch["image"].shape[0]), torch.tensor(0.0), torch.tensor(0.0) + + monkeypatch.setattr(trainer, "_run_training_batch", fake_training_batch) + monkeypatch.setattr(tr, "evaluate", lambda *args, **kwargs: (0.0, 0.0)) + return sizes + + +def test_warmup_covers_the_ragged_train_batch(tiny_trainer, monkeypatch): + # 5 local training samples at local_batch_size 2: every epoch ends with a + # 1-sample batch that the leading warmup batches never present. + trainer = tiny_trainer( + n_train=5, + n_val=4, + config_overrides={"local_batch_size": 2, "warmup_batches": 2}, + ) + sizes = _stub_warmup_steps(trainer, monkeypatch) + + trainer.warmup() + + assert sizes == [2, 2, 1] + + +def test_warmup_adds_no_extra_batch_when_shards_divide_evenly( + tiny_trainer, monkeypatch +): + # local_batch_size 1 can never produce a ragged batch: no extra work. + trainer = tiny_trainer( + n_train=4, + n_val=2, + config_overrides={"local_batch_size": 1, "warmup_batches": 2}, + ) + sizes = _stub_warmup_steps(trainer, monkeypatch) + + trainer.warmup() + + assert sizes == [1, 1] + + +def test_warmup_covers_a_ragged_validation_batch(tiny_trainer, monkeypatch): + # Training divides evenly (4 / 2) but validation does not (3 / 2), so the + # 1-sample shape still has to be warmed. + trainer = tiny_trainer( + n_train=4, + n_val=3, + config_overrides={"local_batch_size": 2, "warmup_batches": 2}, + ) + sizes = _stub_warmup_steps(trainer, monkeypatch) + + trainer.warmup() + + assert sizes == [2, 2, 1] + + +def test_warmup_ragged_sizes_are_agreed_across_ranks(tiny_trainer, monkeypatch): + # Validation is sharded unpadded, so peers can end their epoch with a + # different partial size. Every rank must run the same extra steps or the + # collectives inside them diverge; the peer's remainder (2) is gathered and + # warmed here even though this rank's own is 1. + import torch.distributed as dist + + trainer = tiny_trainer( + n_train=6, + n_val=4, + config_overrides={"local_batch_size": 3, "warmup_batches": 2}, + ) + sizes = _stub_warmup_steps(trainer, monkeypatch) + trainer.world_size = 2 + + def fake_all_gather(tensor_list, tensor, *args, **kwargs): + tensor_list[0].copy_(tensor) + tensor_list[1].fill_(2) + + monkeypatch.setattr(dist, "all_gather", fake_all_gather) + + trainer.warmup() + + assert sizes == [3, 3, 1, 2] + + +def test_warmup_rolls_back_state_including_the_ragged_batch(tiny_trainer, monkeypatch): + # The extra ragged iteration stays inside warmup's snapshot/restore + # envelope, so nothing it touches survives into training. + trainer = tiny_trainer( + n_train=5, + n_val=4, + config_overrides={"local_batch_size": 2, "warmup_batches": 2}, + ) + sizes = _stub_warmup_steps(trainer, monkeypatch, mutate=True) + before = {k: v.detach().clone() for k, v in trainer.model.state_dict().items()} + + trainer.warmup() + + assert sizes == [2, 2, 1] + after = trainer.model.state_dict() + for name, tensor in before.items(): + assert torch.equal(after[name], tensor), name From e6bb5baea0339c832526c38f06a9b0ac149d0406 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 17:43:40 -0700 Subject: [PATCH 48/62] Close the figures standard_viz opens pyplot retained every figure the run-summary plots created, and a sweep calls standard_viz.main in-process once per combination, so figures (and their canvases) piled up for the whole sweep. Each is now closed after its savefig, and unconditionally in a finally since plotting errors are swallowed. R43 --- ScaFFold/viz/standard_viz.py | 19 +++++++++++++--- tests/test_reporting.py | 44 ++++++++++++++++++++++++++++++++++++ 2 files changed, 60 insertions(+), 3 deletions(-) diff --git a/ScaFFold/viz/standard_viz.py b/ScaFFold/viz/standard_viz.py index 50d5beb..d9f2dda 100644 --- a/ScaFFold/viz/standard_viz.py +++ b/ScaFFold/viz/standard_viz.py @@ -27,6 +27,10 @@ def main(config: RunConfig): figures_path = Path(config.run_dir) / "figures" figures_path.mkdir(parents=True, exist_ok=True) + # pyplot keeps a strong reference to every figure until it is closed, and a + # sweep calls this once per combination in the same process, so an unclosed + # figure is retained (canvas included) for the rest of the run. + figures = [] try: epochs = [] train_loss = [] @@ -53,7 +57,7 @@ def main(config: RunConfig): legend_loc = (0, -0.17) # Plot training loss - plt.figure() + figures.append(plt.figure()) plt.plot(epochs, train_loss, label="Train Loss", linewidth=line_thickness) plt.xlabel("Epoch", fontsize=fontsize) plt.ylabel("Train loss", fontsize=fontsize) @@ -63,9 +67,10 @@ def main(config: RunConfig): plt.legend(loc="upper left", bbox_to_anchor=legend_loc, fontsize=legend_fontsize) plt.grid(True, axis="y") plt.savefig(figures_path / "train_loss.png", dpi=300, bbox_inches="tight") + plt.close(figures[-1]) # Plot validation dice - plt.figure() + figures.append(plt.figure()) plt.plot(epochs, val_dice, label="Val Dice Score", linewidth=line_thickness) plt.xlabel("Epoch", fontsize=fontsize) plt.ylabel("Val dice score", fontsize=fontsize) @@ -74,10 +79,11 @@ def main(config: RunConfig): plt.legend(loc="upper left", bbox_to_anchor=legend_loc, fontsize=legend_fontsize) plt.grid(True, axis="y") plt.savefig(figures_path / "val_dice.png", dpi=300, bbox_inches="tight") + plt.close(figures[-1]) # Plot validation loss if available if val_loss: - plt.figure() + figures.append(plt.figure()) plt.plot(epochs, val_loss, label="Val Loss", linewidth=line_thickness) plt.xlabel("Epoch", fontsize=fontsize) plt.ylabel("Val loss", fontsize=fontsize) @@ -86,5 +92,12 @@ def main(config: RunConfig): plt.legend(loc="upper left", bbox_to_anchor=legend_loc, fontsize=legend_fontsize) plt.grid(True, axis="y") plt.savefig(figures_path / "val_loss.png", dpi=300, bbox_inches="tight") + plt.close(figures[-1]) except Exception as e: logger.error(f"Failed to generate figures: {e}") + finally: + # Errors here are logged and swallowed, so the close has to be + # unconditional: a failure between figure() and savefig() would + # otherwise leak exactly the figure nobody goes looking for. + for figure in figures: + plt.close(figure) diff --git a/tests/test_reporting.py b/tests/test_reporting.py index 4a462ca..af50a04 100644 --- a/tests/test_reporting.py +++ b/tests/test_reporting.py @@ -57,6 +57,50 @@ def test_figures_dir_idempotent(self, tmp_path): assert (run_dir / "figures" / "train_loss.png").exists() +class TestFigureLifetime: + """R43: standard_viz closes every figure it opens. + + ``worker.main`` calls ``standard_viz.main`` in-process once per sweep + combination, and pyplot keeps a strong reference to every unclosed figure, + so the canvases (and their Figure/Axes/Line objects) accumulate for the + whole sweep -- 36 live figures / 42 MiB after 12 combinations, plus + matplotlib's max_open_warning from the seventh on. + """ + + def _config(self, tmp_path, name): + run_dir = tmp_path / name + run_dir.mkdir() + (run_dir / "train_stats.csv").write_text( + "epoch,overall_loss,val_dice,val_loss_avg\n1,0.9,0.40,0.8\n2,0.5,0.70,0.4\n" + ) + return SimpleNamespace( + run_dir=str(run_dir), vol_size=32, n_categories=5, unet_layers=2 + ) + + def test_no_figures_left_open(self, tmp_path): + """Repeated calls (a sweep) leave no figure behind.""" + plt.close("all") + for i in range(3): + standard_viz.main(self._config(tmp_path, f"run{i}")) + assert plt.get_fignums() == [], f"figures leaked after call {i}" + + def test_no_figures_left_open_when_plotting_fails(self, tmp_path, monkeypatch): + """A failure between figure() and savefig() does not strand a figure. + + ``main`` logs and swallows plotting errors, so without an unconditional + close the leak survives exactly the case it is hardest to notice. + """ + plt.close("all") + + def boom(*args, **kwargs): + raise OSError("[Errno 28] No space left on device") + + monkeypatch.setattr(plt, "savefig", boom) + standard_viz.main(self._config(tmp_path, "failing_run")) + + assert plt.get_fignums() == [] + + class TestDiceFigure: """F70: Validation Dice figure saved as val_dice.png, not val_loss.png.""" From cb1a22c8ebed031e17279812835485dc223a3453 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 17:43:41 -0700 Subject: [PATCH 49/62] Make activation checkpointing reachable from the config The U-Net's use_checkpointing (made correct by F46) had no config key and no production caller, so the memory/compute trade it offers could not be taken. Per the user's decision to wire the feature rather than delete it, an activation_checkpointing 0/1 flag now enables it in worker.py, before the DDP wrap hides the method behind .module. R44 --- ScaFFold/configs/benchmark_default.yml | 1 + ScaFFold/utils/config_utils.py | 19 +++++++++ ScaFFold/worker.py | 9 ++++ tests/test_config.py | 33 ++++++++++++++ tests/test_worker_dist.py | 59 ++++++++++++++++++++++---- 5 files changed, 112 insertions(+), 9 deletions(-) diff --git a/ScaFFold/configs/benchmark_default.yml b/ScaFFold/configs/benchmark_default.yml index c41f4ed..29be50f 100644 --- a/ScaFFold/configs/benchmark_default.yml +++ b/ScaFFold/configs/benchmark_default.yml @@ -39,5 +39,6 @@ loss_freq: 1 # Number of epochs between logging the overal normalize: 1 # Cateogry search normalization parameter group_norm_groups: 8 # Number of groups used by GroupNorm in the UNet blocks. warmup_batches: 64 # How many warmup batches per rank to run before training. +activation_checkpointing: 0 # If 1, recompute UNet block activations during the backward pass instead of storing them: less memory, more compute. ce_weight_sample_fraction: 0.1 # Fraction of training masks to sample when estimating background vs foreground CE weights. dataset_reuse_enforce_commit_id: 0 # Enforce matching commit IDs for dataset reuse. diff --git a/ScaFFold/utils/config_utils.py b/ScaFFold/utils/config_utils.py index 92af20c..60731fe 100644 --- a/ScaFFold/utils/config_utils.py +++ b/ScaFFold/utils/config_utils.py @@ -27,6 +27,19 @@ def require_positive_int(name: str, value: int) -> int: return value +def require_flag(name: str, value) -> bool: + """Validate an on/off config toggle written as 0/1 (or a YAML boolean). + + ``bool(value)`` would quietly accept ``2``, ``-1`` or ``"no"`` (all true), + so a mistyped toggle would enable the feature it was meant to disable. + """ + if isinstance(value, bool): + return value + if isinstance(value, int) and value in (0, 1): + return bool(value) + raise ValueError(f"{name} must be 0, 1, or a boolean; got {value!r}") + + def validate_unet_dims(problem_scale, unet_bottleneck_dim) -> int: """Check that ``problem_scale``/``unet_bottleneck_dim`` describe a real U-Net. @@ -105,6 +118,7 @@ class Config: "normalize", "group_norm_groups", "warmup_batches", + "activation_checkpointing", "ce_weight_sample_fraction", "dataset_reuse_enforce_commit_id", "target_dice", @@ -163,6 +177,7 @@ class Config: "loss_freq", "group_norm_groups", "warmup_batches", + "activation_checkpointing", "ce_weight_sample_fraction", "target_dice", "checkpoint_interval", @@ -252,6 +267,10 @@ def __init__(self, config_dict, strict=True): self.normalize = config_dict["normalize"] self.group_norm_groups = config_dict.get("group_norm_groups", 8) self.warmup_batches = config_dict.get("warmup_batches") + self.activation_checkpointing = require_flag( + "activation_checkpointing", + config_dict.get("activation_checkpointing", 0), + ) self.ce_weight_sample_fraction = config_dict.get( "ce_weight_sample_fraction", 0.1 ) diff --git a/ScaFFold/worker.py b/ScaFFold/worker.py index 2f36b9a..40458ba 100644 --- a/ScaFFold/worker.py +++ b/ScaFFold/worker.py @@ -227,6 +227,15 @@ def main(kwargs_dict: dict = {}): ) model = model.to(device, memory_format=torch.channels_last_3d) + if config.activation_checkpointing: + # Has to happen before the DDP wrap: afterwards the model is only + # reachable as ``model.module``, and the wrapper does not forward the + # method. + log.info( + "activation_checkpointing TRUE -- recomputing block activations in " + "the backward pass instead of storing them" + ) + model.use_checkpointing() # Wrap with DistConvDDP that corrects gradient scaling for dc submesh model = wrap_model_ddp(model, device, ps) # Store ps for use in the training loop diff --git a/tests/test_config.py b/tests/test_config.py index 28cc15b..086a4f5 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -98,6 +98,39 @@ def test_invalid_type_message_names_type(tmp_path): config_utils.load_config(str(path), "bogus") +def test_activation_checkpointing_is_a_real_option(): + """R44: activation checkpointing is reachable from a config, defaulting off. + + The U-Net has always had ``use_checkpointing``, but with no config key and + no caller it could not be turned on: any attempt was rejected as an unknown + key. + """ + cfg = config_utils.Config(dict(BASE)) + assert cfg.activation_checkpointing is False + assert ( + config_utils.Config({**BASE, "activation_checkpointing": 1}) + ).activation_checkpointing is True + assert ( + config_utils.Config({**BASE, "activation_checkpointing": True}) + ).activation_checkpointing is True + assert ( + config_utils.Config({**BASE, "activation_checkpointing": 0}) + ).activation_checkpointing is False + + +@pytest.mark.parametrize("value", [2, -1, "yes", 0.0]) +def test_activation_checkpointing_rejects_non_flag_values(value): + """Anything that is not a 0/1 (or bool) toggle is rejected by name.""" + with pytest.raises(ValueError, match="activation_checkpointing"): + config_utils.Config({**BASE, "activation_checkpointing": value}) + + +def test_activation_checkpointing_documented_in_the_default_config(): + """The shipped config is the parameter reference, so the key lives there.""" + text = (CONFIG_DIR / "benchmark_default.yml").read_text() + assert "activation_checkpointing:" in text + + def test_async_save_is_real_option(): """async_save is an accepted, defaulted option (consumed by the trainer).""" cfg = config_utils.Config(dict(BASE)) diff --git a/tests/test_worker_dist.py b/tests/test_worker_dist.py index e3441ea..b5c6b32 100644 --- a/tests/test_worker_dist.py +++ b/tests/test_worker_dist.py @@ -172,21 +172,21 @@ def fake_ddp(model, parallel_strategy=None, **kwargs): # --------------------------------------------------------------------------- -def test_worker_singleton_smoke(monkeypatch, tiny_config, tiny_dataset): - """worker.main completes end to end as a one-rank gloo job on CPU. +def _run_singleton_worker( + monkeypatch, tiny_config, tiny_dataset, *, port, config_overrides=None +): + """Run ``worker.main`` as a one-rank gloo job on CPU; return (rc, trainer). - ScaFFold always runs distributed; the supported singleton case is a - one-rank launch. The worker initializes the (gloo) process group itself, - builds a real unsharded ParallelStrategy, and tears the group down before - rank-0 post-processing. + Training itself is stubbed (one synthetic epoch row so post-processing has + data), so what this exercises is the worker's own setup path. """ monkeypatch.setenv("MASTER_ADDR", "127.0.0.1") - monkeypatch.setenv("MASTER_PORT", "29513") + monkeypatch.setenv("MASTER_PORT", str(port)) # Force the CPU path so initialize_dist selects gloo: this test must not # depend on a working GPU/NCCL stack. monkeypatch.setattr(torch.cuda, "is_available", lambda: False) - cfg = tiny_config() + cfg = tiny_config(**(config_overrides or {})) kwargs = dict(vars(cfg)) kwargs.update( { @@ -213,7 +213,20 @@ def fake_train(self, profiler=None): monkeypatch.setattr(worker_mod.PyTorchTrainer, "train", fake_train) result = worker_mod.main(kwargs_dict=kwargs) - trainer = seen.get("trainer") + return result, seen.get("trainer") + + +def test_worker_singleton_smoke(monkeypatch, tiny_config, tiny_dataset): + """worker.main completes end to end as a one-rank gloo job on CPU. + + ScaFFold always runs distributed; the supported singleton case is a + one-rank launch. The worker initializes the (gloo) process group itself, + builds a real unsharded ParallelStrategy, and tears the group down before + rank-0 post-processing. + """ + result, trainer = _run_singleton_worker( + monkeypatch, tiny_config, tiny_dataset, port=29513 + ) assert result == 0 assert trainer is not None @@ -229,6 +242,34 @@ def fake_train(self, profiler=None): assert not torch.distributed.is_initialized() +# --------------------------------------------------------------------------- +# R44: the activation-checkpointing config flag reaches the model +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("flag, expected", [(0, False), (1, True)]) +def test_activation_checkpointing_flag_reaches_the_model( + monkeypatch, tiny_config, tiny_dataset, flag, expected +): + """``activation_checkpointing: 1`` turns the U-Net's flag on. + + ``use_checkpointing`` had no caller at all, and it has to be invoked + *before* the DDP wrap: afterwards the model is only reachable through + ``.module``, which is exactly why this asserts on the wrapped model's + inner module. + """ + _result, trainer = _run_singleton_worker( + monkeypatch, + tiny_config, + tiny_dataset, + port=29520 + flag, + config_overrides={"activation_checkpointing": flag}, + ) + + model = getattr(trainer.model, "module", trainer.model) + assert model.checkpointing is expected + + # --------------------------------------------------------------------------- # Local size detection (R23) # --------------------------------------------------------------------------- From c7e3f5d57a156fe35b2be21413ec644eda18f392 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 17:51:48 -0700 Subject: [PATCH 50/62] Compile the UNet GroupNorm on GPU ATen's GroupNorm computes its per-group statistics with a kernel that launches one workgroup per (batch, group) row. At the benchmark defaults (local_batch_size=1, group_norm_groups=8) that is 8 of an MI300A's 228 CUs, and GroupNorm was 86.98 ms of the 186.97 ms scale-7 training step -- 46.5% of wall time, the single largest cost in the step (R38). FastGroupNorm subclasses nn.GroupNorm and routes its forward through a lazily built, process-wide torch.compile(F.group_norm, dynamic=False, fullgraph=True), which hands the reduction to Inductor to tile across the whole device. The module holds the same weight/bias under the same names with no new buffers, so checkpoints round-trip in both directions against a plain nn.GroupNorm build. The compiled path is used only where it is safe, and every rejection falls back to the stock kernel: CPU tensors (the CPU suite pays no compile latency), tensor subclasses such as DistConv's DCTensor (Dynamo cannot trace __torch_dispatch__ wrappers), an already-compiled enclosing region, an explicit SCAFFOLD_GROUPNORM_COMPILE=0, and -- permanently, with one warning -- any exception out of torch.compile. Dynamo's per-function recompile limit is raised from its stock 8: one UNet needs 10 cache entries (5 distinct GroupNorm shapes, each again under no_grad for evaluation), and past the limit Dynamo gives up and silently reverts every GroupNorm to the slow kernel. Measured on one MI300A with review/round2/repros/perf/step_bench.py --layout cl (1x3x128^3, layers=4, bf16 autocast, GradScaler disabled, warm MIOpen db), median of 20 steps, eager numbers from the same build with SCAFFOLD_GROUPNORM_COMPILE=0: step 184.69 ms -> 100.71 ms (1.83x) forward 104.04 ms -> 30.64 ms backward 75.59 ms -> 64.66 ms peak alloc 9.80 GiB -> 8.22 GiB --batch 2 279.58 ms -> 187.61 ms (1.49x; no regression at B>1) torch.profiler over the same step: GroupNorm 86.98 ms/step (46.5% of wall) -> 6.56 ms/step (6.7% of a 97.91 ms step). Compilation is one-time: the first compiled step takes 15.15 s with a cold Inductor cache and produces 5 graphs, after which 12 steps at 97.4 ms produce none; the first no_grad forward adds the other 5. The default 64 warmup batches absorb it outside every timed epoch. Determinism: no gate needed. Two separate processes running three fwd+bwd+Adam steps of the scale-7 UNet under the more_determinism settings (use_deterministic_algorithms(True, warn_only=True), cudnn.benchmark=False, fixed seeds) hash bitwise identically with the compiled path, exactly as they do with the eager one, so no config flag is plumbed through. Not yet visible in the default configuration: worker.py wraps activations in DCTensor even at dc_num_shards=[1,1,1], and GroupNorm then keeps the eager kernel. Verified through a worker.py-shaped DistConvDDP harness that all three paths (DCTensor, plain tensors, plain + use_checkpointing) train without error and agree on their losses; the speedup lands as soon as the unsharded wrap is skipped (R39). --- ScaFFold/unet/group_norm.py | 203 +++++++++++++++++ ScaFFold/unet/unet_parts.py | 6 +- tests/test_groupnorm.py | 421 ++++++++++++++++++++++++++++++++++++ 3 files changed, 629 insertions(+), 1 deletion(-) create mode 100644 ScaFFold/unet/group_norm.py create mode 100644 tests/test_groupnorm.py diff --git a/ScaFFold/unet/group_norm.py b/ScaFFold/unet/group_norm.py new file mode 100644 index 0000000..1134924 --- /dev/null +++ b/ScaFFold/unet/group_norm.py @@ -0,0 +1,203 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +"""GroupNorm with a ``torch.compile``d fast path on GPU. + +ATen's GroupNorm computes its per-group statistics with a kernel that launches +one workgroup per ``(batch, group)`` row. At this benchmark's defaults +(``local_batch_size=1``, ``group_norm_groups=8``) that is 8 workgroups, so on a +228-CU MI300A the normalization runs at a small fraction of achievable +bandwidth and dominates the step: measured 87 ms of a 187 ms step (47%) at +scale 7. Compiling the same functional GroupNorm hands the reduction to +Inductor, which tiles it across the whole device; the same measurement then +gives a 184.7 ms step at 100.7 ms, with GroupNorm down to ~7% of it. + +``FastGroupNorm`` is a drop-in ``nn.GroupNorm``: same parameters, same names, +same shapes, same numerics -- only the kernel differs, so checkpoints are +interchangeable in both directions with any other GroupNorm-based build. The +compiled path is used only when it is safe and worthwhile, and every rejection +falls back to stock eager ``F.group_norm``: + +* non-CUDA tensors (the CPU test suite never pays compile latency), +* tensor subclasses such as DistConv's ``DCTensor``, whose ``__torch_dispatch__`` + wrapper Dynamo cannot trace, +* an already-compiled enclosing region (the functional call inlines instead), +* an explicit opt-out via ``SCAFFOLD_GROUPNORM_COMPILE=0``, +* any failure inside ``torch.compile`` -- logged once, then eager forever after. + +Determinism: the compiled kernels are bitwise reproducible. Two separate +processes running three fwd+bwd+Adam steps of the scale-7 UNet under +``more_determinism`` (``use_deterministic_algorithms(True, warn_only=True)``, +``cudnn.benchmark=False``, fixed seeds) hash identically with the compiled path, +exactly as they do with the eager one, so no determinism gate is needed. +""" + +import logging +import os + +import torch +import torch.nn as nn +import torch.nn.functional as F + +logger = logging.getLogger(__name__) + +#: Opt-out (``0``/``false``/``off``/``no``) or explicit opt-in (``1``/``true``/ +#: ``on``/``yes``) for the compiled GroupNorm path. Unset means "on wherever it +#: is safe", which is what every production run wants. +COMPILE_ENV_VAR = "SCAFFOLD_GROUPNORM_COMPILE" + +#: Dynamo caches one entry per distinct guard set on the traced function. A +#: UNet presents one entry per distinct activation shape (5 at scale 7) times +#: grad-enabled/no-grad (training vs. evaluation), i.e. 10 -- above the stock +#: limit of 8, which would silently drop the whole model back to eager mid-run. +#: The traced function is a single ``F.group_norm`` call, so the extra entries +#: cost only their one-time compilation. +_MIN_RECOMPILE_LIMIT = 64 + +# Lazily built on the first eligible forward: importing ScaFFold must not drag +# in Dynamo, and a run that never reaches the GPU must not pay for it. +_compiled_group_norm = None + +# Set once if torch.compile raises; the eager path is then used everywhere. +_compile_failed = False + +# None = decide per tensor; True/False = forced by SCAFFOLD_GROUPNORM_COMPILE or +# by set_compile_enabled(). +_compile_override = None + + +def _env_override(): + """Read ``SCAFFOLD_GROUPNORM_COMPILE``; ``None`` when unset or unparsable.""" + raw = os.environ.get(COMPILE_ENV_VAR) + if raw is None: + return None + value = raw.strip().lower() + if value in ("1", "true", "on", "yes"): + return True + if value in ("0", "false", "off", "no"): + return False + logger.warning( + f"Ignoring unrecognized {COMPILE_ENV_VAR}={raw!r}; " + "expected one of 1/0/true/false/on/off/yes/no" + ) + return None + + +_compile_override = _env_override() + + +def set_compile_enabled(enabled): + """Force the compiled path on (``True``) or off (``False``). + + ``None`` restores the default, which is the environment variable if set and + otherwise "compile wherever it is safe". Forcing it on does not override + the device and tensor-subclass checks -- those are correctness conditions, + not preferences. Returns the previous setting so callers (tests) can + restore it. + """ + global _compile_override + previous = _compile_override + _compile_override = _env_override() if enabled is None else bool(enabled) + return previous + + +def _group_norm(input, num_groups, weight, bias, eps): + """The function Dynamo traces: plain functional GroupNorm, nothing else.""" + return F.group_norm(input, num_groups, weight, bias, eps) + + +def _raise_recompile_limit(): + """Lift Dynamo's per-function recompile cap to cover every UNet GN shape. + + Only ever raises it, so a caller that deliberately set a larger limit keeps + theirs. ``cache_size_limit`` is the older spelling of ``recompile_limit``; + set whichever exists. + """ + config = torch._dynamo.config + for name in ("recompile_limit", "cache_size_limit"): + current = getattr(config, name, None) + if isinstance(current, int) and current < _MIN_RECOMPILE_LIMIT: + setattr(config, name, _MIN_RECOMPILE_LIMIT) + + +def _get_compiled_group_norm(): + """Build (once) the compiled functional GroupNorm shared by every module. + + One compiled callable for the whole model, not one per module: the shapes, + not the instances, are what Dynamo specializes on, and sharing keeps the + 18 GroupNorms of a scale-7 UNet down to 5 compilations. ``dynamic=False`` + keeps the specialized kernels (this benchmark runs fixed shapes); + ``fullgraph=True`` turns anything Dynamo cannot handle into an exception we + catch, rather than a silent graph break that reintroduces the slow kernel. + """ + global _compiled_group_norm + if _compiled_group_norm is None: + _raise_recompile_limit() + _compiled_group_norm = torch.compile(_group_norm, dynamic=False, fullgraph=True) + return _compiled_group_norm + + +def _use_compiled(input): + """Whether this particular input should take the compiled path.""" + if _compile_failed or _compile_override is False: + return False + # Tensor subclasses (DistConv's DCTensor) route their ops through + # __torch_dispatch__, which Dynamo cannot trace; eager keeps the wrapper's + # semantics -- including which of its outputs come back wrapped -- exactly + # as they are today. worker.py wraps activations in DCTensor even at + # dc_num_shards=[1,1,1], so this fast path engages once that wrap is + # skipped for the unsharded case (or whenever the model is driven with + # plain tensors, as the tests and the standalone benchmarks do). + if type(input) is not torch.Tensor: + return False + # CPU GroupNorm is not the bottleneck and compiling it would put a + # multi-second C++ build in front of every unit test. + if not input.is_cuda: + return False + # Already inside a compiled region: let the functional call be inlined. + if torch.compiler.is_compiling(): + return False + return True + + +class FastGroupNorm(nn.GroupNorm): + """``nn.GroupNorm`` that runs its GPU forward through ``torch.compile``. + + Identical state: ``weight``/``bias`` of shape ``(num_channels,)``, no + buffers, so state dicts are interchangeable with plain ``nn.GroupNorm`` + in both directions. + """ + + def forward(self, input): + # super().forward() is the stock kernel; deferring to it keeps the eager + # path identical to nn.GroupNorm's by construction. + if not _use_compiled(input): + return super().forward(input) + global _compile_failed + try: + return _get_compiled_group_norm()( + input, self.num_groups, self.weight, self.bias, self.eps + ) + except Exception as e: + # Compilation is an optimization, never a correctness requirement: + # a broken Inductor/Triton install, an unwritable cache directory or + # an untraceable input must degrade to the stock kernel, not kill a + # multi-node run. GroupNorm is pure, so retrying eagerly is safe. + _compile_failed = True + logger.warning( + f"torch.compile of GroupNorm failed ({type(e).__name__}: {e}); " + "falling back to the eager kernel for the rest of this run. " + f"Set {COMPILE_ENV_VAR}=0 to skip this attempt entirely." + ) + return super().forward(input) diff --git a/ScaFFold/unet/unet_parts.py b/ScaFFold/unet/unet_parts.py index 681e44b..c9e6cb0 100644 --- a/ScaFFold/unet/unet_parts.py +++ b/ScaFFold/unet/unet_parts.py @@ -20,6 +20,8 @@ from ScaFFold.utils.perf_measure import annotate +from .group_norm import FastGroupNorm + _doubleconv_annotate = annotate(fmt="DoubleConv.{}") _down_annotate = annotate(fmt="Down.{}") _up_annotate = annotate(fmt="Up.{}") @@ -31,7 +33,9 @@ def _group_norm(num_groups, num_channels): raise ValueError( f"group_norm_groups={num_groups} must evenly divide num_channels={num_channels}" ) - return nn.GroupNorm(num_groups, num_channels) + # FastGroupNorm is nn.GroupNorm plus a compiled GPU kernel; it holds the + # same parameters under the same names, so checkpoints are unaffected. + return FastGroupNorm(num_groups, num_channels) class DoubleConv(nn.Module): diff --git a/tests/test_groupnorm.py b/tests/test_groupnorm.py new file mode 100644 index 0000000..b1e8e6a --- /dev/null +++ b/tests/test_groupnorm.py @@ -0,0 +1,421 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +"""Tests for the compiled GroupNorm fast path (``ScaFFold.unet.group_norm``). + +The optimization must be invisible everywhere except in the profile: the same +state dict as a stock ``nn.GroupNorm`` model (checkpoints stay interchangeable +in both directions), the same numbers within reduction-order noise, and an +eager fallback for every input the compiled kernel cannot or should not take +(CPU, tensor subclasses such as DistConv's ``DCTensor``, a broken compiler). +""" + +from __future__ import annotations + +import logging + +import pytest +import torch +import torch.nn as nn + +from ScaFFold.unet import group_norm as gn_mod +from ScaFFold.unet.group_norm import FastGroupNorm +from ScaFFold.unet.unet_model import UNet + +_N = 16 +_N_CHANNELS = 3 +_N_CLASSES = 2 +_GROUPS = 8 + + +@pytest.fixture(autouse=True) +def _restore_compile_state(): + """Keep per-test overrides of the module-level compile state contained.""" + previous = gn_mod.set_compile_enabled(None) + failed = gn_mod._compile_failed + yield + gn_mod._compile_override = previous + gn_mod._compile_failed = failed + + +def _make_unet(seed: int, group_norm_cls=None): + """Build the worker.py-shaped UNet, optionally with a different norm class.""" + torch.manual_seed(seed) + if group_norm_cls is None: + return UNet( + n_channels=_N_CHANNELS, + n_classes=_N_CLASSES, + trilinear=False, + layers=2, + group_norm_groups=_GROUPS, + ) + import ScaFFold.unet.unet_parts as parts + + original = parts.FastGroupNorm + parts.FastGroupNorm = group_norm_cls + try: + return UNet( + n_channels=_N_CHANNELS, + n_classes=_N_CLASSES, + trilinear=False, + layers=2, + group_norm_groups=_GROUPS, + ) + finally: + parts.FastGroupNorm = original + + +def _make_input(seed: int = 0, channels: int = _N_CHANNELS, size: int = _N): + generator = torch.Generator().manual_seed(seed) + return torch.randn(1, channels, size, size, size, generator=generator) + + +# --------------------------------------------------------------------------- +# state dict compatibility +# --------------------------------------------------------------------------- + + +def test_state_dict_matches_plain_groupnorm_model(): + """Names, shapes and dtypes must be unchanged from the nn.GroupNorm build. + + A checkpoint written before this optimization has to keep loading, so the + parameter inventory of the model may not shift by even one key. + """ + new_model = _make_unet(seed=0) + old_model = _make_unet(seed=0, group_norm_cls=nn.GroupNorm) + + new_sd = new_model.state_dict() + old_sd = old_model.state_dict() + assert list(new_sd.keys()) == list(old_sd.keys()) + for key in old_sd: + assert new_sd[key].shape == old_sd[key].shape, key + assert new_sd[key].dtype == old_sd[key].dtype, key + # The optimization must not have introduced buffers either. + assert [name for name, _ in new_model.named_buffers()] == [ + name for name, _ in old_model.named_buffers() + ] + + +def test_checkpoint_round_trip_both_directions(tmp_path): + """An old checkpoint loads into the new model and vice versa, strict=True. + + Both directions matter: runs resumed onto the new code must accept old + checkpoints, and checkpoints written by the new code must stay readable by + anything still building plain ``nn.GroupNorm`` (e.g. an older analysis + script). After each load the two models must agree bit for bit. + """ + new_model = _make_unet(seed=0) + old_model = _make_unet(seed=1, group_norm_cls=nn.GroupNorm) + x = _make_input(seed=3) + + old_path = tmp_path / "old.pth" + torch.save({"model_state_dict": old_model.state_dict()}, old_path) + loaded = torch.load(old_path, weights_only=True) + missing = new_model.load_state_dict(loaded["model_state_dict"], strict=True) + assert not missing.missing_keys and not missing.unexpected_keys + + new_model.eval() + old_model.eval() + with torch.no_grad(): + assert torch.equal(new_model(x), old_model(x)) + + # ... and the reverse: new checkpoint into the plain-GroupNorm model. + fresh_new = _make_unet(seed=2) + new_path = tmp_path / "new.pth" + torch.save({"model_state_dict": fresh_new.state_dict()}, new_path) + reloaded = torch.load(new_path, weights_only=True) + result = old_model.load_state_dict(reloaded["model_state_dict"], strict=True) + assert not result.missing_keys and not result.unexpected_keys + + fresh_new.eval() + with torch.no_grad(): + assert torch.equal(old_model(x), fresh_new(x)) + + +def test_unet_uses_fast_group_norm(): + """Every norm layer in the model is the fast one -- no half-converted build.""" + model = _make_unet(seed=0) + norms = [m for m in model.modules() if isinstance(m, nn.GroupNorm)] + assert norms, "UNet should contain GroupNorm layers" + assert all(isinstance(m, FastGroupNorm) for m in norms) + + +# --------------------------------------------------------------------------- +# CPU behavior: identical numerics, and no compilation at all +# --------------------------------------------------------------------------- + + +def test_cpu_output_bit_identical_to_eager(): + """On CPU the fast module is literally the stock kernel, so bits must match.""" + fast = FastGroupNorm(_GROUPS, 64) + plain = nn.GroupNorm(_GROUPS, 64) + with torch.no_grad(): + plain.weight.copy_(fast.weight) + plain.bias.copy_(fast.bias) + x = _make_input(seed=5, channels=64, size=8) + assert torch.equal(fast(x), plain(x)) + + +def test_cpu_never_invokes_torch_compile(monkeypatch): + """The CPU unit suite must not pay Inductor's compile latency. + + Guards the ``input.is_cuda`` check: if it ever regresses, a CPU-only test + run would start building C++ kernels for every GroupNorm shape. + """ + calls = [] + + def _boom(*a, **kw): + calls.append(a) + raise AssertionError("torch.compile must not be called for CPU tensors") + + monkeypatch.setattr(torch, "compile", _boom) + monkeypatch.setattr(gn_mod, "_compiled_group_norm", None) + gn_mod.set_compile_enabled(True) # even when explicitly forced on + + model = _make_unet(seed=0) + with torch.no_grad(): + model(_make_input(seed=6)) + assert not calls + + +def test_tensor_subclass_input_stays_eager(): + """DistConv wraps activations in a ``__torch_dispatch__`` tensor subclass. + + Dynamo cannot trace those wrappers, so the predicate must reject anything + that is not exactly ``torch.Tensor`` before a compile is attempted. + """ + + class _Wrapper(torch.Tensor): + pass + + plain = torch.randn(1, 8, 4, 4, 4) + assert gn_mod._use_compiled(plain) is False # CPU + assert gn_mod._use_compiled(plain.as_subclass(_Wrapper)) is False + + +def test_compile_failure_falls_back_to_eager(monkeypatch, caplog): + """A broken compiler degrades to the stock kernel instead of killing the run. + + Simulated by making the compiled callable raise; the module must return the + eager result, warn once, and stop trying for the rest of the process. + """ + + def _raises(*args, **kwargs): + raise RuntimeError("simulated Inductor failure") + + monkeypatch.setattr(gn_mod, "_use_compiled", lambda _input: True) + monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _raises) + gn_mod._compile_failed = False + + fast = FastGroupNorm(_GROUPS, 64) + x = _make_input(seed=7, channels=64, size=8) + with caplog.at_level(logging.WARNING, logger=gn_mod.__name__): + out = fast(x) + assert torch.equal( + out, nn.functional.group_norm(x, _GROUPS, fast.weight, fast.bias) + ) + assert any("falling back to the eager kernel" in r.message for r in caplog.records) + assert gn_mod._compile_failed is True + # Latched off: the predicate now refuses even a would-be eligible tensor. + monkeypatch.undo() + assert gn_mod._use_compiled(torch.randn(1, 8, 4, 4, 4)) is False + + +@pytest.mark.parametrize( + "value,expected", + [ + ("0", False), + ("false", False), + ("OFF", False), + ("no", False), + ("1", True), + ("true", True), + ("On", True), + ("yes", True), + ("maybe", None), + ], +) +def test_env_var_controls_the_fast_path(monkeypatch, value, expected): + """``SCAFFOLD_GROUPNORM_COMPILE`` is the documented run-time opt-out.""" + monkeypatch.setenv(gn_mod.COMPILE_ENV_VAR, value) + gn_mod.set_compile_enabled(None) + assert gn_mod._compile_override is expected + + +def test_env_var_unset_means_auto(monkeypatch): + monkeypatch.delenv(gn_mod.COMPILE_ENV_VAR, raising=False) + gn_mod.set_compile_enabled(None) + assert gn_mod._compile_override is None + + +def test_recompile_limit_is_raised_never_lowered(): + """Dynamo's stock cap of 8 is below what one UNet needs. + + A scale-7 UNet presents 5 distinct GroupNorm shapes, and evaluation runs the + same 5 again under ``no_grad`` -- 10 cache entries (measured). Past the cap + Dynamo gives up and every GroupNorm silently reverts to the slow kernel, so + the module raises the limit; it must never lower one a caller chose. + """ + import torch._dynamo + + config = torch._dynamo.config + name = ( + "recompile_limit" if hasattr(config, "recompile_limit") else "cache_size_limit" + ) + original = getattr(config, name) + try: + setattr(config, name, 8) + gn_mod._raise_recompile_limit() + assert getattr(config, name) >= 10 + setattr(config, name, 4096) + gn_mod._raise_recompile_limit() + assert getattr(config, name) == 4096 + finally: + setattr(config, name, original) + + +# --------------------------------------------------------------------------- +# GPU behavior: numerics, single compile, checkpointing +# --------------------------------------------------------------------------- + + +def _assert_close(actual, expected, tol, what): + diff = (actual.float() - expected.float()).abs().max().item() + scale = expected.float().abs().max().item() + assert diff <= tol * max(scale, 1.0), f"{what}: max|diff|={diff:.3e}" + return diff + + +@pytest.mark.gpu +@pytest.mark.parametrize("shape", [(1, 64, 32, 32, 32), (2, 128, 16, 16, 16)]) +@pytest.mark.parametrize("autocast", [False, True]) +def test_gpu_compiled_matches_eager(shape, autocast): + """Compiled forward and gradients match eager, fp32 and under bf16 autocast. + + ``(1, 64, 32^3)`` is the hot production shape (``[1, 64, 128^3]``) at + reduced size -- same channel count and group count, same reduction + structure, small enough for a unit test. Tolerances are loose enough for + reduction-order differences and tight enough to catch a real numerics bug; + observed maxima on MI300A are ~1e-6 relative. + """ + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(11) + x = torch.randn(*shape, device=device, generator=generator) + grad_out = torch.randn(*shape, device=device, generator=generator) + + fast = FastGroupNorm(_GROUPS, shape[1]).to(device) + with torch.no_grad(): + fast.weight.normal_(1.0, 0.1, generator=generator) + fast.bias.normal_(0.0, 0.1, generator=generator) + + def run(compiled): + gn_mod.set_compile_enabled(compiled) + inp = x.clone().requires_grad_(True) + fast.zero_grad(set_to_none=True) + with torch.autocast("cuda", dtype=torch.bfloat16, enabled=autocast): + out = fast(inp) + out.backward(grad_out.to(out.dtype)) + return ( + out.detach(), + inp.grad.detach(), + fast.weight.grad.detach().clone(), + fast.bias.grad.detach().clone(), + ) + + eager = run(False) + compiled = run(True) + assert gn_mod._compiled_group_norm is not None, "compiled path was not taken" + assert not gn_mod._compile_failed + + _assert_close(compiled[0], eager[0], 1e-5, "output") + _assert_close(compiled[1], eager[1], 1e-4, "d_input") + _assert_close(compiled[2], eager[2], 1e-4, "d_weight") + _assert_close(compiled[3], eager[3], 1e-4, "d_bias") + # Autocast policy must be preserved: GroupNorm is an fp32 op, so the + # compiled path may not quietly hand back bf16 activations. + assert compiled[0].dtype == eager[0].dtype + + +@pytest.mark.gpu +def test_gpu_steady_state_does_not_recompile(): + """Fixed shapes must compile once and then never again. + + A recompile inside a timed epoch would show up as a multi-second outlier in + ``epoch_duration`` (and therefore the FOM), so the guard set has to be + stable across steps. + """ + from torch._dynamo.utils import counters + + gn_mod.set_compile_enabled(True) + device = torch.device("cuda") + fast = FastGroupNorm(_GROUPS, 64).to(device) + x = torch.randn(1, 64, 16, 16, 16, device=device, requires_grad=True) + + def step(): + fast.zero_grad(set_to_none=True) + with torch.autocast("cuda", dtype=torch.bfloat16): + out = fast(x) + out.sum().backward() + + step() # first call: compiles + before = counters["stats"]["unique_graphs"] + for _ in range(5): + step() + assert counters["stats"]["unique_graphs"] == before, "recompiled in steady state" + + +@pytest.mark.gpu +def test_gpu_activation_checkpointing_matches_eager(): + """The compiled kernel must survive recompute under use_checkpointing(). + + Non-reentrant checkpointing replays the block's forward inside the backward + pass; a compiled region has to produce the same activations both times or + the gradients silently change. + + Compared as relative L2 error per gradient tensor, because whole-network + agreement is not bitwise even without this change: with cudnn.benchmark on + and bf16 autocast, two eager runs of this model differ by ~4e-3 relative + (measured), and checkpointing on vs. off differs by the same amount. + Measured here: compiled vs. eager 6.4e-3, i.e. the same order as that noise + floor -- while a genuinely wrong kernel would be O(1). + """ + device = torch.device("cuda") + x = _make_input(seed=9).to(device).requires_grad_(True) + tolerance = 5e-2 + + def grads(compiled, checkpointing): + gn_mod.set_compile_enabled(compiled) + model = _make_unet(seed=0).to(device) + if checkpointing: + model.use_checkpointing() + model.zero_grad(set_to_none=True) + with torch.autocast("cuda", dtype=torch.bfloat16): + out = model(x) + out.float().pow(2).mean().backward() + return {n: p.grad.detach().clone() for n, p in model.named_parameters()} + + def assert_agrees(actual, expected, label): + for name in expected: + reference = expected[name].float() + error = (actual[name].float() - reference).norm().item() + relative = error / max(reference.norm().item(), 1e-12) + assert relative < tolerance, f"{label} {name}: rel L2 {relative:.3e}" + + eager = grads(False, True) + compiled = grads(True, True) + compiled_nockpt = grads(True, False) + assert gn_mod._compiled_group_norm is not None, "compiled path was not taken" + assert not gn_mod._compile_failed + assert_agrees(compiled, eager, "checkpointed grad") + assert_agrees(compiled_nockpt, compiled, "grad") From c669b3f747ecc75e2e30a8e01442f40c2a025cc9 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 18:37:39 -0700 Subject: [PATCH 51/62] Post the ragged-size all_gather before the empty-loader return MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A rank whose warmup fetched no batch returned before the ragged-size all_gather, skipping a collective its peers post — the divergence class this round closes elsewhere. Latent today (padded training shards make loader lengths rank-invariant), but the collective pattern must not depend on local loader state. Found by the final verification pass. --- ScaFFold/utils/trainer.py | 14 +++++++++++--- tests/test_perf_hotpath.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 3 deletions(-) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 9e53ecc..b284708 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -741,11 +741,15 @@ def _warmup_ragged_batches(self, batch): its peers did not, and the collectives inside that step (the gradient all-reduce, the sharded loss reductions) would deadlock. """ - if batch is None: - return local_batch_size = self.config.local_batch_size - available = batch["image"].shape[0] + # Agree on the ragged sizes FIRST. The all_gather below is a + # collective: every rank must post it even when its own warmup fetched + # no batch at all (``batch is None``), or its peers block in the + # gather. Batch availability is rank-invariant in practice (padded + # training shards give every rank the same loader length), but the + # collective pattern must not depend on local loader state, so the + # no-batch early return comes after the gather. ragged_sizes = {len(self.train_sampler) % local_batch_size} local_val_ragged = torch.tensor( [len(self.val_sampler) % local_batch_size], device=self.device @@ -756,6 +760,10 @@ def _warmup_ragged_batches(self, batch): torch.distributed.all_gather(gathered_val_ragged, local_val_ragged) ragged_sizes.update(int(size.item()) for size in gathered_val_ragged) + if batch is None: + return + available = batch["image"].shape[0] + # The batches already run are all ``available`` wide, and ``available`` # is itself rank-invariant (the padded training shards give every rank # the same leading batch size), so this loop is identical on all ranks. diff --git a/tests/test_perf_hotpath.py b/tests/test_perf_hotpath.py index 0eb35d0..a683c59 100644 --- a/tests/test_perf_hotpath.py +++ b/tests/test_perf_hotpath.py @@ -298,6 +298,36 @@ def fake_all_gather(tensor_list, tensor, *args, **kwargs): assert sizes == [3, 3, 1, 2] +def test_warmup_ragged_all_gather_is_posted_even_without_a_batch( + tiny_trainer, monkeypatch +): + # The ragged-size agreement is a collective: a rank whose warmup fetched + # no batch at all (empty train loader) must still post the all_gather, or + # its peers block in it. The no-batch early return has to come after the + # gather, even though such a rank then runs no extra step itself. + import torch.distributed as dist + + trainer = tiny_trainer( + n_train=4, + n_val=3, + config_overrides={"local_batch_size": 2, "warmup_batches": 2}, + ) + sizes = _stub_warmup_steps(trainer, monkeypatch) + gathers = [] + + def fake_all_gather(tensor_list, tensor, *args, **kwargs): + gathers.append(int(tensor.item())) + for out in tensor_list: + out.copy_(tensor) + + monkeypatch.setattr(dist, "all_gather", fake_all_gather) + + trainer._warmup_ragged_batches(None) + + assert gathers, "rank skipped the ragged-size all_gather when batch was None" + assert sizes == [] + + def test_warmup_rolls_back_state_including_the_ragged_batch(tiny_trainer, monkeypatch): # The extra ragged iteration stays inside warmup's snapshot/restore # envelope, so nothing it touches survives into training. From c825cc846a20531df438b6c22aab52e219eef750 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 18:37:39 -0700 Subject: [PATCH 52/62] Correct the best-copy perf claim and widen the copy buffer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The R08 comment claimed the copy halves the bytes pushed at the shared filesystem; measured on Lustre the copy is ~1.02x — the real saving is the serialization CPU. Note the redundancy trade-off (best is now a byte copy of last), use a 16 MiB copy buffer to cut syscall count on parallel filesystems, and document that the Dynamo recompile-limit raise clobbers a deliberately smaller limit. Found by the final verification pass. --- ScaFFold/unet/group_norm.py | 5 +++-- ScaFFold/utils/checkpointing.py | 15 +++++++++++---- 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/ScaFFold/unet/group_norm.py b/ScaFFold/unet/group_norm.py index 1134924..06d7407 100644 --- a/ScaFFold/unet/group_norm.py +++ b/ScaFFold/unet/group_norm.py @@ -121,8 +121,9 @@ def _raise_recompile_limit(): """Lift Dynamo's per-function recompile cap to cover every UNet GN shape. Only ever raises it, so a caller that deliberately set a larger limit keeps - theirs. ``cache_size_limit`` is the older spelling of ``recompile_limit``; - set whichever exists. + theirs -- but note the converse: a limit deliberately set *smaller* than + ours is clobbered up to ``_MIN_RECOMPILE_LIMIT``. ``cache_size_limit`` is + the older spelling of ``recompile_limit``; set whichever exists. """ config = torch._dynamo.config for name in ("recompile_limit", "cache_size_limit"): diff --git a/ScaFFold/utils/checkpointing.py b/ScaFFold/utils/checkpointing.py index 5a04ecd..00dfc14 100644 --- a/ScaFFold/utils/checkpointing.py +++ b/ScaFFold/utils/checkpointing.py @@ -663,7 +663,10 @@ def _atomic_copy(src, dst): tmp_path = dst.with_name(f"{dst.name}.tmp.{os.getpid()}") try: with open(src, "rb") as fsrc, open(tmp_path, "wb") as fdst: - shutil.copyfileobj(fsrc, fdst) + # A large buffer keeps the syscall count low on parallel + # filesystems (the default 64 KiB means ~1k read/write pairs + # per 64 MiB checkpoint). + shutil.copyfileobj(fsrc, fdst, length=16 * 1024 * 1024) fdst.flush() os.fsync(fdst.fileno()) os.replace(tmp_path, dst) @@ -688,9 +691,13 @@ def _write_to_disk(cls, state_dict, last_path, best_path, is_best, log): # Save 'last' atomically. cls._atomic_save(state_dict, last_path) # 'best' is byte-identical to the 'last' just committed, so copy - # that file instead of pickling and fsyncing the same state a - # second time (double the serialization CPU and double the bytes - # pushed at the shared filesystem on every improving epoch). + # that file instead of pickling the same state a second time. The + # saving is the serialization CPU (pickle + zip of the full state + # dict); the filesystem traffic is roughly a wash -- the copy + # writes the same bytes and adds a read (measured ~1.02x faster on + # Lustre). Trade-off: 'best' is now a byte copy of 'last', so a + # silently corrupted 'last' write would propagate into 'best' + # rather than being an independent serialization. # # There is no concurrent writer to race: checkpoint writes are # serialized through a single writer -- the caller's thread in sync From 3b8ead857fbd4e376e10c579dbcc8009fe5dae8f Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 18:52:46 -0700 Subject: [PATCH 53/62] Drop the unreachable packaged fractal library The seed-keyed relayout (R29) made the shipped seed-unknown CSVs at ScaFFold/fractals/var0.15/3DIFS_param unreachable under any configuration. Remove them, their package-data glob, the dead Config.library_root (zero readers), and the README claim; libraries regenerate deterministically from the configured seed. Resolves verification item VB-5 per user decision. --- README.md | 4 +++- ScaFFold/fractals/var0.15/3DIFS_param/000000.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000001.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000002.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000003.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000004.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000005.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000006.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000007.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000008.csv | 3 --- ScaFFold/fractals/var0.15/3DIFS_param/000009.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000010.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000011.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000012.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000013.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000014.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000015.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000016.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000017.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000018.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000019.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000020.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000021.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000022.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000023.csv | 3 --- ScaFFold/fractals/var0.15/3DIFS_param/000024.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000025.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000026.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000027.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000028.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000029.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000030.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000031.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000032.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000033.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000034.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000035.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000036.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000037.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000038.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000039.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000040.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000041.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000042.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000043.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000044.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000045.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000046.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000047.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000048.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000049.csv | 2 -- ScaFFold/utils/config_utils.py | 4 ---- pyproject.toml | 1 - 53 files changed, 3 insertions(+), 108 deletions(-) delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000000.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000001.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000002.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000003.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000004.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000005.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000006.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000007.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000008.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000009.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000010.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000011.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000012.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000013.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000014.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000015.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000016.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000017.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000018.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000019.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000020.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000021.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000022.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000023.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000024.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000025.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000026.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000027.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000028.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000029.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000030.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000031.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000032.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000033.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000034.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000035.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000036.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000037.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000038.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000039.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000040.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000041.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000042.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000043.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000044.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000045.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000046.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000047.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000048.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000049.csv diff --git a/README.md b/README.md index 2442494..53d449c 100644 --- a/README.md +++ b/README.md @@ -54,7 +54,9 @@ The model is trained from a random initialization until convergence, which is de 1. If running the benchmark for the first time, or running with different fractal parameters (`n_categories`, `variance_threshold`) than previously, generate fractal classes and instances: `scaffold generate_fractals -c ScaFFold/configs/benchmark_default.yml` - Note that the benchmark ships with an initial set of 50 fractal classes. + Fractal category libraries are generated deterministically from the + configured seed (under `fract_base_dir/var<...>/seed<...>/`) and reused by + later runs with the same seed. 1. Once fractal generation completes, run the benchmark: `torchrun-hpc -N 1 -n 4 --gpus-per-proc 1 $(which scaffold) benchmark -c ScaFFold/configs/benchmark_default.yml` diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000000.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000000.csv deleted file mode 100644 index 824252c..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000000.csv +++ /dev/null @@ -1,2 +0,0 @@ --3.113880658597081297e-01,1.525788435216699490e-01,-1.677513246653659085e-01,9.450327943163072675e-01,1.437558218378187647e-01,-7.781853097899118499e-01,-7.520371219608632529e-01,4.049719188683833515e-01,3.599506321979530910e-01,-9.209462080360120151e-01,9.797064915552313735e-01,2.046581090479335785e-01,5.952758981070249700e-01 --3.273838635851094025e-01,5.796182968623042608e-01,2.204038860996404559e-02,7.254736836236364006e-02,-3.624484794849551772e-01,3.976647887245436941e-01,-6.272009791693362590e-01,2.300328757717320372e-01,1.439593364137665699e-01,-1.758502593004573900e-01,7.577588977727744979e-01,-4.555583614631009137e-01,4.047241018929750300e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000001.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000001.csv deleted file mode 100644 index fca2ed3..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000001.csv +++ /dev/null @@ -1,2 +0,0 @@ -2.473679574324576524e-01,1.597856635526939684e-02,-1.261210131422938474e-01,2.860947005456315750e-01,8.355771386484407426e-01,-1.930246242945548030e-01,-4.801705885775757743e-01,1.904275728954856195e-01,7.806079367730189844e-01,7.953170851152218113e-01,-5.816986628081313171e-01,6.209542113667643193e-01,1.988928060480311955e-01 --7.041010339743891677e-02,-2.454306473738818717e-01,9.405874516413117448e-01,7.268593945019246050e-01,6.070031387640555387e-01,-9.102525915665182765e-02,6.339356801972018118e-01,-1.532413735258475462e-01,3.654495384649214529e-02,-4.532769270980656628e-01,1.964159205402833397e-01,-5.359870555644301593e-01,8.011071939519688323e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000002.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000002.csv deleted file mode 100644 index c3e5385..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000002.csv +++ /dev/null @@ -1,2 +0,0 @@ --9.413772084350502389e-02,5.254718465280536766e-01,-4.808322979877577286e-01,5.811542374432154823e-01,-4.454009637949891687e-01,-4.249149866229446904e-01,-3.000182472995254201e-01,4.022670562303001240e-01,-2.934854435005793682e-01,-6.605442110513501941e-01,-6.253032129726847632e-01,-5.719885961970654353e-01,4.224941885002395647e-01 --4.407723055180619021e-01,1.325088026302621014e-01,1.840254341100278079e-01,-5.508321197309193895e-01,8.464500652955120330e-01,5.041753392202226181e-01,5.651656783277005935e-01,-8.325090493066080732e-01,-1.381507388162375172e-01,5.571341416171868843e-01,-1.025781698153782617e-01,-8.838988798440101657e-01,5.775058114997604353e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000003.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000003.csv deleted file mode 100644 index 18274fd..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000003.csv +++ /dev/null @@ -1,2 +0,0 @@ --1.294466006455852192e-01,1.974941747383867074e-01,4.427934740152306148e-01,2.336055684661912935e-01,1.300312996679606758e-01,8.725504616000674396e-01,-1.902059536103457571e-01,-3.532649574684778582e-01,-8.430890597933748953e-01,-9.596612485558557726e-01,6.321097711736041180e-01,8.135413458612708038e-01,2.489526692860659640e-01 --1.399185307891788188e-01,-1.602045901301751840e-01,4.335897375838684287e-01,-9.384258169419115170e-01,-1.069737283930594085e-01,2.989241164297768982e-01,-6.355754104429354179e-01,2.746690226736152596e-01,2.732139814068013095e-01,9.765182437472603727e-01,-4.228379364440870702e-02,-5.170594086561124403e-01,7.510473307139340360e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000004.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000004.csv deleted file mode 100644 index ac107cb..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000004.csv +++ /dev/null @@ -1,2 +0,0 @@ --7.074987909298180444e-01,-1.415561527827760013e-01,6.453558418406579733e-02,9.717892764870739164e-01,4.692315433286387005e-02,-9.087912295403599572e-02,-7.306123549338172651e-01,5.199024492355210914e-01,3.242137268086431323e-01,-7.685133956354921470e-01,-2.890512330200536439e-01,-4.414459117103588515e-01,1.767437117859010087e-01 --5.783609543921015561e-01,-6.261288719190640784e-01,6.083280289156212106e-01,5.708229102420392387e-01,6.839066129793194282e-01,-3.730802596414783956e-01,4.320233943214566441e-01,-2.520497672481589735e-01,2.188189615840767654e-01,-5.142628557025761271e-01,9.167783772484443539e-01,-9.404699392000455127e-01,8.232562882140990190e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000005.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000005.csv deleted file mode 100644 index 898a388..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000005.csv +++ /dev/null @@ -1,2 +0,0 @@ --6.016222179722698904e-01,-8.568626391040801149e-01,4.339904630487174675e-01,8.566190336460910437e-02,5.955551394131912701e-01,-7.233426894750305536e-02,-3.439913328223502820e-01,-1.295598008765910247e-01,5.054214022018066466e-01,5.403179943497680160e-01,9.887300162494017108e-01,5.065235980037172681e-01,2.036756775286962806e-01 -4.255811510519635910e-02,-5.861700389763935259e-01,1.539188938105566784e-01,7.464882854730132689e-01,9.985987594821947866e-01,-3.852978585836623893e-01,5.848150331480403974e-01,-8.218568908737176049e-01,-4.786801609318307449e-01,-7.997867917133967275e-01,-3.226694643578802424e-01,-9.055189731557820032e-01,7.963243224713036916e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000006.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000006.csv deleted file mode 100644 index 345f2a4..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000006.csv +++ /dev/null @@ -1,2 +0,0 @@ -4.449557690511229957e-02,1.859345105226184458e-02,-5.492709805712474580e-01,5.002339770143411357e-01,6.022059536539468017e-01,-2.581455709557001210e-01,5.384026026512653829e-01,3.525173161797230392e-01,4.535819387452528773e-01,5.054013100483643051e-01,5.534237431467781132e-01,6.486536701245786407e-01,5.893263018191835512e-01 --4.935786839220561717e-01,-3.483405932567602559e-01,-4.920557235894853498e-01,-3.225857219554804090e-01,-5.369237708506735540e-01,-4.285587807484951828e-01,-2.187430959790792606e-02,8.181072307028691704e-02,4.255774677072714507e-01,-8.957260803321442921e-01,-6.285123937957712847e-01,-7.274955330956049959e-01,4.106736981808163378e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000007.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000007.csv deleted file mode 100644 index fa04c8c..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000007.csv +++ /dev/null @@ -1,2 +0,0 @@ -4.567482715684858530e-01,-2.162812897008825619e-01,8.688925426344686898e-01,9.403035945928239769e-01,-4.923377969407702892e-01,9.039413044894315519e-01,-7.007586459777630505e-01,9.958966466641376858e-01,-2.148840978447630334e-02,-8.820192395688601916e-01,-6.962417306340771272e-01,6.989108727983952551e-01,6.175732443919085268e-01 --1.311082673391426034e-01,2.651764139972054846e-01,2.416531271291553207e-01,2.658826889156997719e-01,5.321835857677645887e-01,8.850640913149898648e-01,-4.146324561488827776e-01,6.016104083259397051e-01,-6.006142415160782289e-01,6.902838974364782221e-01,5.080744400285828188e-01,-5.543689403825666773e-01,3.824267556080915287e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000008.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000008.csv deleted file mode 100644 index b017040..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000008.csv +++ /dev/null @@ -1,3 +0,0 @@ -1.955211818468614027e-01,-1.860012909088504252e-02,1.679370276531577666e-01,-2.474658865690229081e-01,2.002748188638228122e-01,-2.296462069177727106e-01,4.272668474314074150e-01,5.518250090890584048e-01,-1.814099098142296640e-01,-7.759223213118924267e-01,2.396552032441381375e-01,-9.984320127889498853e-01,3.158600947042149998e-01 -5.269496943889822038e-01,5.090567464178141766e-01,3.638063575297016961e-01,-9.409697282110314198e-01,-8.838575707448308449e-01,-9.681971946347656122e-02,-6.464087288109372498e-01,-4.699168427008522109e-01,4.879093108230523335e-01,7.225552560446799610e-01,5.544701020084910059e-01,6.687273760328038552e-01,6.074177711483321751e-01 -3.702243567552572223e-01,2.665339551565091281e-01,-7.302570963874586152e-01,7.128226896410159164e-01,-1.412872752126193010e-01,-5.914873767070649713e-01,4.858707230431502655e-01,2.870662392522644879e-01,-8.632888349796041805e-01,8.908557359947888443e-02,-1.967254792508705830e-01,-6.141181327493672182e-01,7.672213414745278348e-02 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000009.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000009.csv deleted file mode 100644 index f56b11b..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000009.csv +++ /dev/null @@ -1,2 +0,0 @@ -2.948060069505098468e-01,6.887599776806021534e-01,7.611229734423119453e-02,4.932702124290413437e-01,-4.882699298136421451e-01,-3.811733224774245254e-01,7.278710433957895631e-01,6.852985410340497463e-01,-7.967518041925278904e-01,6.365195144826076845e-01,7.940485085563173673e-01,-7.048888688479171272e-02,6.538648345727335887e-01 -2.849341248173462571e-01,-5.989665908310131126e-01,-8.440934355375118159e-01,-5.676180605895972953e-01,2.243341626806083511e-02,-4.393691291202828086e-01,-6.561728033554823369e-01,-1.447073895277080080e-01,-3.027716266953806024e-01,-1.488449066054653436e-01,-7.989435645608391479e-01,6.936025509345438156e-01,3.461351654272663003e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000010.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000010.csv deleted file mode 100644 index af8b628..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000010.csv +++ /dev/null @@ -1,2 +0,0 @@ --2.352521358159322951e-01,-8.804830225066977434e-01,2.445692685925464627e-01,-3.804268259383307704e-01,7.002478757426469080e-01,4.949375238071500593e-01,6.340824702532006363e-01,-1.661009284036518707e-01,6.403120148905094844e-02,-4.544636202226781663e-03,-5.614632388117013484e-01,-8.225067911333789894e-01,7.543884465970274178e-01 --6.205195826406055826e-01,2.082688392648848197e-01,7.051714514820610624e-01,6.615150945575909436e-01,-5.430135870340266901e-01,-7.590298999219038389e-01,-2.484973171764268685e-01,4.507062852245327100e-01,-3.971212226677103274e-01,3.603919868974059249e-01,-6.455363862654992513e-01,6.976679002200263380e-01,2.456115534029726655e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000011.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000011.csv deleted file mode 100644 index 60713bc..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000011.csv +++ /dev/null @@ -1,2 +0,0 @@ --2.647713297675924338e-01,4.532537039449799909e-01,9.126597053670870707e-02,6.297942569420966752e-01,-7.970366903049754814e-01,-5.805081819440571778e-01,1.847828454797417752e-01,1.632357583131875955e-01,1.112052944356816120e-01,5.581182508420681199e-01,-1.235235401936412014e-01,5.525957004344514978e-01,3.917940489256083736e-01 --3.328722698643555855e-01,-6.568958305252938779e-01,9.548663670929014025e-02,2.135325385282085264e-01,-7.994872824743104456e-01,3.120617004916548254e-01,-8.149604985800276147e-01,3.155236614340959367e-01,-5.804561978693489888e-01,-3.465348059373873912e-01,8.779303525751596116e-01,-4.964138499509700431e-01,6.082059510743916819e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000012.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000012.csv deleted file mode 100644 index d75a5f8..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000012.csv +++ /dev/null @@ -1,2 +0,0 @@ --3.732047938669280640e-01,7.457174836228177561e-01,-1.134959963046562326e-01,-2.524155789085107404e-01,6.491113757559063835e-01,5.923873502246719269e-01,-4.163840681869901417e-01,-9.955960643681671662e-02,-4.399149502577506254e-01,-3.298124569267795181e-01,-4.480024355377016931e-01,4.831814457889009873e-01,7.022590685295611035e-01 -3.640511591326804908e-01,-8.138787463613117446e-01,6.945993218897195121e-01,1.561110571982835538e-01,1.420189997595979747e-01,8.921447905553927527e-01,8.983211135903812483e-02,-7.873694636782158085e-02,-1.924595656350303052e-01,3.444934594172788245e-01,1.245560132691707622e-01,-2.441271125304933509e-01,2.977409314704388410e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000013.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000013.csv deleted file mode 100644 index bf897a7..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000013.csv +++ /dev/null @@ -1,2 +0,0 @@ --9.477179726636597579e-02,1.411777244118233021e-01,-3.210166902811444345e-01,8.538383221528667022e-02,5.145040185660725296e-01,5.449889235641092178e-01,1.961008681633682471e-01,5.250846593606550705e-01,2.072609251252632845e-01,8.556312962216905404e-01,7.079983983577677886e-01,-5.970939704307753892e-01,1.986178895167198533e-01 -2.340421417737004184e-02,-2.331838675905175684e-01,8.789698166911528165e-01,4.337374001870073492e-01,-6.744896384287986102e-01,6.995573996808193140e-01,-9.371564747554361752e-02,7.122203806164681961e-01,-3.226197608587253463e-01,5.228286702908606642e-01,6.219275677253508494e-01,6.982539942183398907e-01,8.013821104832802300e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000014.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000014.csv deleted file mode 100644 index 6fd2a78..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000014.csv +++ /dev/null @@ -1,2 +0,0 @@ --5.079619751044526232e-01,-8.448208339950418200e-01,3.379361696842797524e-01,-7.424757815754379209e-01,-2.914860368679528246e-01,7.863175574507479393e-01,-3.913127807203062858e-01,-3.395110295018073376e-01,1.729246237110684259e-01,1.155790784528218929e-01,-1.778336752766045414e-01,-8.237750280597011532e-01,8.128837685573842009e-01 --5.053085469842428790e-01,3.413777808104254685e-01,5.072078013440961541e-02,-2.297771208541388166e-01,1.179085091368414773e-01,-5.353419781682577927e-02,3.033551256644697602e-01,-8.399081644307380135e-01,-2.580876839431038849e-01,-2.915499965121404191e-01,-8.274711651504977894e-01,-3.831424144957116251e-01,1.871162314426157436e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000015.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000015.csv deleted file mode 100644 index 040ab5e..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000015.csv +++ /dev/null @@ -1,2 +0,0 @@ -1.959234370454023821e-01,7.195879279286905295e-01,-8.847434444070787496e-01,-2.025052870905292846e-01,-6.099125369957605347e-01,4.364854998130400787e-01,2.605566328535515730e-01,6.345444901137731186e-02,2.845477814208994261e-01,3.682037052176181380e-01,-5.482600452804842206e-01,8.228596647488930493e-01,2.413865640087159981e-01 --3.105662416993948405e-02,1.541091773938909615e-01,9.832144790547923119e-01,2.617717287788325908e-01,2.964779787430120717e-01,2.867701916196363499e-01,-1.595092390377108593e-01,5.184743556024293820e-01,7.055853027879592787e-01,-2.253650882070545869e-02,3.223158737418712061e-01,-3.130584359547923246e-01,7.586134359912839464e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000016.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000016.csv deleted file mode 100644 index 1262706..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000016.csv +++ /dev/null @@ -1,2 +0,0 @@ --1.300116634675734240e-01,1.252286096526100678e-01,2.499311708856801761e-01,7.830409313847326302e-01,-3.485355694096761159e-01,-2.040759407213834642e-01,3.636277174597402073e-01,-9.857803444259161108e-01,4.953546011072127442e-01,-9.229964481663586184e-01,8.261855069165144894e-01,7.432769447647864514e-01,5.045106210869472196e-01 --9.419389020480599672e-01,-2.909144464059969515e-01,2.371749371359630487e-01,-9.140320234271737121e-02,2.578548696646267846e-01,-6.247044474540388581e-01,5.120075657145519710e-01,7.003880575966952016e-01,-7.351851386341190508e-01,-3.580478535291309328e-02,-3.983862417971921754e-01,-8.667134200938420019e-01,4.954893789130526693e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000017.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000017.csv deleted file mode 100644 index 0a04245..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000017.csv +++ /dev/null @@ -1,2 +0,0 @@ --3.510560689158552794e-01,-2.754404723500034624e-01,-2.684881793088436108e-01,-1.948162421580679204e-01,1.792154998944521793e-01,-3.924494838945669084e-01,-5.433151157278728327e-01,3.050183437518030338e-01,-5.297884663705332287e-01,-4.402434005450059917e-01,3.033450134513224761e-01,5.080005533136775497e-01,1.765071572885704987e-01 --3.296381408504345245e-02,8.213734474661573692e-01,2.884182165116890850e-01,-5.803866935447996589e-01,3.041849216693381930e-01,5.667920446859808781e-01,2.250925284276168448e-01,8.236273827714153395e-01,-4.114514278581935525e-01,2.460381915048910351e-01,-1.767376131099762659e-01,-3.058606496513562867e-01,8.234928427114294180e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000018.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000018.csv deleted file mode 100644 index bca2acf..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000018.csv +++ /dev/null @@ -1,2 +0,0 @@ --5.517066569583584101e-01,6.457030079067409556e-01,1.312452431616923931e-01,-9.654657530459793691e-01,1.001632917413284307e-01,4.808171457391481329e-01,2.887487385246092497e-01,9.356047048047881898e-01,1.042754743685525565e-01,-5.230136305881964986e-01,7.032043120523423507e-01,8.023204324713959501e-01,7.872166104195181813e-01 --5.753978519462528141e-01,7.086092336367590949e-01,8.862996973827064195e-01,-4.508828284143380216e-01,-1.498953015148070111e-01,-1.479318170320651493e-01,-3.642550894021263641e-01,1.915439725996275211e-01,7.097594434528575746e-02,4.751618882704935487e-01,3.779955724380190674e-02,2.814613891706765347e-01,2.127833895804817355e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000019.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000019.csv deleted file mode 100644 index a21ea50..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000019.csv +++ /dev/null @@ -1,2 +0,0 @@ -1.016398917604210972e-01,5.367241283642387728e-01,-6.818635487683848417e-01,-6.492151914242854094e-01,2.951107976835281033e-01,1.210983103608658240e-01,-1.914729327636166545e-01,2.084900726109191194e-01,-9.913674641194303305e-03,-1.318795603716438336e-01,-4.591725358362441778e-01,3.168885180251357347e-02,2.613304272286978147e-01 --2.518998374757852599e-01,-1.545791580965771850e-01,2.591599328760181287e-01,-9.227238644730679784e-01,-6.523143964998687760e-01,1.325715926486024099e-01,-2.634595670714392490e-01,-6.084513786591814188e-01,5.945406152909971098e-01,-7.413295649974527279e-01,-9.836113191580013737e-01,-6.613550934121019687e-01,7.386695727713020743e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000020.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000020.csv deleted file mode 100644 index a24562e..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000020.csv +++ /dev/null @@ -1,2 +0,0 @@ -6.921866875646716100e-01,5.108296431449499408e-03,-8.297196043425478784e-01,-8.782207198751741384e-01,4.881693383540661735e-01,7.602399262614862874e-01,7.929628517423199519e-01,2.747828187373047015e-01,-8.565292578863379358e-01,-8.028470318278093654e-01,-7.666729973480264082e-01,-3.618054675021058486e-01,6.474809588861711873e-01 --4.195656375439105190e-02,3.829378217234820081e-01,-4.738148839593532280e-01,-1.013350142767268647e-01,6.518052790254404982e-01,-2.989039824402659473e-01,3.810369849047898771e-01,-4.761398937049154956e-01,-8.736647575845664093e-01,8.671975808320215862e-01,4.518295413402169114e-01,3.044248270515332866e-01,3.525190411138288682e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000021.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000021.csv deleted file mode 100644 index d861c0d..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000021.csv +++ /dev/null @@ -1,2 +0,0 @@ -1.667868049973333378e-01,-9.628664509718032338e-01,6.014400922608953426e-01,2.972001483931350219e-01,4.941659040388075574e-01,-6.774253220554804500e-01,-1.993299704297657460e-01,6.394927405908403806e-01,-1.095498513905281968e-01,1.797181927692330650e-01,8.966277273316463070e-01,2.586404039871941229e-01,5.552709887608781036e-01 -8.904906062660500332e-01,3.456353981853357293e-01,-5.089078336971060157e-01,5.564474698994368307e-01,1.613474062509465679e-01,-4.536022710624829646e-01,7.605093590363174449e-01,-3.094273061226622268e-01,-6.937981682429052999e-01,8.031386423573720901e-01,-7.976443746091934628e-01,6.529445161208398130e-01,4.447290112391218964e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000022.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000022.csv deleted file mode 100644 index d9cf9f6..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000022.csv +++ /dev/null @@ -1,2 +0,0 @@ --5.460898717556510462e-01,5.858062690874137335e-01,2.194226187473129475e-01,9.188630804699871035e-02,-1.707137636423472493e-01,-3.178585699238225537e-01,3.633095509488837305e-01,-6.887805288505008949e-01,-1.468860672171268256e-02,-2.862987915279111562e-01,-7.723918312056159419e-01,8.126650586523440634e-01,4.458798622363638331e-01 -5.547486983713716402e-01,-5.675784088131794469e-01,-6.187415038981081139e-02,6.835868811760836827e-02,7.617313312413553916e-03,5.521100702462113929e-01,5.867352194847055280e-01,-8.129437845772602422e-01,-1.254568018619570680e-01,8.382074387221385425e-01,9.974232510580696154e-01,1.902631261807281593e-01,5.541201377636362224e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000023.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000023.csv deleted file mode 100644 index c1d0e5a..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000023.csv +++ /dev/null @@ -1,3 +0,0 @@ -2.812878692303713013e-01,-5.380780165501937162e-01,-1.995746226435606285e-01,3.556609488625348536e-01,-3.570058976622436653e-02,3.226211706468127272e-01,7.604367324861782684e-02,-6.726022941928089249e-01,-1.244510164772598682e-01,-5.242674835142189238e-01,6.111367647898253708e-01,5.446881171705224567e-01,1.870729455165833777e-01 --1.804572792799459258e-01,-5.615234318506530098e-01,-5.678667666618428811e-01,-1.973882851907129421e-01,-9.342817705045243226e-01,-2.696256933877165807e-01,-4.754476132734983818e-01,8.379911380637130591e-01,-1.430754625396699620e-01,6.703727340713743210e-01,-7.950803881856465249e-01,1.646446229560936114e-01,5.812129149706145581e-01 --7.303265877923637017e-01,5.163190035467037919e-01,6.970165582377609859e-01,-2.277609099967679018e-01,5.773156478708756367e-01,7.154948998404351279e-01,-4.402130399012762485e-01,-5.815862404831599886e-01,-9.434039458871834594e-01,-3.147493553756923745e-01,-6.215697199671859074e-01,1.823798551890962738e-01,2.317141395128020642e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000024.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000024.csv deleted file mode 100644 index 3eb8c33..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000024.csv +++ /dev/null @@ -1,2 +0,0 @@ -2.166769833429884606e-01,-2.331594788641218052e-01,-8.006870701326018747e-01,-5.266093172746755258e-01,-8.322914134913979023e-01,7.446673507975531958e-01,-2.886575388105372397e-02,-2.990935453031380309e-01,3.936163556018033027e-01,-3.542536932837532238e-01,2.834018169147607402e-01,-9.586592464730825380e-01,4.259584301232757220e-01 -8.416839786286203218e-01,-5.562313049662526154e-02,-3.647897665777326548e-01,-6.475376571556650251e-01,-3.655643797640151238e-01,-3.196437920194410420e-01,5.442135586745493470e-01,5.508198826067678411e-03,-8.528986112361351957e-01,-6.533886080407793617e-01,-8.277524479526696677e-01,-6.784479736983972664e-01,5.740415698767242780e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000025.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000025.csv deleted file mode 100644 index 4f7835f..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000025.csv +++ /dev/null @@ -1,2 +0,0 @@ --3.328441505555097546e-01,-1.387237543865100786e-01,-8.385785443603737122e-01,-6.460892635786552596e-01,-6.206316056012477489e-01,-2.482763796893627806e-02,8.854727682161622759e-02,-3.627820436506015156e-01,-2.149069615655596621e-02,-8.638772932614158240e-01,-9.940866772762460002e-01,7.740573663063279319e-01,4.901549186583971096e-01 --5.533853946343669783e-02,6.783593497875153311e-01,-3.470947229869234540e-01,-8.081860602696091522e-01,-2.234790784515090500e-01,5.246894583705352666e-02,4.184989961414606885e-01,5.085439272585972059e-01,2.229705198326268345e-01,-2.452705816636990832e-02,-3.598764438670341015e-01,-4.548214145779025941e-01,5.098450813416028904e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000026.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000026.csv deleted file mode 100644 index cf18c74..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000026.csv +++ /dev/null @@ -1,2 +0,0 @@ -9.215997820884049840e-01,-8.886305956631088687e-01,7.918854213071213621e-02,8.549195867901333568e-01,-7.539338040902159310e-01,2.127998942252871117e-01,9.334583063583679063e-02,2.117336123070212572e-01,-3.905267778142720303e-01,4.922987444042437044e-01,3.538552884251142672e-01,-4.995810348035092385e-01,3.171922691986440168e-01 --2.710991105815863111e-01,6.963869144978018788e-02,-2.614443785391928898e-01,-1.745302413226517135e-01,2.582270480554498260e-01,-7.700432645988855018e-01,-9.312241172350603780e-01,5.326596688321914019e-01,6.785270353740011640e-01,-6.480593296020313865e-01,-5.862245859094883382e-01,-8.798502128158780522e-02,6.828077308013559277e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000027.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000027.csv deleted file mode 100644 index 5645631..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000027.csv +++ /dev/null @@ -1,2 +0,0 @@ --3.363859180192507736e-01,-2.005356922417877996e-02,4.105299683991425752e-03,4.520536430696797670e-01,-2.397748696508128496e-01,-9.149591317297329773e-01,-2.687516636255817826e-01,2.502854147579025579e-02,-5.658580150798908637e-01,-3.444933603162623204e-01,3.980110762078155062e-01,6.019507953845388837e-01,2.306736239861526538e-01 --9.924270668037067367e-01,1.880015741363572079e-01,7.482014409903954277e-01,-8.717841019029326510e-01,3.178806298249037265e-01,4.369798231563537527e-02,-4.363922340879544670e-01,-3.301960725800534568e-01,5.876506713161555595e-01,-6.176395243356269660e-02,-3.698941006853062596e-01,-6.461432233004753556e-01,7.693263760138474572e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000028.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000028.csv deleted file mode 100644 index 3901f6d..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000028.csv +++ /dev/null @@ -1,2 +0,0 @@ --9.488858097444059503e-01,4.079256280033416449e-01,-7.273191618538961123e-01,-2.164528040135893505e-02,1.801964698567355416e-01,2.110643357786530228e-01,-6.481561009612546442e-02,-2.415201336393544285e-01,8.209213192744986287e-01,1.342520688379920113e-01,6.654275517850556376e-01,4.674644534794492046e-02,2.727174515041750347e-01 -5.889383911906385105e-01,-4.716661362499952048e-01,8.398891628230311657e-01,9.444296451726641450e-01,6.852416895301041144e-01,-2.718291944059443299e-01,9.759285577328178363e-01,4.221736051403957024e-01,-5.860661824356676597e-01,6.789765242522334265e-01,-6.864599914320155261e-01,8.076484878207652596e-01,7.272825484958250764e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000029.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000029.csv deleted file mode 100644 index 7f0e048..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000029.csv +++ /dev/null @@ -1,2 +0,0 @@ --3.156004044367652916e-01,4.539226489835086475e-01,3.386522834986691599e-01,-6.897275360788317489e-01,-6.389720151715840846e-01,-3.773349166157233814e-01,1.813606744935740700e-01,-6.205184082106425247e-01,-3.558681169774229325e-01,3.970835926210039002e-01,-8.915932500008973971e-01,-7.113662644290583703e-01,2.064773053549207871e-01 --2.222304390465441593e-01,2.761428589612251461e-01,-4.556405012454298742e-01,1.006421622526911808e-01,1.163138470867763896e-01,4.306829086925703098e-01,7.112480073546407766e-01,8.255028592978328472e-01,-9.373835503173011396e-02,-7.714024572181317208e-01,9.021782591265152806e-01,-1.140412270850570398e-01,7.935226946450792962e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000030.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000030.csv deleted file mode 100644 index 76e9192..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000030.csv +++ /dev/null @@ -1,2 +0,0 @@ --3.054766017168615289e-01,3.285035685711270581e-01,3.339186103540698891e-01,8.138242613802961767e-01,-5.205905000696835483e-01,1.529228105404809579e-01,8.466677815415111219e-01,-7.061057292640156025e-01,-1.216890214920127722e-01,-1.716434030659630405e-01,-3.234074144154277519e-01,-5.208430114476063633e-01,5.154903134260301334e-01 -8.269501238092979989e-02,9.228403821803805585e-01,-9.202930276956799993e-01,-1.265069840482921926e-01,2.493516975241008016e-01,-3.998067562089782090e-01,-4.105033927816157391e-01,2.800337569818500683e-02,-5.976660197456356016e-01,4.214440523544955575e-01,4.705808164701692498e-01,4.943614828825371177e-01,4.845096865739698111e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000031.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000031.csv deleted file mode 100644 index 297ec32..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000031.csv +++ /dev/null @@ -1,2 +0,0 @@ -5.931682635522148583e-01,8.699251727039116755e-01,-8.264376444320893356e-01,6.777316299673374900e-03,1.259812425166662031e-02,7.526947568173611991e-01,-6.456449467182574509e-01,-1.867494885666358684e-01,-5.751692444575797758e-01,1.762422487332968579e-01,-4.008003102356285652e-01,-8.546965786953328870e-01,6.116071270069189936e-01 -9.061805329925218810e-01,5.725102286958729803e-01,6.505301386644584127e-01,-5.634820490409915283e-01,1.322069399068586115e-01,-5.320988916994693341e-01,-8.018914495696449762e-02,-1.672922951958961679e-01,-5.236940947370711807e-01,2.913933672990698387e-01,-7.670309932568115663e-01,2.266688501822122781e-01,3.883928729930810619e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000032.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000032.csv deleted file mode 100644 index 7b3c85b..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000032.csv +++ /dev/null @@ -1,2 +0,0 @@ --4.029727807509135840e-01,2.221234395356765923e-01,1.653311729326147894e-01,2.854384277332262521e-01,-5.299668975252573855e-01,2.277161429752445621e-01,-8.605748587143935424e-01,2.679358058329555092e-01,-1.753822281114920667e-01,-4.080365861113419701e-01,-7.990197638852070128e-02,6.451913142929026623e-01,2.646000967988433872e-01 --5.825951636626720553e-01,9.933789450712708913e-01,6.017632717440635215e-01,1.285660707793123692e-01,9.225675787278662110e-01,7.629667460674474100e-01,-1.269325586916481008e-01,-2.399724811514538647e-01,2.244383569074661633e-01,3.864334357836947120e-01,1.407009021060205978e-01,-4.470113212371948919e-01,7.353999032011566683e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000033.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000033.csv deleted file mode 100644 index a6af9d6..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000033.csv +++ /dev/null @@ -1,2 +0,0 @@ --3.343344504186707145e-01,-1.788377734180790490e-02,5.238121179021124618e-01,-1.801466014877619592e-01,2.444788960837480651e-01,5.831808836713883171e-01,-8.022784620348435425e-01,-4.295701503377884478e-01,-3.972178837057684930e-01,7.118887696693811939e-01,-8.639924687935303105e-01,4.474523280618651899e-01,4.561726115084100419e-01 --3.159104803788210791e-01,-8.391006409409682565e-01,-3.390316252365024319e-01,1.553039413054588813e-01,8.170194759350031255e-01,7.527728133817603862e-01,6.978584447967217663e-01,1.128938492491866619e-01,-8.250572721858402403e-01,-5.512768790771354066e-01,3.630337106170709038e-01,-2.222138950067416019e-01,5.438273884915899581e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000034.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000034.csv deleted file mode 100644 index d18f1b5..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000034.csv +++ /dev/null @@ -1,2 +0,0 @@ -5.697263627900581717e-01,9.724449544795059630e-01,-1.674757567980476036e-01,-1.372771123535918569e-01,-7.764922940367462445e-01,5.410911337820014655e-01,3.028425898014133200e-01,4.025386950980982537e-01,-1.476033977383572893e-01,3.294617484232353899e-01,-5.024261635732010234e-01,-4.464667963982544840e-01,6.101539033152794111e-01 --7.758966548598598134e-01,6.989084956917288594e-01,7.492372266276172699e-01,1.570862193960709252e-01,-2.601786362671802966e-01,2.727571671873869619e-01,-1.349546389088733811e-01,1.024561812582318598e-01,5.506075183810597018e-01,2.805526623632508265e-01,-9.364615058079310828e-01,5.565384469523662059e-01,3.898460966847205333e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000035.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000035.csv deleted file mode 100644 index f61ede0..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000035.csv +++ /dev/null @@ -1,2 +0,0 @@ --7.898235741040993130e-01,5.427817216628392227e-01,6.300274118070616769e-01,-5.211653741616515401e-01,-6.981622668479592342e-02,9.578781805862026655e-01,1.666377264587282081e-01,1.637757194655065085e-01,7.571393872230314237e-02,9.020579858164654574e-01,8.210889142502095783e-01,9.377974534116928496e-01,8.615315349321857052e-01 -2.679807526537674178e-01,-1.111544236701951238e-01,-5.024699421507621278e-01,7.457314846733458236e-01,-4.885664083857288453e-02,-7.673089298447215434e-01,-8.883324858657668521e-01,-2.464007149940794505e-01,6.129545902906914367e-01,2.592414698046310306e-01,-3.794121113471968787e-01,1.341081142088484945e-02,1.384684650678143225e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000036.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000036.csv deleted file mode 100644 index bcb6f88..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000036.csv +++ /dev/null @@ -1,2 +0,0 @@ --7.654273192322054165e-01,-4.094450069738897469e-02,-7.595383743431962653e-01,-8.557505909970739566e-01,9.064524921829164583e-01,1.517470713028805651e-01,9.075466163718142187e-01,-4.142547321400837923e-01,4.595383450523824465e-01,2.242844526762681756e-01,6.565226571017350743e-01,-6.406292457846078925e-01,2.164925796749257170e-01 -4.134224368111192316e-01,-1.279193807354295220e-01,1.395348624965011552e-01,-7.113928155897839556e-01,-2.037598375484022117e-01,-1.397999201487876153e-01,-8.100399318887028244e-01,-3.540527934775026253e-01,-8.142623919357250273e-01,-7.886095433794413356e-01,-3.674068192061692439e-01,-8.141558656567047247e-01,7.835074203250743663e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000037.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000037.csv deleted file mode 100644 index d666294..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000037.csv +++ /dev/null @@ -1,2 +0,0 @@ -6.234462065020718313e-01,1.300245200362042386e-01,-2.842597130938973038e-01,1.595438141453187075e-01,-4.713993691785938189e-01,1.486436548511567146e-01,-6.366613328949060069e-01,7.675794196907781419e-02,-9.937538509582277690e-02,5.037179007533050257e-01,-6.402726075739686440e-01,-3.221377519083947760e-01,4.313968476057338797e-01 --2.673181139834390763e-03,-2.286479913327021940e-02,6.059211548471905573e-01,-3.648636787800783043e-01,-4.414015900975392093e-01,2.960060241178898988e-01,-3.576760461529826518e-01,-9.644156958072134245e-01,-5.959032460738362680e-01,-7.214674024314973177e-01,5.674627493986834637e-01,3.261734992328386706e-01,5.686031523942661758e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000038.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000038.csv deleted file mode 100644 index df4a712..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000038.csv +++ /dev/null @@ -1,2 +0,0 @@ --7.420103884864390764e-01,-6.381524612883886505e-01,-1.613639618245801266e-01,7.199764450024330742e-01,-1.006210277087187244e-01,5.736088536664056825e-01,-2.078207707668398019e-01,3.313608111485661922e-01,-1.567612276534600113e-01,3.527694725218057936e-01,-1.767897433287017872e-01,-3.553205709558269199e-01,5.623580600962118092e-01 -1.803325284453796140e-01,1.231843219533304001e-01,5.739878271111997776e-03,-8.697875645645707365e-01,4.277980606243028117e-01,1.010399067634688564e-01,9.441666464685340987e-01,-2.209014237470598996e-01,3.359090300130838092e-01,5.940980517447476128e-01,2.482757629902709873e-01,-1.834052354989281763e-01,4.376419399037881908e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000039.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000039.csv deleted file mode 100644 index a71f835..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000039.csv +++ /dev/null @@ -1,2 +0,0 @@ -7.369773596082351830e-02,-2.482646733587001719e-01,5.057638234463792681e-01,-6.242771609481161388e-01,-7.429818940347492351e-01,-4.707267465509237248e-01,-4.626361468138358024e-01,-6.467146033806137062e-02,2.131345959536605772e-01,7.725903761561427885e-01,-4.240286243640691843e-01,-4.218771823875742122e-01,5.823322236392094453e-01 --6.554490572494988676e-01,-7.815506609888773770e-01,-7.332898854928275867e-01,6.612085576900739170e-02,2.178726193278017753e-01,5.569210460230202830e-01,-2.319012715747397202e-01,-9.022663389078762197e-01,-4.304867213238519064e-01,-5.467178901348748177e-01,5.657685803999605856e-01,7.731699809585836913e-01,4.176677763607904992e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000040.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000040.csv deleted file mode 100644 index 2012dad..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000040.csv +++ /dev/null @@ -1,2 +0,0 @@ -3.730533296626739048e-01,1.433183679154126366e-01,3.637273449949744997e-01,2.796358796157163429e-01,9.971687232829262726e-03,7.480311521939866370e-01,3.937355475477484212e-01,4.560000728934507919e-01,-4.984654853379433259e-01,-3.455550319734554954e-01,-9.835584381128927856e-01,-2.633451200187155727e-01,1.651205815987711323e-01 -8.112653261883815414e-01,-3.950454378848380355e-01,3.176817652040930806e-01,-5.239337818191858176e-03,4.250105481704846699e-01,3.604098546798490954e-01,5.490989594952853103e-01,-7.353484188252294995e-01,1.375370159687769878e-01,3.739887895489748537e-01,7.194665959008399447e-01,-2.052330554944428176e-01,8.348794184012288122e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000041.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000041.csv deleted file mode 100644 index 9697871..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000041.csv +++ /dev/null @@ -1,2 +0,0 @@ -1.985151568679008882e-01,-4.118089649714407052e-02,-1.173690980919479543e-01,-6.147768674282323431e-01,5.384205158174693029e-01,-8.092057756629023046e-02,2.779284151732872576e-01,2.862361042030712177e-01,3.840843723060081150e-01,9.898121724234101304e-01,9.713830508604506253e-01,7.822167263660317893e-01,7.111855449058056555e-01 -5.892885167675354641e-01,-1.963827275432317165e-03,-7.606511330012744043e-01,5.586488854728632880e-01,7.846991525938085132e-02,-6.184691865438018965e-01,2.366672318317677437e-01,-4.646883212217289838e-02,2.801530524023061464e-01,-6.722047100717012391e-01,-7.126605435777322306e-01,-6.314786757874539802e-01,2.888144550941944000e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000042.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000042.csv deleted file mode 100644 index 9760c67..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000042.csv +++ /dev/null @@ -1,2 +0,0 @@ --5.511430473636560468e-01,2.075233910672833471e-01,9.276713288167097726e-01,-1.822552515865676348e-01,4.964535495465318693e-01,4.198110192190658285e-01,-7.110023362885946607e-01,3.517890700491155265e-02,1.069361824193393318e-02,-7.884513208221513025e-01,-8.560972160329762826e-01,-9.047373606248456657e-01,8.266997873488395321e-01 --5.807999610356651132e-02,-1.039348906028856323e-01,3.141343076567044701e-01,-2.056624804791262751e-01,5.792891943323019710e-02,-9.509068945320611199e-01,-3.764789001983093186e-01,2.488096127203558439e-01,-1.853125052175255139e-01,8.301118240605835918e-01,-1.216569088549470656e-01,8.391547337468705514e-01,1.733002126511604402e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000043.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000043.csv deleted file mode 100644 index a0119e8..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000043.csv +++ /dev/null @@ -1,2 +0,0 @@ --5.633361199129609531e-01,-1.242868737984672567e-01,-8.018904487385585256e-01,-5.010751601456076010e-02,-5.054966895260570858e-02,4.270084989655040797e-01,5.613658733821358382e-02,1.067847504111867352e-01,9.715247219631861775e-02,-5.884270629564716248e-01,2.869301460815778526e-01,-8.916215668672899941e-01,1.127259639768470878e-01 -5.813484676405529239e-01,-7.608577159641982668e-02,-1.726841049851548515e-01,1.495963314716997061e-01,9.282798648258450136e-02,-9.208945702831410340e-01,-3.570382956573512345e-01,-3.065376051478874153e-01,-3.776733142035426649e-01,-9.376120443254172265e-01,-1.646485576867517953e-01,4.424834047857026942e-01,8.872740360231529122e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000044.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000044.csv deleted file mode 100644 index 088d30f..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000044.csv +++ /dev/null @@ -1,2 +0,0 @@ -1.245129250727008419e-01,-6.636582935874648648e-01,-7.095021996374384354e-01,-5.708048087285482186e-01,-2.613952915421429157e-01,4.360332099195141087e-01,4.386695896700476549e-01,2.591668816930605690e-01,2.710449407828230406e-01,-1.994249360309625629e-01,-9.884196752400828956e-01,-8.948557108504355817e-01,4.008892167512856930e-01 --7.977101440753309181e-01,-5.891213181154428824e-01,-5.394841756954262824e-01,6.845658522320718919e-01,5.996325324640972010e-01,-7.950515319880713250e-02,1.314403032344408917e-01,8.974954145234985692e-01,3.185168035203989056e-02,-8.571054947316427697e-01,8.942330986453337349e-01,4.800334580863008238e-01,5.991107832487143625e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000045.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000045.csv deleted file mode 100644 index 379affb..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000045.csv +++ /dev/null @@ -1,2 +0,0 @@ -4.156697549963126459e-01,1.938812896637998051e-02,3.762020023825551895e-02,-4.543409658210944002e-01,-4.160769028245498991e-01,-4.404768450621054932e-02,-1.226606756646233531e-01,2.246644739370420307e-01,-1.930408370461424994e-01,4.195049700021282746e-01,-7.252898605016289135e-01,-1.186045724865631978e-01,2.376483080486445632e-01 --5.048426707707873717e-01,2.678319927728600724e-01,-2.276997135948761741e-01,5.838793463819900165e-01,-5.592597689481533241e-01,5.818759639934625305e-01,-5.955380441213122822e-01,-6.677174966534085154e-01,2.195016551675315064e-01,-5.308958776483587716e-01,8.413556421996504220e-02,6.724915832844733377e-01,7.623516919513554368e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000046.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000046.csv deleted file mode 100644 index 674f6e6..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000046.csv +++ /dev/null @@ -1,2 +0,0 @@ --1.237880297177609279e-01,-5.389232504861669604e-01,8.351615604828057648e-01,-5.197299868564686509e-01,-6.477376740352802642e-02,6.235481432700507032e-01,8.648963276371262054e-01,-4.513602462196104614e-01,1.121472463841981515e-01,9.872965050461273151e-01,9.679965448742415823e-01,-1.147346194826741606e-01,3.912812329938487044e-01 --5.999286838715911507e-02,-2.097115124782031881e-01,3.409859382911564207e-01,-1.638644314851909201e-01,-5.392371544653988824e-01,1.453514884790074735e-01,-9.158051602213537201e-01,7.348874336021522513e-01,5.552381537384321053e-01,-2.804504387202546578e-01,7.872454148265923823e-01,2.773799831190322251e-01,6.087187670061512401e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000047.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000047.csv deleted file mode 100644 index 4580b5a..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000047.csv +++ /dev/null @@ -1,2 +0,0 @@ --4.781135907766911330e-01,-1.276623027908179164e-01,-5.244677607605363612e-02,7.793736542701723558e-02,-4.965349935613236898e-01,-7.450674808538539917e-01,2.215427033187817862e-01,-2.835553192890256646e-01,-5.879019088656525227e-01,-8.655409501489235158e-01,1.484038914439216317e-01,2.340787049439099210e-01,4.441034483063370786e-01 -8.918907925254648816e-02,9.583522210758066429e-02,3.380023039545712038e-01,-2.757160640620373027e-01,9.491265801938739699e-01,-2.949240243853041843e-01,2.368000271881480767e-01,-3.274535743006037336e-01,2.320714265769052709e-01,6.681697133141641931e-01,-8.151860272732713852e-02,-5.223786728288988268e-01,5.558965516936629214e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000048.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000048.csv deleted file mode 100644 index 4b22744..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000048.csv +++ /dev/null @@ -1,2 +0,0 @@ --2.430657731263985433e-01,-5.883176849573654721e-01,-1.051727327384854860e-02,-5.613318623862584289e-01,-3.654238252241008844e-01,-2.903699040052056812e-01,-3.376162129993560690e-01,-1.937379177653244522e-01,1.447762165448007732e-01,9.132031396935695877e-01,2.309821948964190241e-01,4.969710849235846606e-01,3.556727315444812021e-01 --5.700025402755233284e-01,3.962163276669257161e-01,4.106860990482503748e-01,-4.445226161030975121e-01,2.216226823069946672e-01,5.337581703209348660e-01,-5.759837441605704100e-01,-3.921802841679722373e-01,-5.148726436855111110e-01,-4.398568392874997457e-01,-8.504659419589648550e-01,6.585902602061193267e-01,6.443272684555187979e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000049.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000049.csv deleted file mode 100644 index e66cf87..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000049.csv +++ /dev/null @@ -1,2 +0,0 @@ --3.200427851155407399e-01,3.576359420223551577e-01,-2.513016740058249265e-01,-7.546934815443087086e-01,-2.786933152965342941e-01,2.327058657559735178e-01,7.113022836300710861e-01,6.680266543764468157e-01,2.330623308024013518e-01,-6.132416434718550580e-01,-1.253524191738015769e-01,4.190854187019641408e-01,6.428379861042887722e-01 --2.051650609723685292e-01,-4.086715839711074771e-02,1.337484780156206199e-02,-3.217966923911506072e-01,2.645280631801716353e-01,-5.515626691943973370e-01,-7.300788769231552067e-01,-9.944065418179377502e-01,-6.942904879158731113e-01,3.461188780593014158e-01,-2.042033665456883806e-01,-6.116735230775620646e-01,3.571620138957111168e-01 diff --git a/ScaFFold/utils/config_utils.py b/ScaFFold/utils/config_utils.py index 60731fe..3d7ebe3 100644 --- a/ScaFFold/utils/config_utils.py +++ b/ScaFFold/utils/config_utils.py @@ -18,8 +18,6 @@ import yaml -import ScaFFold.paths - def require_positive_int(name: str, value: int) -> int: if not isinstance(value, int) or isinstance(value, bool) or value < 1: @@ -89,7 +87,6 @@ class Config: "dataset_dir", "fract_base_dir", "job_name", - "library_root", "n_categories", "problem_scale", "unet_bottleneck_dim", @@ -209,7 +206,6 @@ def _validate_keys(cls, config_dict, strict): def __init__(self, config_dict, strict=True): self._validate_keys(config_dict, strict) - self.library_root = str(ScaFFold.paths.scaffold_root).rstrip("/") + "/ScaFFold/" self.base_run_dir = str(Path(config_dict["base_run_dir"]).resolve()) self.dataset_dir = str( Path(config_dict.get("dataset_dir", "datasets/")).resolve() diff --git a/pyproject.toml b/pyproject.toml index c0e6c3e..80b5aa7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,6 @@ build-backend = "setuptools.build_meta" [tool.setuptools] package-data = { "ScaFFold" = [ "package_data/weights_ins145.csv", - "fractals/var0.15/3DIFS_param/*", "configs/*", ] } include-package-data = true From 199251867f2dc8242190d0e5240769668b03bb34 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Sat, 1 Aug 2026 00:07:05 -0700 Subject: [PATCH 54/62] Serve DistConv's DCTensor from the compiled GroupNorm worker.py wraps every activation in a DCTensor even at dc_num_shards=1, and FastGroupNorm rejected tensor subclasses outright, so production never took the compiled path: the 1.84x GroupNorm win measured in round 2 was latent. DistConv's __torch_dispatch__ has no GroupNorm-specific handling -- it unwraps to the local shard, runs the stock aten kernel and rewraps -- so there is no distributed GroupNorm semantics to preserve, only a wrapper Dynamo cannot trace. forward() now does that same unwrap itself, in front of the compiled kernel. It cannot use dispatch's mechanism: dispatch runs below autograd where a bare _tensor read is safe, while forward runs above it and must go through DistConv's _ToTensor/_FromTensor pair or the graph back to the producing convolution is severed. Measured on MI300A at the scale-8 GroupNorm shapes, fwd+bwd, DCTensor-wrapped: [1,64,256^3] 129.56 -> 11.82 ms (10.96x), [1,128,128^3] 31.28 -> 3.19 ms, [1,256,64^3] 7.85 -> 1.11 ms; the two smallest shapes are launch-bound and within noise. Values and gradients are bitwise identical to the eager wrapped route, including a 2-rank sharded run. Statistics stay per-shard at every shard count, exactly as DistConv computes them today; making them global is an upstream question (R32/R39), not one this change touches. --- ScaFFold/unet/group_norm.py | 64 +++- .../rank_scripts/groupnorm_shards_2rank.py | 123 +++++++ tests/test_groupnorm.py | 300 +++++++++++++++++- 3 files changed, 471 insertions(+), 16 deletions(-) create mode 100644 tests/helpers/rank_scripts/groupnorm_shards_2rank.py diff --git a/ScaFFold/unet/group_norm.py b/ScaFFold/unet/group_norm.py index 06d7407..33d3323 100644 --- a/ScaFFold/unet/group_norm.py +++ b/ScaFFold/unet/group_norm.py @@ -30,8 +30,9 @@ falls back to stock eager ``F.group_norm``: * non-CUDA tensors (the CPU test suite never pays compile latency), -* tensor subclasses such as DistConv's ``DCTensor``, whose ``__torch_dispatch__`` - wrapper Dynamo cannot trace, +* tensor subclasses, whose ``__torch_dispatch__`` wrappers Dynamo cannot trace + -- except DistConv's ``DCTensor``, which is unwrapped to its local shard + around the compiled kernel instead (see ``FastGroupNorm.forward``), * an already-compiled enclosing region (the functional call inlines instead), * an explicit opt-out via ``SCAFFOLD_GROUPNORM_COMPILE=0``, * any failure inside ``torch.compile`` -- logged once, then eager forever after. @@ -45,6 +46,7 @@ import logging import os +import sys import torch import torch.nn as nn @@ -153,13 +155,11 @@ def _use_compiled(input): """Whether this particular input should take the compiled path.""" if _compile_failed or _compile_override is False: return False - # Tensor subclasses (DistConv's DCTensor) route their ops through - # __torch_dispatch__, which Dynamo cannot trace; eager keeps the wrapper's - # semantics -- including which of its outputs come back wrapped -- exactly - # as they are today. worker.py wraps activations in DCTensor even at - # dc_num_shards=[1,1,1], so this fast path engages once that wrap is - # skipped for the unsharded case (or whenever the model is driven with - # plain tensors, as the tests and the standalone benchmarks do). + # Tensor subclasses route their ops through __torch_dispatch__, which + # Dynamo cannot trace. DistConv's DCTensor never reaches this check -- + # forward() peeks at its local shard instead -- so anything rejected here + # is an unknown wrapper, and eager keeps its semantics exactly as they + # are today. if type(input) is not torch.Tensor: return False # CPU GroupNorm is not the bottleneck and compiling it would put a @@ -172,21 +172,61 @@ def _use_compiled(input): return True +def _dctensor_ops(input): + """The ``distconv.distconv`` module when ``input`` is a DCTensor, else None. + + Resolved through ``sys.modules`` instead of an import: a DCTensor can only + exist if DistConv is already imported, and this module must stay importable + (and the CPU suite runnable) without DistConv installed. + """ + distconv = sys.modules.get("distconv.distconv") + if distconv is not None and isinstance(input, distconv.DCTensor): + return distconv + return None + + class FastGroupNorm(nn.GroupNorm): """``nn.GroupNorm`` that runs its GPU forward through ``torch.compile``. Identical state: ``weight``/``bias`` of shape ``(num_channels,)``, no buffers, so state dicts are interchangeable with plain ``nn.GroupNorm`` in both directions. + + DistConv's ``DCTensor`` gets the compiled kernel too: its generic + ``__torch_dispatch__`` has no GroupNorm-specific handling -- it unwraps to + the local shard, runs the stock aten kernels, and rewraps the outputs, so + statistics are per-shard and no communication happens at any shard count. + ``forward`` moves that same unwrap up in front of the compiled kernel, + preserving those semantics exactly (DCTensor in -> DCTensor out) while + keeping the fast kernel Dynamo's inability to trace the wrapper would + otherwise forfeit. It cannot copy dispatch's *mechanism*, though: dispatch + runs below autograd, where reading ``_tensor`` directly is safe, whereas + this runs above it, so the unwrap has to go through DistConv's + ``_ToTensor``/``_FromTensor`` autograd pair or the graph back to the + producing convolution is severed. """ def forward(self, input): - # super().forward() is the stock kernel; deferring to it keeps the eager - # path identical to nn.GroupNorm's by construction. - if not _use_compiled(input): + distconv = _dctensor_ops(input) + # The eligibility checks look at the local shard for a DCTensor (the + # peek is a plain attribute read, no autograd involvement) and at the + # tensor itself otherwise. + local_view = input._tensor if distconv is not None else input + if not _use_compiled(local_view): + # super().forward() is the stock kernel; deferring to it keeps the + # eager path identical to nn.GroupNorm's by construction. return super().forward(input) global _compile_failed try: + if distconv is not None: + # _ToTensor is the autograd-aware unwrap DistConv itself uses; + # DCTensor.from_shard is the public spelling of _FromTensor. + # (There is no public unwrap yet -- upstream ask.) + local = distconv._ToTensor.apply(input) + out = _get_compiled_group_norm()( + local, self.num_groups, self.weight, self.bias, self.eps + ) + return distconv.DCTensor.from_shard(out, input._parallel_strategy) return _get_compiled_group_norm()( input, self.num_groups, self.weight, self.bias, self.eps ) diff --git a/tests/helpers/rank_scripts/groupnorm_shards_2rank.py b/tests/helpers/rank_scripts/groupnorm_shards_2rank.py new file mode 100644 index 0000000..a69bcfd --- /dev/null +++ b/tests/helpers/rank_scripts/groupnorm_shards_2rank.py @@ -0,0 +1,123 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +"""Two-rank check that FastGroupNorm's DCTensor route is shard-count agnostic. + +Run under ``torchrun --nproc_per_node=2`` with the gloo backend (see +``tests/test_groupnorm.py``). Every other DCTensor test uses +``num_shards=(1, 1, 1)``, where sharding is a no-op and the local shard is the +whole tensor; this one shards a spatial dim across two ranks so the claim the +fast path actually rests on -- GroupNorm statistics are per-shard, and the +route does not add communication or change which elements are reduced together +-- is exercised where it can fail. + +Each rank prints one ``RESULT ...`` line plus ``DONE``; the parent asserts. +""" + +import os +import sys + +import torch +import torch.distributed as dist +import torch.nn.functional as F + +sys.path.insert(0, os.environ.get("SCAFFOLD_ROOT", "/usr/WS1/dryden1/ScaFFold")) + +from ScaFFold.unet import group_norm as gn_mod # noqa: E402 +from ScaFFold.unet.group_norm import FastGroupNorm # noqa: E402 + +GROUPS = 8 +CHANNELS = 16 +SIZE = 8 # dim 2 is split into two shards of 4 + + +def run(): + dist.init_process_group(backend="gloo") + rank = dist.get_rank() + + import distconv + + ps = distconv.ParallelStrategy(num_shards=(2,), shard_dim=(2,), device_type="cpu") + + # The same global volume on both ranks; each takes its own slab. + generator = torch.Generator().manual_seed(41) + volume = torch.randn(1, CHANNELS, SIZE, SIZE, SIZE, generator=generator) + half = SIZE // 2 + local = volume.narrow(2, rank * half, half).contiguous() + + norm = FastGroupNorm(GROUPS, CHANNELS) + param_generator = torch.Generator().manual_seed(97) + with torch.no_grad(): + norm.weight.normal_(1.0, 0.1, generator=param_generator) + norm.bias.normal_(0.0, 0.1, generator=param_generator) + + def forward(compiled): + """Run the wrapped GroupNorm with the compiled route on or off. + + The compiled route is forced on CPU by standing in the stock functional + kernel for the compiled callable: what is under test here is the + unwrap/rewrap plumbing at shard counts > 1, not Inductor. + """ + original_use, original_get = ( + gn_mod._use_compiled, + gn_mod._get_compiled_group_norm, + ) + if compiled: + gn_mod._use_compiled = lambda t: type(t) is torch.Tensor + gn_mod._get_compiled_group_norm = lambda: F.group_norm + else: + gn_mod._use_compiled = lambda t: False + try: + norm.zero_grad(set_to_none=True) + x = local.clone().requires_grad_(True) + out = norm(distconv.DCTensor.from_shard(x, ps)) + assert isinstance(out, distconv.DCTensor), type(out) + distconv.distconv._ToTensor.apply(out).pow(2).sum().backward() + weight_grad = norm.weight.grad + if isinstance(weight_grad, distconv.DCTensor): + weight_grad = weight_grad._tensor + return out._tensor.detach().clone(), x.grad.clone(), weight_grad.clone() + finally: + gn_mod._use_compiled = original_use + gn_mod._get_compiled_group_norm = original_get + + eager = forward(compiled=False) + compiled = forward(compiled=True) + identical = all(torch.equal(a, b) for a, b in zip(eager, compiled)) + + # Per-shard statistics: this rank's output normalizes its own slab only. + with torch.no_grad(): + per_shard = F.group_norm(local, GROUPS, norm.weight, norm.bias, norm.eps) + global_slice = F.group_norm( + volume, GROUPS, norm.weight, norm.bias, norm.eps + ).narrow(2, rank * half, half) + + # No spaces in any field: torchrun interleaves the ranks' stdout and a + # line can arrive without its trailing newline, so the parent's regex has + # to be able to tell two RESULT lines apart when they run together. + shape = "x".join(str(dim) for dim in compiled[0].shape) + print( + f"RESULT rank={rank} shape={shape} " + f"identical={identical} " + f"per_shard={torch.equal(compiled[0], per_shard)} " + f"global={torch.allclose(compiled[0], global_slice, atol=1e-6)}", + flush=True, + ) + print("DONE", flush=True) + dist.barrier() + dist.destroy_process_group() + + +if __name__ == "__main__": + run() diff --git a/tests/test_groupnorm.py b/tests/test_groupnorm.py index b1e8e6a..dec8b70 100644 --- a/tests/test_groupnorm.py +++ b/tests/test_groupnorm.py @@ -18,12 +18,21 @@ state dict as a stock ``nn.GroupNorm`` model (checkpoints stay interchangeable in both directions), the same numbers within reduction-order noise, and an eager fallback for every input the compiled kernel cannot or should not take -(CPU, tensor subclasses such as DistConv's ``DCTensor``, a broken compiler). +(CPU, unknown tensor subclasses, a broken compiler). DistConv's ``DCTensor`` +is not in that list: ``forward`` unwraps it to its local shard around the +compiled kernel, so the wrapped production path is served too. That unwrap is +*not* the bare attribute read DistConv's own dispatch does -- dispatch runs +below autograd, where a bare read is safe, while ``forward`` runs above it and +must go through DistConv's ``_ToTensor``/``_FromTensor`` pair to keep the graph +connected. """ from __future__ import annotations import logging +import os +import re +from pathlib import Path import pytest import torch @@ -32,6 +41,7 @@ from ScaFFold.unet import group_norm as gn_mod from ScaFFold.unet.group_norm import FastGroupNorm from ScaFFold.unet.unet_model import UNet +from tests.helpers import mpi_runner _N = 16 _N_CHANNELS = 3 @@ -190,10 +200,13 @@ def _boom(*a, **kw): def test_tensor_subclass_input_stays_eager(): - """DistConv wraps activations in a ``__torch_dispatch__`` tensor subclass. + """Unknown tensor subclasses must stay on the eager path. - Dynamo cannot trace those wrappers, so the predicate must reject anything - that is not exactly ``torch.Tensor`` before a compile is attempted. + Dynamo cannot trace ``__torch_dispatch__`` wrappers, so the predicate must + reject anything that is not exactly ``torch.Tensor`` before a compile is + attempted. DistConv's ``DCTensor`` is handled separately -- ``forward`` + unwraps it before consulting the predicate -- but any other wrapper has + unknown semantics and keeps the stock kernel. """ class _Wrapper(torch.Tensor): @@ -285,6 +298,285 @@ def test_recompile_limit_is_raised_never_lowered(): setattr(config, name, original) +# --------------------------------------------------------------------------- +# DCTensor routing: unwrap -> compiled kernel -> rewrap +# --------------------------------------------------------------------------- + + +@pytest.fixture +def dc_cpu(gloo_group_1rank): + """DistConv package plus a CPU ParallelStrategy over the 1-rank group. + + ``num_shards=(1, 1, 1)`` on dims (2, 3, 4) is exactly what worker.py builds + for a single-device run; a process group must exist even then. + """ + import distconv + + ps = distconv.ParallelStrategy( + num_shards=(1, 1, 1), shard_dim=(2, 3, 4), device_type="cpu" + ) + return distconv, ps + + +def _seeded_norm(channels=16): + """A FastGroupNorm with non-default affine params (the defaults are 1/0, + which would let a kernel that drops weight/bias slip through).""" + fast = FastGroupNorm(_GROUPS, channels) + generator = torch.Generator().manual_seed(97) + with torch.no_grad(): + fast.weight.normal_(1.0, 0.1, generator=generator) + fast.bias.normal_(0.0, 0.1, generator=generator) + return fast + + +def test_dctensor_routes_through_compiled_kernel(monkeypatch, dc_cpu): + """A DCTensor input reaches the compiled callable as its plain local shard. + + Verified with a recording stand-in for the compiled callable: it must see + exactly ``torch.Tensor`` (Dynamo cannot trace the wrapper), and the caller + must get a DCTensor back with the same values the stock kernel produces. + """ + distconv, ps = dc_cpu + seen = [] + + def _recording(input, num_groups, weight, bias, eps): + seen.append(type(input)) + return nn.functional.group_norm(input, num_groups, weight, bias, eps) + + monkeypatch.setattr(gn_mod, "_use_compiled", lambda t: type(t) is torch.Tensor) + monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _recording) + + fast = _seeded_norm() + x = _make_input(seed=21, channels=16, size=4) + out = fast(distconv.DCTensor.from_shard(x, ps)) + + assert seen == [torch.Tensor] + assert isinstance(out, distconv.DCTensor) + reference = nn.functional.group_norm(x, _GROUPS, fast.weight, fast.bias, fast.eps) + assert torch.equal(out._tensor, reference) + + +def test_dctensor_gradients_reach_the_layer_upstream(monkeypatch, dc_cpu): + """Gradients must flow past GroupNorm into the layer that produced its input. + + The unwrap has to go through DistConv's autograd pair rather than a bare + ``input._tensor`` read. The distinction is invisible when the DCTensor + wraps a leaf -- the leaf *is* ``_tensor``, so even a bare read reaches it -- + which is why this test puts a producer in front, as production does + (``conv -> GroupNorm`` in every block). With a bare read the graph is + severed there: the input and the producer's weight get no gradient at all + while GroupNorm's own weight/bias still look healthy. + """ + distconv, ps = dc_cpu + monkeypatch.setattr(gn_mod, "_use_compiled", lambda t: type(t) is torch.Tensor) + monkeypatch.setattr( + gn_mod, "_get_compiled_group_norm", lambda: nn.functional.group_norm + ) + + fast = _seeded_norm() + reference = nn.GroupNorm(_GROUPS, 16) + producer = nn.Conv3d(16, 16, 1, bias=False) + reference_producer = nn.Conv3d(16, 16, 1, bias=False) + with torch.no_grad(): + reference.weight.copy_(fast.weight) + reference.bias.copy_(fast.bias) + reference_producer.weight.copy_(producer.weight) + + x_fast = _make_input(seed=22, channels=16, size=4).requires_grad_(True) + x_ref = x_fast.detach().clone().requires_grad_(True) + + # A 1x1x1 conv on a DCTensor takes DistConv's convolution path, so the + # DCTensor handed to GroupNorm is a genuine intermediate, not a leaf. + out = fast(producer(distconv.DCTensor.from_shard(x_fast, ps))) + assert isinstance(out, distconv.DCTensor) + # Unwrap the way a downstream consumer would (autograd-aware) and drive a + # scalar backward through it. + distconv.distconv._ToTensor.apply(out).pow(2).sum().backward() + reference(reference_producer(x_ref)).pow(2).sum().backward() + + assert x_fast.grad is not None, "gradient never reached the input" + assert producer.weight.grad is not None, "gradient never reached the producer" + assert torch.equal(x_fast.grad, x_ref.grad) + assert torch.equal(producer.weight.grad, reference_producer.weight.grad) + assert torch.equal(fast.weight.grad, reference.weight.grad) + assert torch.equal(fast.bias.grad, reference.bias.grad) + + +def test_dctensor_on_cpu_never_invokes_torch_compile(monkeypatch, dc_cpu): + """The unwrap route obeys the same CPU guard as plain tensors. + + On CPU the local shard fails the ``is_cuda`` check, so a DCTensor must fall + through to the stock eager dispatch -- and still come back wrapped. The + stand-in records instead of only raising: ``forward`` catches ``Exception`` + to fall back, so a raise alone would be swallowed by the very code path + under test and the assertion would never fire. + """ + distconv, ps = dc_cpu + calls = [] + + def _boom(*a, **kw): + calls.append(a) + raise AssertionError("torch.compile must not be called for CPU tensors") + + monkeypatch.setattr(torch, "compile", _boom) + monkeypatch.setattr(gn_mod, "_compiled_group_norm", None) + gn_mod.set_compile_enabled(True) # even when explicitly forced on + + fast = _seeded_norm() + x = _make_input(seed=23, channels=16, size=4) + out = fast(distconv.DCTensor.from_shard(x, ps)) + + assert not calls + assert isinstance(out, distconv.DCTensor) + reference = nn.functional.group_norm(x, _GROUPS, fast.weight, fast.bias, fast.eps) + assert torch.equal(out._tensor, reference) + + +def test_dctensor_compile_failure_falls_back_to_eager(monkeypatch, caplog, dc_cpu): + """A broken compiler degrades the wrapped path to eager, like the plain one.""" + distconv, ps = dc_cpu + + def _raises(*args, **kwargs): + raise RuntimeError("simulated Inductor failure") + + monkeypatch.setattr(gn_mod, "_use_compiled", lambda t: type(t) is torch.Tensor) + monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _raises) + gn_mod._compile_failed = False + + fast = _seeded_norm() + x = _make_input(seed=24, channels=16, size=4) + with caplog.at_level(logging.WARNING, logger=gn_mod.__name__): + out = fast(distconv.DCTensor.from_shard(x, ps)) + + assert isinstance(out, distconv.DCTensor) + reference = nn.functional.group_norm(x, _GROUPS, fast.weight, fast.bias, fast.eps) + assert torch.equal(out._tensor, reference) + assert gn_mod._compile_failed is True + assert any("falling back to the eager kernel" in r.message for r in caplog.records) + + +def test_dctensor_two_shards_matches_eager_and_normalizes_per_shard(): + """Shard count > 1: same values as the eager route, still per-shard stats. + + Every other DCTensor test here runs ``num_shards=(1, 1, 1)``, where the + local shard is the whole tensor and the sharding is a no-op -- so none of + them can catch a fast path that quietly reduced over the wrong set of + elements. This one splits a spatial dim over two ranks and asserts both + halves of the claim the fast path rests on: bit-identical to the eager + wrapped route, and normalizing the local shard rather than the global + volume (which is DistConv's existing semantics, not something this change + introduces -- see the R32/R39 notes for the upstream discussion). + """ + script = ( + Path(__file__).resolve().parent + / "helpers" + / "rank_scripts" + / "groupnorm_shards_2rank.py" + ) + rc, out, err = mpi_runner.torchrun_gloo(str(script), n=2, timeout=180) + assert rc == 0, f"2-rank job failed rc={rc}\nstdout:\n{out}\nstderr:\n{err[-3000:]}" + + results = { + rank: match + for rank, *match in re.findall( + # Bounded alternatives, not \S+: the ranks' lines can arrive + # concatenated, so a greedy final field would swallow the next + # line's "RESULT". + r"RESULT rank=(\d+) shape=(\S+) identical=(True|False) " + r"per_shard=(True|False) global=(True|False)", + out, + ) + } + assert set(results) == {"0", "1"}, f"missing ranks\nstdout:\n{out}" + for rank, (shape, identical, per_shard, matches_global) in results.items(): + assert identical == "True", f"rank {rank}: compiled route differs from eager" + assert per_shard == "True", f"rank {rank}: not per-shard statistics" + assert matches_global == "False", f"rank {rank}: matched global statistics" + # Each rank holds half of the sharded dim. + assert shape == "1x16x4x8x8", f"rank {rank}: unexpected shard shape {shape}" + + +@pytest.fixture +def dc_cuda(): + """DistConv package plus a CUDA ParallelStrategy over a 1-rank NCCL group.""" + import distconv + import torch.distributed as dist + + created = False + if not dist.is_initialized(): + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29517") + dist.init_process_group(backend="nccl", rank=0, world_size=1) + created = True + ps = distconv.ParallelStrategy( + num_shards=(1, 1, 1), shard_dim=(2, 3, 4), device_type="cuda" + ) + yield distconv, ps + if created and dist.is_initialized(): + dist.destroy_process_group() + + +@pytest.mark.gpu +@pytest.mark.parametrize("autocast", [False, True]) +@pytest.mark.parametrize("layout", ["contiguous", "channels_last_3d"]) +def test_gpu_dctensor_matches_eager_dctensor(dc_cuda, autocast, layout): + """The compiled unwrap path matches today's eager wrapped path on GPU. + + This is the production configuration: worker.py wraps every activation in + a DCTensor (even at dc_num_shards=[1,1,1]), which used to force the eager + kernel. Values and gradients must agree within reduction-order noise, the + output must still be a DCTensor, and the compile must actually engage. + + Both layouts are covered because production requests ``channels_last_3d`` + (worker.py) and, with ``PYTORCH_MIOPEN_SUGGEST_NHWC=1`` set as it is there, + the convolutions really do hand GroupNorm channels-last activations. Only + parity is asserted, not the output layout: both routes return contiguous + today regardless of the input layout. + """ + distconv, ps = dc_cuda + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(31) + x = torch.randn(1, 64, 32, 32, 32, device=device, generator=generator) + if layout == "channels_last_3d": + x = x.to(memory_format=torch.channels_last_3d) + grad_out = torch.randn(*x.shape, device=device, generator=generator) + + fast = FastGroupNorm(_GROUPS, 64).to(device) + with torch.no_grad(): + fast.weight.normal_(1.0, 0.1, generator=generator) + fast.bias.normal_(0.0, 0.1, generator=generator) + + def plain(t): + return t._tensor if isinstance(t, distconv.DCTensor) else t + + def run(compiled): + gn_mod.set_compile_enabled(compiled) + inp = x.clone().requires_grad_(True) + fast.zero_grad(set_to_none=True) + with torch.autocast("cuda", dtype=torch.bfloat16, enabled=autocast): + out = fast(distconv.DCTensor.from_shard(inp, ps)) + assert isinstance(out, distconv.DCTensor) + local = distconv.distconv._ToTensor.apply(out) + local.backward(grad_out.to(local.dtype)) + return ( + local.detach(), + inp.grad.detach().clone(), + plain(fast.weight.grad).detach().clone(), + plain(fast.bias.grad).detach().clone(), + ) + + eager = run(False) + compiled = run(True) + assert gn_mod._compiled_group_norm is not None, "compiled path was not taken" + assert not gn_mod._compile_failed + + _assert_close(compiled[0], eager[0], 1e-5, "output") + _assert_close(compiled[1], eager[1], 1e-4, "d_input") + _assert_close(compiled[2], eager[2], 1e-4, "d_weight") + _assert_close(compiled[3], eager[3], 1e-4, "d_bias") + assert compiled[0].dtype == eager[0].dtype + + # --------------------------------------------------------------------------- # GPU behavior: numerics, single compile, checkpointing # --------------------------------------------------------------------------- From 76d30dfb4879fc54702b5dd51e8871b62e9857a9 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Sat, 1 Aug 2026 01:25:03 -0700 Subject: [PATCH 55/62] Add a channels-last-native Triton GroupNorm Nothing imports this yet; it is inert until a later commit wires it into FastGroupNorm. With PYTORCH_MIOPEN_SUGGEST_NHWC=1, which production sets, every convolution in the UNet emits channels_last_3d but GroupNorm consumes it and emits contiguous -- 22 layout breaks per forward -- and torch's compiled GroupNorm is 6.5x slower on channels-last input than on contiguous. The loss is pure access pattern: in NDHWC one program reads a dense (BLOCK_S, C) run and reshapes the inner axis to (G, C/G), so all groups' statistics come out of one coalesced pass, where Inductor walks the logical NCDHW order over channels-last memory as a strided gather. Measured on MI300A, fp32, 22 GroupNorm sites of a scale-8 UNet, fwd+bwd: 442.8 -> 69.0 ms/step (6.4x), which is a dead heat with compiled GroupNorm on contiguous input while additionally preserving the layout. 95-98% of measured streaming bandwidth at the large shapes. Statistics use a corrected two-pass per tile combined by Chan's parallel formula rather than E[x^2]-E[x]^2, which costs 0.8% at the dominant shape and is 4-23x more accurate than ATen at large input means (at mu/sigma=1e5 the naive form loses the variance entirely). Reductions are register-only with a fixed order and no float atomics, so results are bitwise reproducible. Registered as scaffold_gn::group_norm with a fake kernel and autograd, so it traces under torch.compile(fullgraph=True) without a graph break and composes with DistConv's DCTensor. Only 5-D channels_last_3d input takes the kernel; is_supported() returns False for contiguous input so callers keep their own compiled fallback, which is already at 89-92% of roofline there. Also: fp32/bf16/fp16 with fp32 statistics and torch's autocast contract, an int64 tile-base path for volumes past INT_MAX (measured free, verified at 2.16e9 elements), and an optional fused ReLU that is bit-exact against unfused+F.relu and takes 38% off the forward. --- ScaFFold/unet/triton_group_norm.py | 1385 ++++++++++++++++++++++++++++ tests/test_triton_group_norm.py | 868 +++++++++++++++++ 2 files changed, 2253 insertions(+) create mode 100644 ScaFFold/unet/triton_group_norm.py create mode 100644 tests/test_triton_group_norm.py diff --git a/ScaFFold/unet/triton_group_norm.py b/ScaFFold/unet/triton_group_norm.py new file mode 100644 index 0000000..cad62ad --- /dev/null +++ b/ScaFFold/unet/triton_group_norm.py @@ -0,0 +1,1385 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +"""Channels-last-native Triton GroupNorm (NDHWC in, NDHWC out). + +Why this exists +=============== +With ``PYTORCH_MIOPEN_SUGGEST_NHWC=1`` -- which every production ScaFFold run +sets -- every convolution in the UNet emits ``channels_last_3d`` activations, +but *every* stock GroupNorm variant (eager or Inductor-compiled) consumes them +through the *logical* NCDHW iteration order, which over a channels-last tensor +is a strided gather, and then emits a **contiguous** tensor. Measured on one +MI300A at scale 8 that costs 6.4x on GroupNorm itself (443 ms/step of GN +fwd+bwd against 69 ms/step here) *and* breaks the channels-last chain 22 times +per forward, forcing the following convolution to convert back. + +A channels-last-3d contiguous ``(N, C, D, H, W)`` tensor is *physically* a dense +``(N, S, C)`` array with ``S = D*H*W``; group ``g`` owns a contiguous run of +``C/G`` channels *inside every voxel*. The kernels below therefore let **one +program handle all groups at once** for a chunk of voxels: they read a dense +``(BLOCK_S, C)`` run, reshape the inner axis to ``(G, C/G)``, and get perfectly +coalesced loads and stores with the group axis costing nothing. Measured at +95-98% of this device's streaming roofline at the two largest UNet shapes. + +Measured on one MI300A (228 CUs), fp32, ``num_groups=8``, median of 20, the six +scale-8 UNet GroupNorm shapes, fwd / fwd+bwd in ms:: + + shape this compiled-CL compiled-CONT eager-CL fwd + [1,64,256^3] 4.28/11.61 19.51/77.01 4.77/11.85 151.2 + [1,128,128^3] 1.19/ 3.20 7.23/24.64 1.22/ 3.07 36.5 + [1,256,64^3] 0.35/ 0.97 2.27/ 7.42 0.34/ 0.84 9.1 + [1,512,32^3] 0.14/ 0.59 0.19/ 1.03 0.11/ 0.33 1.7 + [1,1024,16^3] 0.13/ 0.59 0.10/ 0.39 0.07/ 0.33 0.4 + [1,2048,8^3] 0.13/ 0.58 0.08/ 0.40 0.07/ 0.33 0.1 + +Over the 22 scale-8 call sites that is **442.8 -> 69.0 ms/step** of GroupNorm +fwd+bwd against today's production path (compiled GroupNorm on channels-last +input), i.e. **374 ms/step recovered**, and a dead heat with compiled GroupNorm +on *contiguous* input (66.4 ms/step) while additionally not breaking the +layout chain. The three smallest shapes lose on host dispatch, not on GPU +work -- see :func:`select_strategy`. + +Public API +========== +``triton_group_norm(input, num_groups, weight=None, bias=None, eps=1e-5, +activation=None)`` + Drop-in for ``F.group_norm`` (plus an optionally fused ReLU) with full + autograd support. Accepts *anything* ``F.group_norm`` accepts; inputs the + Triton kernel cannot serve fall back to ``F.group_norm`` internally (see + "Layouts" below). + +``is_supported(input, num_groups, weight=None, bias=None, activation=None)`` + Cheap, side-effect-free predicate: ``True`` exactly when the native Triton + kernel will run. Callers that already have a good fallback (e.g. a + ``torch.compile``d GroupNorm) should test this and route rejects + themselves; ``triton_group_norm``'s own fallback is plain eager + ``F.group_norm``. + +Contract +======== +For every input ``is_supported`` accepts, the result matches ``F.group_norm`` +to within fp32 reduction-order noise, with: + +* **dtype** -- output dtype is exactly ``F.group_norm``'s. Verified + empirically on this build (torch 2.13.0+rocm7.2): without autocast the output + dtype is the input dtype (fp32/bf16/fp16); under ``torch.autocast("cuda", + ...)`` GroupNorm is an fp32-policy op, so the output is **fp32** for any + input dtype. This module reproduces that rule (see ``_autocast_out_dtype``) + without materializing the fp32 copy of the input that autocast's cast would + create: the kernels read the input at its native width and accumulate in + fp32, which is bit-for-bit the same computation as upcasting first, but reads + half the bytes. Gradients follow the same rule: ``d_input`` has the input's + dtype, ``d_weight``/``d_bias`` have the parameter's dtype. Forward time at + ``[1,64,256^3]`` / ``[1,128,128^3]`` / ``[1,256,64^3]`` in ms: fp32 + 4.28/1.20/0.35, bf16 2.42/0.68/0.22, fp16 2.43/0.66/0.22, and bf16-in with + the fp32-out autocast contract 3.08/0.84/0.28 -- so honouring autocast's + fp32 output still buys 1.4x over fp32 end to end, because only the read side + narrows. +* **statistics** -- always accumulated in fp32, never in the input dtype. +* **memory format** -- the output has the *input's* memory format. This is the + one deliberate difference from stock GroupNorm, which returns a contiguous + tensor for every input layout (measured, see the table in + ``review/gn-dctensor/triton/RESULTS.md``); preserving channels-last is the + entire point of the kernel. +* **autograd** -- ``d_input``, ``d_weight``, ``d_bias``; ``weight=None`` and/or + ``bias=None`` supported. +* **determinism** -- bitwise reproducible run to run and process to process. + There are no float atomics anywhere, and the grid, split count and tile sizes + are pure functions of the shape (the tuning table is frozen in this file for + exactly that reason -- a *runtime* autotuner would break reproducibility by + changing the reduction order between runs). + +Reduction strategy +================== +Group statistics span ``S * C/G`` elements (134M at the largest UNet shape), so +one pass cannot produce them. Split-K partial reductions land at a fixed +scratch index and are combined by a fixed-order tree:: + + fwd: stats_partial -> stats_finalize -> normalize (3 kernels) + bwd: bwd_partial -> bwd_finalize -> dwdb_reduce -> dx (4 kernels) + +Traffic (``B = numel * itemsize``): 3B forward, 5B backward. + +Numerics: Welford, not ``E[x^2]-E[x]^2`` +======================================== +The prototype accumulated ``sum(x)`` and ``sum(x*x)`` and formed +``var = E[x^2] - E[x]^2``. That is split-friendly and cheap but cancels +catastrophically once ``mean >> std``, because it subtracts two nearly equal +large numbers to recover a small one. + +Here each tile instead produces ``(count, mean, M2)`` via a *corrected* +two-pass over registers -- ``mean0 = sum(x)/n``, then ``corr = sum(x-mean0)/n`` +to recover the digits the first sum lost, then ``M2 = sum((x-mean0-corr)^2)`` +-- and tiles and splits are merged with Chan's parallel combine. Every step is +register-only (the tile is read from HBM exactly once either way) and +atomic-free, so neither the traffic model nor determinism changes. + +Measured at ``[1,256,64^3]``, ``num_groups=8``, affine, relative error of the +*output* against a float64 reference computed from the same fp32 samples: + + x ~ N(mu, sigma) this kernel ATen fp32 E[x^2]-E[x]^2 + mu=0, sigma=1 1.6e-07 4.1e-07 1.8e-07 + mu=10, sigma=1 3.0e-07 6.9e-07 9.9e-06 + mu=100, sigma=1 8.0e-07 4.2e-06 5.6e-04 + mu=1e3, sigma=1e-2 1.1e-04 2.5e-03 2.3e+00 + +At ``mu/sigma = 1e5`` the old formulation has lost the variance outright (the +difference of the two ~1e6-sized fp32 terms is below one ulp, so ``rstd`` +saturates on ``eps`` and the output is meaningless), while this kernel is still +good to 1.1e-04 -- and is 4-23x *more* accurate than ATen's own fp32 GroupNorm +at every non-trivial mean. The residual 1.1e-04 is the fp32 representation +floor rather than an algorithm defect: a mean of 1e3 held in fp32 is quantized +to ~6e-5, which is 6e-3 of a standard deviation here, and both kernels sit on +that floor. + +Cost of the rewrite, isolated by timing the kernels alone against the +prototype's: **+0.8%** on the forward at ``[1,64,256^3]`` (the shape that +dominates the step), +3-8% at the middle shapes, and **0%** on the backward, +which does not compute a variance. A cheaper shifted-mean variant (two tile +reductions instead of three, shift taken from a peeled first tile) would +recover most of that; it was not worth the extra failure mode for ~4 ms/step. + +Layouts +======= +* ``channels_last_3d`` 5-D input -> **native Triton kernel**, channels-last + output. This is the fast path and the only one ``is_supported`` accepts. +* Plain contiguous NCDHW (and every other layout/rank) -> ``triton_group_norm`` + falls back to ``F.group_norm``, which returns a contiguous tensor, so the + input's memory format is still preserved. ``is_supported`` returns ``False`` + so that callers keep their own (probably compiled) fallback rather than + silently dropping to the eager kernel. + + This is a deliberate scope decision, not an oversight. A native NCDHW kernel + would need a *different* tiling -- with C outermost the fast axis is spatial, + so a program must own one group and stream S, rather than owning all groups + and streaming voxels -- i.e. a second family of four kernels. The payoff is + small: on contiguous input Inductor's compiled GroupNorm already reaches + 89-92% of this device's measured streaming roofline (RESULTS.md 4) -- and the + table above confirms it, 66.4 ms/step against this kernel's 69.0 -- so a + native NCDHW kernel could win ~10% there, against the 6.4x it wins on + channels-last input. If a mixed-layout model ever makes that 10% matter, the + place to add it is the strategy hook below. + +Addressing +========== +``[2, 64, 256^3]`` is *exactly* 2^31 elements, so int32 linear offsets block +batch>1 at the largest UNet shape and every shape above it. The kernels widen +only the **scalar tile base** to int64 (``INT64`` is a ``tl.constexpr``, so +shapes that fit still emit pure 32-bit code); the vector offsets inside a tile +span at most ``BLOCK_S*C + C`` elements and stay int32 either way. That is why +the wide path is free: forcing int64 on every scale-8 shape moves fwd+bwd by +-0.8% to +0.8% and forward by -3.7% to +4% (a 0.02 ms swing on the two smallest, +launch-bound shapes) -- noise in both directions. The switch is kept anyway +because it costs one constexpr and documents where the boundary is; correctness +above 2^31 elements is covered by a test at ``[2, 64, 256, 256, 257]`` +(2_155_872_256 elements, 2.5e-05 relative error on *both* batch items, i.e. the +same reduction noise a 134M-element fp32 reduction has anywhere). + +Fused activation +================ +``activation="relu"`` folds the ReLU into the forward store. In a store-bound +kernel that is free (one ``tl.maximum``) and it removes an entire 2B streaming +pass. Measured against ``F.relu(triton_group_norm(x))``: 38% off the forward +and 35% off fwd+bwd at ``[1,64,256^3]`` (6.94 -> 4.29 ms and 18.27 -> 11.80 +ms), 37%/33% at ``[1,128,128^3]``, tapering to ~11% at the launch-bound +shapes. + +The backward gates the incoming gradient on the sign of the **pre-activation** +value, which it *recomputes* from the saved ``(x, mean, rstd, weight, bias)`` +using the identical expression the forward used. Recomputation costs two FLOPs +on values already in registers and is bit-exact -- same inputs, same operation +order, same fp32 rounding -- so the sign always agrees with the forward's. The +alternative, testing ``y > 0`` on the saved output, would need the output kept +alive *in addition to* ``x`` (which the GroupNorm backward needs regardless), +and in bf16/fp16 it would also mis-gate any element whose positive +pre-activation rounded to zero on the store. + +Composition +=========== +Registered as real dispatcher ops (``scaffold_gn::group_norm`` / +``scaffold_gn::group_norm_backward``) via ``torch.library.custom_op``, with a +fake/meta kernel and ``register_autograd``. Consequences: + +* ``torch.compile(..., fullgraph=True)`` traces through without a graph break. +* Tensor subclasses that dispatch via ``__torch_dispatch__`` -- notably + DistConv's ``DCTensor`` -- intercept the op, unwrap to the local shard, run + it, and rewrap, so a DCTensor goes in and a DCTensor comes out with the graph + intact. As with the rest of DistConv today, statistics are per-shard. + +Going through the dispatcher costs ~35 us of host time per forward call +(measured against calling ``_forward``/``_backward`` directly), which is +invisible at the three largest shapes and is roughly two kernel launches at the +three smallest. That is the price of composing, and it is the same order as +the launch overhead those shapes already pay; see :func:`select_strategy`. + +Triton is imported lazily, on the first call that actually reaches the kernel, +so importing this module (or running the CPU test suite) costs nothing. +""" + +import functools +import importlib.util +from typing import Optional, Tuple + +import torch +import torch.nn.functional as F + +__all__ = [ + "triton_group_norm", + "is_supported", + "select_strategy", + "GNConfig", + "default_config", + "SUPPORTED_ACTIVATIONS", +] + +#: The activations that may be fused into the forward store. +SUPPORTED_ACTIVATIONS = (None, "relu") + +#: Input dtypes the kernels read directly (statistics are always fp32). +SUPPORTED_DTYPES = (torch.float32, torch.bfloat16, torch.float16) + +#: Largest linear element index representable in int32. +_INT32_MAX = 2**31 - 1 + + +# --------------------------------------------------------------------------- # +# tiling configuration +# --------------------------------------------------------------------------- # +class GNConfig: + """Tiling knobs. A pure function of the shape => bitwise determinism. + + ``stats_tile``/``elem_tile`` are *element* budgets per program (the spatial + block is ``tile // channels_per_voxel``, rounded down to a power of two); + ``nsplit_target`` is the total number of split-K partials wanted across the + batch, so the per-sample split count is ``nsplit_target // N``. + """ + + __slots__ = ( + "stats_tile", + "stats_warps", + "nsplit_target", + "elem_tile", + "elem_warps", + ) + + def __init__( + self, + stats_tile=8192, + stats_warps=4, + nsplit_target=2048, + elem_tile=8192, + elem_warps=4, + ): + self.stats_tile = stats_tile + self.stats_warps = stats_warps + self.nsplit_target = nsplit_target + self.elem_tile = elem_tile + self.elem_warps = elem_warps + + def key(self): + return ( + self.stats_tile, + self.stats_warps, + self.nsplit_target, + self.elem_tile, + self.elem_warps, + ) + + def __eq__(self, other): + return isinstance(other, GNConfig) and self.key() == other.key() + + def __hash__(self): + return hash(self.key()) + + def __repr__(self): + return ( + "GNConfig(stats_tile=%d, stats_warps=%d, nsplit_target=%d, " + "elem_tile=%d, elem_warps=%d)" % self.key() + ) + + +#: Frozen tuning table, produced by coordinate descent on fwd+bwd time on one +#: MI300A (228 CUs) at fp32 with ``num_groups=8``, keyed by the +#: ``(num_channels, cube-root spatial extent)`` of the scale-8 ScaFFold UNet +#: GroupNorm sites. Frozen -- never autotuned at run time -- because the split +#: count fixes the reduction order and therefore the bits of the result. +_TUNED = { + (64, 256): GNConfig(8192, 4, 2048, 8192, 4), + (128, 128): GNConfig(16384, 4, 512, 16384, 4), + (256, 64): GNConfig(4096, 4, 512, 8192, 4), + (512, 32): GNConfig(32768, 4, 8192, 8192, 4), + (1024, 16): GNConfig(16384, 8, 512, 8192, 4), + (2048, 8): GNConfig(32768, 4, 512, 8192, 4), +} + +_DEFAULT_CONFIG = GNConfig() + + +def default_config(num_channels: int, spatial: int) -> GNConfig: + """Tiling for ``num_channels`` channels and ``spatial = D*H*W`` voxels.""" + edge = round(spatial ** (1.0 / 3.0)) + if edge**3 != spatial: + edge = None + return _TUNED.get((num_channels, edge), _DEFAULT_CONFIG) + + +# --------------------------------------------------------------------------- # +# small-shape dispatch hook +# --------------------------------------------------------------------------- # +#: Every strategy name ``select_strategy`` may return. Only ``"split_k"`` is +#: implemented; anything else raises rather than silently doing the wrong +#: thing. +STRATEGIES = ("split_k",) + +#: Spatial extent (``D*H*W``) below which the split-K chain is expected to be +#: host-dispatch bound rather than bandwidth bound. Measured on MI300A: at +#: ``[1,2048,8^3]`` the seven kernels do 0.031 ms of GPU work behind 0.600 ms +#: of Python/autograd/launch cost, a 19x overhead tax (RESULTS.md 3). Purely +#: informational today -- ``select_strategy`` does not use it yet. +SMALL_SPATIAL_THRESHOLD = 4096 + + +def select_strategy(n: int, num_channels: int, spatial: int, num_groups: int) -> str: + """### SMALL-SHAPE DISPATCH HOOK ### -- the single point where a different + kernel strategy is chosen for a shape. + + Returns a name from :data:`STRATEGIES`. Today it always returns + ``"split_k"``: three forward and four backward kernels with split-K partial + reductions, which is bandwidth-optimal for the large shapes but pays seven + kernel launches (~17 us each on this node) plus autograd overhead + regardless of size -- so below roughly ``SMALL_SPATIAL_THRESHOLD`` voxels + the whole call is host bound and a *single-program-per-(n, group)* kernel + that never leaves registers would win. + + That regime is under active investigation; when a second strategy lands, + add its name to :data:`STRATEGIES`, return it from here on a rule that is a + **pure function of the shape** (determinism depends on it), and branch on + it in ``_dispatch`` -- which is the only caller, sits in front of the + memoized tiling plan, and is itself called by both ``_forward`` and + ``_backward``. Nothing else in this file needs to change. + """ + return "split_k" + + +# --------------------------------------------------------------------------- # +# planning helpers +# --------------------------------------------------------------------------- # +def _prev_pow2(x: int) -> int: + p = 1 + while p * 2 <= x: + p *= 2 + return p + + +def _next_pow2(x: int) -> int: + p = 1 + while p < x: + p *= 2 + return p + + +def _cdiv(a: int, b: int) -> int: + return -(-a // b) + + +class _Plan: + """Everything the launcher needs, derived only from the shape + config.""" + + __slots__ = ( + "n", + "channels", + "spatial", + "groups", + "group_channels", + "groups_p2", + "group_channels_p2", + "masked_c", + "int64", + "block_s_stats", + "nsplit", + "chunk", + "block_s_elem", + "nblk_elem", + "cfg", + ) + + def __init__(self, n, channels, spatial, groups, cfg, numel): + self.n = n + self.channels = channels + self.spatial = spatial + self.groups = groups + self.group_channels = channels // groups + self.groups_p2 = _next_pow2(groups) + self.group_channels_p2 = _next_pow2(self.group_channels) + # Only power-of-two group/channel counts tile the (G, C/G) axes exactly; + # anything else is rounded up and masked, which is correct but reads a + # few lanes it throws away. + self.masked_c = ( + self.groups_p2 != groups or self.group_channels_p2 != self.group_channels + ) + # int64 addressing is needed once a linear element index can exceed + # INT32_MAX. [2,64,256^3] is exactly 2^31 elements, so this is not + # hypothetical at scale. Only the *scalar* tile base is widened (see + # the kernels), which measurement showed to be free. + self.int64 = numel + channels > _INT32_MAX + self.cfg = cfg + + voxel = self.groups_p2 * self.group_channels_p2 + self.block_s_stats = max(1, _prev_pow2(cfg.stats_tile // max(1, voxel))) + ntiles = max(1, spatial // self.block_s_stats) + self.nsplit = _prev_pow2( + max(1, min(ntiles, max(1, cfg.nsplit_target // max(1, n)))) + ) + self.chunk = _cdiv(spatial, self.nsplit) + self.block_s_elem = max(1, _prev_pow2(cfg.elem_tile // max(1, voxel))) + self.nblk_elem = _cdiv(spatial, self.block_s_elem) + + @property + def elements_per_group(self) -> float: + return float(self.spatial * self.group_channels) + + +@functools.lru_cache(maxsize=256) +def _plan(n, channels, spatial, groups, numel) -> _Plan: + """Memoized: a UNet presents a handful of shapes, and the three smallest + GroupNorm sites are host-dispatch bound, so rebuilding the plan (two + power-of-two loops and a dict lookup) on every call is measurable there. + Memoization cannot affect results -- the plan is a pure function of its + arguments, which is also what makes the kernels bitwise reproducible.""" + return _Plan(n, channels, spatial, groups, default_config(channels, spatial), numel) + + +def _dispatch(n, channels, spatial, groups, numel) -> _Plan: + """Consult the strategy hook, then build (or reuse) the tiling plan.""" + strategy = select_strategy(n, channels, spatial, groups) + if strategy != "split_k": + raise NotImplementedError( + f"kernel strategy {strategy!r} selected by select_strategy() is not " + f"implemented; known strategies are {STRATEGIES}" + ) + return _plan(n, channels, spatial, groups, numel) + + +# --------------------------------------------------------------------------- # +# Triton kernels (built lazily -- importing this module must not import triton) +# --------------------------------------------------------------------------- # +triton = None +tl = None +_welford_combine = None +_stats_partial_kernel = None +_stats_finalize_kernel = None +_normalize_kernel = None +_bwd_partial_kernel = None +_bwd_finalize_kernel = None +_dwdb_reduce_kernel = None +_dx_kernel = None + + +_TRITON_AVAILABLE = None + + +def triton_available() -> bool: + """Whether Triton is importable, without importing it. + + ``find_spec`` on a top-level name only touches the finders, so this stays + side-effect free and is safe to call from :func:`is_supported`. The answer + is memoized in a plain global rather than an ``lru_cache`` because Dynamo + warns (loudly, once per process) when it traces through a cache wrapper. + """ + global _TRITON_AVAILABLE + if _TRITON_AVAILABLE is None: + try: + _TRITON_AVAILABLE = importlib.util.find_spec("triton") is not None + except (ImportError, ValueError): + _TRITON_AVAILABLE = False + return _TRITON_AVAILABLE + + +def _build_kernels(): + """Import Triton and install the JIT kernels into this module's globals. + + The kernels are defined inside a function purely so that ``import triton`` + is deferred to the first GPU call; they are written into ``globals()`` so + Triton's name resolution (which reads ``fn.__globals__``) sees them. + """ + global triton, tl + import triton as _triton + import triton.language as _tl + + triton = _triton + tl = _tl + + # ---------------------------------------------------------------- stats -- + @_triton.jit + def _welford_combine(cnt_a, mean_a, m2_a, cnt_b, mean_b, m2_b): + """Chan's parallel merge of two (count, mean, M2) triples. + + Exact for empty partials on either side (``cnt == 0`` leaves the other + triple untouched), which matters because the last split of a shape + whose spatial extent is not a multiple of the chunk size can be empty. + """ + cnt = cnt_a + cnt_b + denom = tl.where(cnt == 0.0, 1.0, cnt) + delta = mean_b - mean_a + mean = mean_a + delta * (cnt_b / denom) + m2 = m2_a + m2_b + delta * delta * (cnt_a * cnt_b / denom) + return cnt, mean, m2 + + @_triton.jit + def _stats_partial_kernel( + X, + PCNT, + PMEAN, + PM2, + S, + CHUNK, + C: tl.constexpr, + G: tl.constexpr, + CG: tl.constexpr, + GP: tl.constexpr, + CGP: tl.constexpr, + NSPLIT: tl.constexpr, + BLOCK_S: tl.constexpr, + MASKED_C: tl.constexpr, + INT64: tl.constexpr, + ): + """One program per ``(split, n)``: Welford partials for every group. + + Reads a dense ``(BLOCK_S, C)`` run of memory per step -- perfectly + coalesced -- and produces the statistics of all G groups at once, the + group axis being the inner channel axis reshaped to ``(G, C/G)``. + """ + sp = tl.program_id(0) + n = tl.program_id(1) + + offs_g = tl.arange(0, GP) + offs_j = tl.arange(0, CGP) + offs_s = tl.arange(0, BLOCK_S) + inner = offs_g[None, :, None] * CG + offs_j[None, None, :] + cmask = (offs_g[None, :, None] < G) & (offs_j[None, None, :] < CG) + off = offs_s[:, None, None] * C + inner + + s_begin = sp * CHUNK + s_end = tl.minimum(s_begin + CHUNK, S) + + cnt = tl.zeros((GP,), dtype=tl.float32) + mean = tl.zeros((GP,), dtype=tl.float32) + m2 = tl.zeros((GP,), dtype=tl.float32) + + for s0 in range(s_begin, s_end, BLOCK_S): + # Only the scalar tile base is ever widened to int64; the vector + # offsets stay int32 because they span at most BLOCK_S*C+C + # elements. That keeps the wide arithmetic off the hot path. + if INT64: + base = (n.to(tl.int64) * S + s0) * C + else: + base = (n * S + s0) * C + nvalid = tl.minimum(BLOCK_S, s_end - s0) + m = tl.broadcast_to((offs_s < nvalid)[:, None, None], (BLOCK_S, GP, CGP)) + if MASKED_C: + m = m & cmask + x = tl.load(X + base + off, mask=m, other=0.0).to(tl.float32) + + # Corrected two-pass within the tile: the first mean loses digits + # to the magnitude of the data, `corr` puts them back, and the + # centred squares are then accurate to fp32 roundoff. Everything + # here is register traffic; the tile is read from HBM exactly once. + cnt_t = (nvalid * CG).to(tl.float32) + mean0 = tl.sum(tl.sum(x, 2), 0) / cnt_t + d = tl.where(m, x - mean0[None, :, None], 0.0) + corr = tl.sum(tl.sum(d, 2), 0) / cnt_t + dd = tl.where(m, d - corr[None, :, None], 0.0) + m2_t = tl.sum(tl.sum(dd * dd, 2), 0) + mean_t = mean0 + corr + + new_cnt = cnt + cnt_t + delta = mean_t - mean + mean = mean + delta * (cnt_t / new_cnt) + m2 = m2 + m2_t + delta * delta * (cnt * cnt_t / new_cnt) + cnt = new_cnt + + o = (n * NSPLIT + sp) * G + offs_g + gm = offs_g < G + tl.store(PCNT + o, cnt, mask=gm) + tl.store(PMEAN + o, mean, mask=gm) + tl.store(PM2 + o, m2, mask=gm) + + @_triton.jit + def _stats_finalize_kernel( + PCNT, + PMEAN, + PM2, + MEAN, + RSTD, + M, + eps, + G: tl.constexpr, + NSPLIT: tl.constexpr, + ): + """Merge the NSPLIT partials of one ``(n, g)`` into mean and 1/std.""" + pid = tl.program_id(0) # n * G + g + n = pid // G + g = pid % G + offs = tl.arange(0, NSPLIT) + idx = (n * NSPLIT + offs) * G + g + cnt, mean, m2 = tl.reduce( + (tl.load(PCNT + idx), tl.load(PMEAN + idx), tl.load(PM2 + idx)), + 0, + _welford_combine, + ) + # `cnt` equals M by construction; M is passed in so the divisor is the + # exact element count rather than a float accumulated from partials. + var = m2 / M + tl.store(MEAN + pid, mean) + tl.store(RSTD + pid, 1.0 / tl.sqrt(var + eps)) + + # ------------------------------------------------------------ normalize -- + @_triton.jit + def _normalize_kernel( + X, + Y, + MEAN, + RSTD, + W, + B, + S, + C: tl.constexpr, + G: tl.constexpr, + CG: tl.constexpr, + GP: tl.constexpr, + CGP: tl.constexpr, + BLOCK_S: tl.constexpr, + RELU: tl.constexpr, + HAS_W: tl.constexpr, + HAS_B: tl.constexpr, + MASKED_C: tl.constexpr, + INT64: tl.constexpr, + ): + blk = tl.program_id(0) + n = tl.program_id(1) + + offs_g = tl.arange(0, GP) + offs_j = tl.arange(0, CGP) + offs_s = tl.arange(0, BLOCK_S) + inner = offs_g[None, :, None] * CG + offs_j[None, None, :] + off = offs_s[:, None, None] * C + inner + wb = offs_g[:, None] * CG + offs_j[None, :] + wbm = (offs_g[:, None] < G) & (offs_j[None, :] < CG) + + s0 = blk * BLOCK_S + m = tl.broadcast_to((offs_s < S - s0)[:, None, None], (BLOCK_S, GP, CGP)) + if MASKED_C: + m = m & ((offs_g[None, :, None] < G) & (offs_j[None, None, :] < CG)) + if INT64: + base = (n.to(tl.int64) * S + s0) * C + else: + base = (n * S + s0) * C + + gm = offs_g < G + mean = tl.load(MEAN + n * G + offs_g, mask=gm, other=0.0)[None, :, None] + rstd = tl.load(RSTD + n * G + offs_g, mask=gm, other=0.0)[None, :, None] + if HAS_W: + w = tl.load(W + wb, mask=wbm, other=0.0).to(tl.float32)[None, :, :] + else: + w = tl.full((1, GP, CGP), 1.0, tl.float32) + if HAS_B: + b = tl.load(B + wb, mask=wbm, other=0.0).to(tl.float32)[None, :, :] + else: + b = tl.zeros((1, GP, CGP), dtype=tl.float32) + + x = tl.load(X + base + off, mask=m, other=0.0).to(tl.float32) + xhat = (x - mean) * rstd + y = xhat * w + b + if RELU: + y = tl.maximum(y, 0.0) + tl.store(Y + base + off, y.to(Y.dtype.element_ty), mask=m) + + # ------------------------------------------------------------- backward -- + @_triton.jit + def _bwd_partial_kernel( + X, + DY, + MEAN, + RSTD, + W, + B, + PS1, + PS2, + PDW, + PDB, + S, + CHUNK, + C: tl.constexpr, + G: tl.constexpr, + CG: tl.constexpr, + GP: tl.constexpr, + CGP: tl.constexpr, + NSPLIT: tl.constexpr, + BLOCK_S: tl.constexpr, + RELU: tl.constexpr, + HAS_W: tl.constexpr, + HAS_B: tl.constexpr, + MASKED_C: tl.constexpr, + INT64: tl.constexpr, + ): + """Partials for the two per-``(n, g)`` reductions used by dx, and for + the per-channel dweight / dbias reductions.""" + sp = tl.program_id(0) + n = tl.program_id(1) + + offs_g = tl.arange(0, GP) + offs_j = tl.arange(0, CGP) + offs_s = tl.arange(0, BLOCK_S) + inner = offs_g[None, :, None] * CG + offs_j[None, None, :] + cmask = (offs_g[None, :, None] < G) & (offs_j[None, None, :] < CG) + off = offs_s[:, None, None] * C + inner + wb = offs_g[:, None] * CG + offs_j[None, :] + wbm = (offs_g[:, None] < G) & (offs_j[None, :] < CG) + + s_begin = sp * CHUNK + s_end = tl.minimum(s_begin + CHUNK, S) + + gm = offs_g < G + mean = tl.load(MEAN + n * G + offs_g, mask=gm, other=0.0)[None, :, None] + rstd = tl.load(RSTD + n * G + offs_g, mask=gm, other=0.0)[None, :, None] + if HAS_W: + w = tl.load(W + wb, mask=wbm, other=0.0).to(tl.float32)[None, :, :] + else: + w = tl.full((1, GP, CGP), 1.0, tl.float32) + if HAS_B: + b = tl.load(B + wb, mask=wbm, other=0.0).to(tl.float32)[None, :, :] + else: + b = tl.zeros((1, GP, CGP), dtype=tl.float32) + + acc1 = tl.zeros((GP,), dtype=tl.float32) + acc2 = tl.zeros((GP,), dtype=tl.float32) + accdw = tl.zeros((GP, CGP), dtype=tl.float32) + accdb = tl.zeros((GP, CGP), dtype=tl.float32) + + for s0 in range(s_begin, s_end, BLOCK_S): + if INT64: + base = (n.to(tl.int64) * S + s0) * C + else: + base = (n * S + s0) * C + nvalid = tl.minimum(BLOCK_S, s_end - s0) + m = tl.broadcast_to((offs_s < nvalid)[:, None, None], (BLOCK_S, GP, CGP)) + if MASKED_C: + m = m & cmask + x = tl.load(X + base + off, mask=m, other=0.0).to(tl.float32) + dy = tl.load(DY + base + off, mask=m, other=0.0).to(tl.float32) + xhat = (x - mean) * rstd + if RELU: + # Identical expression (and therefore identical rounding) to + # the forward's pre-activation, so the sign test agrees with + # the forward bit for bit. Masked lanes carry dy == 0, so + # gating cannot resurrect them. + dy = tl.where(xhat * w + b > 0.0, dy, 0.0) + dyw = dy * w + acc1 += tl.sum(tl.sum(dyw, 2), 0) + acc2 += tl.sum(tl.sum(dyw * xhat, 2), 0) + accdw += tl.sum(dy * xhat, 0) + accdb += tl.sum(dy, 0) + + o = (n * NSPLIT + sp) * G + offs_g + tl.store(PS1 + o, acc1, mask=gm) + tl.store(PS2 + o, acc2, mask=gm) + row = (n * NSPLIT + sp) * C + wb + tl.store(PDW + row, accdw, mask=wbm) + tl.store(PDB + row, accdb, mask=wbm) + + @_triton.jit + def _bwd_finalize_kernel( + PS1, + PS2, + C1, + C2, + M, + G: tl.constexpr, + NSPLIT: tl.constexpr, + ): + pid = tl.program_id(0) # n * G + g + n = pid // G + g = pid % G + offs = tl.arange(0, NSPLIT) + idx = (n * NSPLIT + offs) * G + g + tl.store(C1 + pid, tl.sum(tl.load(PS1 + idx)) / M) + tl.store(C2 + pid, tl.sum(tl.load(PS2 + idx)) / M) + + @_triton.jit + def _dwdb_reduce_kernel( + PDW, + PDB, + DW, + DB, + ROWS, + C, + BLOCK_C: tl.constexpr, + BLOCK_R: tl.constexpr, + ): + pid = tl.program_id(0) + offs_c = pid * BLOCK_C + tl.arange(0, BLOCK_C) + mc = offs_c < C + accw = tl.zeros((BLOCK_C,), dtype=tl.float32) + accb = tl.zeros((BLOCK_C,), dtype=tl.float32) + for r0 in range(0, ROWS, BLOCK_R): + offs_r = r0 + tl.arange(0, BLOCK_R) + m = (offs_r[:, None] < ROWS) & mc[None, :] + off = offs_r[:, None] * C + offs_c[None, :] + accw += tl.sum(tl.load(PDW + off, mask=m, other=0.0), 0) + accb += tl.sum(tl.load(PDB + off, mask=m, other=0.0), 0) + tl.store(DW + offs_c, accw, mask=mc) + tl.store(DB + offs_c, accb, mask=mc) + + @_triton.jit + def _dx_kernel( + X, + DY, + DX, + MEAN, + RSTD, + W, + B, + C1, + C2, + S, + C: tl.constexpr, + G: tl.constexpr, + CG: tl.constexpr, + GP: tl.constexpr, + CGP: tl.constexpr, + BLOCK_S: tl.constexpr, + RELU: tl.constexpr, + HAS_W: tl.constexpr, + HAS_B: tl.constexpr, + MASKED_C: tl.constexpr, + INT64: tl.constexpr, + ): + blk = tl.program_id(0) + n = tl.program_id(1) + + offs_g = tl.arange(0, GP) + offs_j = tl.arange(0, CGP) + offs_s = tl.arange(0, BLOCK_S) + inner = offs_g[None, :, None] * CG + offs_j[None, None, :] + off = offs_s[:, None, None] * C + inner + wb = offs_g[:, None] * CG + offs_j[None, :] + wbm = (offs_g[:, None] < G) & (offs_j[None, :] < CG) + + s0 = blk * BLOCK_S + m = tl.broadcast_to((offs_s < S - s0)[:, None, None], (BLOCK_S, GP, CGP)) + if MASKED_C: + m = m & ((offs_g[None, :, None] < G) & (offs_j[None, None, :] < CG)) + if INT64: + base = (n.to(tl.int64) * S + s0) * C + else: + base = (n * S + s0) * C + + gm = offs_g < G + mean = tl.load(MEAN + n * G + offs_g, mask=gm, other=0.0)[None, :, None] + rstd = tl.load(RSTD + n * G + offs_g, mask=gm, other=0.0)[None, :, None] + c1 = tl.load(C1 + n * G + offs_g, mask=gm, other=0.0)[None, :, None] + c2 = tl.load(C2 + n * G + offs_g, mask=gm, other=0.0)[None, :, None] + if HAS_W: + w = tl.load(W + wb, mask=wbm, other=0.0).to(tl.float32)[None, :, :] + else: + w = tl.full((1, GP, CGP), 1.0, tl.float32) + if HAS_B: + b = tl.load(B + wb, mask=wbm, other=0.0).to(tl.float32)[None, :, :] + else: + b = tl.zeros((1, GP, CGP), dtype=tl.float32) + + x = tl.load(X + base + off, mask=m, other=0.0).to(tl.float32) + dy = tl.load(DY + base + off, mask=m, other=0.0).to(tl.float32) + xhat = (x - mean) * rstd + if RELU: + dy = tl.where(xhat * w + b > 0.0, dy, 0.0) + dyw = dy * w + dx = rstd * (dyw - c1 - xhat * c2) + tl.store(DX + base + off, dx.to(DX.dtype.element_ty), mask=m) + + globals().update( + _welford_combine=_welford_combine, + _stats_partial_kernel=_stats_partial_kernel, + _stats_finalize_kernel=_stats_finalize_kernel, + _normalize_kernel=_normalize_kernel, + _bwd_partial_kernel=_bwd_partial_kernel, + _bwd_finalize_kernel=_bwd_finalize_kernel, + _dwdb_reduce_kernel=_dwdb_reduce_kernel, + _dx_kernel=_dx_kernel, + ) + + +def _ensure_kernels(): + if _stats_partial_kernel is None: + _build_kernels() + + +# --------------------------------------------------------------------------- # +# python drivers +# --------------------------------------------------------------------------- # +_CL_FORMAT = torch.channels_last_3d + + +def _shape_of(input: torch.Tensor): + n, channels = input.shape[0], input.shape[1] + spatial = 1 + for d in input.shape[2:]: + spatial *= d + return n, channels, spatial + + +def _forward(input, num_groups, weight, bias, eps, activation, out_dtype): + _ensure_kernels() + n, channels, spatial = _shape_of(input) + plan = _dispatch(n, channels, spatial, num_groups, input.numel()) + groups = num_groups + device = input.device + + pcnt = torch.empty(n * plan.nsplit * groups, device=device, dtype=torch.float32) + pmean = torch.empty_like(pcnt) + pm2 = torch.empty_like(pcnt) + mean = torch.empty((n, groups), device=device, dtype=torch.float32) + rstd = torch.empty_like(mean) + + _stats_partial_kernel[(plan.nsplit, n)]( + input, + pcnt, + pmean, + pm2, + spatial, + plan.chunk, + C=channels, + G=groups, + CG=plan.group_channels, + GP=plan.groups_p2, + CGP=plan.group_channels_p2, + NSPLIT=plan.nsplit, + BLOCK_S=plan.block_s_stats, + MASKED_C=plan.masked_c, + INT64=plan.int64, + num_warps=plan.cfg.stats_warps, + ) + _stats_finalize_kernel[(n * groups,)]( + pcnt, + pmean, + pm2, + mean, + rstd, + plan.elements_per_group, + eps, + G=groups, + NSPLIT=plan.nsplit, + num_warps=4, + ) + + out = torch.empty_like(input, dtype=out_dtype, memory_format=_CL_FORMAT) + _normalize_kernel[(plan.nblk_elem, n)]( + input, + out, + mean, + rstd, + weight, + bias, + spatial, + C=channels, + G=groups, + CG=plan.group_channels, + GP=plan.groups_p2, + CGP=plan.group_channels_p2, + BLOCK_S=plan.block_s_elem, + RELU=activation == "relu", + HAS_W=weight is not None, + HAS_B=bias is not None, + MASKED_C=plan.masked_c, + INT64=plan.int64, + num_warps=plan.cfg.elem_warps, + ) + return out, mean, rstd + + +def _backward(grad_out, input, weight, bias, mean, rstd, num_groups, activation): + _ensure_kernels() + n, channels, spatial = _shape_of(input) + plan = _dispatch(n, channels, spatial, num_groups, input.numel()) + groups = num_groups + device = input.device + + ps1 = torch.empty(n * plan.nsplit * groups, device=device, dtype=torch.float32) + ps2 = torch.empty_like(ps1) + pdw = torch.empty(n * plan.nsplit * channels, device=device, dtype=torch.float32) + pdb = torch.empty_like(pdw) + + _bwd_partial_kernel[(plan.nsplit, n)]( + input, + grad_out, + mean, + rstd, + weight, + bias, + ps1, + ps2, + pdw, + pdb, + spatial, + plan.chunk, + C=channels, + G=groups, + CG=plan.group_channels, + GP=plan.groups_p2, + CGP=plan.group_channels_p2, + NSPLIT=plan.nsplit, + BLOCK_S=plan.block_s_stats, + RELU=activation == "relu", + HAS_W=weight is not None, + HAS_B=bias is not None, + MASKED_C=plan.masked_c, + INT64=plan.int64, + num_warps=plan.cfg.stats_warps, + ) + + c1 = torch.empty(n * groups, device=device, dtype=torch.float32) + c2 = torch.empty_like(c1) + _bwd_finalize_kernel[(n * groups,)]( + ps1, + ps2, + c1, + c2, + plan.elements_per_group, + G=groups, + NSPLIT=plan.nsplit, + num_warps=4, + ) + + d_weight = torch.empty(channels, device=device, dtype=torch.float32) + d_bias = torch.empty_like(d_weight) + rows = n * plan.nsplit + block_c = min(256, max(64, _next_pow2(channels))) + block_r = 32 if rows >= 32 else 1 + _dwdb_reduce_kernel[(_cdiv(channels, block_c),)]( + pdw, + pdb, + d_weight, + d_bias, + rows, + channels, + BLOCK_C=block_c, + BLOCK_R=block_r, + num_warps=4, + ) + + d_input = torch.empty_like(input, memory_format=_CL_FORMAT) + _dx_kernel[(plan.nblk_elem, n)]( + input, + grad_out, + d_input, + mean, + rstd, + weight, + bias, + c1, + c2, + spatial, + C=channels, + G=groups, + CG=plan.group_channels, + GP=plan.groups_p2, + CGP=plan.group_channels_p2, + BLOCK_S=plan.block_s_elem, + RELU=activation == "relu", + HAS_W=weight is not None, + HAS_B=bias is not None, + MASKED_C=plan.masked_c, + INT64=plan.int64, + num_warps=plan.cfg.elem_warps, + ) + return d_input, d_weight, d_bias + + +# --------------------------------------------------------------------------- # +# torch.library registration +# --------------------------------------------------------------------------- # +def _validate(input, num_groups, weight, bias, activation): + if activation not in SUPPORTED_ACTIVATIONS: + raise ValueError( + f"activation must be one of {SUPPORTED_ACTIVATIONS}, got {activation!r}" + ) + if input.dim() != 5: + raise ValueError(f"expected a 5-D NCDHW tensor, got {tuple(input.shape)}") + if num_groups <= 0 or input.shape[1] % num_groups != 0: + raise ValueError( + f"num_channels={input.shape[1]} is not divisible by num_groups={num_groups}" + ) + if input.dtype not in SUPPORTED_DTYPES: + raise ValueError(f"unsupported input dtype {input.dtype}") + if not input.is_contiguous(memory_format=_CL_FORMAT): + # Required, not converted: the fake kernel promises the *input's* + # memory format for the output, so silently converting here would make + # the traced and eager results disagree on strides. The public + # ``triton_group_norm`` routes non-channels-last input to + # ``F.group_norm`` before it ever reaches this op. + raise ValueError( + "input must be channels_last_3d-contiguous; use triton_group_norm() " + "which falls back to F.group_norm for other layouts" + ) + for name, t in (("weight", weight), ("bias", bias)): + if t is not None and t.numel() != input.shape[1]: + raise ValueError( + f"{name} has {t.numel()} elements, expected {input.shape[1]}" + ) + + +@torch.library.custom_op( + "scaffold_gn::group_norm", mutates_args=(), device_types="cuda" +) +def _group_norm_op( + input: torch.Tensor, + num_groups: int, + weight: Optional[torch.Tensor], + bias: Optional[torch.Tensor], + eps: float, + activation: Optional[str], + out_dtype: Optional[torch.dtype], +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Channels-last GroupNorm forward: returns ``(output, mean, rstd)``. + + ``mean``/``rstd`` are ``(N, num_groups)`` fp32 tensors kept for the + backward; they are *not* differentiable (nothing produces a gradient for + them) and callers should treat them as opaque. + """ + _validate(input, num_groups, weight, bias, activation) + weight = None if weight is None else weight.contiguous() + bias = None if bias is None else bias.contiguous() + out, mean, rstd = _forward( + input, num_groups, weight, bias, eps, activation, out_dtype or input.dtype + ) + return out, mean, rstd + + +@_group_norm_op.register_fake +def _(input, num_groups, weight, bias, eps, activation, out_dtype): + # empty_like preserves the input's memory format, which is the contract. + out = torch.empty_like(input, dtype=out_dtype or input.dtype) + mean = input.new_empty((input.shape[0], num_groups), dtype=torch.float32) + rstd = input.new_empty((input.shape[0], num_groups), dtype=torch.float32) + return out, mean, rstd + + +@torch.library.custom_op( + "scaffold_gn::group_norm_backward", mutates_args=(), device_types="cuda" +) +def _group_norm_backward_op( + grad_out: torch.Tensor, + input: torch.Tensor, + weight: Optional[torch.Tensor], + bias: Optional[torch.Tensor], + mean: torch.Tensor, + rstd: torch.Tensor, + num_groups: int, + activation: Optional[str], +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Returns ``(d_input, d_weight, d_bias)``. + + ``d_weight``/``d_bias`` are zero-element tensors when the corresponding + parameter is ``None``. ``d_input`` always has the input's dtype and + channels-last memory format. + """ + if not grad_out.is_contiguous(memory_format=_CL_FORMAT): + grad_out = grad_out.contiguous(memory_format=_CL_FORMAT) + if not input.is_contiguous(memory_format=_CL_FORMAT): + input = input.contiguous(memory_format=_CL_FORMAT) + weight = None if weight is None else weight.contiguous() + bias = None if bias is None else bias.contiguous() + d_input, d_weight, d_bias = _backward( + grad_out, input, weight, bias, mean, rstd, num_groups, activation + ) + d_input = d_input.to(input.dtype) + if weight is None: + d_weight = d_weight.new_empty(0) + else: + d_weight = d_weight.to(weight.dtype) + if bias is None: + d_bias = d_bias.new_empty(0) + else: + d_bias = d_bias.to(bias.dtype) + return d_input, d_weight, d_bias + + +@_group_norm_backward_op.register_fake +def _(grad_out, input, weight, bias, mean, rstd, num_groups, activation): + channels = input.shape[1] + d_input = torch.empty_like(input) + d_weight = input.new_empty( + channels if weight is not None else 0, + dtype=weight.dtype if weight is not None else torch.float32, + ) + d_bias = input.new_empty( + channels if bias is not None else 0, + dtype=bias.dtype if bias is not None else torch.float32, + ) + return d_input, d_weight, d_bias + + +def _setup_context(ctx, inputs, output): + input, num_groups, weight, bias, eps, activation, out_dtype = inputs + _out, mean, rstd = output + ctx.save_for_backward(input, weight, bias, mean, rstd) + ctx.num_groups = num_groups + ctx.activation = activation + ctx.needs = ( + ctx.needs_input_grad[0], + ctx.needs_input_grad[2], + ctx.needs_input_grad[3], + ) + + +def _autograd_backward(ctx, grad_out, grad_mean, grad_rstd): + input, weight, bias, mean, rstd = ctx.saved_tensors + need_x, need_w, need_b = ctx.needs + if not (need_x or need_w or need_b): + return None, None, None, None, None, None, None + d_input, d_weight, d_bias = torch.ops.scaffold_gn.group_norm_backward( + grad_out, input, weight, bias, mean, rstd, ctx.num_groups, ctx.activation + ) + return ( + d_input if need_x else None, + None, # num_groups + d_weight if need_w else None, + d_bias if need_b else None, + None, # eps + None, # activation + None, # out_dtype + ) + + +torch.library.register_autograd( + "scaffold_gn::group_norm", _autograd_backward, setup_context=_setup_context +) + + +# --------------------------------------------------------------------------- # +# public API +# --------------------------------------------------------------------------- # +def _autocast_active(input: torch.Tensor) -> bool: + """Whether autocast is enabled for this tensor's device type.""" + try: + return bool(torch.is_autocast_enabled(input.device.type)) + except (RuntimeError, TypeError): # device type autocast does not know + return False + + +def _autocast_out_dtype(input: torch.Tensor) -> Optional[torch.dtype]: + """``F.group_norm``'s output dtype for this input, or None for "unchanged". + + ``at::group_norm`` carries autocast's ``fp32`` cast policy, so under an + enabled autocast region it upcasts its input and returns fp32 whatever came + in. Verified empirically on torch 2.13.0+rocm7.2 for fp32/bf16/fp16 input + and both bf16 and fp16 autocast dtypes. + """ + if input.dtype is not torch.float32 and _autocast_active(input): + return torch.float32 + return None + + +def is_supported( + input, + num_groups: int, + weight=None, + bias=None, + activation: Optional[str] = None, +) -> bool: + """Whether the native channels-last Triton kernel can serve this call. + + Cheap (a handful of attribute reads and one stride check) and side-effect + free -- in particular it does not import Triton, allocate, or launch. + ``False`` means "use ``F.group_norm``": the fast path needs a 5-D CUDA + tensor that is ``channels_last_3d``-contiguous, an fp32/bf16/fp16 dtype, a + channel count divisible by ``num_groups``, and affine parameters whose + dtype ``F.group_norm`` would itself accept for this input (equal to the + input's, or fp32 under autocast, which is what autocast would produce). + """ + if activation not in SUPPORTED_ACTIVATIONS: + return False + if not isinstance(input, torch.Tensor): + return False + if input.device.type != "cuda" or not triton_available(): + return False + if input.dim() != 5 or input.dtype not in SUPPORTED_DTYPES: + return False + if not isinstance(num_groups, int) or num_groups <= 0: + return False + channels = input.shape[1] + if channels % num_groups != 0 or input.numel() == 0: + return False + if not input.is_contiguous(memory_format=_CL_FORMAT): + return False + autocast = None + for t in (weight, bias): + if t is None: + continue + if not isinstance(t, torch.Tensor): + return False + if t.dim() != 1 or t.numel() != channels: + return False + if t.device != input.device: + return False + if t.dtype is not input.dtype: + if t.dtype is not torch.float32: + return False + if autocast is None: + autocast = _autocast_active(input) + if not autocast: + # F.group_norm would raise "expected scalar type ..." here; + # reject so the caller reproduces that behaviour exactly. + return False + return True + + +def triton_group_norm( + input, + num_groups: int, + weight=None, + bias=None, + eps: float = 1e-5, + activation: Optional[str] = None, +): + """GroupNorm with an optionally fused activation, channels-last native. + + A drop-in replacement for ``F.group_norm(input, num_groups, weight, bias, + eps)`` (followed by ``F.relu`` when ``activation="relu"``). Inputs that + :func:`is_supported` rejects are served by ``F.group_norm`` itself, which + keeps this function total but means such calls get the *eager* kernel -- + callers with a faster fallback should branch on :func:`is_supported` + themselves. + + The output has the input's memory format and ``F.group_norm``'s dtype; see + the module docstring for the full contract. + """ + if activation not in SUPPORTED_ACTIVATIONS: + raise ValueError( + f"activation must be one of {SUPPORTED_ACTIVATIONS}, got {activation!r}" + ) + if not is_supported(input, num_groups, weight, bias, activation): + out = F.group_norm(input, num_groups, weight, bias, eps) + return F.relu(out) if activation == "relu" else out + out, _mean, _rstd = torch.ops.scaffold_gn.group_norm( + input, + num_groups, + weight, + bias, + float(eps), + activation, + _autocast_out_dtype(input), + ) + return out diff --git a/tests/test_triton_group_norm.py b/tests/test_triton_group_norm.py new file mode 100644 index 0000000..ff4d0ee --- /dev/null +++ b/tests/test_triton_group_norm.py @@ -0,0 +1,868 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +"""Tests for the channels-last Triton GroupNorm (``ScaFFold.unet.triton_group_norm``). + +The kernel replaces a stock op, so almost every test here is a *parity* test: +values and gradients against ``F.group_norm``, but with the reference computed +in **float64** rather than against another fp32 result -- an fp32-vs-fp32 +comparison cannot tell a correct kernel from one that has merely made the same +mistake, and it cannot see the variance-formula failure the Welford +implementation exists to fix (``test_welford_survives_large_mean``). Each +parity test reports the measured relative error so a regression shows up as a +number, not just a boolean. + +The other three things being pinned down: + +* the **contract** -- output dtype exactly matches ``F.group_norm``'s + (including its fp32 autocast policy), output *memory format* matches the + input's (which is where the kernel deliberately differs from stock, and the + entire reason it exists), and ``is_supported`` accepts exactly the inputs the + native kernel serves; +* **determinism** -- the same call twice is bitwise identical, forward and + backward, because the split count and tiling are pure functions of the shape; +* **composition** -- the op is a real dispatcher op, so it must survive + ``torch.compile(fullgraph=True)`` without a graph break and a ``DCTensor`` + round trip through ``__torch_dispatch__`` with the autograd graph intact. + +CPU runs never touch Triton: the module defers ``import triton`` to the first +call that reaches a kernel, which ``test_import_does_not_pull_in_triton`` +checks in a fresh interpreter. +""" + +import os +import subprocess +import sys + +import pytest +import torch +import torch.nn.functional as F + +from ScaFFold.unet import triton_group_norm as tgn +from ScaFFold.unet.triton_group_norm import is_supported, triton_group_norm + +CL = torch.channels_last_3d +GROUPS = 8 +EPS = 1e-5 + +#: Relative-error ceilings against a float64 reference, by input dtype. The +#: fp32 numbers observed on MI300A at these (small) test shapes are ~2e-07 for +#: y/dx/dweight/dbias; the ceiling leaves room for the 3e-05 that a 134M-element +#: fp32 reduction shows at the largest production shape. The low-precision +#: ceilings are set just above the output's own rounding: 2^-8 for bf16 and +#: 2^-11 for fp16, measured 3.3e-03 and 3.8e-04. +_TOL = { + torch.float32: 1e-4, + torch.bfloat16: 2e-2, + torch.float16: 3e-3, +} + + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + + +def _rel(actual, expected): + """max|actual - expected| / max|expected|, computed in float64.""" + a = actual.detach().double() + e = expected.detach().double() + scale = e.abs().max().clamp_min(1e-30) + return ((a - e).abs().max() / scale).item() + + +def _tensors(shape, dtype, device, affine=True, seed=0, mean=0.0, std=1.0): + """Channels-last input plus (optionally) affine parameters and a cotangent.""" + gen = torch.Generator(device=device).manual_seed(seed) + x = torch.empty(shape, device=device, dtype=dtype, memory_format=CL) + x.normal_(mean, std, generator=gen) + channels = shape[1] + if affine: + weight = torch.empty(channels, device=device, dtype=dtype) + weight.normal_(1.0, 0.25, generator=gen) + bias = torch.empty(channels, device=device, dtype=dtype) + bias.normal_(0.0, 0.25, generator=gen) + else: + weight = bias = None + grad_out = torch.empty(shape, device=device, dtype=dtype, memory_format=CL) + grad_out.normal_(generator=gen) + return x, weight, bias, grad_out + + +def _run(fn, x, weight, bias, grad_out, activation=None, eps=EPS): + """Forward + backward through ``fn``, returning detached results.""" + x = x.detach().clone().requires_grad_(True) + weight = None if weight is None else weight.detach().clone().requires_grad_(True) + bias = None if bias is None else bias.detach().clone().requires_grad_(True) + out = fn(x, GROUPS, weight, bias, eps, activation) + out.backward(grad_out.to(out.dtype)) + return ( + out.detach(), + x.grad, + None if weight is None else weight.grad, + None if bias is None else bias.grad, + ) + + +def _reference(x, weight, bias, grad_out, activation=None, eps=EPS): + """``F.group_norm`` (+ optional ReLU) evaluated entirely in float64.""" + + def fn(x, groups, weight, bias, eps, activation): + out = F.group_norm(x, groups, weight, bias, eps) + return F.relu(out) if activation == "relu" else out + + return _run( + fn, + x.double(), + None if weight is None else weight.double(), + None if bias is None else bias.double(), + grad_out.double(), + activation, + eps, + ) + + +def _assert_parity(got, ref, dtype, label, tol=None): + """Compare (y, dx, dweight, dbias) against the float64 reference.""" + tol = _TOL[dtype] if tol is None else tol + errors = {} + for name, a, e in zip(("y", "dx", "dweight", "dbias"), got, ref): + if a is None: + assert e is None or True # no parameter -> no gradient to compare + continue + errors[name] = _rel(a, e) + print( + f"[{label}] " + + " ".join(f"{k}={v:.3e}" for k, v in errors.items()) + + f" (tol {tol:.1e})" + ) + for name, err in errors.items(): + assert err <= tol, f"{label}: {name} relative error {err:.3e} > {tol:.1e}" + return errors + + +def _cuda_shapes(): + """Shapes covering N>1, non-power-of-two extents and a wide channel count.""" + return [ + (1, 64, 8, 8, 8), # the canonical UNet shape, shrunk + (2, 64, 9, 7, 5), # N>1, all three extents non-power-of-two + (1, 128, 5, 6, 7), + (3, 256, 4, 4, 4), + (1, 2048, 6, 6, 6), # widest UNet channel count + ] + + +# --------------------------------------------------------------------------- +# CPU-only behaviour (no Triton, no GPU) +# --------------------------------------------------------------------------- + + +def test_import_does_not_pull_in_triton(): + """Importing the module must not import Triton. + + Run in a fresh interpreter because any earlier GPU test in this session + would already have built the kernels. The guarantee matters twice over: a + CPU-only unit run must not pay Triton's import, and the module must stay + importable on a build that has no Triton at all. + """ + script = ( + "import sys; import ScaFFold.unet.triton_group_norm as m; " + "assert m.tl is None, 'kernels built at import time'; " + "print('triton' in sys.modules)" + ) + repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + cwd=repo_root, + timeout=300, + ) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == "False", ( + f"triton was imported at module import time: {result.stdout!r}" + ) + + +def test_cpu_input_falls_back_bitwise(): + """A CPU tensor is not supported, and the fallback is the stock kernel.""" + gen = torch.Generator().manual_seed(0) + x = torch.randn(1, 64, 4, 4, 4, generator=gen).requires_grad_(True) + weight = torch.randn(64, generator=gen).requires_grad_(True) + bias = torch.randn(64, generator=gen).requires_grad_(True) + assert is_supported(x, GROUPS, weight, bias) is False + + out = triton_group_norm(x, GROUPS, weight, bias, EPS) + assert torch.equal(out, F.group_norm(x, GROUPS, weight, bias, EPS)) + out.pow(2).sum().backward() + assert x.grad is not None and weight.grad is not None and bias.grad is not None + + +def test_cpu_fused_relu_falls_back_bitwise(): + gen = torch.Generator().manual_seed(1) + x = torch.randn(2, 32, 3, 4, 5, generator=gen) + weight = torch.randn(32, generator=gen) + bias = torch.randn(32, generator=gen) + got = triton_group_norm(x, 8, weight, bias, EPS, "relu") + assert torch.equal(got, F.relu(F.group_norm(x, 8, weight, bias, EPS))) + + +def test_unknown_activation_raises(): + x = torch.randn(1, 8, 2, 2, 2) + with pytest.raises(ValueError, match="activation"): + triton_group_norm(x, 2, activation="gelu") + assert is_supported(x, 2, activation="gelu") is False + + +def test_select_strategy_is_a_pure_function_of_shape(): + """The dispatch hook must be deterministic -- the reduction order, and so + the bits of the result, depend on it.""" + for args in ((1, 64, 8**3, 8), (2, 2048, 16**3, 8), (1, 128, 7 * 5 * 3, 4)): + first = tgn.select_strategy(*args) + assert first in tgn.STRATEGIES + assert all(tgn.select_strategy(*args) == first for _ in range(3)) + + +def test_tuning_table_covers_the_scale8_shapes(): + """The frozen table is what makes the kernel reproducible; keep it honest.""" + for channels, edge in ((64, 256), (128, 128), (256, 64), (512, 32), (1024, 16)): + assert tgn.default_config(channels, edge**3) is tgn._TUNED[(channels, edge)] + # An unlisted shape falls back to the generic config rather than failing. + assert tgn.default_config(96, 11**3) == tgn.GNConfig() + + +def test_plan_depends_only_on_shape(): + """Two plans for the same shape must be identical objects of identical + content, or the split count could drift between calls and break bitwise + reproducibility.""" + a = tgn._plan(2, 128, 32**3, 8, 2 * 128 * 32**3) + b = tgn._plan(2, 128, 32**3, 8, 2 * 128 * 32**3) + assert (a.nsplit, a.chunk, a.block_s_stats, a.block_s_elem, a.int64) == ( + b.nsplit, + b.chunk, + b.block_s_stats, + b.block_s_elem, + b.int64, + ) + # int64 addressing turns on exactly when a linear index can overflow int32. + small = tgn._plan(1, 64, 128**3, 8, 64 * 128**3) + big = tgn._plan(2, 64, 256**3, 8, 2 * 64 * 256**3) + assert small.int64 is False + assert big.int64 is True + + +# --------------------------------------------------------------------------- +# is_supported +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +def test_is_supported_accepts_the_fast_path(): + device = torch.device("cuda") + x = torch.empty(1, 64, 6, 6, 6, device=device, memory_format=CL).normal_() + weight = torch.randn(64, device=device) + assert is_supported(x, GROUPS, weight, weight) is True + assert is_supported(x, GROUPS) is True + assert is_supported(x, GROUPS, activation="relu") is True + + +@pytest.mark.gpu +def test_is_supported_rejections(): + """Everything ``is_supported`` rejects must be something a caller can hand + to ``F.group_norm`` instead -- so the rejections are the contract's edge.""" + device = torch.device("cuda") + cl = torch.empty(1, 64, 6, 6, 6, device=device, memory_format=CL).normal_() + cases = { + "cpu tensor": (torch.randn(1, 64, 6, 6, 6), GROUPS, None, None, None), + "contiguous (NCDHW)": ( + torch.randn(1, 64, 6, 6, 6, device=device), + GROUPS, + None, + None, + None, + ), + "float64": ( + torch.empty( + 1, 64, 6, 6, 6, device=device, dtype=torch.float64, memory_format=CL + ), + GROUPS, + None, + None, + None, + ), + "4-D": (torch.randn(1, 64, 6, 6, device=device), GROUPS, None, None, None), + "channels not divisible": (cl, 7, None, None, None), + "num_groups=0": (cl, 0, None, None, None), + "bad activation": (cl, GROUPS, None, None, "gelu"), + "weight wrong size": (cl, GROUPS, torch.randn(32, device=device), None, None), + "weight on cpu": (cl, GROUPS, torch.randn(64), None, None), + "weight dtype mismatch": ( + cl, + GROUPS, + torch.randn(64, device=device, dtype=torch.bfloat16), + None, + None, + ), + "sliced (non-contiguous)": ( + torch.empty(1, 64, 6, 6, 12, device=device, memory_format=CL)[..., ::2], + GROUPS, + None, + None, + None, + ), + "not a tensor": (None, GROUPS, None, None, None), + } + for label, args in cases.items(): + assert is_supported(*args) is False, f"{label} should be rejected" + + +@pytest.mark.gpu +def test_rejected_inputs_still_produce_stock_results(): + """``triton_group_norm`` stays total: rejects go to ``F.group_norm``.""" + device = torch.device("cuda") + gen = torch.Generator(device=device).manual_seed(4) + x = torch.randn(2, 64, 5, 6, 7, device=device, generator=gen).requires_grad_(True) + weight = torch.randn(64, device=device, generator=gen).requires_grad_(True) + bias = torch.randn(64, device=device, generator=gen).requires_grad_(True) + assert is_supported(x, GROUPS, weight, bias) is False + got = triton_group_norm(x, GROUPS, weight, bias, EPS) + assert torch.equal(got, F.group_norm(x, GROUPS, weight, bias, EPS)) + + +# --------------------------------------------------------------------------- +# value / gradient parity against float64 +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +@pytest.mark.parametrize("shape", _cuda_shapes()) +def test_parity_fp32(shape): + device = torch.device("cuda") + x, weight, bias, grad_out = _tensors( + shape, torch.float32, device, seed=hash(shape) % 1000 + ) + assert is_supported(x, GROUPS, weight, bias) + got = _run(triton_group_norm, x, weight, bias, grad_out) + ref = _reference(x, weight, bias, grad_out) + _assert_parity(got, ref, torch.float32, f"fp32 {shape}") + assert got[0].is_contiguous(memory_format=CL) + assert got[1].is_contiguous(memory_format=CL) + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("activation", [None, "relu"]) +def test_parity_every_dtype_and_activation(dtype, activation): + device = torch.device("cuda") + shape = (2, 128, 7, 6, 5) + x, weight, bias, grad_out = _tensors(shape, dtype, device, seed=11) + assert is_supported(x, GROUPS, weight, bias, activation) + got = _run(triton_group_norm, x, weight, bias, grad_out, activation) + ref = _reference(x, weight, bias, grad_out, activation) + _assert_parity(got, ref, dtype, f"{dtype} act={activation}") + assert got[0].dtype == dtype + assert got[1].dtype == dtype + assert got[2].dtype == dtype and got[3].dtype == dtype + + +@pytest.mark.gpu +@pytest.mark.parametrize("affine", ["both", "weight_only", "bias_only", "neither"]) +def test_parity_without_affine_parameters(affine): + """``weight=None`` / ``bias=None`` are separate kernel constexpr paths.""" + device = torch.device("cuda") + shape = (2, 64, 5, 5, 5) + x, weight, bias, grad_out = _tensors(shape, torch.float32, device, seed=19) + if affine in ("bias_only", "neither"): + weight = None + if affine in ("weight_only", "neither"): + bias = None + assert is_supported(x, GROUPS, weight, bias) + got = _run(triton_group_norm, x, weight, bias, grad_out) + + reference_weight = weight + if weight is None and bias is not None: + # Upstream limitation, not a difference in this kernel: + # ``F.group_norm(x, g, None, bias).backward()`` raises "tensor does not + # have a device" on both CPU and CUDA (torch 2.13.0+rocm7.2), so the + # float64 reference has to spell the same computation with weight=1. + # This kernel handles the combination directly. + reference_weight = torch.ones_like(bias) + ref = _reference(x, reference_weight, bias, grad_out) + # ``_assert_parity`` skips outputs this configuration does not produce. + _assert_parity(got, ref, torch.float32, f"affine={affine}") + + +@pytest.mark.gpu +def test_partial_gradient_requirements(): + """Only some inputs requiring grad must not change the ones that do.""" + device = torch.device("cuda") + x, weight, bias, grad_out = _tensors( + (1, 64, 5, 5, 5), torch.float32, device, seed=23 + ) + full = _run(triton_group_norm, x, weight, bias, grad_out) + + frozen_w = weight.detach().clone() + frozen_b = bias.detach().clone() + xi = x.detach().clone().requires_grad_(True) + out = triton_group_norm(xi, GROUPS, frozen_w, frozen_b, EPS) + out.backward(grad_out) + assert torch.equal(xi.grad, full[1]) + assert frozen_w.grad is None and frozen_b.grad is None + + # ... and the mirror image: parameters only. + xn = x.detach().clone() + wn = weight.detach().clone().requires_grad_(True) + bn = bias.detach().clone().requires_grad_(True) + triton_group_norm(xn, GROUPS, wn, bn, EPS).backward(grad_out) + assert torch.equal(wn.grad, full[2]) + assert torch.equal(bn.grad, full[3]) + + +# --------------------------------------------------------------------------- +# fused activation +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_fused_relu_matches_unfused(dtype): + """The fused store and ``F.relu`` on the unfused output must agree bitwise, + forward *and* backward -- the backward recomputes the pre-activation rather + than reading it back, so this is the test that recomputation is exact.""" + device = torch.device("cuda") + x, weight, bias, grad_out = _tensors((2, 64, 6, 7, 8), dtype, device, seed=29) + fused = _run(triton_group_norm, x, weight, bias, grad_out, "relu") + + def unfused(x, groups, weight, bias, eps, _activation): + return F.relu(triton_group_norm(x, groups, weight, bias, eps)) + + separate = _run(unfused, x, weight, bias, grad_out) + for name, a, b in zip(("y", "dx", "dweight", "dbias"), fused, separate): + assert torch.equal(a, b), f"fused vs unfused+relu differ in {name}" + # A ReLU that never fires would make this test vacuous. + assert (fused[0] == 0).any() and (fused[0] > 0).any() + + +# --------------------------------------------------------------------------- +# determinism +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +@pytest.mark.parametrize("activation", [None, "relu"]) +def test_bitwise_determinism(activation): + """Same input twice => bitwise-equal output and gradients. + + Guaranteed by construction (no float atomics; grid, split count and tile + sizes are pure functions of the shape) and asserted here because a future + run-time autotuner would silently break it. + """ + device = torch.device("cuda") + x, weight, bias, grad_out = _tensors( + (2, 128, 9, 11, 13), torch.float32, device, seed=31 + ) + first = _run(triton_group_norm, x, weight, bias, grad_out, activation) + second = _run(triton_group_norm, x, weight, bias, grad_out, activation) + for name, a, b in zip(("y", "dx", "dweight", "dbias"), first, second): + assert torch.equal(a, b), f"{name} is not bitwise reproducible" + + +# --------------------------------------------------------------------------- +# numerics: Welford vs E[x^2] - E[x]^2 +# --------------------------------------------------------------------------- + + +def _naive_group_norm(x, num_groups, weight, bias, eps): + """The prototype's variance formula, reproduced in fp32 torch ops. + + ``var = E[x^2] - E[x]^2`` is split-friendly and cheap, and it is what the + kernel used before the Welford rewrite; this is the thing the test below + must show is broken so that "the new one passes" means something. + """ + n, channels = x.shape[0], x.shape[1] + flat = x.reshape(n, num_groups, -1) + mean = flat.mean(-1) + mean_sq = (flat * flat).mean(-1) + var = mean_sq - mean * mean + rstd = 1.0 / torch.sqrt(var + eps) + out = (flat - mean[..., None]) * rstd[..., None] + out = out.reshape(x.shape) + shape = (1, channels) + (1,) * (x.dim() - 2) + return out * weight.reshape(shape) + bias.reshape(shape) + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "mean,std,naive_floor", + [ + (0.0, 1.0, None), # both formulations are fine here + (100.0, 1.0, 1e-4), # E[x^2]-E[x]^2 already an order of magnitude off + (1e3, 1e-2, 1e-1), # ... and here it has lost the variance entirely + ], +) +def test_welford_survives_large_mean(mean, std, naive_floor): + """Large-mean / small-variance input: the regression case for the rewrite. + + Measured here on MI300A at ``[1, 256, 24^3]`` (relative error of the output + against a float64 reference computed from the same fp32 samples), with the + production shape ``[1, 256, 64^3]`` in parentheses: + + mean=0, std=1 this 1.3e-07 (1.6e-07) E[x^2]-E[x]^2 1.4e-07 (1.8e-07) + mean=1e2, std=1 this 1.1e-06 (8.0e-07) E[x^2]-E[x]^2 9.5e-04 (5.6e-04) + mean=1e3, std=1e-2 this 4.4e-04 (1.1e-04) E[x^2]-E[x]^2 2.3e+00 (2.3e+00) + + i.e. at ``mean/std = 1e5`` the old formulation loses the variance outright + (the difference of the two ~1e6-sized fp32 terms is below one ulp, so + ``rstd`` saturates on ``eps`` and the output is meaningless) while Welford + is still correct to 4.4e-04 -- itself *5x better* than ATen's own fp32 + GroupNorm on the same input (2.2e-03), and dominated by the fp32 + representation of a mean of 1e3 rather than by anything the kernel does. + """ + device = torch.device("cuda") + x, weight, bias, grad_out = _tensors( + (1, 256, 24, 24, 24), torch.float32, device, seed=37, mean=mean, std=std + ) + got = triton_group_norm(x, GROUPS, weight, bias, EPS) + ref = F.group_norm(x.double(), GROUPS, weight.double(), bias.double(), EPS) + stock = F.group_norm(x, GROUPS, weight, bias, EPS) + naive = _naive_group_norm(x, GROUPS, weight, bias, EPS) + + err = _rel(got, ref) + err_stock = _rel(stock, ref) + err_naive = _rel(naive, ref) + print( + f"[welford mean={mean:g} std={std:g}] triton={err:.3e} " + f"aten_fp32={err_stock:.3e} naive_Ex2={err_naive:.3e}" + ) + # Never worse than ATen's own fp32 kernel by more than a small factor. + assert err <= max(4.0 * err_stock, 1e-5), ( + f"triton {err:.3e} vs aten fp32 {err_stock:.3e}" + ) + if naive_floor is not None: + assert err_naive > naive_floor, ( + "the naive formulation was expected to fail here " + f"but only reached {err_naive:.3e}" + ) + assert err < err_naive / 10.0, ( + f"Welford ({err:.3e}) is not clearly better than " + f"E[x^2]-E[x]^2 ({err_naive:.3e})" + ) + + +# --------------------------------------------------------------------------- +# dtypes, layouts, autocast +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("autocast_dtype", [None, torch.bfloat16, torch.float16]) +def test_output_dtype_matches_stock(dtype, autocast_dtype): + """The dtype contract, including autocast's fp32 policy for GroupNorm. + + Stock behaviour on this build (measured, not assumed): without autocast the + output dtype is the input dtype; inside *any* enabled CUDA autocast region + it is fp32, because ``at::group_norm`` carries the fp32 cast policy. + """ + device = torch.device("cuda") + x, weight, bias, _ = _tensors((1, 64, 5, 5, 5), dtype, device, seed=41) + if autocast_dtype is not None: + # Autocast casts the parameters itself, and production keeps them fp32. + weight = weight.float() + bias = bias.float() + ctx = ( + torch.autocast("cuda", dtype=autocast_dtype) + if autocast_dtype is not None + else torch.autocast("cuda", enabled=False) + ) + with ctx: + assert is_supported(x, GROUPS, weight, bias) + stock = F.group_norm(x, GROUPS, weight, bias, EPS) + got = triton_group_norm(x, GROUPS, weight, bias, EPS) + expected = torch.float32 if autocast_dtype is not None else dtype + assert stock.dtype == expected, "assumption about stock GroupNorm broke" + assert got.dtype == stock.dtype + # ... and the one deliberate difference: stock always returns contiguous. + assert stock.is_contiguous() and not stock.is_contiguous(memory_format=CL) + assert got.is_contiguous(memory_format=CL) + print( + f"[dtype in={dtype} autocast={autocast_dtype}] " + f"stock={stock.dtype}/CONT mine={got.dtype}/CL rel={_rel(got, stock):.3e}" + ) + assert _rel(got.float(), stock.float()) <= _TOL[stock.dtype] + + +@pytest.mark.gpu +def test_autocast_gradient_dtypes_match_stock(): + """Under autocast, ``d_input`` keeps the input's dtype and the parameter + gradients stay fp32 -- exactly what the cast nodes around stock GroupNorm + produce.""" + device = torch.device("cuda") + x, _, _, grad_out = _tensors((1, 64, 5, 5, 5), torch.bfloat16, device, seed=43) + weight = torch.randn(64, device=device, requires_grad=True) + bias = torch.randn(64, device=device, requires_grad=True) + + def run(fn): + xi = x.detach().clone().requires_grad_(True) + w = weight.detach().clone().requires_grad_(True) + b = bias.detach().clone().requires_grad_(True) + with torch.autocast("cuda", dtype=torch.bfloat16): + out = fn(xi, GROUPS, w, b, EPS) + out.backward(grad_out.to(out.dtype)) + return out, xi.grad, w.grad, b.grad + + stock = run(F.group_norm) + got = run(triton_group_norm) + for name, a, b in zip(("y", "dx", "dweight", "dbias"), got, stock): + assert a.dtype == b.dtype, f"{name}: {a.dtype} != {b.dtype}" + assert _rel(a.float(), b.float()) <= 5e-2, name + + +@pytest.mark.gpu +@pytest.mark.parametrize("layout", ["channels_last_3d", "contiguous"]) +def test_memory_format_is_preserved(layout): + """Both layouts round-trip their own format; contiguous input takes the + documented ``F.group_norm`` fallback rather than silently changing layout.""" + device = torch.device("cuda") + gen = torch.Generator(device=device).manual_seed(47) + x = torch.randn(2, 64, 5, 6, 7, device=device, generator=gen) + if layout == "channels_last_3d": + x = x.contiguous(memory_format=CL) + weight = torch.randn(64, device=device, generator=gen) + bias = torch.randn(64, device=device, generator=gen) + grad_out = torch.randn(2, 64, 5, 6, 7, device=device, generator=gen) + if layout == "channels_last_3d": + grad_out = grad_out.contiguous(memory_format=CL) + + assert is_supported(x, GROUPS, weight, bias) is (layout == "channels_last_3d") + got = _run(triton_group_norm, x, weight, bias, grad_out) + ref = _reference(x, weight, bias, grad_out) + _assert_parity(got, ref, torch.float32, f"layout={layout}") + if layout == "channels_last_3d": + assert got[0].is_contiguous(memory_format=CL) + assert got[1].is_contiguous(memory_format=CL) + else: + assert got[0].is_contiguous() + assert got[1].is_contiguous() + + +# --------------------------------------------------------------------------- +# int64 offsets +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +def test_int64_switch_flips_at_int32_max(): + """The switch is a pure function of the element count, so pin the boundary. + + ``[2, 64, 256^3]`` is *exactly* 2^31 elements: the shape that made an + int64 path mandatory before batch>1 or scale 16. + """ + assert tgn._plan(1, 64, 255**3, 8, 64 * 255**3).int64 is False + assert tgn._plan(2, 64, 256**3, 8, 2 * 64 * 256**3).int64 is True + + +@pytest.mark.gpu +@pytest.mark.slow +def test_correct_above_int32_max_elements(): + """Correctness at a shape whose linear element count exceeds INT32_MAX. + + ``[2, 64, 256, 256, 257]`` is 2_155_872_256 elements -- 8.4M past 2^31, and + non-power-of-two in the fastest spatial dimension so a truncated offset + cannot accidentally land on the right address. fp32 (8.03 GiB per tensor) + keeps the comparison sharp; the reference needs an NCDHW copy, so the peak + is ~48 GiB and the test skips, loudly, if the device cannot hold that. + """ + device = torch.device("cuda") + shape = (2, 64, 256, 256, 257) + numel = 1 + for dim in shape: + numel *= dim + assert numel > 2**31 - 1 + needed = 6 * numel * 4 # x, y, x_contig, reference, and slack for the diff + free, total = torch.cuda.mem_get_info() + if free < needed: + pytest.skip( + f"needs ~{needed / 2**30:.0f} GiB free, device has " + f"{free / 2**30:.0f} GiB of {total / 2**30:.0f} GiB" + ) + + gen = torch.Generator(device=device).manual_seed(53) + x = torch.empty(shape, device=device, memory_format=CL) + x.normal_(generator=gen) + weight = torch.randn(64, device=device, generator=gen) + bias = torch.randn(64, device=device, generator=gen) + + assert tgn._plan( + shape[0], shape[1], shape[2] * shape[3] * shape[4], GROUPS, numel + ).int64 + got = triton_group_norm(x, GROUPS, weight, bias, EPS) + assert got.is_contiguous(memory_format=CL) + + contiguous = x.contiguous() + del x + torch.cuda.empty_cache() + reference = F.group_norm(contiguous, GROUPS, weight, bias, EPS) + del contiguous + torch.cuda.empty_cache() + + # Compare both batch items separately: a truncated 32-bit offset wraps + # partway through, so the second half would be wrong while the first is not. + errors = [_rel(got[i], reference[i]) for i in range(shape[0])] + print(f"[int64 {shape}] per-sample relative error {errors}") + for i, err in enumerate(errors): + assert err < 1e-4, f"sample {i}: relative error {err:.3e}" + del got, reference + torch.cuda.empty_cache() + + +# --------------------------------------------------------------------------- +# composition: torch.compile and DCTensor +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +def test_custom_op_is_registered_with_a_fake_kernel(): + """A meta/fake kernel is what lets Dynamo trace the op without running it.""" + from torch._subclasses.fake_tensor import FakeTensorMode + + assert hasattr(torch.ops.scaffold_gn, "group_norm") + assert hasattr(torch.ops.scaffold_gn, "group_norm_backward") + with FakeTensorMode(): + x = torch.empty(2, 64, 5, 6, 7, device="cuda", memory_format=CL) + weight = torch.empty(64, device="cuda") + out, mean, rstd = torch.ops.scaffold_gn.group_norm( + x, GROUPS, weight, weight, EPS, "relu", None + ) + assert out.shape == x.shape and out.dtype == x.dtype + assert out.is_contiguous(memory_format=CL) + assert mean.shape == (2, GROUPS) and rstd.dtype == torch.float32 + # ... and the dtype override autocast uses. + bf16 = torch.empty( + 2, 64, 5, 6, 7, device="cuda", dtype=torch.bfloat16, memory_format=CL + ) + out32, _, _ = torch.ops.scaffold_gn.group_norm( + bf16, GROUPS, None, None, EPS, None, torch.float32 + ) + assert out32.dtype == torch.float32 + assert out32.is_contiguous(memory_format=CL) + + +@pytest.mark.gpu +@pytest.mark.parametrize("activation", [None, "relu"]) +def test_torch_compile_fullgraph(activation): + """``fullgraph=True`` raises on a graph break, so this *is* the no-break + test; the compiled result must additionally be bitwise equal to eager, + because the op is opaque to Inductor and so cannot be re-associated.""" + device = torch.device("cuda") + x, weight, bias, grad_out = _tensors( + (2, 64, 6, 6, 6), torch.float32, device, seed=59 + ) + + def fn(x, weight, bias): + return triton_group_norm(x, GROUPS, weight, bias, EPS, activation) * 2.0 + + def wrapped(x, groups, weight, bias, eps, _activation): + return fn(x, weight, bias) + + eager = _run(wrapped, x, weight, bias, grad_out) + + compiled_fn = torch.compile(fn, fullgraph=True, dynamic=False) + + def wrapped_compiled(x, groups, weight, bias, eps, _activation): + return compiled_fn(x, weight, bias) + + compiled = _run(wrapped_compiled, x, weight, bias, grad_out) + for name, a, b in zip(("y", "dx", "dweight", "dbias"), compiled, eager): + assert torch.equal(a, b), f"compiled and eager differ in {name}" + assert compiled[0].is_contiguous(memory_format=CL) + + +@pytest.fixture +def dc_cuda(): + """DistConv package plus a CUDA ParallelStrategy over a 1-rank NCCL group. + + Mirrors ``tests/test_groupnorm.py``'s fixture (``num_shards=(1, 1, 1)`` on + dims (2, 3, 4) is what worker.py builds for a single-device run) on its own + rendezvous port so the two suites can run in one session. + """ + import torch.distributed as dist + + distconv = pytest.importorskip("distconv") + + created = False + if not dist.is_initialized(): + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29519") + dist.init_process_group(backend="nccl", rank=0, world_size=1) + created = True + strategy = distconv.ParallelStrategy( + num_shards=(1, 1, 1), shard_dim=(2, 3, 4), device_type="cuda" + ) + yield distconv, strategy + if created and dist.is_initialized(): + dist.destroy_process_group() + + +@pytest.mark.gpu +def test_dctensor_round_trip(dc_cuda): + """A DCTensor must go in and come out, with the graph back to its producer + intact. + + The op is a real dispatcher op, so DistConv's generic + ``__torch_dispatch__`` unwraps to the local shard, runs it, and rewraps -- + no GroupNorm-specific handling needed on either side. The producer in + front matters: with a bare ``input._tensor`` read the graph would be severed + there and only GroupNorm's own parameters would see gradients. + """ + distconv, strategy = dc_cuda + device = torch.device("cuda") + gen = torch.Generator(device=device).manual_seed(61) + x = torch.empty(1, 64, 6, 6, 6, device=device, memory_format=CL) + x.normal_(generator=gen) + grad_out = torch.empty_like(x) + grad_out.normal_(generator=gen) + weight = torch.randn(64, device=device, generator=gen) + bias = torch.randn(64, device=device, generator=gen) + producer = torch.nn.Conv3d(64, 64, 1, bias=False).to(device) + + def run(fn, wrap): + xi = x.detach().clone().requires_grad_(True) + conv = torch.nn.Conv3d(64, 64, 1, bias=False).to(device) + with torch.no_grad(): + conv.weight.copy_(producer.weight) + w = weight.detach().clone().requires_grad_(True) + b = bias.detach().clone().requires_grad_(True) + inp = distconv.DCTensor.from_shard(xi, strategy) if wrap else xi + # The conv is the producer; the explicit channels-last conversion is + # what PYTORCH_MIOPEN_SUGGEST_NHWC=1 gives production for free. + hidden = conv(inp).contiguous(memory_format=CL) + out = fn(hidden, GROUPS, w, b, EPS, "relu") + if wrap: + assert isinstance(out, distconv.DCTensor) + assert out.is_contiguous(memory_format=CL) + out = distconv.distconv._ToTensor.apply(out) + out.backward(grad_out) + return out.detach(), xi.grad, conv.weight.grad, w.grad, b.grad + + def stock(x, groups, weight, bias, eps, _activation): + return F.relu(F.group_norm(x, groups, weight, bias, eps)) + + got = run(triton_group_norm, wrap=True) + ref = run(stock, wrap=False) + for name, a, b in zip(("y", "dx", "dconv", "dweight", "dbias"), got, ref): + assert a is not None, f"{name} never received a gradient" + err = _rel(a, b) + print(f"[dctensor] {name}={err:.3e}") + assert err <= 1e-4, f"{name}: relative error {err:.3e}" From c3020c227d8ae2549d577b25d96a64c85105d252 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Sat, 1 Aug 2026 03:33:00 -0700 Subject: [PATCH 56/62] Fix six defects found reviewing the Triton GroupNorm An independent review of 1a8aced attacked the kernel and its suite; these are its findings, plus the edge suite it wrote to pin them (153 tests, none of which the original 51 subsumed). The serious one: every kernel launched on torch.cuda.current_device() rather than the input's device, so a tensor on cuda:1 while cuda:0 was current took a memory access fault and dumped core, where ATen's group_norm handles the same call. One rank per GPU hides this in production, but nothing guaranteed it. The guard costs 0.51 us and is skipped outright when the device is already current. The rest were quieter. The backward fake kernel promised contiguous d_input strides while the real one returns channels-last, which eager never notices and torch.compile miscompiles. mean/rstd were differentiable and answered grad(mean.sum(), x) with zeros -- a plausible wrong number rather than an error -- after materializing a full-size cotangent and running the entire backward for it. (1,C,1,1,1) input was accepted and returned bias where stock raises. Single-element groups returned d_input = 2.2e-05 instead of the exact zero, because fma(dy, w, -c1) contracts and leaves the product's rounding error multiplied by rstd = 316; that case now answers analytically. Double backward and subnormal eps are documented rather than fixed. The review also proposed deleting the Welford correction term, having measured that it does not move the output error. It does not: the fp32 mean rounds the correction away in y. It moves rstd, which is what it computes -- 1472x in a single-tile reduction at mu/sigma=1e6, 3.8x at [1,512,32^3] in the production configuration -- so it stays, now priced honestly at 0.9% of fwd+bwd (the docstring's old 0.8% was the cost of the whole two-pass rewrite) and pinned by a test that measures rstd rather than y. Mutation coverage over both suites: 38/38 killed, up from 30/33. Interleaved A/B over 4 runs shows no fwd+bwd regression at any of the six shapes. --- ScaFFold/unet/triton_group_norm.py | 489 +++++++--- tests/test_triton_group_norm_edge.py | 1339 ++++++++++++++++++++++++++ 2 files changed, 1676 insertions(+), 152 deletions(-) create mode 100644 tests/test_triton_group_norm_edge.py diff --git a/ScaFFold/unet/triton_group_norm.py b/ScaFFold/unet/triton_group_norm.py index cad62ad..8b4e9c8 100644 --- a/ScaFFold/unet/triton_group_norm.py +++ b/ScaFFold/unet/triton_group_norm.py @@ -55,10 +55,10 @@ ========== ``triton_group_norm(input, num_groups, weight=None, bias=None, eps=1e-5, activation=None)`` - Drop-in for ``F.group_norm`` (plus an optionally fused ReLU) with full - autograd support. Accepts *anything* ``F.group_norm`` accepts; inputs the - Triton kernel cannot serve fall back to ``F.group_norm`` internally (see - "Layouts" below). + Drop-in for ``F.group_norm`` (plus an optionally fused ReLU) with + first-order autograd support. Accepts *anything* ``F.group_norm`` accepts; + inputs the Triton kernel cannot serve fall back to ``F.group_norm`` + internally (see "Layouts" below). ``is_supported(input, num_groups, weight=None, bias=None, activation=None)`` Cheap, side-effect-free predicate: ``True`` exactly when the native Triton @@ -94,12 +94,35 @@ ``review/gn-dctensor/triton/RESULTS.md``); preserving channels-last is the entire point of the kernel. * **autograd** -- ``d_input``, ``d_weight``, ``d_bias``; ``weight=None`` and/or - ``bias=None`` supported. + ``bias=None`` supported. **First order only**: the backward is itself a + custom op with no autograd formula of its own, so a second + ``torch.autograd.grad`` through this op raises ``RuntimeError: Trying to + backward through scaffold_gn.group_norm_backward.default but no autograd + formula was registered``. Stock ``F.group_norm`` *does* support double + backward, so a gradient penalty or a Hessian-vector product must route + around this kernel (``is_supported`` says nothing about second derivatives; + it is documented there too). It fails loudly rather than returning garbage. +* **device** -- the kernels run on the *input's* device, whatever device is + current, matching ATen's ``DeviceGuard`` behaviour; see ``_device_guard``. * **determinism** -- bitwise reproducible run to run and process to process. There are no float atomics anywhere, and the grid, split count and tile sizes are pure functions of the shape (the tuning table is frozen in this file for exactly that reason -- a *runtime* autotuner would break reproducibility by changing the reduction order between runs). +* **rejections** -- every shape/dtype/parameter combination ``F.group_norm`` + raises on is one ``is_supported`` returns ``False`` for, including the + degenerate "1 value per channel" shape (``N*(C/G)*D*H*W == 1``), so a caller + that branches on ``is_supported`` never gets an answer where the op this + replaces would have raised. +* **eps** -- one deliberate divergence, at a value no run uses: for a + *subnormal* fp32 ``eps`` (``< 1.18e-38``) on a zero-variance group the GPU + flushes ``var + eps`` to zero, so ``rstd`` is ``inf`` and ``y`` is ``NaN`` + where ATen stays finite. The boundary is exactly the normal/subnormal one + (``eps=1.2e-38`` gives ``rstd=9.1e18``, ``eps=1e-38`` gives ``inf``); at + ``eps == 0`` both implementations produce non-finite output identically. + Left as is rather than clamped because clamping would perturb every + ordinary call to defend a value nine orders of magnitude below the smallest + plausible one. Reduction strategy ================== @@ -151,6 +174,54 @@ reductions instead of three, shift taken from a peeled first tile) would recover most of that; it was not worth the extra failure mode for ~4 ms/step. +What the third pass (``corr``) is worth, separately +--------------------------------------------------- +The accuracy above is mostly the *two-pass* structure; the ``corr`` term is a +third reduction on top of it and deserves its own accounting. Deleting it +outright (keeping ``mean_t = mean0``, ``M2 = sum((x-mean0)^2)``) and comparing +both against float64 on the same fp32 samples, relative error of ``rstd``, +10 seeds each: + + regime with corr without ratio + one tile per group reduction (nsplit=1): + [2,64,8,4,4] G=4, mu/sigma=1e6 1.1e-07 1.6e-04 1472x + [2,64,8,4,4] G=2, mu/sigma=1e6 8.4e-08 4.9e-05 580x + [2,64,8,4,4] G=1, mu/sigma=1e6 9.8e-08 2.1e-05 216x + [2,64,8,4,4] G=4, mu/sigma=1e5 9.9e-08 8.4e-06 85x + many tiles and splits (the production configs): + [1,512,32^3] G=8, mu/sigma=1e5 1.2e-05 4.4e-05 3.8x + [1,1024,16^3] G=8, mu/sigma=1e5 8.6e-06 2.5e-05 2.9x + [1,256,24^3] G=8, mu/sigma=1e5 3.0e-05 3.7e-05 1.3x + [1,256,24^3] G=8, mu/sigma=1e7 5.8e-04 2.3e-06 0.004x + +So it is decisively load-bearing exactly where the tile mean is formed from +many large values -- up to 1472x on ``rstd`` -- and worth a steady 1.3-4x in +the multi-split configs the tuning table actually picks, at ``mu/sigma = 1e5``. +Past ``mu/sigma ~ 1e6`` with many splits it can go the *other* way (last row): +there the true spread between tile means is smaller than one ulp of the means +themselves, so Chan's between-tile term is computed from quantization noise +either way and the uncorrected version's inflated ``M2`` partly cancels it. +That regime is past fp32's floor for this computation (a mean of 1e6 held in +fp32 quantizes to 0.06, i.e. 6% of a standard deviation at sigma=1) and no +production input is near it. + +The *output* error is nearly unmoved by any of this -- at most ~1.4x in either +direction -- which is why the term looks free to delete if you only measure +``y``: the output is dominated by the fp32 representation of the mean, which +``corr`` cannot improve (``mean0 + corr`` rounds straight back to ``mean0`` +once the mean is large). It is ``rstd`` that carries the benefit. + +Price, measured the same way (median of 20 forwards, correction removed +outright rather than zeroed): **+2.3%** of the forward at ``[1,64,256^3]``, ++3.0% at ``[1,128,128^3]``, +4.7% at ``[1,256,64^3]``, and nothing measurable +(-1.2% to +0.5%, i.e. noise) at the three launch-bound shapes. At the shape +that dominates the step that is +0.10 ms of a 11.6 ms fwd+bwd, i.e. +0.9%. +Kept: a 1.3-1472x accuracy factor on the statistic the whole rewrite exists to +protect is worth ~1% of GroupNorm time. The load-bearing case is pinned by +``test_welford_correction_recovers_rstd_in_a_single_tile_reduction`` in +``tests/test_triton_group_norm_edge.py``, so deleting the term now fails the +suite instead of passing it silently. + Layouts ======= * ``channels_last_3d`` 5-D input -> **native Triton kernel**, channels-last @@ -228,6 +299,7 @@ so importing this module (or running the CPU test suite) costs nothing. """ +import contextlib import functools import importlib.util from typing import Optional, Tuple @@ -590,12 +662,26 @@ def _stats_partial_kernel( m = tl.broadcast_to((offs_s < nvalid)[:, None, None], (BLOCK_S, GP, CGP)) if MASKED_C: m = m & cmask + # `other` is not load-bearing for any *bounded* value: `cnt_t` + # counts only the valid lanes, so `corr` below evaluates to + # `sum_valid(x)/cnt_t - mean0` and `mean_t = mean0 + corr` is the + # true mean whatever the masked lanes contributed, while `d`/`dd` + # are re-masked before they reach `m2_t`. (Bounded: `other=1e30` + # would swamp `mean0` and the correction with it, and `other=inf` + # or `nan` would poison it outright.) 0.0 is kept because it is + # the value that survives all three of those, not because the + # cancellation is something to rely on. x = tl.load(X + base + off, mask=m, other=0.0).to(tl.float32) # Corrected two-pass within the tile: the first mean loses digits # to the magnitude of the data, `corr` puts them back, and the # centred squares are then accurate to fp32 roundoff. Everything # here is register traffic; the tile is read from HBM exactly once. + # `corr` is worth 1.3x to 1472x on `rstd` once the data's mean + # dominates its spread -- see "What the third pass is worth" in the + # module docstring for the measurements, for the one regime where + # it goes the other way, and for why the *output* error barely + # moves even where `rstd` improves by three orders of magnitude. cnt_t = (nvalid * CG).to(tl.float32) mean0 = tl.sum(tl.sum(x, 2), 0) / cnt_t d = tl.where(m, x - mean0[None, :, None], 0.0) @@ -931,6 +1017,34 @@ def _ensure_kernels(): # --------------------------------------------------------------------------- # _CL_FORMAT = torch.channels_last_3d +#: Reused so the common (already-current device) path allocates nothing. +_NO_GUARD = contextlib.nullcontext() + + +def _device_guard(device: torch.device): + """Make ``device`` current for the kernel launches inside the ``with``. + + A Triton launch goes to whatever device is *current*, not to the device the + argument tensors live on, so without this a tensor on ``cuda:1`` while + ``cuda:0`` is current makes the kernel dereference another device's pointers + and the process dies with ``Memory access fault by GPU node-N``. ATen ops + (including ``F.group_norm``) carry a ``DeviceGuard`` and handle the same + call, so this is required for the drop-in contract, not a nicety. + + The ``current_device()`` test is not about correctness but about *cost*. + Measured on this node (median of 200k calls, torch 2.13.0+rocm7.2): + ``with torch.cuda.device(t.device)`` is **1.55 us** of host time per call, + ``with torch.cuda._DeviceGuard(t.device.index)`` **0.61 us**, and this + helper **0.51 us** when the tensor is already on the current device -- + which it is on every ScaFFold call, since ScaFFold pins one device per + rank. Two of those (forward + backward) against the 0.65 ms fwd+bwd of the + two smallest scale-8 shapes, which are host-dispatch bound, is 0.16% + instead of 0.48%. + """ + if device.index == torch.cuda.current_device(): + return _NO_GUARD + return torch.cuda.device(device) + def _shape_of(input: torch.Tensor): n, channels = input.shape[0], input.shape[1] @@ -947,65 +1061,66 @@ def _forward(input, num_groups, weight, bias, eps, activation, out_dtype): groups = num_groups device = input.device - pcnt = torch.empty(n * plan.nsplit * groups, device=device, dtype=torch.float32) - pmean = torch.empty_like(pcnt) - pm2 = torch.empty_like(pcnt) - mean = torch.empty((n, groups), device=device, dtype=torch.float32) - rstd = torch.empty_like(mean) - - _stats_partial_kernel[(plan.nsplit, n)]( - input, - pcnt, - pmean, - pm2, - spatial, - plan.chunk, - C=channels, - G=groups, - CG=plan.group_channels, - GP=plan.groups_p2, - CGP=plan.group_channels_p2, - NSPLIT=plan.nsplit, - BLOCK_S=plan.block_s_stats, - MASKED_C=plan.masked_c, - INT64=plan.int64, - num_warps=plan.cfg.stats_warps, - ) - _stats_finalize_kernel[(n * groups,)]( - pcnt, - pmean, - pm2, - mean, - rstd, - plan.elements_per_group, - eps, - G=groups, - NSPLIT=plan.nsplit, - num_warps=4, - ) + with _device_guard(device): + pcnt = torch.empty(n * plan.nsplit * groups, device=device, dtype=torch.float32) + pmean = torch.empty_like(pcnt) + pm2 = torch.empty_like(pcnt) + mean = torch.empty((n, groups), device=device, dtype=torch.float32) + rstd = torch.empty_like(mean) + + _stats_partial_kernel[(plan.nsplit, n)]( + input, + pcnt, + pmean, + pm2, + spatial, + plan.chunk, + C=channels, + G=groups, + CG=plan.group_channels, + GP=plan.groups_p2, + CGP=plan.group_channels_p2, + NSPLIT=plan.nsplit, + BLOCK_S=plan.block_s_stats, + MASKED_C=plan.masked_c, + INT64=plan.int64, + num_warps=plan.cfg.stats_warps, + ) + _stats_finalize_kernel[(n * groups,)]( + pcnt, + pmean, + pm2, + mean, + rstd, + plan.elements_per_group, + eps, + G=groups, + NSPLIT=plan.nsplit, + num_warps=4, + ) - out = torch.empty_like(input, dtype=out_dtype, memory_format=_CL_FORMAT) - _normalize_kernel[(plan.nblk_elem, n)]( - input, - out, - mean, - rstd, - weight, - bias, - spatial, - C=channels, - G=groups, - CG=plan.group_channels, - GP=plan.groups_p2, - CGP=plan.group_channels_p2, - BLOCK_S=plan.block_s_elem, - RELU=activation == "relu", - HAS_W=weight is not None, - HAS_B=bias is not None, - MASKED_C=plan.masked_c, - INT64=plan.int64, - num_warps=plan.cfg.elem_warps, - ) + out = torch.empty_like(input, dtype=out_dtype, memory_format=_CL_FORMAT) + _normalize_kernel[(plan.nblk_elem, n)]( + input, + out, + mean, + rstd, + weight, + bias, + spatial, + C=channels, + G=groups, + CG=plan.group_channels, + GP=plan.groups_p2, + CGP=plan.group_channels_p2, + BLOCK_S=plan.block_s_elem, + RELU=activation == "relu", + HAS_W=weight is not None, + HAS_B=bias is not None, + MASKED_C=plan.masked_c, + INT64=plan.int64, + num_warps=plan.cfg.elem_warps, + ) return out, mean, rstd @@ -1016,100 +1131,133 @@ def _backward(grad_out, input, weight, bias, mean, rstd, num_groups, activation) groups = num_groups device = input.device - ps1 = torch.empty(n * plan.nsplit * groups, device=device, dtype=torch.float32) - ps2 = torch.empty_like(ps1) - pdw = torch.empty(n * plan.nsplit * channels, device=device, dtype=torch.float32) - pdb = torch.empty_like(pdw) - - _bwd_partial_kernel[(plan.nsplit, n)]( - input, - grad_out, - mean, - rstd, - weight, - bias, - ps1, - ps2, - pdw, - pdb, - spatial, - plan.chunk, - C=channels, - G=groups, - CG=plan.group_channels, - GP=plan.groups_p2, - CGP=plan.group_channels_p2, - NSPLIT=plan.nsplit, - BLOCK_S=plan.block_s_stats, - RELU=activation == "relu", - HAS_W=weight is not None, - HAS_B=bias is not None, - MASKED_C=plan.masked_c, - INT64=plan.int64, - num_warps=plan.cfg.stats_warps, - ) + with _device_guard(device): + ps1 = torch.empty(n * plan.nsplit * groups, device=device, dtype=torch.float32) + ps2 = torch.empty_like(ps1) + pdw = torch.empty( + n * plan.nsplit * channels, device=device, dtype=torch.float32 + ) + pdb = torch.empty_like(pdw) + + _bwd_partial_kernel[(plan.nsplit, n)]( + input, + grad_out, + mean, + rstd, + weight, + bias, + ps1, + ps2, + pdw, + pdb, + spatial, + plan.chunk, + C=channels, + G=groups, + CG=plan.group_channels, + GP=plan.groups_p2, + CGP=plan.group_channels_p2, + NSPLIT=plan.nsplit, + BLOCK_S=plan.block_s_stats, + RELU=activation == "relu", + HAS_W=weight is not None, + HAS_B=bias is not None, + MASKED_C=plan.masked_c, + INT64=plan.int64, + num_warps=plan.cfg.stats_warps, + ) - c1 = torch.empty(n * groups, device=device, dtype=torch.float32) - c2 = torch.empty_like(c1) - _bwd_finalize_kernel[(n * groups,)]( - ps1, - ps2, - c1, - c2, - plan.elements_per_group, - G=groups, - NSPLIT=plan.nsplit, - num_warps=4, - ) + c1 = torch.empty(n * groups, device=device, dtype=torch.float32) + c2 = torch.empty_like(c1) + _bwd_finalize_kernel[(n * groups,)]( + ps1, + ps2, + c1, + c2, + plan.elements_per_group, + G=groups, + NSPLIT=plan.nsplit, + num_warps=4, + ) - d_weight = torch.empty(channels, device=device, dtype=torch.float32) - d_bias = torch.empty_like(d_weight) - rows = n * plan.nsplit - block_c = min(256, max(64, _next_pow2(channels))) - block_r = 32 if rows >= 32 else 1 - _dwdb_reduce_kernel[(_cdiv(channels, block_c),)]( - pdw, - pdb, - d_weight, - d_bias, - rows, - channels, - BLOCK_C=block_c, - BLOCK_R=block_r, - num_warps=4, - ) + d_weight = torch.empty(channels, device=device, dtype=torch.float32) + d_bias = torch.empty_like(d_weight) + rows = n * plan.nsplit + block_c = min(256, max(64, _next_pow2(channels))) + block_r = 32 if rows >= 32 else 1 + _dwdb_reduce_kernel[(_cdiv(channels, block_c),)]( + pdw, + pdb, + d_weight, + d_bias, + rows, + channels, + BLOCK_C=block_c, + BLOCK_R=block_r, + num_warps=4, + ) - d_input = torch.empty_like(input, memory_format=_CL_FORMAT) - _dx_kernel[(plan.nblk_elem, n)]( - input, - grad_out, - d_input, - mean, - rstd, - weight, - bias, - c1, - c2, - spatial, - C=channels, - G=groups, - CG=plan.group_channels, - GP=plan.groups_p2, - CGP=plan.group_channels_p2, - BLOCK_S=plan.block_s_elem, - RELU=activation == "relu", - HAS_W=weight is not None, - HAS_B=bias is not None, - MASKED_C=plan.masked_c, - INT64=plan.int64, - num_warps=plan.cfg.elem_warps, - ) + d_input = torch.empty_like(input, memory_format=_CL_FORMAT) + if plan.group_channels * spatial == 1: + # One element per group: mean == x and var == 0 identically, so + # xhat is the constant 0 and y does not depend on x at all -- the + # exact d_input is zero everywhere. _dx_kernel would instead + # return rstd * (dy*w - c1), and since the compiler contracts that + # to fma(dy, w, -c1) while c1 was accumulated from the *rounded* + # product, what survives is the product's rounding error amplified + # by rstd = 1/sqrt(eps) ~ 316 (2.2e-05 at eps=1e-5). Answering + # with the exact zero costs one integer test per backward call. + d_input.zero_() + else: + _dx_kernel[(plan.nblk_elem, n)]( + input, + grad_out, + d_input, + mean, + rstd, + weight, + bias, + c1, + c2, + spatial, + C=channels, + G=groups, + CG=plan.group_channels, + GP=plan.groups_p2, + CGP=plan.group_channels_p2, + BLOCK_S=plan.block_s_elem, + RELU=activation == "relu", + HAS_W=weight is not None, + HAS_B=bias is not None, + MASKED_C=plan.masked_c, + INT64=plan.int64, + num_warps=plan.cfg.elem_warps, + ) return d_input, d_weight, d_bias # --------------------------------------------------------------------------- # # torch.library registration # --------------------------------------------------------------------------- # +def _one_value_per_channel(input, num_groups: int) -> bool: + """Whether ``F.group_norm`` would reject this shape as degenerate. + + ``F.group_norm`` runs ``_verify_batch_size([N*C//G, G, *spatial])``, which + raises ``ValueError("Expected more than 1 value per channel when + training")`` exactly when ``N * (C/G) * D*H*W == 1``. All three factors are + positive, so that holds iff ``N == 1``, ``C == num_groups`` and the spatial + extent is 1 -- i.e. iff ``numel == C == num_groups``, which is the cheap + form used here (``numel`` is wanted by the caller anyway). + + Rejected rather than served: the kernel *can* compute it (it returns + ``bias``, since every group has zero variance), but a caller that branches + on :func:`is_supported` would then get a result where the op this replaces + raises, which is a worse failure than being slower. + """ + channels = input.shape[1] + return channels == num_groups and input.numel() == channels + + def _validate(input, num_groups, weight, bias, activation): if activation not in SUPPORTED_ACTIVATIONS: raise ValueError( @@ -1121,6 +1269,13 @@ def _validate(input, num_groups, weight, bias, activation): raise ValueError( f"num_channels={input.shape[1]} is not divisible by num_groups={num_groups}" ) + if _one_value_per_channel(input, num_groups): + # Same rejection, and the same exception type, as F.group_norm's + # _verify_batch_size; see _one_value_per_channel. + raise ValueError( + f"Expected more than 1 value per channel when training, got input " + f"size {tuple(input.shape)} with num_groups={num_groups}" + ) if input.dtype not in SUPPORTED_DTYPES: raise ValueError(f"unsupported input dtype {input.dtype}") if not input.is_contiguous(memory_format=_CL_FORMAT): @@ -1155,8 +1310,10 @@ def _group_norm_op( """Channels-last GroupNorm forward: returns ``(output, mean, rstd)``. ``mean``/``rstd`` are ``(N, num_groups)`` fp32 tensors kept for the - backward; they are *not* differentiable (nothing produces a gradient for - them) and callers should treat them as opaque. + backward; they are marked non-differentiable in ``_setup_context`` + (nothing produces a gradient for them), so they come back with + ``requires_grad=False`` and differentiating through them raises rather than + returning zeros. Callers should treat them as opaque. """ _validate(input, num_groups, weight, bias, activation) weight = None if weight is None else weight.contiguous() @@ -1204,7 +1361,9 @@ def _group_norm_backward_op( d_input, d_weight, d_bias = _backward( grad_out, input, weight, bias, mean, rstd, num_groups, activation ) - d_input = d_input.to(input.dtype) + # No `d_input.to(input.dtype)`: `_backward` allocates it with + # `empty_like(input)` and `_dx_kernel` stores through `DX.dtype.element_ty`, + # so it already *is* the input's dtype. if weight is None: d_weight = d_weight.new_empty(0) else: @@ -1219,7 +1378,12 @@ def _group_norm_backward_op( @_group_norm_backward_op.register_fake def _(grad_out, input, weight, bias, mean, rstd, num_groups, activation): channels = input.shape[1] - d_input = torch.empty_like(input) + # channels_last_3d, *not* the input's own format: the real op relayouts a + # non-channels-last `input` and always returns a channels-last `d_input`, + # so promising `empty_like(input)` here would hand torch.compile the wrong + # strides for any contiguous NCDHW input -- silently, since eager never + # consults the fake kernel. + d_input = torch.empty_like(input, memory_format=_CL_FORMAT) d_weight = input.new_empty( channels if weight is not None else 0, dtype=weight.dtype if weight is not None else torch.float32, @@ -1234,6 +1398,12 @@ def _(grad_out, input, weight, bias, mean, rstd, num_groups, activation): def _setup_context(ctx, inputs, output): input, num_groups, weight, bias, eps, activation, out_dtype = inputs _out, mean, rstd = output + # Outputs 1 and 2 are backward state, not results: nothing produces a + # gradient for them. Without this they come back requiring grad, and + # differentiating through them *succeeds* -- autograd materializes an + # all-zero cotangent for the unused `out` and runs the whole backward to + # return zeros, which is a plausible wrong answer rather than an error. + ctx.mark_non_differentiable(mean, rstd) ctx.save_for_backward(input, weight, bias, mean, rstd) ctx.num_groups = num_groups ctx.activation = activation @@ -1308,6 +1478,19 @@ def is_supported( channel count divisible by ``num_groups``, and affine parameters whose dtype ``F.group_norm`` would itself accept for this input (equal to the input's, or fp32 under autocast, which is what autocast would produce). + Shapes ``F.group_norm`` itself rejects are rejected here too, so that + branching on this predicate can never turn a stock ``ValueError`` into an + answer (see :func:`_one_value_per_channel`). + + ``True`` promises the *first* derivative only: the backward is itself a + custom op with no autograd formula, so a second ``torch.autograd.grad`` + raises where stock ``F.group_norm`` would succeed. Callers that need a + gradient penalty or a Hessian-vector product must not take this path. + + Note that this is a capability predicate, not a layout classifier: for + shapes whose spatial *and* channel extents make the contiguous and + channels-last-3d stride patterns coincide (e.g. ``(N, C, 1, 1, 1)``), a + plain contiguous tensor is accepted, correctly -- it is the same bytes. """ if activation not in SUPPORTED_ACTIVATIONS: return False @@ -1322,6 +1505,8 @@ def is_supported( channels = input.shape[1] if channels % num_groups != 0 or input.numel() == 0: return False + if _one_value_per_channel(input, num_groups): + return False if not input.is_contiguous(memory_format=_CL_FORMAT): return False autocast = None diff --git a/tests/test_triton_group_norm_edge.py b/tests/test_triton_group_norm_edge.py new file mode 100644 index 0000000..65fe148 --- /dev/null +++ b/tests/test_triton_group_norm_edge.py @@ -0,0 +1,1339 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +"""Adversarial edge-case tests for the channels-last Triton GroupNorm. + +Companion to ``tests/test_triton_group_norm.py``, written independently during +an audit of the kernel. It covers the ground the author's suite does not, and +pins the divergences that audit found. See +``review/gn-dctensor/KERNEL_REVIEW.md`` for the full write-up. + +The two structural gaps this file closes: + +* **The masked channel axis is never exercised upstream.** Every GPU test in + ``test_triton_group_norm.py`` uses ``num_groups=8`` with a channel count of + 64/128/256/2048, so ``G`` and ``C/G`` are *always* powers of two and + ``_Plan.masked_c`` is always ``False``. The entire ``MASKED_C=True`` code + path -- the ``cmask``/``wbm`` predicates in all four kernels, and the + ``inner`` offsets that deliberately run past the end of a voxel -- ships + untested. :func:`test_masked_channel_axis_parity` and friends run it. + +* **Uninitialised split-K scratch is never checked.** ``_forward`` and + ``_backward`` allocate their partial buffers with ``torch.empty``, so a slot + that is read before it is written would surface as *plausible* numbers, not + as a crash. :func:`test_scratch_slots_are_all_written` poisons every + ``torch.empty`` with NaN for the duration of the call, which turns that class + of bug into a hard failure. + +The audit's six findings -- no device guard, the backward fake kernel's stride +promise, silently-differentiable ``mean``/``rstd``, accepting a shape +``F.group_norm`` rejects, a non-zero ``d_input`` for single-element groups, and +undocumented double backward -- were tested here as ``xfail(strict=True)`` +first and fixed afterwards; the tests remain, without the markers, as the +regression pins. The last section adds the coverage a mutation sweep of the +kernels found thinnest: the ``INT64=True`` branch (which a default run never +compiled), the split-K Welford merge on *unequal* split counts, ``eps`` +placement, and the tile-mean correction term. +""" + +import contextlib +import os +import subprocess +import sys +import textwrap + +import pytest +import torch +import torch.nn.functional as F + +from ScaFFold.unet import triton_group_norm as tgn +from ScaFFold.unet.triton_group_norm import is_supported, triton_group_norm + +CL = torch.channels_last_3d +EPS = 1e-5 + +#: Relative-error ceiling against the float64 reference below. fp32 parity at +#: these (small) shapes measures ~1e-07; the ceiling leaves room for the +#: reduction noise a production-sized split-K reduction shows. +FP32_TOL = 1e-4 + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + + +# --------------------------------------------------------------------------- +# independent float64 reference (deliberately *not* F.group_norm, and +# deliberately not the helper the author's suite uses) +# --------------------------------------------------------------------------- + + +def _ref64(x, groups, weight, bias, eps, activation=None): + """GroupNorm written from scratch in float64, in the (N, G, ...) view.""" + xd = x.double() + n, channels = xd.shape[0], xd.shape[1] + flat = xd.reshape(n, groups, -1) + mu = flat.mean(-1, keepdim=True) + var = ((flat - mu) ** 2).mean(-1, keepdim=True) + y = ((flat - mu) / torch.sqrt(var + eps)).reshape(xd.shape) + shape = (1, channels) + (1,) * (xd.dim() - 2) + if weight is not None: + y = y * weight.double().reshape(shape) + if bias is not None: + y = y + bias.double().reshape(shape) + return torch.relu(y) if activation == "relu" else y + + +def _rel(actual, expected): + a = actual.detach().double() + e = expected.detach().double() + return ((a - e).abs().max() / e.abs().max().clamp_min(1e-300)).item() + + +def _make(shape, groups, dtype=torch.float32, affine=True, seed=0, mean=0.0, std=1.0): + device = torch.device("cuda") + gen = torch.Generator(device=device).manual_seed(seed) + x = torch.empty(shape, device=device, dtype=dtype, memory_format=CL) + x.normal_(mean, std, generator=gen) + channels = shape[1] + if affine: + weight = torch.empty(channels, device=device, dtype=dtype) + weight.normal_(1.0, 0.25, generator=gen) + bias = torch.empty(channels, device=device, dtype=dtype) + bias.normal_(0.0, 0.25, generator=gen) + else: + weight = bias = None + grad_out = torch.empty(shape, device=device, dtype=dtype, memory_format=CL) + grad_out.normal_(generator=gen) + return x, weight, bias, grad_out + + +def _parity(shape, groups, activation=None, affine=True, seed=0, eps=EPS, label=""): + """Forward + backward against the float64 reference. Returns the errors.""" + x, weight, bias, grad_out = _make(shape, groups, affine=affine, seed=seed) + assert is_supported(x, groups, weight, bias, activation), ( + f"is_supported rejected {shape} groups={groups}" + ) + + xi = x.detach().clone().requires_grad_(True) + wi = None if weight is None else weight.detach().clone().requires_grad_(True) + bi = None if bias is None else bias.detach().clone().requires_grad_(True) + triton_group_norm(xi, groups, wi, bi, eps, activation).backward(grad_out) + + xd = x.detach().clone().double().requires_grad_(True) + wd = ( + None + if weight is None + else weight.detach().clone().double().requires_grad_(True) + ) + bd = None if bias is None else bias.detach().clone().double().requires_grad_(True) + _ref64(xd, groups, wd, bd, eps, activation).backward(grad_out.double()) + + errors = {"dx": _rel(xi.grad, xd.grad)} + if wi is not None: + errors["dw"] = _rel(wi.grad, wd.grad) + if bi is not None: + errors["db"] = _rel(bi.grad, bd.grad) + xj = x.detach().clone() + errors["y"] = _rel( + triton_group_norm(xj, groups, weight, bias, eps, activation), + _ref64(x, groups, weight, bias, eps, activation), + ) + print(f"[{label or shape}] " + " ".join(f"{k}={v:.2e}" for k, v in errors.items())) + for name, err in errors.items(): + assert err <= FP32_TOL, f"{label or shape}: {name} rel err {err:.3e}" + return errors + + +# --------------------------------------------------------------------------- +# 1. the masked channel axis (MASKED_C=True) -- never reached upstream +# --------------------------------------------------------------------------- + +#: ``(shape, num_groups)`` pairs for which ``_Plan.masked_c`` is True, i.e. +#: ``num_groups`` and/or ``num_channels // num_groups`` is not a power of two, +#: so ``GP``/``CGP`` over-cover the channel axis and every load, store and +#: reduction in all four kernels has to be predicated. +_MASKED_CASES = [ + ((1, 6, 4, 4, 4), 3), # G=3 -> GP=4, CG=2 + ((1, 15, 5, 5, 5), 3), # G=3, CG=5 -> both padded + ((2, 12, 7, 5, 3), 3), # N>1 with a padded group axis + ((1, 24, 9, 9, 9), 6), # G=6 -> GP=8, CG=4 + ((1, 20, 4, 4, 4), 5), # G=5 -> GP=8, CG=4 + ((1, 20, 4, 4, 4), 4), # G=4, CG=5 -> only the inner axis padded + ((3, 20, 3, 5, 7), 5), + ((1, 63, 5, 5, 5), 7), # G=7, CG=9 + ((2, 63, 5, 5, 5), 7), + ((1, 96, 5, 5, 5), 6), # G=6, CG=16 + ((1, 10, 3, 3, 3), 10), # G == C, neither a power of two +] + + +@pytest.mark.gpu +@pytest.mark.parametrize("shape,groups", _MASKED_CASES) +@pytest.mark.parametrize("activation", [None, "relu"]) +def test_masked_channel_axis_parity(shape, groups, activation): + """``MASKED_C=True``: the padded (G, C/G) tile must be fully predicated. + + ``inner = g * CG + j`` deliberately runs past the end of a voxel for the + padding lanes, so a wrong ``cmask``/``wbm`` predicate reads (or writes) the + *next* voxel's channels, and a wrong ``other=`` poisons the Welford sums. + Neither shows up anywhere in the author's suite, which only ever runs + ``num_groups=8`` over 64/128/256/2048 channels. + """ + plan = tgn._plan(shape[0], shape[1], shape[2] * shape[3] * shape[4], groups, 0) + assert plan.masked_c, "case is supposed to exercise the padded channel axis" + _parity(shape, groups, activation, seed=abs(hash((shape, groups))) % 997) + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "shape,groups", + [ + ((1, 64, 4, 4, 4), 64), # instance norm, C/G == 1 + ((1, 7, 3, 3, 3), 7), # instance norm, prime channel count + ((1, 64, 4, 4, 4), 1), # layer norm, G == 1 + ((1, 7, 3, 3, 3), 1), # layer norm, prime channel count + ((1, 1, 4, 4, 4), 1), # single channel + ((2, 1, 4, 4, 4), 1), + ((1, 3, 5, 5, 5), 3), + ], +) +def test_extreme_group_counts(shape, groups): + """``num_groups == num_channels`` (instance norm) and ``== 1`` (layer norm). + + Both collapse one axis of the ``(BLOCK_S, GP, CGP)`` tile to length 1 and + are accepted by ``is_supported``; neither appears upstream. + """ + _parity(shape, groups, seed=abs(hash((shape, groups))) % 997) + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "shape", + [ + (1, 64, 1, 1, 1), # S == 1: a single voxel, far below one tile + (2, 64, 1, 1, 1), + (4, 64, 2, 1, 1), + (1, 64, 1, 1, 127), # just under a 128-voxel stats tile + (1, 64, 1, 1, 128), # exactly one tile + (1, 64, 1, 1, 129), # just over + (1, 64, 1, 1, 257), + (1, 64, 13, 17, 19), # three primes + (5, 64, 3, 3, 3), # N not a power of two + (1, 2048, 1, 1, 2), # widest channel count, two voxels + ], +) +def test_ragged_spatial_tails(shape): + """Spatial extents that are prime, or sit just either side of a tile edge. + + The ragged tail is where ``offs_s < S - s0`` in ``_normalize_kernel`` / + ``_dx_kernel`` and ``nvalid = min(BLOCK_S, s_end - s0)`` in the two partial + kernels have to agree; ``cnt_t = nvalid * CG`` also has to be the *valid* + lane count or the Welford mean is scaled wrong. + """ + _parity(shape, 8, seed=abs(hash(shape)) % 997) + + +# --------------------------------------------------------------------------- +# 2. split-K scratch +# --------------------------------------------------------------------------- + + +def _empty_split_count(n, channels, spatial, groups): + plan = tgn._plan(n, channels, spatial, groups, n * channels * spatial) + return sum(1 for sp in range(plan.nsplit) if sp * plan.chunk >= spatial), plan + + +#: Shapes whose ``ceil(S / nsplit)`` chunking leaves at least one split with +#: ``s_begin >= S``, i.e. a program that writes an all-zero ``(cnt, mean, M2)`` +#: partial that the finalize tree then has to absorb. Found by search over the +#: plan; ``_welford_combine``'s ``cnt == 0`` guard is what makes them harmless. +_EMPTY_SPLIT_CASES = [ + ((1, 64, 1, 1, 32775), 8), + ((2, 64, 1, 1, 32775), 8), + ((1, 128, 1, 1, 8198), 8), + ((1, 256, 1, 1, 2049), 8), + ((1, 2048, 1, 1, 33), 8), +] + + +@pytest.mark.gpu +@pytest.mark.parametrize("shape,groups", _EMPTY_SPLIT_CASES) +def test_empty_split_slots(shape, groups): + """A split whose whole chunk lies past ``S`` still has to combine cleanly. + + ``chunk = ceil(S / nsplit)`` can leave trailing splits entirely empty; that + program's loop never runs, so it stores ``(0, 0, 0)``. Chan's combine is + only exact for those because of its ``cnt == 0`` guard, and no upstream + shape produces one. + """ + empties, plan = _empty_split_count( + shape[0], shape[1], shape[2] * shape[3] * shape[4], groups + ) + assert empties > 0, ( + f"expected an empty split for {shape}; plan has nsplit={plan.nsplit} " + f"chunk={plan.chunk}" + ) + _parity(shape, groups, seed=abs(hash(shape)) % 997) + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "shape,groups", + [ + ((1, 64, 8, 8, 8), 8), + ((2, 64, 8, 8, 8), 8), + ((3, 15, 5, 5, 5), 3), + ((1, 64, 1, 1, 32775), 8), # has an empty split + ((2, 2048, 1, 1, 1), 8), + ], +) +def test_scratch_slots_are_all_written(shape, groups): + """Every split-K partial slot must be written before it is read. + + ``_forward``/``_backward`` allocate ``pcnt/pmean/pm2`` and + ``ps1/ps2/pdw/pdb`` with ``torch.empty``. A slot that is read but never + written would inherit whatever the caching allocator last left there -- + usually finite, plausible numbers, which no parity test can be relied on to + catch. Poisoning every ``torch.empty``/``empty_like`` with NaN for the + duration of the call turns that into a hard failure, and also proves the + output buffer itself is fully covered by the store masks. + """ + x, weight, bias, grad_out = _make(shape, groups, seed=5) + real_empty, real_empty_like = torch.empty, torch.empty_like + + def poisoned_empty(*args, **kwargs): + t = real_empty(*args, **kwargs) + return t.fill_(float("nan")) if t.is_floating_point() else t + + def poisoned_empty_like(*args, **kwargs): + t = real_empty_like(*args, **kwargs) + return t.fill_(float("nan")) if t.is_floating_point() else t + + xi = x.detach().clone().requires_grad_(True) + wi = weight.detach().clone().requires_grad_(True) + bi = bias.detach().clone().requires_grad_(True) + torch.empty, torch.empty_like = poisoned_empty, poisoned_empty_like + try: + out = triton_group_norm(xi, groups, wi, bi, EPS) + out.backward(grad_out) + finally: + torch.empty, torch.empty_like = real_empty, real_empty_like + + for name, t in (("y", out), ("dx", xi.grad), ("dw", wi.grad), ("db", bi.grad)): + assert torch.isfinite(t).all(), ( + f"{name} contains NaN with poisoned scratch: a split-K slot (or an " + f"output element) is read/returned without ever being written" + ) + ref = _ref64(x, groups, weight, bias, EPS) + assert _rel(out, ref) <= FP32_TOL + + +# --------------------------------------------------------------------------- +# 3. numerics +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +@pytest.mark.parametrize("value", [0.0, 3.0, 1e3]) +@pytest.mark.parametrize("eps", [1e-5, 1e-12]) +def test_all_equal_input_has_exactly_zero_variance(value, eps): + """Variance exactly 0 => ``rstd = 1/sqrt(eps)`` and ``xhat`` exactly 0. + + This is the sharpest possible statement of the Welford claim: with + ``weight=1, bias=0`` the output must be *identically* zero, with no + tolerance at all. ATen's fp32 GroupNorm does not manage it (it forms the + variance by cancellation and leaves ~1e-05 of noise at ``value=3`` and + ~3e-03 at ``value=1e3``), which is asserted here so the comparison stays + honest if ATen ever changes. + """ + device = torch.device("cuda") + shape = (2, 64, 8, 8, 8) + x = torch.full(shape, value, device=device).contiguous(memory_format=CL) + weight = torch.ones(64, device=device) + bias = torch.zeros(64, device=device) + + got = triton_group_norm(x, 8, weight, bias, eps) + assert torch.equal(got, torch.zeros_like(got)), ( + f"all-equal input must normalise to exactly 0, got max " + f"{got.abs().max().item():.3e}" + ) + # ... and rstd really is 1/sqrt(eps), which only the output scale can show. + _out, _mean, rstd = torch.ops.scaffold_gn.group_norm( + x, 8, None, None, eps, None, None + ) + assert torch.allclose(rstd, torch.full_like(rstd, 1.0 / eps**0.5), rtol=1e-6), ( + f"rstd={rstd.flatten()[0].item():.6e} != 1/sqrt(eps)={1.0 / eps**0.5:.6e}" + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "mean,std,naive_floor", + [(1e4, 1e-2, 1e-1), (1e6, 1.0, 1e1)], +) +def test_welford_at_extreme_mean_to_std_ratio(mean, std, naive_floor): + """Past the ratios the author's suite tests (mu/sigma up to 1e5). + + At ``mu/sigma = 1e6`` the ``E[x^2]-E[x]^2`` formulation is off by ~2e0 to + ~3e2 relative while this kernel holds ~5e-04, and ATen's own fp32 kernel is + 50-70x worse than this one. Measured on MI300A at ``[1, 256, 24^3]``. + """ + x, weight, bias, _ = _make((1, 256, 24, 24, 24), 8, seed=37, mean=mean, std=std) + got = triton_group_norm(x, 8, weight, bias, EPS) + ref = _ref64(x, 8, weight, bias, EPS) + stock = F.group_norm(x, 8, weight, bias, EPS) + + flat = x.reshape(1, 8, -1) + mu = flat.mean(-1) + var = (flat * flat).mean(-1) - mu * mu + naive = (flat - mu[..., None]) / torch.sqrt(var + EPS)[..., None] + naive = naive.reshape(x.shape) * weight.reshape(1, 256, 1, 1, 1) + bias.reshape( + 1, 256, 1, 1, 1 + ) + + err, err_stock, err_naive = _rel(got, ref), _rel(stock, ref), _rel(naive, ref) + print( + f"[mu={mean:g} sd={std:g}] triton={err:.3e} aten={err_stock:.3e} " + f"naive={err_naive:.3e}" + ) + assert err_naive > naive_floor, "the naive formulation was supposed to fail here" + assert err < err_naive / 100.0 + assert err <= err_stock, ( + f"triton {err:.3e} is worse than ATen fp32 {err_stock:.3e} at mu/sigma=" + f"{mean / std:g}" + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize("eps", [1e-5, 1e-8, 1e-12]) +def test_tiny_eps_with_tiny_variance(eps): + """``eps`` far below the default with data whose std is ~1e-4. + + ``rstd = 1/sqrt(var + eps)`` reaches ~1e4 here, so any error in the + variance is amplified by that factor before it reaches the output. + """ + device = torch.device("cuda") + gen = torch.Generator(device=device).manual_seed(7) + x = torch.empty(1, 64, 16, 16, 16, device=device, memory_format=CL) + x.normal_(0.0, 1e-4, generator=gen) + weight = torch.ones(64, device=device) + bias = torch.zeros(64, device=device) + got = triton_group_norm(x, 8, weight, bias, eps) + assert _rel(got, _ref64(x, 8, weight, bias, eps)) <= FP32_TOL + + +# --------------------------------------------------------------------------- +# 4. autograd plumbing +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "kind", ["contiguous", "sliced", "expanded", "transposed", "channels_last"] +) +def test_grad_out_layout_variants(kind): + """A cotangent that is not channels-last-contiguous. + + ``_group_norm_backward_op`` relayouts it; the kernels index it with the + *input's* channels-last stride pattern, so a missed relayout silently + permutes the gradient rather than raising. The author's suite only ever + feeds a channels-last-contiguous cotangent to the fast path. + """ + shape = (2, 64, 5, 6, 7) + x, weight, bias, _ = _make(shape, 8, seed=71) + device = torch.device("cuda") + gen = torch.Generator(device=device).manual_seed(77) + base = torch.empty(shape, device=device) + base.normal_(generator=gen) + if kind == "contiguous": + grad_out = base.contiguous() + elif kind == "channels_last": + grad_out = base.contiguous(memory_format=CL) + elif kind == "sliced": + wide = torch.empty( + (shape[0], shape[1], shape[2], shape[3], shape[4] * 2), device=device + ) + wide.normal_(generator=gen) + grad_out = wide.contiguous(memory_format=CL)[..., ::2] + elif kind == "expanded": + col = torch.empty((shape[0], shape[1], shape[2], shape[3], 1), device=device) + col.normal_(generator=gen) + grad_out = col.expand(shape) + else: # transposed + swapped = torch.empty( + (shape[0], shape[1], shape[2], shape[4], shape[3]), device=device + ) + swapped.normal_(generator=gen) + grad_out = swapped.contiguous(memory_format=CL).transpose(3, 4) + + xi = x.detach().clone().requires_grad_(True) + wi = weight.detach().clone().requires_grad_(True) + bi = bias.detach().clone().requires_grad_(True) + triton_group_norm(xi, 8, wi, bi, EPS).backward(grad_out) + + xd = x.detach().clone().double().requires_grad_(True) + wd = weight.detach().clone().double().requires_grad_(True) + bd = bias.detach().clone().double().requires_grad_(True) + _ref64(xd, 8, wd, bd, EPS).backward(grad_out.double()) + + assert _rel(xi.grad, xd.grad) <= FP32_TOL + assert _rel(wi.grad, wd.grad) <= FP32_TOL + assert _rel(bi.grad, bd.grad) <= FP32_TOL + # d_input keeps the *input's* format regardless of the cotangent's. + assert xi.grad.is_contiguous(memory_format=CL) + + +@pytest.mark.gpu +def test_affine_parameters_may_be_non_contiguous_views(): + """``weight``/``bias`` sliced out of a bigger parameter tensor. + + ``is_supported`` only checks rank, numel, device and dtype, so a strided or + offset 1-D parameter reaches the op, which is why it calls ``.contiguous()`` + on both. Nothing upstream tests that. + """ + x, _weight, _bias, _ = _make((1, 64, 4, 4, 4), 8, seed=83) + device = torch.device("cuda") + gen = torch.Generator(device=device).manual_seed(83) + strided = torch.empty(128, device=device) + strided.normal_(1.0, 0.25, generator=gen) + weight = strided[::2] # stride 2 + pack = torch.empty(4, 64, device=device) + pack.normal_(0.0, 0.25, generator=gen) + bias = pack[2] # storage offset + assert not weight.is_contiguous() + assert is_supported(x, 8, weight, bias) + got = triton_group_norm(x, 8, weight, bias, EPS) + assert _rel(got, _ref64(x, 8, weight, bias, EPS)) <= FP32_TOL + + +@pytest.mark.gpu +@pytest.mark.parametrize("lo,hi", [(0, 2), (1, 3), (3, 4)]) +def test_channels_last_views_with_a_storage_offset(lo, hi): + """A batch slice of a bigger channels-last tensor stays channels-last + contiguous but has a non-zero storage offset -- the kernels must address + from ``data_ptr()``, not from the storage base.""" + device = torch.device("cuda") + gen = torch.Generator(device=device).manual_seed(89) + big = torch.empty((4, 64, 5, 6, 7), device=device, memory_format=CL) + big.normal_(generator=gen) + weight = torch.empty(64, device=device) + weight.normal_(1.0, 0.25, generator=gen) + bias = torch.empty(64, device=device) + bias.normal_(0.0, 0.25, generator=gen) + view = big[lo:hi] + assert view.is_contiguous(memory_format=CL) and is_supported(view, 8, weight, bias) + got = triton_group_norm(view, 8, weight, bias, EPS) + assert _rel(got, _ref64(view, 8, weight, bias, EPS)) <= FP32_TOL + + +@pytest.mark.gpu +def test_double_backward_raises_instead_of_returning_garbage(): + """Higher-order gradients are *not* supported, and must say so. + + ``scaffold_gn::group_norm_backward`` has no autograd formula of its own, so + a second ``torch.autograd.grad`` through the kernel raises. Stock + ``F.group_norm`` supports double backward, so this is a real (if narrow) + behavioural difference from the op it replaces -- anything that needs a + gradient penalty or a Hessian-vector product cannot use this kernel. The + test pins "raises loudly", which is the safe half of the story. + """ + x, weight, bias, grad_out = _make((1, 64, 4, 4, 4), 8, seed=97) + xi = x.detach().clone().requires_grad_(True) + wi = weight.detach().clone().requires_grad_(True) + bi = bias.detach().clone().requires_grad_(True) + y = triton_group_norm(xi, 8, wi, bi, EPS) + (gx,) = torch.autograd.grad(y, xi, grad_out, create_graph=True) + with pytest.raises(RuntimeError, match="no autograd formula was registered"): + torch.autograd.grad(gx.sum(), xi) + + # ... and stock really does support it, so this is a divergence not a law. + xr = x.detach().clone().requires_grad_(True) + yr = F.group_norm(xr, 8, weight, bias, EPS) + (gxr,) = torch.autograd.grad(yr, xr, grad_out, create_graph=True) + (ggr,) = torch.autograd.grad(gxr.sum(), xr) + assert torch.isfinite(ggr).all() + + +# --------------------------------------------------------------------------- +# 5. fake / meta kernel +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("out_dtype", [None, torch.float32]) +@pytest.mark.parametrize("has_w,has_b", [(True, True), (False, False), (True, False)]) +def test_fake_forward_matches_real_in_every_branch(dtype, out_dtype, has_w, has_b): + """The fake kernel must promise the real shape, dtype, stride *and* device. + + A meta mismatch is invisible in eager and silently corrupts + ``torch.compile``; the author's suite spot-checks two combinations, this + walks the whole cross product of dtype x out_dtype override x affine. + """ + from torch._subclasses.fake_tensor import FakeTensorMode + + device = torch.device("cuda") + shape = (2, 64, 5, 6, 7) + x = torch.empty(shape, device=device, dtype=dtype, memory_format=CL).normal_() + weight = torch.randn(64, device=device, dtype=dtype) if has_w else None + bias = torch.randn(64, device=device, dtype=dtype) if has_b else None + + real = torch.ops.scaffold_gn.group_norm(x, 8, weight, bias, EPS, "relu", out_dtype) + with FakeTensorMode() as mode: + args = [None if t is None else mode.from_tensor(t) for t in (x, weight, bias)] + fake = torch.ops.scaffold_gn.group_norm( + args[0], 8, args[1], args[2], EPS, "relu", out_dtype + ) + for i, (r, f) in enumerate(zip(real, fake)): + assert r.shape == f.shape, f"out[{i}] shape" + assert r.dtype == f.dtype, f"out[{i}] dtype {r.dtype} != {f.dtype}" + assert r.stride() == f.stride(), f"out[{i}] stride {r.stride()} != {f.stride()}" + assert r.device.type == f.device.type, f"out[{i}] device" + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "layout", ["contiguous", "channels_last", "sliced", "degenerate"] +) +@pytest.mark.parametrize("has_w,has_b", [(True, True), (False, False), (True, False)]) +def test_fake_backward_matches_real_in_every_branch(layout, has_w, has_b): + """The backward's fake kernel must promise what the real op returns. + + The real op relayouts a non-channels-last ``input`` and *always* returns a + channels-last ``d_input``; ``torch.empty_like(input)`` would instead + preserve the input's own format, so for a plain contiguous NCDHW input the + two disagree ((13440, 1, 2688, 448, 64) against (13440, 210, 42, 7, 1)). A + meta mismatch is invisible in eager and silently corrupts + ``torch.compile``, so every branch of the promise -- both layouts, a + non-contiguous view, the shape where the two formats coincide, and each + affine combination -- is checked here rather than only the CL case. + """ + from torch._subclasses.fake_tensor import FakeTensorMode + + device = torch.device("cuda") + if layout == "degenerate": + shape = (2, 64, 1, 1, 1) # contiguous *is* channels_last_3d here + x = torch.randn(shape, device=device) + else: + shape = (2, 64, 5, 6, 7) + if layout == "contiguous": + x = torch.randn(shape, device=device) + elif layout == "channels_last": + x = torch.randn(shape, device=device).contiguous(memory_format=CL) + else: # sliced: neither contiguous nor channels-last contiguous + x = torch.randn((2, 64, 5, 6, 14), device=device)[..., ::2] + grad_out = torch.randn(shape, device=device).contiguous(memory_format=CL) + weight = torch.randn(64, device=device) if has_w else None + bias = torch.randn(64, device=device) if has_b else None + mean = torch.zeros(2, 8, device=device) + rstd = torch.ones(2, 8, device=device) + + real = torch.ops.scaffold_gn.group_norm_backward( + grad_out, x, weight, bias, mean, rstd, 8, None + ) + with FakeTensorMode() as mode: + a = [ + None if t is None else mode.from_tensor(t) + for t in (grad_out, x, weight, bias, mean, rstd) + ] + fake = torch.ops.scaffold_gn.group_norm_backward( + a[0], a[1], a[2], a[3], a[4], a[5], 8, None + ) + names = ("d_input", "d_weight", "d_bias") + for name, r, f in zip(names, real, fake): + assert r.shape == f.shape, f"{name} shape {r.shape} != {f.shape}" + assert r.dtype == f.dtype, f"{name} dtype {r.dtype} != {f.dtype}" + assert r.stride() == f.stride(), f"{name} stride {r.stride()} != {f.stride()}" + assert r.device.type == f.device.type, f"{name} device" + assert real[0].is_contiguous(memory_format=CL) + + +@pytest.mark.gpu +def test_mean_and_rstd_are_not_silently_differentiable(): + """``mean``/``rstd`` are backward state, so they must refuse, not lie. + + They are documented as "not differentiable". Before they were marked as + such, they came back with ``requires_grad=True`` and differentiating + through them *succeeded*: autograd materialised an all-zero cotangent for + the unused ``out``, ran the entire backward (a full-size zeros allocation + plus four kernels) and returned zeros -- a plausible wrong answer where the + true value is ~6e-04. ``ctx.mark_non_differentiable`` turns that into an + error, which is the only safe outcome short of a real formula. + """ + device = torch.device("cuda") + shape = (2, 64, 5, 6, 7) + x = torch.empty(shape, device=device, memory_format=CL).normal_() + xi = x.clone().requires_grad_(True) + out, mean, rstd = torch.ops.scaffold_gn.group_norm( + xi, 8, None, None, EPS, None, None + ) + assert out.requires_grad, "the forward output must still be differentiable" + assert not mean.requires_grad, "mean must be marked non-differentiable" + assert not rstd.requires_grad, "rstd must be marked non-differentiable" + for name, t in (("mean", mean), ("rstd", rstd)): + with pytest.raises(RuntimeError, match="does not require grad"): + torch.autograd.grad(t.sum(), xi) + assert xi.grad is None, f"differentiating {name} left a gradient behind" + # The value that used to come back silently wrong is genuinely non-zero, + # so "returns zeros" was never defensible as an answer. + xd = x.clone().double().requires_grad_(True) + (want,) = torch.autograd.grad(xd.reshape(2, 8, -1).mean(-1).sum(), xd) + assert want.abs().max() > 0 + + +# --------------------------------------------------------------------------- +# 6. contract / drop-in divergences +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "shape,groups", + [((1, 64, 1, 1, 1), 8), ((2, 64, 1, 1, 1), 8), ((1, 1, 4, 5, 6), 1)], +) +def test_is_supported_accepts_layout_ambiguous_contiguous_input(shape, groups): + """``is_supported`` is *not* simply "False for contiguous input". + + For shapes whose spatial or channel extents are all 1 the contiguous and + channels-last-3d stride patterns coincide, so a plain ``torch.randn`` + tensor is accepted by the fast path. That is benign -- the two layouts are + the same bytes -- but it means callers cannot use ``is_supported`` as a + layout *classifier*. Pinned here so the behaviour is deliberate. + """ + device = torch.device("cuda") + x = torch.randn(shape, device=device) # never asked for channels_last + assert x.is_contiguous() + assert x.is_contiguous(memory_format=CL) + assert is_supported(x, groups) is True + got = triton_group_norm(x, groups, None, None, EPS) + assert _rel(got, _ref64(x, groups, None, None, EPS)) <= FP32_TOL + + +@pytest.mark.gpu +@pytest.mark.parametrize("shape,groups", [((1, 8, 1, 1, 1), 8), ((1, 1, 1, 1, 1), 1)]) +def test_one_value_per_channel_matches_stock_rejection(shape, groups): + """``N*(C/G)*D*H*W == 1`` is a shape ``F.group_norm`` refuses to run. + + The kernel *can* compute it (every group has zero variance, so the answer + is ``bias``), and it used to: ``is_supported`` returned True and + ``triton_group_norm`` returned a value where the op it is a drop-in for + raises ``ValueError``. A caller branching on ``is_supported`` would then + get a different answer from the reference path, which is worse than being + slower, so all three of ``is_supported``, the public wrapper and the raw op + now reject it the same way stock does. + """ + device = torch.device("cuda") + x = torch.empty(shape, device=device, memory_format=CL).normal_() + with pytest.raises(ValueError, match="more than 1 value per channel"): + F.group_norm(x, groups, None, None, EPS) + assert is_supported(x, groups) is False, ( + "is_supported accepts a shape F.group_norm rejects" + ) + # The public wrapper reaches the same rejection through its fallback... + with pytest.raises(ValueError, match="more than 1 value per channel"): + triton_group_norm(x, groups, None, None, EPS) + # ... and the op itself refuses too, for anyone calling it directly. + with pytest.raises(ValueError, match="more than 1 value per channel"): + torch.ops.scaffold_gn.group_norm(x, groups, None, None, EPS, None, None) + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "shape,groups", + [((2, 8, 1, 1, 1), 8), ((1, 8, 2, 1, 1), 8), ((1, 16, 1, 1, 1), 8)], +) +def test_neighbours_of_the_one_value_per_channel_shape_are_still_served(shape, groups): + """The rejection must be exactly stock's, not a shape family around it. + + ``_verify_batch_size`` rejects ``N*(C/G)*spatial == 1`` and nothing else, so + bumping *any one* of N, C/G or the spatial extent to 2 has to come back to + the fast path -- including ``(2, 8, 1, 1, 1)``, which still has a single + element per group. + """ + device = torch.device("cuda") + x = torch.empty(shape, device=device, memory_format=CL).normal_() + F.group_norm(x, groups, None, None, EPS) # stock accepts it + assert is_supported(x, groups) is True + got = triton_group_norm(x, groups, None, None, EPS) + assert _rel(got, _ref64(x, groups, None, None, EPS)) <= FP32_TOL + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "shape,groups,activation", + [ + ((2, 8, 1, 1, 1), 8, None), + ((2, 8, 1, 1, 1), 8, "relu"), + ((1, 8, 2, 1, 1), 8, None), # 2 elements per group: NOT the degenerate case + ((3, 16, 1, 1, 1), 16, None), + ], +) +def test_single_element_group_gradient_is_exactly_zero(shape, groups, activation): + """One element per group => y is constant in x => dx must be identically 0. + + ``mean == x`` and ``var == 0`` identically, so ``xhat`` is the constant 0 + and nothing downstream depends on ``x``. ``_dx_kernel`` used to answer + 2.2e-05 instead: the compiler contracts ``dy*w - c1`` to + ``fma(dy, w, -c1)`` while ``c1`` was accumulated from the *rounded* + product, so what survives is the product's rounding error (7.0e-08, well + under one ulp of ``dyw``), amplified by ``rstd = 1/sqrt(eps) = 316``. + ``_backward`` now recognises the degenerate case and returns the exact + zero; ATen, on the shapes where it will run at all, leaves ~3e-05 there. + + The ``(1, 8, 2, 1, 1)`` case is the control: two elements per group, so the + gradient is *not* identically zero and the kernel must not zero it. + """ + device = torch.device("cuda") + gen = torch.Generator(device=device).manual_seed(3) + channels = shape[1] + x = torch.empty(shape, device=device, memory_format=CL) + x.normal_(generator=gen) + weight = torch.empty(channels, device=device) + weight.normal_(1.0, 0.25, generator=gen) + bias = torch.empty(channels, device=device) + bias.normal_(0.0, 0.25, generator=gen) + grad_out = torch.empty(shape, device=device, memory_format=CL) + grad_out.normal_(generator=gen) + + assert is_supported(x, groups, weight, bias, activation) + xi = x.clone().requires_grad_(True) + wi = weight.clone().requires_grad_(True) + bi = bias.clone().requires_grad_(True) + triton_group_norm(xi, groups, wi, bi, EPS, activation).backward(grad_out) + + xd = x.clone().double().requires_grad_(True) + wd = weight.clone().double().requires_grad_(True) + bd = bias.clone().double().requires_grad_(True) + _ref64(xd, groups, wd, bd, EPS, activation).backward(grad_out.double()) + + if channels // groups * shape[2] * shape[3] * shape[4] == 1: + assert torch.equal(xi.grad, torch.zeros_like(xi.grad)), ( + f"dx should be exactly 0, got {xi.grad.abs().max().item():.3e}" + ) + assert xd.grad.abs().max() == 0, "the float64 reference disagrees" + # d_weight is exactly 0 too (xhat is exactly 0); d_bias is not. + assert torch.equal(wi.grad, torch.zeros_like(wi.grad)) + assert _rel(bi.grad, bd.grad) <= FP32_TOL + else: + assert xd.grad.abs().max() > 0, "control case is supposed to be non-trivial" + assert xi.grad.abs().max() > 0, "the kernel zeroed a non-degenerate gradient" + # A *two*-element group is merely ill-conditioned, not degenerate: + # xhat is +-1/sqrt(1+eps/var) and dx is a difference of near-equal + # terms, so every fp32 implementation loses digits here -- 4.9e-04 + # relative for this kernel and 1.4e-04 for ATen on this input. The + # bound is therefore loose against float64, and tight against ATen, + # which suffers the same cancellation. + assert _rel(xi.grad, xd.grad) <= 1e-3 + xa = x.clone().requires_grad_(True) + F.group_norm(xa, groups, weight, bias, EPS).backward(grad_out) + assert _rel(xi.grad, xa.grad) <= 1e-3 + + +# --------------------------------------------------------------------------- +# 7. composition +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +@pytest.mark.parametrize("activation", [None, "relu"]) +def test_torch_compile_with_dynamic_shapes(activation): + """``dynamic=True`` as well as the author's ``dynamic=False``. + + With dynamic shapes the fake kernel is invoked on *symbolic* sizes, so a + shape/stride promise that only happens to hold for a concrete size shows up + here and nowhere else. ``fullgraph=True`` is the no-graph-break assertion. + """ + x, weight, bias, grad_out = _make((2, 64, 6, 6, 6), 8, seed=59) + + def fn(x, weight, bias): + return triton_group_norm(x, 8, weight, bias, EPS, activation) * 2.0 + + def run(f): + xi = x.detach().clone().requires_grad_(True) + wi = weight.detach().clone().requires_grad_(True) + bi = bias.detach().clone().requires_grad_(True) + out = f(xi, wi, bi) + out.backward(grad_out) + return out.detach(), xi.grad, wi.grad, bi.grad + + torch._dynamo.reset() + eager = run(fn) + compiled = run(torch.compile(fn, fullgraph=True, dynamic=True)) + for name, a, b in zip(("y", "dx", "dweight", "dbias"), compiled, eager): + assert torch.equal(a, b), f"dynamic-shape compile differs from eager in {name}" + assert compiled[0].is_contiguous(memory_format=CL) + + +# --------------------------------------------------------------------------- +# 8. determinism, across processes +# --------------------------------------------------------------------------- + +_DETERMINISM_SCRIPT = textwrap.dedent( + """ + import hashlib, sys, torch + from ScaFFold.unet.triton_group_norm import triton_group_norm + CL = torch.channels_last_3d + + def h(t): + b = t.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes() + return hashlib.sha256(b).hexdigest() + + if sys.argv[1] == "warm": + # Different JIT order, different lru_cache occupancy, different + # allocator state and different free memory before the real work. + junk = [] + for shape, g in (((3, 128, 7, 7, 7), 8), ((1, 15, 5, 5, 5), 3)): + a = torch.empty(shape, device="cuda", memory_format=CL).normal_() + triton_group_norm(a, g, None, None, 1e-5, "relu") + junk.append(torch.empty(1 << 25, device="cuda")) + del junk + torch.cuda.empty_cache() + + for shape, groups, act, dtype in ( + ((2, 128, 9, 11, 13), 8, None, torch.float32), + ((2, 64, 6, 7, 8), 8, "relu", torch.bfloat16), + ((3, 15, 5, 5, 5), 3, None, torch.float32), + ): + gen = torch.Generator(device="cuda").manual_seed(31) + x = torch.empty(shape, device="cuda", dtype=dtype, memory_format=CL) + x.normal_(generator=gen) + w = torch.empty(shape[1], device="cuda", dtype=dtype) + w.normal_(1.0, 0.25, generator=gen) + b = torch.empty(shape[1], device="cuda", dtype=dtype) + b.normal_(0.0, 0.25, generator=gen) + go = torch.empty(shape, device="cuda", dtype=dtype, memory_format=CL) + go.normal_(generator=gen) + xi = x.clone().requires_grad_(True) + wi = w.clone().requires_grad_(True) + bi = b.clone().requires_grad_(True) + y = triton_group_norm(xi, groups, wi, bi, 1e-5, act) + y.backward(go) + print(shape, groups, act, dtype, + h(y), h(xi.grad), h(wi.grad), h(bi.grad), flush=True) + """ +) + + +@pytest.mark.gpu +@pytest.mark.slow +@pytest.mark.timeout(900) +def test_bitwise_determinism_across_processes(): + """Process-to-process bitwise reproducibility, which is half the claim. + + ``test_bitwise_determinism`` upstream only calls the kernel twice in *one* + process, where the plan is already memoised and the JIT cache already warm. + This runs three fresh interpreters -- one of which first JITs other shapes, + churns the caching allocator and changes how much memory is free -- and + compares SHA-256 of the raw output bytes. Anything that made the split + count, tile size or launch geometry depend on device state rather than on + the shape would show up only here. + """ + outputs = [] + for mode in ("plain", "plain", "warm"): + result = subprocess.run( + [sys.executable, "-c", _DETERMINISM_SCRIPT, mode], + capture_output=True, + text=True, + cwd=REPO_ROOT, + timeout=600, + ) + assert result.returncode == 0, result.stderr[-3000:] + outputs.append(result.stdout) + assert outputs[0] == outputs[1], "two identical processes disagree" + assert outputs[0] == outputs[2], ( + "a process that JITted other shapes first disagrees:\n" + f"{outputs[0]}\n--- vs ---\n{outputs[2]}" + ) + assert outputs[0].count("\n") >= 3 + + +# --------------------------------------------------------------------------- +# 9. multi-device +# --------------------------------------------------------------------------- + +_DEVICE_GUARD_SCRIPT = textwrap.dedent( + """ + import sys, torch + import torch.nn.functional as F + from ScaFFold.unet.triton_group_norm import triton_group_norm + CL = torch.channels_last_3d + torch.cuda.set_device(0) # current device = 0 + other = "cuda:1" + g = torch.Generator(device=other).manual_seed(1) + x = torch.empty((1, 64, 4, 4, 4), device=other, memory_format=CL) + x.normal_(generator=g) + w = torch.empty(64, device=other); w.normal_(1.0, 0.25, generator=g) + b = torch.empty(64, device=other); b.normal_(0.0, 0.25, generator=g) + go = torch.empty((1, 64, 4, 4, 4), device=other, memory_format=CL) + go.normal_(generator=g) + + def rel(a, e): + return ((a.double() - e.double()).abs().max() + / e.double().abs().max().clamp_min(1e-30)).item() + + # ATen carries a DeviceGuard, so this is the behaviour to match. + xr = x.clone().requires_grad_(True) + wr = w.clone().requires_grad_(True) + br = b.clone().requires_grad_(True) + F.group_norm(xr, 8, wr, br, 1e-5).backward(go) + + xi = x.clone().requires_grad_(True) + wi = w.clone().requires_grad_(True) + bi = b.clone().requires_grad_(True) + y = triton_group_norm(xi, 8, wi, bi, 1e-5) # tensors on 1, current is 0 + y.backward(go) # ... and so is the backward + torch.cuda.synchronize() + assert torch.cuda.current_device() == 0, "the guard leaked the device" + for name, got, want in (("y", y, F.group_norm(xr.detach(), 8, w, b, 1e-5)), + ("dx", xi.grad, xr.grad), + ("dw", wi.grad, wr.grad), + ("db", bi.grad, br.grad)): + assert got.device == torch.device(other), f"{name} on {got.device}" + e = rel(got, want) + assert e < 1e-4, f"{name}: rel err {e}" + print("OK") + """ +) + + +@pytest.mark.gpu +@pytest.mark.slow +@pytest.mark.timeout(600) +def test_kernel_runs_on_the_inputs_device_not_the_current_one(): + """Tensors on cuda:1 while cuda:0 is current, forward *and* backward. + + A Triton launch goes to whatever device is *current*, so without a device + guard the kernel dereferences another device's pointers and the process + dies with ``Memory access fault by GPU node-N``. ``F.group_norm`` carries + ATen's ``DeviceGuard`` and handles the identical call, so this is a + divergence from the op being replaced, not a PyTorch limitation. + + Run in a subprocess because the failure mode is an unrecoverable GPU memory + fault, which would take the whole pytest session with it. + """ + if torch.cuda.device_count() < 2: + pytest.skip("needs 2 visible CUDA devices") + result = subprocess.run( + [sys.executable, "-c", _DEVICE_GUARD_SCRIPT], + capture_output=True, + text=True, + cwd=REPO_ROOT, + timeout=480, + ) + assert result.returncode == 0 and "OK" in result.stdout, ( + f"returncode={result.returncode}\nstdout={result.stdout}\n" + f"stderr={result.stderr[-2000:]}" + ) + + +@pytest.mark.gpu +def test_device_guard_helper_is_a_no_op_on_the_current_device(): + """The guard must be free on the hot path and real off it. + + ``_device_guard`` skips ``torch.cuda.device`` when the tensor already lives + on the current device (1.55 us against 0.51 us of host time per call, which + is 0.5% of the two smallest scale-8 shapes' 0.65 ms fwd+bwd because they + are host-dispatch bound). This pins both halves of that shortcut so a + future edit cannot quietly turn it into "no guard at all"; the multi-device + behaviour itself is covered by the subprocess test above. + """ + device = torch.device("cuda", torch.cuda.current_device()) + guard = tgn._device_guard(device) + assert guard is tgn._NO_GUARD, "should not build a guard for the current device" + # Constructing a guard for another index does not touch that device. + other = torch.device("cuda", device.index + 1) + assert tgn._device_guard(other) is not tgn._NO_GUARD, ( + "a foreign device must get a real guard" + ) + + +# --------------------------------------------------------------------------- +# 10. addressing at the int32 boundary +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +@pytest.mark.slow +@pytest.mark.timeout(1800) +def test_int32_addressing_at_its_documented_maximum(): + """``numel = INT32_MAX - 127`` with N=2, i.e. ``plan.int64 is False``. + + The author's suite tests the shape *above* the switch + (``test_correct_above_int32_max_elements``) but never the largest shape the + **int32** path itself has to serve, which is where a missing term in the + ``numel + channels > INT32_MAX`` guard would bite. ``65 * 63 * 4097`` is + ``2^24 - 1`` voxels, so nothing about the extents is a power of two. + + Verified without materialising an NCDHW reference: the statistics are + checked against a chunked float64 reduction over the physical (N, S, C) + view, and the output against an elementwise recomputation done per batch + item (a truncated offset wraps partway through, so sample 1 would break + while sample 0 did not). + """ + device = torch.device("cuda") + shape = (2, 64, 65, 63, 4097) + n, channels = shape[0], shape[1] + spatial = shape[2] * shape[3] * shape[4] + numel = n * channels * spatial + assert numel == 2**31 - 128, numel + + plan = tgn._plan(n, channels, spatial, 8, numel) + assert plan.int64 is False, "this shape is supposed to use the int32 path" + + free, total = torch.cuda.mem_get_info() + needed = 4 * numel * 4 + if free < needed: + pytest.skip( + f"needs ~{needed / 2**30:.0f} GiB free, device has " + f"{free / 2**30:.0f} GiB of {total / 2**30:.0f} GiB" + ) + + gen = torch.Generator(device=device).manual_seed(53) + x = torch.empty(shape, device=device, memory_format=CL) + x.normal_(generator=gen) + weight = torch.randn(channels, device=device, generator=gen) + bias = torch.randn(channels, device=device, generator=gen) + out, mean, rstd = torch.ops.scaffold_gn.group_norm( + x, 8, weight, bias, EPS, None, None + ) + + group_channels = channels // 8 + flat = x.permute(0, 2, 3, 4, 1).reshape(n, spatial, channels) # no copy + chunk = 1 << 20 + for i in range(n): + acc = torch.zeros(8, dtype=torch.float64, device=device) + for s in range(0, spatial, chunk): + acc += ( + flat[i, s : s + chunk] + .double() + .reshape(-1, 8, group_channels) + .sum(dim=(0, 2)) + ) + mu = acc / (spatial * group_channels) + acc2 = torch.zeros(8, dtype=torch.float64, device=device) + for s in range(0, spatial, chunk): + d = ( + flat[i, s : s + chunk].double().reshape(-1, 8, group_channels) + - mu[None, :, None] + ) + acc2 += (d * d).sum(dim=(0, 2)) + var = acc2 / (spatial * group_channels) + assert _rel(mean[i], mu) <= 1e-5, f"sample {i} mean" + assert _rel(rstd[i], 1.0 / torch.sqrt(var + EPS)) <= 1e-5, f"sample {i} rstd" + del flat + + mv = ( + mean.reshape(n, 8, 1).expand(n, 8, group_channels).reshape(n, channels, 1, 1, 1) + ) + rv = ( + rstd.reshape(n, 8, 1).expand(n, 8, group_channels).reshape(n, channels, 1, 1, 1) + ) + for i in range(n): + recomputed = (x[i : i + 1] - mv[i : i + 1]) * rv[i : i + 1] * weight.reshape( + 1, channels, 1, 1, 1 + ) + bias.reshape(1, channels, 1, 1, 1) + # fp32 subtraction of two near-equal fp32 values is exact. + err = (out[i : i + 1] - recomputed).abs().max().item() + scale = recomputed.abs().max().item() + print(f"[int32-max sample {i}] elementwise rel err {err / scale:.3e}") + assert err / scale < 1e-5, f"sample {i}" + del recomputed + del x, out, mean, rstd, mv, rv + torch.cuda.empty_cache() + + +# --------------------------------------------------------------------------- +# 11. coverage the mutation sweep found thin +# --------------------------------------------------------------------------- + + +@contextlib.contextmanager +def _force_int64_addressing(): + """Make every plan take the int64 tile-base path, whatever the shape. + + ``_Plan`` sets ``int64 = numel + channels > _INT32_MAX``, so dropping the + threshold turns the wide path on for a shape that fits in a few MiB. The + plan cache is keyed on the shape, not on the threshold, so it has to be + cleared on the way in *and* on the way out. + """ + real = tgn._INT32_MAX + tgn._plan.cache_clear() + tgn._INT32_MAX = -1 + try: + yield + finally: + tgn._INT32_MAX = real + tgn._plan.cache_clear() + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "shape,groups", + [ + ((2, 64, 5, 6, 7), 8), # ragged tail, several splits + ((1, 2048, 6, 6, 6), 8), # widest channel count + ((3, 15, 5, 5, 5), 3), # masked channel axis as well + ((1, 64, 1, 1, 32775), 8), # has an empty split + ], +) +@pytest.mark.parametrize("activation", [None, "relu"]) +def test_int64_addressing_path_is_behaviourally_correct(shape, groups, activation): + """Run the ``INT64=True`` branch of all seven kernels on a small shape. + + ``INT64`` is a ``tl.constexpr``, so the wide and narrow paths are *different + compiled kernels*; only shapes above 2^31 elements reach the wide one + naturally, and the one test that does is ``@pytest.mark.slow`` and needs + 8 GiB. In a default ``-m "not slow"`` run the int64 branch therefore has no + behavioural coverage at all -- forcing ``self.int64`` gives it some for the + price of a few MiB. + + The two paths differ only in the *type* of the scalar tile base, so the + results must be **bitwise** identical, which is a far sharper assertion than + a tolerance and would catch a widened offset that lost or duplicated a tile. + """ + x, weight, bias, grad_out = _make(shape, groups, seed=abs(hash(shape)) % 997) + + def run(): + xi = x.detach().clone().requires_grad_(True) + wi = weight.detach().clone().requires_grad_(True) + bi = bias.detach().clone().requires_grad_(True) + out = triton_group_norm(xi, groups, wi, bi, EPS, activation) + out.backward(grad_out) + return out.detach(), xi.grad, wi.grad, bi.grad + + plan32 = tgn._plan( + shape[0], shape[1], shape[2] * shape[3] * shape[4], groups, x.numel() + ) + assert plan32.int64 is False, "shape is supposed to fit the int32 path" + narrow = run() + + with _force_int64_addressing(): + plan64 = tgn._plan( + shape[0], shape[1], shape[2] * shape[3] * shape[4], groups, x.numel() + ) + assert plan64.int64 is True, "the int64 path was not forced on" + wide = run() + + for name, a, b in zip(("y", "dx", "dweight", "dbias"), wide, narrow): + assert torch.equal(a, b), f"int64 path differs from int32 in {name}" + # ... and both are actually right, not identically wrong. + ref = _ref64(x, groups, weight, bias, EPS, activation) + assert _rel(wide[0], ref) <= FP32_TOL + + +#: ``(shape, groups)`` whose split-K partials have *unequal* counts, because +#: ``chunk = ceil(S / nsplit)`` does not divide ``S``. Chan's combine weights +#: the delta by ``cnt_b / (cnt_a + cnt_b)``; with equal counts every level of +#: the reduction tree has ``cnt_a == cnt_b``, so weighting by the wrong one is +#: invisible. Only a ragged (or empty) trailing split exposes it -- which is +#: why the mutation sweep killed that bug with exactly two parametrizations of +#: one test upstream. +_UNEVEN_SPLIT_CASES = [ + ((2, 64, 9, 7, 5), 8), + ((1, 2048, 6, 6, 6), 8), + ((1, 64, 1, 1, 32775), 8), + ((1, 256, 1, 1, 2049), 8), + ((2, 128, 11, 13, 17), 8), + ((2, 15, 9, 9, 9), 3), # masked channel axis as well + ((1, 20, 17, 17, 17), 5), # 16 splits, trailing split 15 voxels short +] + + +@pytest.mark.gpu +@pytest.mark.parametrize("shape,groups", _UNEVEN_SPLIT_CASES) +@pytest.mark.parametrize("eps", [1e-5, 0.5]) +def test_group_statistics_match_float64_with_uneven_splits(shape, groups, eps): + """Assert ``mean``/``rstd`` themselves, not just the output they feed. + + Two things hide inside the output's 1e-4 tolerance and show up here: + + * **the Welford merge.** The shapes above all have at least one split with + a different element count from its neighbours, which is the only + configuration in which mis-weighting Chan's delta changes the answer. + * **where ``eps`` goes.** Every parity test in both files uses + ``eps=1e-5`` against a variance of ~1, where ``1/sqrt(var+eps)`` and + ``1/(sqrt(var)+eps)`` agree to ~1e-5 -- inside that tolerance. At + ``eps=0.5`` they are 0.816 and 0.667, a 22% difference that no tolerance + can absorb. + """ + spatial = shape[2] * shape[3] * shape[4] + plan = tgn._plan(shape[0], shape[1], spatial, groups, shape[0] * shape[1] * spatial) + counts = { + max(0, min(sp * plan.chunk + plan.chunk, spatial) - sp * plan.chunk) + for sp in range(plan.nsplit) + } + assert plan.nsplit > 1 and len(counts) > 1, ( + f"{shape} was supposed to give unequal split counts; nsplit=" + f"{plan.nsplit} chunk={plan.chunk} counts={sorted(counts)}" + ) + + x, _weight, _bias, _ = _make(shape, groups, seed=abs(hash(shape)) % 997) + _out, mean, rstd = torch.ops.scaffold_gn.group_norm( + x, groups, None, None, eps, None, None + ) + flat = x.double().reshape(shape[0], groups, -1) + mean64 = flat.mean(-1) + var64 = ((flat - mean64[..., None]) ** 2).mean(-1) + rstd64 = 1.0 / torch.sqrt(var64 + eps) + assert _rel(mean, mean64) <= 1e-5, "group mean" + assert _rel(rstd, rstd64) <= 1e-5, "group rstd (eps placement / Welford merge)" + + +@pytest.mark.gpu +@pytest.mark.parametrize("groups", [1, 2, 4]) +@pytest.mark.parametrize("seed", [0, 1, 2, 3]) +def test_welford_correction_recovers_rstd_in_a_single_tile_reduction(groups, seed): + """The third reduction pass (``corr``) is load-bearing, and here is where. + + ``mean0 = sum(x)/n`` loses digits in proportion to the tile's element count + times ``mu/sigma``; ``corr = sum(x-mean0)/n`` recovers them, and ``M2`` is + then formed around the corrected mean. The effect is largest when one tile + carries a whole group's reduction, which is this shape: ``block_s_stats`` + covers all 128 voxels and ``nsplit == 1``, so 8192/``G`` elements per group + go through a single ``mean0``. + + At ``mu/sigma = 1e6`` the correction is worth **216x** (G=1), **580x** + (G=2) and **1472x** (G=4) on the relative error of ``rstd`` -- measured by + running a copy of this module with the term deleted. Corrected lands at + ~1e-07 for every seed and group count; without it, at 2.1e-05 to 1.6e-04. + The 1e-06 ceiling below sits an order of magnitude above the first and an + order of magnitude below the second. + + The *output* is not a witness for this: ``y`` moves by at most ~1.4x with + or without the term, because it is dominated by the fp32 representation of + the mean. That is why this asserts ``rstd`` directly. + """ + device = torch.device("cuda") + shape = (2, 64, 8, 4, 4) + spatial = shape[2] * shape[3] * shape[4] + plan = tgn._plan(shape[0], shape[1], spatial, groups, shape[0] * shape[1] * spatial) + assert plan.nsplit == 1 and plan.block_s_stats >= spatial, ( + f"case is supposed to be a single-tile reduction; nsplit={plan.nsplit} " + f"block_s_stats={plan.block_s_stats} spatial={spatial}" + ) + + gen = torch.Generator(device=device).manual_seed(seed) + x = torch.empty(shape, device=device, memory_format=CL) + x.normal_(1e4, 1e-2, generator=gen) # mu/sigma = 1e6 + _out, mean, rstd = torch.ops.scaffold_gn.group_norm( + x, groups, None, None, EPS, None, None + ) + flat = x.double().reshape(shape[0], groups, -1) + mean64 = flat.mean(-1) + var64 = ((flat - mean64[..., None]) ** 2).mean(-1) + err = _rel(rstd, 1.0 / torch.sqrt(var64 + EPS)) + print(f"[corr G={groups} seed={seed}] rstd rel err {err:.3e}") + assert err <= 1e-6, ( + f"rstd rel err {err:.3e} at mu/sigma=1e6 with G={groups}: the tile mean " + f"correction is not doing its job" + ) From c64c67b61bbb478a34d8ee907f1aea2099511a18 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Sat, 1 Aug 2026 05:13:23 -0700 Subject: [PATCH 57/62] Fold the GroupNorm finalize kernels into their consumers Seven launches per fwd+bwd become four: _stats_finalize is folded into _normalize and _bwd_finalize + _dwdb_reduce into _dx, each elementwise program recomputing the tiny finalize redundantly rather than round-tripping through a launch. Folding alone regresses the dominant shape by 34%, which is the part worth recording: at [1,64,256^3] the elementwise grid is 131072 programs and each would re-read all 2048 partials, so the "tiny" finalize becomes 25.7 GB of reads against a 4.3 GB tensor. Lowering the split count cannot fix it -- the stats kernel needs ~456 splits at that shape just to fill the device. So the elementwise grid is now capped separately (GNConfig.elem_progs) with each program striding over its tiles, which makes the redundant read nprog x nsplit instead of nblk x nsplit: 0.1-403 MB, L2-resident, 1-9% of the tensor. Tiles and split counts were then retuned jointly with the folding, per shape, since the two are not independent. Every candidate was measured interleaved against the incumbent in one process; measured sequentially, this node's background load drifts 35% and swamps the effect entirely. fwd+bwd, min of 7 medians-of-20: [1,64,256^3] -1.7%, [1,128,128^3] -6.8%, [1,256,64^3] -4.9%, and -7.7% to -9.5% at the six small and N>1 shapes. No shape regresses. Scale-8 rollup over 22 sites: 69.5 -> 67.1 ms/step. Results differ bitwise from the previous table -- a different split count is a different reduction order -- but accuracy is unchanged (1.1-1.9e-07 against float64 for both, on all six shapes) and determinism is undiminished: grid, split count, tile sizes and reduction order all remain pure functions of the shape, verified run-to-run and across three interpreters. A new test pins that elem_progs, the one plan field that is a free parameter, cannot move a bit. Also checked and deliberately not changed: the tl.sum(tl.sum(x, 2), 0) reduction in the partial kernels. The documented LDS-staging hazard does not reproduce here (5 shapes x 5 tile sizes x 3 warp counts all compile both ways), and the equivalent spelling measures 0.84% slower, so changing every result's bits buys nothing. --- ScaFFold/unet/triton_group_norm.py | 660 +++++++++++++++++---------- tests/test_triton_group_norm_edge.py | 183 ++++++++ 2 files changed, 590 insertions(+), 253 deletions(-) diff --git a/ScaFFold/unet/triton_group_norm.py b/ScaFFold/unet/triton_group_norm.py index 8b4e9c8..2ba2220 100644 --- a/ScaFFold/unet/triton_group_norm.py +++ b/ScaFFold/unet/triton_group_norm.py @@ -37,20 +37,27 @@ scale-8 UNet GroupNorm shapes, fwd / fwd+bwd in ms:: shape this compiled-CL compiled-CONT eager-CL fwd - [1,64,256^3] 4.28/11.61 19.51/77.01 4.77/11.85 151.2 - [1,128,128^3] 1.19/ 3.20 7.23/24.64 1.22/ 3.07 36.5 + [1,64,256^3] 4.23/11.39 19.51/77.01 4.77/11.85 151.2 + [1,128,128^3] 1.15/ 3.00 7.23/24.64 1.22/ 3.07 36.5 [1,256,64^3] 0.35/ 0.97 2.27/ 7.42 0.34/ 0.84 9.1 - [1,512,32^3] 0.14/ 0.59 0.19/ 1.03 0.11/ 0.33 1.7 - [1,1024,16^3] 0.13/ 0.59 0.10/ 0.39 0.07/ 0.33 0.4 - [1,2048,8^3] 0.13/ 0.58 0.08/ 0.40 0.07/ 0.33 0.1 + [1,512,32^3] 0.14/ 0.57 0.19/ 1.03 0.11/ 0.33 1.7 + [1,1024,16^3] 0.11/ 0.57 0.10/ 0.39 0.07/ 0.33 0.4 + [1,2048,8^3] 0.11/ 0.57 0.08/ 0.40 0.07/ 0.33 0.1 -Over the 22 scale-8 call sites that is **442.8 -> 69.0 ms/step** of GroupNorm +Over the 22 scale-8 call sites that is **442.8 -> 67.1 ms/step** of GroupNorm fwd+bwd against today's production path (compiled GroupNorm on channels-last -input), i.e. **374 ms/step recovered**, and a dead heat with compiled GroupNorm +input), i.e. **376 ms/step recovered**, and a dead heat with compiled GroupNorm on *contiguous* input (66.4 ms/step) while additionally not breaking the layout chain. The three smallest shapes lose on host dispatch, not on GPU work -- see :func:`select_strategy`. +The ``this`` column and the rollup were re-measured after the launch folding +below (2+2 kernels instead of 3+4, retuned jointly); the same measurement of +the unfused chain in the same process gives 4.26/11.59, 1.19/3.21, 0.35/1.01, +0.14/0.63, 0.13/0.63, 0.13/0.62 and a 69.5 ms/step rollup, i.e. **-3.4% over +the 22 sites** and -1.7% to -9.5% per shape. The other three columns are from +the earlier sweep and are unchanged. + Public API ========== ``triton_group_norm(input, num_groups, weight=None, bias=None, eps=1e-5, @@ -130,11 +137,54 @@ one pass cannot produce them. Split-K partial reductions land at a fixed scratch index and are combined by a fixed-order tree:: - fwd: stats_partial -> stats_finalize -> normalize (3 kernels) - bwd: bwd_partial -> bwd_finalize -> dwdb_reduce -> dx (4 kernels) + fwd: stats_partial -> normalize (2 kernels) + bwd: bwd_partial -> dx (2 kernels) Traffic (``B = numel * itemsize``): 3B forward, 5B backward. +Each pass is **two** launches, not the three and four an unfused split-K chain +needs: the two finalize passes and the dweight/dbias row reduction are folded +into the elementwise kernel that consumes them. ``_normalize_kernel`` +re-derives ``mean``/``rstd`` from the split-K partials itself (and program 0 +stores them for the backward); ``_dx_kernel`` re-derives ``c1``/``c2`` the same +way and its first ``ceil(C/BLOCK_C)`` programs also do the dweight/dbias +reduction. Folding plus the retuning below is worth 8.0-9.5% of fwd+bwd at the +four smallest shapes, which are host-dispatch bound, 4.9-6.8% at the two middle +ones and 1.7% at the largest. + +The catch, and the reason the tuning table was re-derived rather than inherited: +**the fusion and the tiling are one problem, not two.** A fused finalize is +recomputed by every elementwise *program*, so its cost is +``nprog_elem * nsplit`` triples of redundant (L2-resident) traffic. Keeping the +unfused table's ``nsplit_target=2048`` at ``[1,64,256^3]``, whose flat +elementwise grid is 131072 programs, asks for 25.7 GB of redundant reads +against a 4.3 GB tensor and costs **+34% of fwd+bwd** (+70% of the forward). +Two things fix it, both of them in ``GNConfig``: the elementwise grid is capped +at ``elem_progs`` programs which then stride over the tiles (so the redundancy +is bounded by the *grid*, not by the tile count), and ``nsplit_target`` is +retuned per shape against that cap. With both, the same shape is 1-2% *faster* +than the unfused chain. Do not change one without re-running the other; the +coordinate-descent tuner is ``review/gn-dctensor/kernel-opt/tune.py``. + +Why not one launch per pass +--------------------------- +A device-scope software barrier (int32 atomics with volatile loads, no float +atomics, so still bitwise deterministic) collapses each pass to a single +launch and was measured at 2.3-2.6x on the four smallest shapes. It is +deliberately **not** used. A grid barrier requires every workgroup to be +co-resident, which caps the grid at the CU count (228 here); the kernel then +tops out at 0.5-0.9 TB/s against split-K's 2.7, so it loses catastrophically +the moment the shape is bandwidth-bound rather than dispatch-bound -- +**18.0 ms against 3.0 ms at [1,128,128^3]**, and it does not compile at all at +``[1,64,256^3]``. Serving both regimes therefore means shipping two kernel +families plus a crossover rule, for a whole-model gain of ~3% (65.7 -> 63.6 +ms/step at scale 8); and under CUDA-graph capture, where launch count is free, +the ten small-shape sites are already only 1.05 ms of a 64.1 ms/step total, so +the gain is zero. Hand-rolled inter-workgroup synchronisation is not a good +trade for 3% in a benchmark whose value depends on being trustworthy and +reproducible. The measurements are in +``review/gn-dctensor/triton-small/RESULTS.md``. + Numerics: Welford, not ``E[x^2]-E[x]^2`` ======================================== The prototype accumulated ``sum(x)`` and ``sum(x*x)`` and formed @@ -238,7 +288,7 @@ and streaming voxels -- i.e. a second family of four kernels. The payoff is small: on contiguous input Inductor's compiled GroupNorm already reaches 89-92% of this device's measured streaming roofline (RESULTS.md 4) -- and the - table above confirms it, 66.4 ms/step against this kernel's 69.0 -- so a + table above confirms it, 66.4 ms/step against this kernel's 67.1 -- so a native NCDHW kernel could win ~10% there, against the 6.4x it wins on channels-last input. If a mixed-layout model ever makes that 10% matter, the place to add it is the strategy hook below. @@ -262,10 +312,11 @@ ================ ``activation="relu"`` folds the ReLU into the forward store. In a store-bound kernel that is free (one ``tl.maximum``) and it removes an entire 2B streaming -pass. Measured against ``F.relu(triton_group_norm(x))``: 38% off the forward -and 35% off fwd+bwd at ``[1,64,256^3]`` (6.94 -> 4.29 ms and 18.27 -> 11.80 -ms), 37%/33% at ``[1,128,128^3]``, tapering to ~11% at the launch-bound -shapes. +pass. Measured against ``F.relu(triton_group_norm(x))``: 39% off the forward +and 35% off fwd+bwd at ``[1,64,256^3]`` (6.83 -> 4.20 ms and 17.82 -> 11.61 +ms), 38%/35% at ``[1,128,128^3]``, 34%/30% at ``[1,256,64^3]``, tapering to +21%/9% at ``[1,512,32^3]`` and below, where the call is host bound and there is +less streaming pass to remove. The backward gates the incoming gradient on the sign of the **pre-activation** value, which it *recomputes* from the saved ``(x, mean, rstd, weight, bias)`` @@ -289,11 +340,49 @@ it, and rewrap, so a DCTensor goes in and a DCTensor comes out with the graph intact. As with the rest of DistConv today, statistics are per-shard. -Going through the dispatcher costs ~35 us of host time per forward call -(measured against calling ``_forward``/``_backward`` directly), which is -invisible at the three largest shapes and is roughly two kernel launches at the -three smallest. That is the price of composing, and it is the same order as -the launch overhead those shapes already pay; see :func:`select_strategy`. +Where the host time goes, and what is left +========================================== +Composing has a price, and at the launch-bound shapes it is now the *dominant* +cost. Peeling the layers at ``[1,2048,8^3]``, steady-state wall clock per +fwd+bwd (median of 200, min of 7 rounds; GPU work is 0.030 ms):: + + kernels + this file's Python (_forward/_backward called directly) 0.145 ms + + torch.library dispatcher (both custom ops) +0.054 ms + + autograd (register_autograd node, save_for_backward, ctx) +0.338 ms + = triton_group_norm(x).backward(dy) 0.537 ms + + for scale: an *empty* python torch.autograd.Function, fwd+bwd 0.065 ms + +So **63% of the call is the autograd layer** and 10% is the dispatcher -- both +of them the cost of being a real dispatcher op that ``torch.compile`` and +``DCTensor`` can see, which is the whole point of registering it that way. Of +the 0.145 ms this file is responsible for, 0.030 ms is GPU and the remaining +0.115 ms is four launches plus the allocations, plan lookup and argument +binding around them: launching every kernel twice measures the marginal cost of +a whole invocation site at **35 us**, so the four of them are ~0.14 ms of host +work that the GPU work does not cover. + +Two things were considered for that 0.14 ms and rejected: + +* **Bypassing ``JITFunction.run`` for a cached ``CompiledKernel`` handle** + (8.68 us -> 3.94 us per launch on this node) would recover ~19 us, i.e. 3.3% + of the call. It buys that by asserting that Triton's specialization key -- + including 16-byte pointer alignment -- is a pure function of the shape. It is + not: this module accepts channels-last *views with a storage offset* and + non-contiguous affine parameters, both of which the test suite exercises, and + a stale specialization there is a wrong answer rather than a crash. 3.3% on + the host-bound shapes only is not worth a silent-miscompute failure mode. +* **Caching the scratch buffers** across calls saves ~1.7 us per allocation, + ~20 us here, and makes the buffers shared mutable state across call sites -- + correct on one stream, wrong on two, and this module has no way to know. + +What that leaves: the kernels are at 95-98% of the streaming roofline at the +two largest shapes and the fused chain is 1.7-6.8% faster than the unfused one +there, so there is no meaningful GPU headroom left. At the launch-bound +shapes the remaining 0.4 ms is torch's own plumbing, and the two ways to remove +it are both outside this file: CUDA-graph capture of the training step (which +takes the ten small scale-8 sites to ~1.05 ms/step of GPU time in total), or a +C++ autograd node. Triton is imported lazily, on the first call that actually reaches the kernel, so importing this module (or running the CPU test suite) costs nothing. @@ -335,7 +424,17 @@ class GNConfig: ``stats_tile``/``elem_tile`` are *element* budgets per program (the spatial block is ``tile // channels_per_voxel``, rounded down to a power of two); ``nsplit_target`` is the total number of split-K partials wanted across the - batch, so the per-sample split count is ``nsplit_target // N``. + batch, so the per-sample split count is ``nsplit_target // N``; + ``elem_progs`` caps the elementwise grid, each program then striding over + ``ceil(nblk_elem / elem_progs)`` tiles (0 = one program per tile). + + These are **not** independent knobs, and in particular they stopped being + independent when the finalize passes were folded into the elementwise + kernels: each elementwise *program* now re-reads all ``nsplit`` split-K + partials, so the redundant traffic is ``min(nblk_elem, elem_progs) * + nsplit`` triples. Raising ``nsplit_target`` for stats-kernel occupancy and + lowering ``elem_progs`` for redundancy pull against each other and were + tuned together; see the module docstring. """ __slots__ = ( @@ -344,6 +443,7 @@ class GNConfig: "nsplit_target", "elem_tile", "elem_warps", + "elem_progs", ) def __init__( @@ -353,12 +453,14 @@ def __init__( nsplit_target=2048, elem_tile=8192, elem_warps=4, + elem_progs=2048, ): self.stats_tile = stats_tile self.stats_warps = stats_warps self.nsplit_target = nsplit_target self.elem_tile = elem_tile self.elem_warps = elem_warps + self.elem_progs = elem_progs def key(self): return ( @@ -367,6 +469,7 @@ def key(self): self.nsplit_target, self.elem_tile, self.elem_warps, + self.elem_progs, ) def __eq__(self, other): @@ -378,22 +481,33 @@ def __hash__(self): def __repr__(self): return ( "GNConfig(stats_tile=%d, stats_warps=%d, nsplit_target=%d, " - "elem_tile=%d, elem_warps=%d)" % self.key() + "elem_tile=%d, elem_warps=%d, elem_progs=%d)" % self.key() ) #: Frozen tuning table, produced by coordinate descent on fwd+bwd time on one #: MI300A (228 CUs) at fp32 with ``num_groups=8``, keyed by the #: ``(num_channels, cube-root spatial extent)`` of the scale-8 ScaFFold UNet -#: GroupNorm sites. Frozen -- never autotuned at run time -- because the split -#: count fixes the reduction order and therefore the bits of the result. +#: GroupNorm sites plus the ``[1,4096,4^3]`` tail. Frozen -- never autotuned at +#: run time -- because the split count fixes the reduction order and therefore +#: the bits of the result. +#: +#: Re-derived for the fused (2+2 launch) kernels: every candidate was timed +#: **interleaved against the incumbent** in one process, because this node moves +#: the host-bound shapes by +-35% over the minutes a sweep takes. The table is +#: keyed on ``(C, edge)`` and not on ``N``: ``nsplit_target`` is a target for +#: the split count *summed over the batch* (the per-sample count is +#: ``nsplit_target // N``), so the same entry serves ``N > 1`` with the same +#: total number of stats programs. Verified at ``[2,1024,16^3]``, +#: ``[4,2048,8^3]`` and ``[2,256,64^3]``. _TUNED = { - (64, 256): GNConfig(8192, 4, 2048, 8192, 4), - (128, 128): GNConfig(16384, 4, 512, 16384, 4), - (256, 64): GNConfig(4096, 4, 512, 8192, 4), - (512, 32): GNConfig(32768, 4, 8192, 8192, 4), - (1024, 16): GNConfig(16384, 8, 512, 8192, 4), - (2048, 8): GNConfig(32768, 4, 512, 8192, 4), + (64, 256): GNConfig(16384, 4, 2048, 8192, 4, 2048), + (128, 128): GNConfig(16384, 4, 2048, 16384, 4, 912), + (256, 64): GNConfig(16384, 4, 512, 16384, 8, 912), + (512, 32): GNConfig(32768, 8, 4096, 16384, 4, 0), + (1024, 16): GNConfig(65536, 4, 8192, 16384, 4, 0), + (2048, 8): GNConfig(16384, 8, 1024, 8192, 8, 228), + (4096, 4): GNConfig(16384, 4, 32, 4096, 8, 0), } _DEFAULT_CONFIG = GNConfig() @@ -415,11 +529,13 @@ def default_config(num_channels: int, spatial: int) -> GNConfig: #: thing. STRATEGIES = ("split_k",) -#: Spatial extent (``D*H*W``) below which the split-K chain is expected to be -#: host-dispatch bound rather than bandwidth bound. Measured on MI300A: at -#: ``[1,2048,8^3]`` the seven kernels do 0.031 ms of GPU work behind 0.600 ms -#: of Python/autograd/launch cost, a 19x overhead tax (RESULTS.md 3). Purely -#: informational today -- ``select_strategy`` does not use it yet. +#: Spatial extent (``D*H*W``) below which the split-K chain is host-dispatch +#: bound rather than bandwidth bound. Measured on MI300A: at ``[1,2048,8^3]`` +#: the kernels do 0.030 ms of GPU work behind ~0.58 ms of +#: Python/autograd/launch cost, and 0.086 ms of that is the *empty* +#: ``torch.autograd.Function`` wrapper -- i.e. 68% of the remaining call is +#: torch's plumbing, not this file's. Purely informational -- +#: ``select_strategy`` does not use it. SMALL_SPATIAL_THRESHOLD = 4096 @@ -428,19 +544,22 @@ def select_strategy(n: int, num_channels: int, spatial: int, num_groups: int) -> kernel strategy is chosen for a shape. Returns a name from :data:`STRATEGIES`. Today it always returns - ``"split_k"``: three forward and four backward kernels with split-K partial - reductions, which is bandwidth-optimal for the large shapes but pays seven - kernel launches (~17 us each on this node) plus autograd overhead - regardless of size -- so below roughly ``SMALL_SPATIAL_THRESHOLD`` voxels - the whole call is host bound and a *single-program-per-(n, group)* kernel - that never leaves registers would win. - - That regime is under active investigation; when a second strategy lands, - add its name to :data:`STRATEGIES`, return it from here on a rule that is a - **pure function of the shape** (determinism depends on it), and branch on - it in ``_dispatch`` -- which is the only caller, sits in front of the - memoized tiling plan, and is itself called by both ``_forward`` and - ``_backward``. Nothing else in this file needs to change. + ``"split_k"``: two forward and two backward kernels with split-K partial + reductions and the finalize passes fused into their consumers. That is + bandwidth-optimal for the large shapes and, after the fusion, within + ~0.09 ms of the floor a Python ``autograd.Function`` can reach at the small + ones -- so there is much less left here than there looks. Below roughly + ``SMALL_SPATIAL_THRESHOLD`` voxels the call is host bound, but the host cost + is now dominated by autograd and the dispatcher rather than by launches: + see "Why not one launch per pass" in the module docstring for the one + strategy that *would* cut it further and why it is not here. + + If a second strategy ever lands, add its name to :data:`STRATEGIES`, return + it from here on a rule that is a **pure function of the shape** + (determinism depends on it), and branch on it in ``_dispatch`` -- which is + the only caller, sits in front of the memoized tiling plan, and is itself + called by both ``_forward`` and ``_backward``. Nothing else in this file + needs to change. """ return "split_k" @@ -484,6 +603,14 @@ class _Plan: "chunk", "block_s_elem", "nblk_elem", + "nprog_elem", + "elements_per_group", + "dwdb_rows", + "dwdb_block_c", + "dwdb_block_r", + "dwdb_progs", + "grid_dx", + "zero_dx", "cfg", ) @@ -517,10 +644,26 @@ def __init__(self, n, channels, spatial, groups, cfg, numel): self.chunk = _cdiv(spatial, self.nsplit) self.block_s_elem = max(1, _prev_pow2(cfg.elem_tile // max(1, voxel))) self.nblk_elem = _cdiv(spatial, self.block_s_elem) - - @property - def elements_per_group(self) -> float: - return float(self.spatial * self.group_channels) + # Grid cap for the two elementwise kernels; each program then strides + # over its share of the tiles. Bounds the cost of the fused finalize, + # which every *program* pays once. + self.nprog_elem = ( + self.nblk_elem + if cfg.elem_progs <= 0 + else min(self.nblk_elem, cfg.elem_progs) + ) + self.elements_per_group = float(spatial * self.group_channels) + # Everything the fused dweight/dbias reduction in _dx_kernel needs. + # Precomputed rather than derived per call: the launch-bound shapes pay + # every Python statement in _backward, and _next_pow2 is a loop. + self.dwdb_rows = n * self.nsplit + self.dwdb_block_c = min(256, max(64, _next_pow2(channels))) + self.dwdb_block_r = 32 if self.dwdb_rows >= 32 else 1 + self.dwdb_progs = _cdiv(channels, self.dwdb_block_c) + # Programs past nprog_elem run no elementwise loop iterations; they + # exist only when there are more dweight/dbias blocks than tiles. + self.grid_dx = max(self.nprog_elem, self.dwdb_progs) + self.zero_dx = self.group_channels * spatial == 1 @functools.lru_cache(maxsize=256) @@ -551,11 +694,8 @@ def _dispatch(n, channels, spatial, groups, numel) -> _Plan: tl = None _welford_combine = None _stats_partial_kernel = None -_stats_finalize_kernel = None _normalize_kernel = None _bwd_partial_kernel = None -_bwd_finalize_kernel = None -_dwdb_reduce_kernel = None _dx_kernel = None @@ -702,58 +842,56 @@ def _stats_partial_kernel( tl.store(PMEAN + o, mean, mask=gm) tl.store(PM2 + o, m2, mask=gm) - @_triton.jit - def _stats_finalize_kernel( - PCNT, - PMEAN, - PM2, - MEAN, - RSTD, - M, - eps, - G: tl.constexpr, - NSPLIT: tl.constexpr, - ): - """Merge the NSPLIT partials of one ``(n, g)`` into mean and 1/std.""" - pid = tl.program_id(0) # n * G + g - n = pid // G - g = pid % G - offs = tl.arange(0, NSPLIT) - idx = (n * NSPLIT + offs) * G + g - cnt, mean, m2 = tl.reduce( - (tl.load(PCNT + idx), tl.load(PMEAN + idx), tl.load(PM2 + idx)), - 0, - _welford_combine, - ) - # `cnt` equals M by construction; M is passed in so the divisor is the - # exact element count rather than a float accumulated from partials. - var = m2 / M - tl.store(MEAN + pid, mean) - tl.store(RSTD + pid, 1.0 / tl.sqrt(var + eps)) - # ------------------------------------------------------------ normalize -- @_triton.jit def _normalize_kernel( X, Y, + PCNT, + PMEAN, + PM2, MEAN, RSTD, W, B, S, + M, + eps, C: tl.constexpr, G: tl.constexpr, CG: tl.constexpr, GP: tl.constexpr, CGP: tl.constexpr, + NSPLIT: tl.constexpr, BLOCK_S: tl.constexpr, + NBLK: tl.constexpr, + NPROG: tl.constexpr, RELU: tl.constexpr, HAS_W: tl.constexpr, HAS_B: tl.constexpr, MASKED_C: tl.constexpr, INT64: tl.constexpr, ): - blk = tl.program_id(0) + """Finalize the split-K statistics, then normalize NBLK/NPROG tiles. + + The finalize is *recomputed by every program* rather than round-tripped + through its own kernel launch: merging NSPLIT Welford triples is a few + KB of L2-resident traffic and a tree reduction over a ``(NSPLIT, GP)`` + tile, which is cheaper than the ~9 us launch it replaces. What it is + *not* cheap enough for is being paid once per tile at the largest + shapes, where the flat grid is 10^5 programs: the grid is therefore + capped at ``NPROG`` and each program strides over its share of the + ``NBLK`` tiles, so the redundant read costs ``NPROG * NSPLIT`` and not + ``NBLK * NSPLIT``. See :class:`GNConfig` -- ``nsplit_target``, + ``elem_tile`` and ``elem_progs`` are one joint tuning problem, not + three independent knobs. + + Every program reads the same partials with the same tile shape, so they + all get bit-identical ``mean``/``rstd``; program 0 stores them for the + backward. The loop carries nothing across iterations, so the striding + cannot affect the result. + """ + pid = tl.program_id(0) n = tl.program_id(1) offs_g = tl.arange(0, GP) @@ -764,18 +902,33 @@ def _normalize_kernel( wb = offs_g[:, None] * CG + offs_j[None, :] wbm = (offs_g[:, None] < G) & (offs_j[None, :] < CG) - s0 = blk * BLOCK_S - m = tl.broadcast_to((offs_s < S - s0)[:, None, None], (BLOCK_S, GP, CGP)) - if MASKED_C: - m = m & ((offs_g[None, :, None] < G) & (offs_j[None, None, :] < CG)) - if INT64: - base = (n.to(tl.int64) * S + s0) * C - else: - base = (n * S + s0) * C - gm = offs_g < G - mean = tl.load(MEAN + n * G + offs_g, mask=gm, other=0.0)[None, :, None] - rstd = tl.load(RSTD + n * G + offs_g, mask=gm, other=0.0)[None, :, None] + offs_p = tl.arange(0, NSPLIT) + pidx = (n * NSPLIT + offs_p[:, None]) * G + offs_g[None, :] + pm = tl.broadcast_to(gm[None, :], (NSPLIT, GP)) + # Padded group lanes load cnt == 0, which _welford_combine treats as the + # identity, so they merge to (0, 0, 0) and are masked off on the store. + # Reduced over axis 0 -- the *slowest* axis -- deliberately: reducing + # the fastest axis of a 2-D tile makes Triton stage the whole tile + # through LDS, which for a (G, NSPLIT) tile is 64 KB per array. + cnt_p = tl.load(PCNT + pidx, mask=pm, other=0.0) + _cnt, mu, m2 = tl.reduce( + ( + cnt_p, + tl.load(PMEAN + pidx, mask=pm, other=0.0), + tl.load(PM2 + pidx, mask=pm, other=0.0), + ), + 0, + _welford_combine, + ) + # `_cnt` equals M by construction; M is passed in so the divisor is the + # exact element count rather than a float accumulated from partials. + rs = 1.0 / tl.sqrt(m2 / M + eps) + if pid == 0: + tl.store(MEAN + n * G + offs_g, mu, mask=gm) + tl.store(RSTD + n * G + offs_g, rs, mask=gm) + mean = mu[None, :, None] + rstd = rs[None, :, None] if HAS_W: w = tl.load(W + wb, mask=wbm, other=0.0).to(tl.float32)[None, :, :] else: @@ -785,12 +938,21 @@ def _normalize_kernel( else: b = tl.zeros((1, GP, CGP), dtype=tl.float32) - x = tl.load(X + base + off, mask=m, other=0.0).to(tl.float32) - xhat = (x - mean) * rstd - y = xhat * w + b - if RELU: - y = tl.maximum(y, 0.0) - tl.store(Y + base + off, y.to(Y.dtype.element_ty), mask=m) + for blk in tl.range(pid, NBLK, NPROG): + s0 = blk * BLOCK_S + m = tl.broadcast_to((offs_s < S - s0)[:, None, None], (BLOCK_S, GP, CGP)) + if MASKED_C: + m = m & ((offs_g[None, :, None] < G) & (offs_j[None, None, :] < CG)) + if INT64: + base = (n.to(tl.int64) * S + s0) * C + else: + base = (n * S + s0) * C + x = tl.load(X + base + off, mask=m, other=0.0).to(tl.float32) + xhat = (x - mean) * rstd + y = xhat * w + b + if RELU: + y = tl.maximum(y, 0.0) + tl.store(Y + base + off, y.to(Y.dtype.element_ty), mask=m) # ------------------------------------------------------------- backward -- @_triton.jit @@ -885,49 +1047,6 @@ def _bwd_partial_kernel( tl.store(PDW + row, accdw, mask=wbm) tl.store(PDB + row, accdb, mask=wbm) - @_triton.jit - def _bwd_finalize_kernel( - PS1, - PS2, - C1, - C2, - M, - G: tl.constexpr, - NSPLIT: tl.constexpr, - ): - pid = tl.program_id(0) # n * G + g - n = pid // G - g = pid % G - offs = tl.arange(0, NSPLIT) - idx = (n * NSPLIT + offs) * G + g - tl.store(C1 + pid, tl.sum(tl.load(PS1 + idx)) / M) - tl.store(C2 + pid, tl.sum(tl.load(PS2 + idx)) / M) - - @_triton.jit - def _dwdb_reduce_kernel( - PDW, - PDB, - DW, - DB, - ROWS, - C, - BLOCK_C: tl.constexpr, - BLOCK_R: tl.constexpr, - ): - pid = tl.program_id(0) - offs_c = pid * BLOCK_C + tl.arange(0, BLOCK_C) - mc = offs_c < C - accw = tl.zeros((BLOCK_C,), dtype=tl.float32) - accb = tl.zeros((BLOCK_C,), dtype=tl.float32) - for r0 in range(0, ROWS, BLOCK_R): - offs_r = r0 + tl.arange(0, BLOCK_R) - m = (offs_r[:, None] < ROWS) & mc[None, :] - off = offs_r[:, None] * C + offs_c[None, :] - accw += tl.sum(tl.load(PDW + off, mask=m, other=0.0), 0) - accb += tl.sum(tl.load(PDB + off, mask=m, other=0.0), 0) - tl.store(DW + offs_c, accw, mask=mc) - tl.store(DB + offs_c, accb, mask=mc) - @_triton.jit def _dx_kernel( X, @@ -937,24 +1056,64 @@ def _dx_kernel( RSTD, W, B, - C1, - C2, + PS1, + PS2, + PDW, + PDB, + DW, + DB, + ROWS, S, + M, C: tl.constexpr, G: tl.constexpr, CG: tl.constexpr, GP: tl.constexpr, CGP: tl.constexpr, + NSPLIT: tl.constexpr, BLOCK_S: tl.constexpr, + NBLK: tl.constexpr, + NPROG: tl.constexpr, + NDW: tl.constexpr, + BLOCK_C: tl.constexpr, + BLOCK_R: tl.constexpr, RELU: tl.constexpr, HAS_W: tl.constexpr, HAS_B: tl.constexpr, MASKED_C: tl.constexpr, INT64: tl.constexpr, + ZERO_DX: tl.constexpr, ): - blk = tl.program_id(0) + """The whole backward tail: dweight/dbias, the c1/c2 finalize, and dx. + + Two reductions that used to be their own launches ride along here. The + per-channel dweight/dbias row reduction is done by the first ``NDW`` + programs of ``n == 0`` (a single pass over an ``(n*nsplit, C)`` scratch, + i.e. a few hundred KB); the per-``(n, g)`` c1/c2 finalize is recomputed + redundantly by every program, once, before the tile loop -- exactly as + in ``_normalize_kernel``, and capped the same way. The grid is + ``(max(NPROG, NDW), n)``; programs past ``NPROG`` exist only to cover + the dweight/dbias rows and run no loop iterations. + """ + pid = tl.program_id(0) n = tl.program_id(1) + # ---- dweight / dbias: rows of the split-K scratch, once per channel -- + if n == 0: + if pid < NDW: + offs_c = pid * BLOCK_C + tl.arange(0, BLOCK_C) + mc = offs_c < C + accw = tl.zeros((BLOCK_C,), dtype=tl.float32) + accb = tl.zeros((BLOCK_C,), dtype=tl.float32) + for r0 in range(0, ROWS, BLOCK_R): + offs_r = r0 + tl.arange(0, BLOCK_R) + rm = (offs_r[:, None] < ROWS) & mc[None, :] + roff = offs_r[:, None] * C + offs_c[None, :] + accw += tl.sum(tl.load(PDW + roff, mask=rm, other=0.0), 0) + accb += tl.sum(tl.load(PDB + roff, mask=rm, other=0.0), 0) + tl.store(DW + offs_c, accw, mask=mc) + tl.store(DB + offs_c, accb, mask=mc) + offs_g = tl.arange(0, GP) offs_j = tl.arange(0, CGP) offs_s = tl.arange(0, BLOCK_S) @@ -963,46 +1122,72 @@ def _dx_kernel( wb = offs_g[:, None] * CG + offs_j[None, :] wbm = (offs_g[:, None] < G) & (offs_j[None, :] < CG) - s0 = blk * BLOCK_S - m = tl.broadcast_to((offs_s < S - s0)[:, None, None], (BLOCK_S, GP, CGP)) - if MASKED_C: - m = m & ((offs_g[None, :, None] < G) & (offs_j[None, None, :] < CG)) - if INT64: - base = (n.to(tl.int64) * S + s0) * C - else: - base = (n * S + s0) * C - - gm = offs_g < G - mean = tl.load(MEAN + n * G + offs_g, mask=gm, other=0.0)[None, :, None] - rstd = tl.load(RSTD + n * G + offs_g, mask=gm, other=0.0)[None, :, None] - c1 = tl.load(C1 + n * G + offs_g, mask=gm, other=0.0)[None, :, None] - c2 = tl.load(C2 + n * G + offs_g, mask=gm, other=0.0)[None, :, None] - if HAS_W: - w = tl.load(W + wb, mask=wbm, other=0.0).to(tl.float32)[None, :, :] - else: - w = tl.full((1, GP, CGP), 1.0, tl.float32) - if HAS_B: - b = tl.load(B + wb, mask=wbm, other=0.0).to(tl.float32)[None, :, :] + if ZERO_DX: + # One element per group: mean == x and var == 0 identically, so + # xhat is the constant 0 and y does not depend on x at all -- the + # exact d_input is zero everywhere. The expression below would + # instead return rstd * (dy*w - c1), and since the compiler + # contracts that to fma(dy, w, -c1) while c1 was accumulated from + # the *rounded* product, what survives is the product's rounding + # error amplified by rstd = 1/sqrt(eps) ~ 316 (2.2e-05 at + # eps=1e-5). Answering with the exact zero costs one constexpr. + zero = tl.zeros((BLOCK_S, GP, CGP), dtype=tl.float32) + for blk in tl.range(pid, NBLK, NPROG): + s0 = blk * BLOCK_S + m = tl.broadcast_to( + (offs_s < S - s0)[:, None, None], (BLOCK_S, GP, CGP) + ) + if MASKED_C: + m = m & ((offs_g[None, :, None] < G) & (offs_j[None, None, :] < CG)) + if INT64: + base = (n.to(tl.int64) * S + s0) * C + else: + base = (n * S + s0) * C + tl.store(DX + base + off, zero.to(DX.dtype.element_ty), mask=m) else: - b = tl.zeros((1, GP, CGP), dtype=tl.float32) - - x = tl.load(X + base + off, mask=m, other=0.0).to(tl.float32) - dy = tl.load(DY + base + off, mask=m, other=0.0).to(tl.float32) - xhat = (x - mean) * rstd - if RELU: - dy = tl.where(xhat * w + b > 0.0, dy, 0.0) - dyw = dy * w - dx = rstd * (dyw - c1 - xhat * c2) - tl.store(DX + base + off, dx.to(DX.dtype.element_ty), mask=m) + gm = offs_g < G + offs_p = tl.arange(0, NSPLIT) + pidx = (n * NSPLIT + offs_p[:, None]) * G + offs_g[None, :] + pm = tl.broadcast_to(gm[None, :], (NSPLIT, GP)) + c1 = (tl.sum(tl.load(PS1 + pidx, mask=pm, other=0.0), 0) / M)[None, :, None] + c2 = (tl.sum(tl.load(PS2 + pidx, mask=pm, other=0.0), 0) / M)[None, :, None] + + mean = tl.load(MEAN + n * G + offs_g, mask=gm, other=0.0)[None, :, None] + rstd = tl.load(RSTD + n * G + offs_g, mask=gm, other=0.0)[None, :, None] + if HAS_W: + w = tl.load(W + wb, mask=wbm, other=0.0).to(tl.float32)[None, :, :] + else: + w = tl.full((1, GP, CGP), 1.0, tl.float32) + if HAS_B: + b = tl.load(B + wb, mask=wbm, other=0.0).to(tl.float32)[None, :, :] + else: + b = tl.zeros((1, GP, CGP), dtype=tl.float32) + + for blk in tl.range(pid, NBLK, NPROG): + s0 = blk * BLOCK_S + m = tl.broadcast_to( + (offs_s < S - s0)[:, None, None], (BLOCK_S, GP, CGP) + ) + if MASKED_C: + m = m & ((offs_g[None, :, None] < G) & (offs_j[None, None, :] < CG)) + if INT64: + base = (n.to(tl.int64) * S + s0) * C + else: + base = (n * S + s0) * C + x = tl.load(X + base + off, mask=m, other=0.0).to(tl.float32) + dy = tl.load(DY + base + off, mask=m, other=0.0).to(tl.float32) + xhat = (x - mean) * rstd + if RELU: + dy = tl.where(xhat * w + b > 0.0, dy, 0.0) + dyw = dy * w + dx = rstd * (dyw - c1 - xhat * c2) + tl.store(DX + base + off, dx.to(DX.dtype.element_ty), mask=m) globals().update( _welford_combine=_welford_combine, _stats_partial_kernel=_stats_partial_kernel, - _stats_finalize_kernel=_stats_finalize_kernel, _normalize_kernel=_normalize_kernel, _bwd_partial_kernel=_bwd_partial_kernel, - _bwd_finalize_kernel=_bwd_finalize_kernel, - _dwdb_reduce_kernel=_dwdb_reduce_kernel, _dx_kernel=_dx_kernel, ) @@ -1067,6 +1252,7 @@ def _forward(input, num_groups, weight, bias, eps, activation, out_dtype): pm2 = torch.empty_like(pcnt) mean = torch.empty((n, groups), device=device, dtype=torch.float32) rstd = torch.empty_like(mean) + out = torch.empty_like(input, dtype=out_dtype, memory_format=_CL_FORMAT) _stats_partial_kernel[(plan.nsplit, n)]( input, @@ -1086,34 +1272,28 @@ def _forward(input, num_groups, weight, bias, eps, activation, out_dtype): INT64=plan.int64, num_warps=plan.cfg.stats_warps, ) - _stats_finalize_kernel[(n * groups,)]( + _normalize_kernel[(plan.nprog_elem, n)]( + input, + out, pcnt, pmean, pm2, mean, rstd, - plan.elements_per_group, - eps, - G=groups, - NSPLIT=plan.nsplit, - num_warps=4, - ) - - out = torch.empty_like(input, dtype=out_dtype, memory_format=_CL_FORMAT) - _normalize_kernel[(plan.nblk_elem, n)]( - input, - out, - mean, - rstd, weight, bias, spatial, + plan.elements_per_group, + eps, C=channels, G=groups, CG=plan.group_channels, GP=plan.groups_p2, CGP=plan.group_channels_p2, + NSPLIT=plan.nsplit, BLOCK_S=plan.block_s_elem, + NBLK=plan.nblk_elem, + NPROG=plan.nprog_elem, RELU=activation == "relu", HAS_W=weight is not None, HAS_B=bias is not None, @@ -1167,72 +1347,46 @@ def _backward(grad_out, input, weight, bias, mean, rstd, num_groups, activation) num_warps=plan.cfg.stats_warps, ) - c1 = torch.empty(n * groups, device=device, dtype=torch.float32) - c2 = torch.empty_like(c1) - _bwd_finalize_kernel[(n * groups,)]( - ps1, - ps2, - c1, - c2, - plan.elements_per_group, - G=groups, - NSPLIT=plan.nsplit, - num_warps=4, - ) - d_weight = torch.empty(channels, device=device, dtype=torch.float32) d_bias = torch.empty_like(d_weight) - rows = n * plan.nsplit - block_c = min(256, max(64, _next_pow2(channels))) - block_r = 32 if rows >= 32 else 1 - _dwdb_reduce_kernel[(_cdiv(channels, block_c),)]( + d_input = torch.empty_like(input, memory_format=_CL_FORMAT) + _dx_kernel[(plan.grid_dx, n)]( + input, + grad_out, + d_input, + mean, + rstd, + weight, + bias, + ps1, + ps2, pdw, pdb, d_weight, d_bias, - rows, - channels, - BLOCK_C=block_c, - BLOCK_R=block_r, - num_warps=4, + plan.dwdb_rows, + spatial, + plan.elements_per_group, + C=channels, + G=groups, + CG=plan.group_channels, + GP=plan.groups_p2, + CGP=plan.group_channels_p2, + NSPLIT=plan.nsplit, + BLOCK_S=plan.block_s_elem, + NBLK=plan.nblk_elem, + NPROG=plan.nprog_elem, + NDW=plan.dwdb_progs, + BLOCK_C=plan.dwdb_block_c, + BLOCK_R=plan.dwdb_block_r, + RELU=activation == "relu", + HAS_W=weight is not None, + HAS_B=bias is not None, + MASKED_C=plan.masked_c, + INT64=plan.int64, + ZERO_DX=plan.zero_dx, + num_warps=plan.cfg.elem_warps, ) - - d_input = torch.empty_like(input, memory_format=_CL_FORMAT) - if plan.group_channels * spatial == 1: - # One element per group: mean == x and var == 0 identically, so - # xhat is the constant 0 and y does not depend on x at all -- the - # exact d_input is zero everywhere. _dx_kernel would instead - # return rstd * (dy*w - c1), and since the compiler contracts that - # to fma(dy, w, -c1) while c1 was accumulated from the *rounded* - # product, what survives is the product's rounding error amplified - # by rstd = 1/sqrt(eps) ~ 316 (2.2e-05 at eps=1e-5). Answering - # with the exact zero costs one integer test per backward call. - d_input.zero_() - else: - _dx_kernel[(plan.nblk_elem, n)]( - input, - grad_out, - d_input, - mean, - rstd, - weight, - bias, - c1, - c2, - spatial, - C=channels, - G=groups, - CG=plan.group_channels, - GP=plan.groups_p2, - CGP=plan.group_channels_p2, - BLOCK_S=plan.block_s_elem, - RELU=activation == "relu", - HAS_W=weight is not None, - HAS_B=bias is not None, - MASKED_C=plan.masked_c, - INT64=plan.int64, - num_warps=plan.cfg.elem_warps, - ) return d_input, d_weight, d_bias diff --git a/tests/test_triton_group_norm_edge.py b/tests/test_triton_group_norm_edge.py index 65fe148..be88cd4 100644 --- a/tests/test_triton_group_norm_edge.py +++ b/tests/test_triton_group_norm_edge.py @@ -1337,3 +1337,186 @@ def test_welford_correction_recovers_rstd_in_a_single_tile_reduction(groups, see f"rstd rel err {err:.3e} at mu/sigma=1e6 with G={groups}: the tile mean " f"correction is not doing its job" ) + + +# --------------------------------------------------------------------------- +# 12. the fused finalize and the capped elementwise grid +# --------------------------------------------------------------------------- +# +# ``_stats_finalize``/``_bwd_finalize``/``_dwdb_reduce`` are no longer their own +# launches: each is recomputed inside the elementwise kernel that consumes it. +# Two consequences need pinning. +# +# * The elementwise grid is capped at ``GNConfig.elem_progs`` and each program +# *strides* over its share of the tiles, so that the fused finalize costs +# ``nprog_elem * nsplit`` and not ``nblk_elem * nsplit`` reads. No scale-8 +# shape and no shape in either suite reaches that path with the shipped +# table -- ``nprog_elem == nblk_elem`` at every small shape -- so it has to be +# reached deliberately. +# * The cap is a *performance* knob. If it could change a single bit of the +# output it would break the module's reproducibility contract, since it is +# the one plan field that does not follow from the shape alone. + + +@contextlib.contextmanager +def _forced_config(channels, spatial, **overrides): + """Temporarily install a tiling config for one ``(channels, spatial)`` key. + + ``default_config`` keys the frozen table by ``(num_channels, cube-root + spatial extent)``, so the spatial extent has to be a perfect cube here. + Restores the previous entry (or its absence) and clears the plan cache on + the way out, so no other test can see it. + """ + edge = round(spatial ** (1.0 / 3.0)) + assert edge**3 == spatial, "forced configs need a cube spatial extent" + key = (channels, edge) + cfg = tgn.GNConfig(*tgn.default_config(channels, spatial).key()) + for name, value in overrides.items(): + assert hasattr(cfg, name), name + setattr(cfg, name, value) + sentinel = object() + saved = tgn._TUNED.get(key, sentinel) + tgn._TUNED[key] = cfg + tgn._plan.cache_clear() + try: + yield cfg + finally: + if saved is sentinel: + del tgn._TUNED[key] + else: + tgn._TUNED[key] = saved + tgn._plan.cache_clear() + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "shape,groups", + [ + ((1, 64, 16, 16, 16), 8), + ((2, 64, 16, 16, 16), 8), # N > 1: the stride is per (blk, n) program + ((1, 20, 8, 8, 8), 5), # capped grid *and* a padded channel axis + ], +) +@pytest.mark.parametrize("elem_progs", [1, 3, 8]) +@pytest.mark.parametrize("activation", [None, "relu"]) +def test_capped_elementwise_grid_strides_over_its_tiles( + shape, groups, elem_progs, activation +): + """Fewer elementwise programs than tiles: each must cover several tiles. + + A grid-stride loop that got its start, stride or trip count wrong leaves + part of the output (and of ``d_input``) unwritten -- which, since both are + ``torch.empty``, surfaces as plausible stale numbers rather than as a + crash. ``elem_progs=1`` is the extreme: one program per sample walks every + tile, so it also pins that the fused statistics are hoisted out of the loop + correctly rather than being recomputed per iteration from stale state. + """ + spatial = shape[2] * shape[3] * shape[4] + with _forced_config(shape[1], spatial, elem_tile=1024, elem_progs=elem_progs): + plan = tgn._plan(shape[0], shape[1], spatial, groups, 0) + assert plan.nprog_elem == min(plan.nblk_elem, elem_progs) + assert plan.nprog_elem < plan.nblk_elem, ( + f"the cap has to actually bite: nprog={plan.nprog_elem} " + f"nblk={plan.nblk_elem}" + ) + _parity( + shape, + groups, + activation, + seed=abs(hash((shape, elem_progs))) % 997, + label=f"{shape} elem_progs={elem_progs}", + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "shape,groups", [((2, 64, 16, 16, 16), 8), ((1, 20, 8, 8, 8), 5)] +) +def test_elementwise_grid_cap_is_bitwise_neutral(shape, groups): + """``elem_progs`` may not change a single bit of any output. + + It is the only field of ``_Plan`` that is a free parameter rather than a + consequence of the shape, and the module promises bitwise reproducibility. + That promise holds only because the elementwise kernels carry nothing + across loop iterations: the fused finalize is computed once per program + from the *same* partials with the *same* tile shape, and the tile bodies + are pure elementwise. If tuning this knob ever moved a result, the frozen + table would have become part of the numerical contract. + """ + spatial = shape[2] * shape[3] * shape[4] + x, weight, bias, grad_out = _make(shape, groups, seed=11) + results = [] + for elem_progs in (0, 1, 5, 64, 4096): + with _forced_config(shape[1], spatial, elem_tile=1024, elem_progs=elem_progs): + xi = x.detach().clone().requires_grad_(True) + wi = weight.detach().clone().requires_grad_(True) + bi = bias.detach().clone().requires_grad_(True) + y = triton_group_norm(xi, groups, wi, bi, EPS) + y.backward(grad_out) + results.append( + (y.detach().clone(), xi.grad.clone(), wi.grad.clone(), bi.grad.clone()) + ) + for elem_progs, got in zip((1, 5, 64, 4096), results[1:]): + for name, a, b in zip(("y", "dx", "dweight", "dbias"), got, results[0]): + assert torch.equal(a, b), ( + f"elem_progs={elem_progs} changed {name} bitwise; the grid cap " + f"is supposed to be a pure performance knob" + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "shape,groups,elem_progs", + [ + ((1, 64, 16, 16, 16), 8, 0), + ((1, 64, 16, 16, 16), 8, 3), + ((2, 96, 8, 8, 8), 6, 2), # padded channel axis, N > 1, capped grid + ], +) +def test_fused_finalize_publishes_the_statistics(shape, groups, elem_progs): + """``mean``/``rstd`` are published by program 0 of the *normalize* kernel. + + There is no separate finalize launch any more: every elementwise program + re-derives the statistics from the split-K Welford partials, and program 0 + is the one that stores them for the backward pass. A wrong publishing + program, a wrong partials index, or a group-mask slip in that fused + reduction would hand the backward garbage while leaving the forward -- which + uses its own locally computed copy -- perfectly correct. So check the + published tensors directly against float64. + """ + spatial = shape[2] * shape[3] * shape[4] + with _forced_config(shape[1], spatial, elem_tile=1024, elem_progs=elem_progs): + x, weight, bias, _grad = _make(shape, groups, seed=3) + _y, mean, rstd = torch.ops.scaffold_gn.group_norm( + x, groups, weight, bias, EPS, None, None + ) + flat = x.double().reshape(shape[0], groups, -1) + mean64 = flat.mean(-1) + var64 = ((flat - mean64[..., None]) ** 2).mean(-1) + assert _rel(mean, mean64) <= FP32_TOL + assert _rel(rstd, 1.0 / torch.sqrt(var64 + EPS)) <= FP32_TOL + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "shape,groups", [((2, 2048, 1, 1, 1), 8), ((1, 1024, 2, 1, 1), 8)] +) +def test_dweight_blocks_are_covered_when_there_are_more_of_them_than_tiles( + shape, groups +): + """The dweight/dbias reduction rides in ``_dx_kernel``'s first NDW programs. + + Those blocks are per-*channel*, the elementwise tiles are per-*voxel*, and + nothing makes the first outnumber the second: at ``(2, 2048, 1, 1, 1)`` + there is one elementwise tile and eight dweight blocks. The grid is + ``max(nprog_elem, dwdb_progs)`` for exactly that reason, and a grid of + ``nprog_elem`` alone would silently leave 7/8 of ``d_weight`` unwritten. + """ + spatial = shape[2] * shape[3] * shape[4] + plan = tgn._plan(shape[0], shape[1], spatial, groups, 0) + assert plan.dwdb_progs > plan.nprog_elem, ( + f"case is supposed to have more dweight blocks ({plan.dwdb_progs}) than " + f"elementwise programs ({plan.nprog_elem})" + ) + assert plan.grid_dx == plan.dwdb_progs + _parity(shape, groups, seed=13) From 38f2fc90744924e37c39230113271cda9969c7be Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Sat, 1 Aug 2026 12:32:09 -0700 Subject: [PATCH 58/62] Route the UNet GroupNorm through the Triton kernel FastGroupNorm now tries Triton, then the compiled kernel, then eager, and a failure at any rung latches that rung off and drops to the next rather than all the way to the bottom. SCAFFOLD_GROUPNORM_TRITON gates it exactly as SCAFFOLD_GROUPNORM_COMPILE gates the rung below. DCTensor is unwrapped explicitly for both fast rungs rather than dispatched through __torch_dispatch__, which keeps the tensor-subclass policy in one place: the kernel's is_supported() only asks isinstance, so an unknown wrapper would otherwise be routed into it where today it stays eager. The ReLU that followed every GroupNorm moves inside it. FastGroupNorm gains an activation argument and always applies the ReLU -- fused into the Triton store, explicit and in-place on the other two rungs -- and DoubleConv keeps an nn.Identity in the vacated Sequential slot so the state dict does not move by one key. That is asserted on the serialized bytes, not just the key list. Scale 7, DCTensor, channels-last, bf16 autocast: 193.25 -> 91.41 ms/step (2.12x), peak allocated 8.33 -> 7.21 GiB. Most of that is not GroupNorm. Its own kernels only go 18.69 -> 7.53 ms; the rest is elementwise and copy time, 99.68 -> 13.63 ms, which was the conv-side layout conversions GroupNorm was forcing by emitting contiguous into a channels-last model. Counted together, GroupNorm and the layout tax it imposed go 118.4 -> 21.2 ms/step, from 61.8% of device time to 23.9%. Whole-model gradient deviation from eager is 4.83e-03 median / 6.73e-02 max over 64 parameters, against 4.94e-03 / 6.64e-02 for the compiled rung that production already ships; triton-vs-compiled is smaller than either. Two processes under more_determinism produce identical losses, activation hashes and parameter hashes, with a counter confirming the Triton rung ran. One fix beyond the wiring: the ladder now re-raises _StopRecomputationError. Non-reentrant checkpointing ends its recompute by raising that from a saved-tensor pack hook, i.e. from inside whichever op is saving a tensor. Absorbing the trailing ReLU moved that boundary inside the try, so the ladder caught it, decided the kernel was broken and dropped the whole model to eager for the rest of the run. The hazard predates this commit; it was simply unreachable while the ReLU sat outside forward(). --- ScaFFold/unet/group_norm.py | 369 +++++++++++++++---- ScaFFold/unet/unet_parts.py | 35 +- tests/test_groupnorm.py | 714 +++++++++++++++++++++++++++++++++--- 3 files changed, 1003 insertions(+), 115 deletions(-) diff --git a/ScaFFold/unet/group_norm.py b/ScaFFold/unet/group_norm.py index 33d3323..9d7e95f 100644 --- a/ScaFFold/unet/group_norm.py +++ b/ScaFFold/unet/group_norm.py @@ -12,36 +12,73 @@ # # SPDX-License-Identifier: (Apache-2.0) -"""GroupNorm with a ``torch.compile``d fast path on GPU. - -ATen's GroupNorm computes its per-group statistics with a kernel that launches -one workgroup per ``(batch, group)`` row. At this benchmark's defaults -(``local_batch_size=1``, ``group_norm_groups=8``) that is 8 workgroups, so on a -228-CU MI300A the normalization runs at a small fraction of achievable -bandwidth and dominates the step: measured 87 ms of a 187 ms step (47%) at -scale 7. Compiling the same functional GroupNorm hands the reduction to -Inductor, which tiles it across the whole device; the same measurement then -gives a 184.7 ms step at 100.7 ms, with GroupNorm down to ~7% of it. +"""GroupNorm with a Triton fast path and a ``torch.compile``d one behind it. + +Three kernels, tried in order, all of them producing the same numbers: + +1. **Native channels-last Triton** (:mod:`ScaFFold.unet.triton_group_norm`), + whenever that module's ``is_supported`` accepts the input. Production runs + set ``PYTORCH_MIOPEN_SUGGEST_NHWC=1``, under which every convolution emits + ``channels_last_3d`` -- and every *stock* GroupNorm (eager or Inductor) reads + that layout through the logical NCDHW order, a strided gather, and returns a + contiguous tensor, breaking the layout chain at all 22 call sites of the + forward. The Triton kernel is NDHWC in and NDHWC out and is 6.5x faster than + the compiled kernel on that input; with the ReLU fused it takes another 38% + off the forward. See that module's docstring for the measurements. +2. **``torch.compile``d ``F.group_norm``**, for inputs the Triton kernel does + not serve (contiguous NCDHW, non-5-D, unsupported dtypes) and as the landing + place if the Triton path ever raises. ATen's own kernel launches one + workgroup per ``(batch, group)`` row -- 8 of them at this benchmark's + defaults -- so on a 228-CU MI300A it runs at a small fraction of achievable + bandwidth: measured 87 ms of a 187 ms step (47%) at scale 7, against 7% of a + 184.7 ms step once Inductor tiles the reduction across the device. +3. **Stock eager ``F.group_norm``**, which is what every rejection falls back + to and what defines the semantics the other two must match. ``FastGroupNorm`` is a drop-in ``nn.GroupNorm``: same parameters, same names, same shapes, same numerics -- only the kernel differs, so checkpoints are interchangeable in both directions with any other GroupNorm-based build. The -compiled path is used only when it is safe and worthwhile, and every rejection -falls back to stock eager ``F.group_norm``: - -* non-CUDA tensors (the CPU test suite never pays compile latency), -* tensor subclasses, whose ``__torch_dispatch__`` wrappers Dynamo cannot trace - -- except DistConv's ``DCTensor``, which is unwrapped to its local shard - around the compiled kernel instead (see ``FastGroupNorm.forward``), -* an already-compiled enclosing region (the functional call inlines instead), -* an explicit opt-out via ``SCAFFOLD_GROUPNORM_COMPILE=0``, -* any failure inside ``torch.compile`` -- logged once, then eager forever after. - -Determinism: the compiled kernels are bitwise reproducible. Two separate -processes running three fwd+bwd+Adam steps of the scale-7 UNet under -``more_determinism`` (``use_deterministic_algorithms(True, warn_only=True)``, -``cudnn.benchmark=False``, fixed seeds) hash identically with the compiled path, -exactly as they do with the eager one, so no determinism gate is needed. +one addition is the optional fused ``activation`` (see below), which adds no +state either. + +Routing rejections, in the order they are tested: + +* an explicit opt-out via ``SCAFFOLD_GROUPNORM_TRITON=0`` / + ``SCAFFOLD_GROUPNORM_COMPILE=0``, +* non-CUDA tensors -- the CPU test suite pays neither compile latency nor the + Triton import, +* tensor subclasses, whose ``__torch_dispatch__`` wrappers have unknown + semantics -- except DistConv's ``DCTensor``, which is unwrapped to its local + shard around both fast kernels (see ``FastGroupNorm``), +* for the Triton kernel, anything its ``is_supported`` rejects (a layout, dtype, + degenerate shape or affine-parameter dtype it does not serve); for the + compiled one, an already-compiled enclosing region (the functional call + inlines instead), +* any failure inside either kernel -- logged once, latched off for the rest of + the process, and retried on the next path down. Both are optimizations, never + correctness requirements: a broken Triton install must degrade a multi-node + run, not kill it. + +Determinism: all three kernels are bitwise reproducible. Two separate processes +running the scale-7 UNet under ``more_determinism`` +(``use_deterministic_algorithms(True, warn_only=True)``, ``cudnn.benchmark=False``, +fixed seeds) hash identically with the Triton path, the compiled path and the +eager one alike, so no determinism gate is needed. The Triton kernel's grid, +split count and tile sizes are pure functions of the shape and it uses no float +atomics, which is what buys that. + +Fused activation +================ +Every GroupNorm in the UNet is immediately followed by a ReLU, and the Triton +kernel can fold that into its forward store for free (it is store-bound) while +removing a whole streaming pass -- 38% of the forward at the shapes that +dominate. ``FastGroupNorm(..., activation="relu")`` therefore *always* applies +the ReLU: fused inside the Triton kernel where that path is taken, and as an +explicit in-place ``F.relu`` on the compiled and eager paths. ``DoubleConv`` +consequently holds an ``nn.Identity`` where its ``nn.ReLU`` used to be, so the +positional keys of its ``nn.Sequential`` -- and therefore every checkpoint -- +are unchanged (neither module has parameters or buffers). Correctness holds on +every path; the fusion is purely an optimization inside the module. """ import logging @@ -51,6 +88,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +import torch.utils.checkpoint # for _CONTROL_FLOW_EXCEPTIONS; torch loads it anyway logger = logging.getLogger(__name__) @@ -59,6 +97,17 @@ #: is safe", which is what every production run wants. COMPILE_ENV_VAR = "SCAFFOLD_GROUPNORM_COMPILE" +#: The same, for the native channels-last Triton kernel, which is tried first. +#: Same spellings, same "unset means on wherever it is safe" default -- the +#: whole point of the kernel is that production takes it. +TRITON_ENV_VAR = "SCAFFOLD_GROUPNORM_TRITON" + +#: Activations this module can apply after normalizing. Must stay a subset of +#: ``triton_group_norm.SUPPORTED_ACTIVATIONS`` (pinned by a test); spelled out +#: here rather than imported so that constructing a module -- or running the +#: whole CPU suite -- never imports the kernel module. +SUPPORTED_ACTIVATIONS = (None, "relu") + #: Dynamo caches one entry per distinct guard set on the traced function. A #: UNet presents one entry per distinct activation shape (5 at scale 7) times #: grad-enabled/no-grad (training vs. evaluation), i.e. 10 -- above the stock @@ -78,10 +127,22 @@ # by set_compile_enabled(). _compile_override = None +# The triton_group_norm module, imported on the first CUDA forward. Importing +# it registers two dispatcher ops, and a CPU-only run must pay neither that nor +# the `triton` import the module itself defers to its first launch. +_triton_module = None + +# Set once if the Triton kernel raises; the compiled path is used from then on. +_triton_failed = False + +# None = decide per tensor; True/False = forced by SCAFFOLD_GROUPNORM_TRITON or +# by set_triton_enabled(). +_triton_override = None -def _env_override(): - """Read ``SCAFFOLD_GROUPNORM_COMPILE``; ``None`` when unset or unparsable.""" - raw = os.environ.get(COMPILE_ENV_VAR) + +def _env_override(name): + """Read boolean env var ``name``; ``None`` when unset or unparsable.""" + raw = os.environ.get(name) if raw is None: return None value = raw.strip().lower() @@ -90,13 +151,14 @@ def _env_override(): if value in ("0", "false", "off", "no"): return False logger.warning( - f"Ignoring unrecognized {COMPILE_ENV_VAR}={raw!r}; " + f"Ignoring unrecognized {name}={raw!r}; " "expected one of 1/0/true/false/on/off/yes/no" ) return None -_compile_override = _env_override() +_compile_override = _env_override(COMPILE_ENV_VAR) +_triton_override = _env_override(TRITON_ENV_VAR) def set_compile_enabled(enabled): @@ -110,7 +172,26 @@ def set_compile_enabled(enabled): """ global _compile_override previous = _compile_override - _compile_override = _env_override() if enabled is None else bool(enabled) + _compile_override = ( + _env_override(COMPILE_ENV_VAR) if enabled is None else bool(enabled) + ) + return previous + + +def set_triton_enabled(enabled): + """Force the Triton path on (``True``) or off (``False``). + + The exact counterpart of :func:`set_compile_enabled`: ``None`` restores the + default (``SCAFFOLD_GROUPNORM_TRITON`` if set, otherwise "wherever + ``is_supported`` accepts"), forcing it on does not override the device, + subclass or ``is_supported`` checks -- those are correctness conditions -- + and the previous setting is returned so tests can restore it. + """ + global _triton_override + previous = _triton_override + _triton_override = ( + _env_override(TRITON_ENV_VAR) if enabled is None else bool(enabled) + ) return previous @@ -151,6 +232,68 @@ def _get_compiled_group_norm(): return _compiled_group_norm +def _get_triton_module(): + """Import (once) :mod:`ScaFFold.unet.triton_group_norm`. + + Deferred rather than imported at the top of this file: that module registers + two dispatcher ops and builds an autograd formula at import time, and a run + that never reaches the GPU (the whole CPU unit suite) must not pay for it. + Only ever called after the input has been shown to be a CUDA tensor, which + is also what keeps ``import triton`` -- which that module defers again, to + its first kernel launch -- out of a CPU-only process entirely. + """ + global _triton_module + if _triton_module is None: + from . import triton_group_norm + + _triton_module = triton_group_norm + return _triton_module + + +#: Exceptions torch raises *through* this module as control flow rather than as +#: a kernel failure, and which the fallback ladder must therefore re-raise. +#: +#: ``torch.utils.checkpoint``'s non-reentrant recompute stops itself early by +#: raising ``_StopRecomputationError`` from its saved-tensor *pack hook* -- i.e. +#: from inside whichever op happens to be saving a tensor when the recompute has +#: produced everything the backward needs. In a ``DoubleConv`` that op is this +#: module (GroupNorm saves its input and statistics, and the fused ReLU means +#: nothing follows it in the block), so the exception surfaces inside the +#: ``try``. Swallowing it would latch the fast kernel off, silently drop the +#: model to eager mid-run, and leave the checkpoint machinery waiting for a stop +#: that never came. Private API, so tolerate its absence rather than importing +#: it by name. +_CONTROL_FLOW_EXCEPTIONS = tuple( + exception + for exception in (getattr(torch.utils.checkpoint, "_StopRecomputationError", None),) + if isinstance(exception, type) and issubclass(exception, BaseException) +) + + +def _use_triton(input, num_groups, weight, bias, activation): + """Whether this particular input should take the native Triton kernel. + + Ordered so that the cheap local tests come first and the module import last: + a CPU tensor is rejected before ``_get_triton_module`` is ever called. + """ + if _triton_failed or _triton_override is False: + return False + # Same policy as _use_compiled: an unknown __torch_dispatch__ wrapper has + # unknown semantics and keeps the stock kernel. is_supported() would accept + # one (it only asks isinstance), so this check is load-bearing here, not a + # copy for symmetry. DistConv's DCTensor never reaches it -- forward() + # unwraps to the local shard first. + if type(input) is not torch.Tensor: + return False + if not input.is_cuda: + return False + # is_supported() is cheap and side-effect free: a handful of attribute reads + # and one stride check, no allocation, no launch, no triton import. + return _get_triton_module().is_supported( + input, num_groups, weight, bias, activation + ) + + def _use_compiled(input): """Whether this particular input should take the compiled path.""" if _compile_failed or _compile_override is False: @@ -185,51 +328,149 @@ def _dctensor_ops(input): return None +def _run_local(input, distconv, kernel): + """Run ``kernel`` on a plain tensor, DCTensor in -> DCTensor out. + + ``distconv`` is ``None`` for a plain tensor, where this is just + ``kernel(input)``. For a ``DCTensor`` the unwrap goes through DistConv's + ``_ToTensor``/``_FromTensor`` autograd pair (``DCTensor.from_shard`` is the + public spelling of the latter; there is no public unwrap yet -- upstream + ask) rather than a bare ``input._tensor`` read: DistConv's own dispatch may + read ``_tensor`` directly because it runs *below* autograd, while this runs + above it and a bare read would sever the graph back to the producing + convolution. + """ + if distconv is None: + return kernel(input) + local = distconv._ToTensor.apply(input) + return distconv.DCTensor.from_shard(kernel(local), input._parallel_strategy) + + class FastGroupNorm(nn.GroupNorm): - """``nn.GroupNorm`` that runs its GPU forward through ``torch.compile``. + """``nn.GroupNorm`` with a Triton GPU kernel and an optional fused ReLU. Identical state: ``weight``/``bias`` of shape ``(num_channels,)``, no buffers, so state dicts are interchangeable with plain ``nn.GroupNorm`` - in both directions. - - DistConv's ``DCTensor`` gets the compiled kernel too: its generic - ``__torch_dispatch__`` has no GroupNorm-specific handling -- it unwraps to - the local shard, runs the stock aten kernels, and rewraps the outputs, so - statistics are per-shard and no communication happens at any shard count. - ``forward`` moves that same unwrap up in front of the compiled kernel, - preserving those semantics exactly (DCTensor in -> DCTensor out) while - keeping the fast kernel Dynamo's inability to trace the wrapper would - otherwise forfeit. It cannot copy dispatch's *mechanism*, though: dispatch - runs below autograd, where reading ``_tensor`` directly is safe, whereas - this runs above it, so the unwrap has to go through DistConv's - ``_ToTensor``/``_FromTensor`` autograd pair or the graph back to the - producing convolution is severed. + in both directions. ``activation`` is a plain Python attribute, not a + submodule or a buffer, so setting it does not add a key either. + + ``activation="relu"`` makes this module's forward *always* apply a ReLU -- + fused into the Triton kernel's store where that path is taken, and as an + explicit in-place ``F.relu`` on the compiled and eager paths. The + correctness of the model therefore does not depend on which kernel runs; + only the number of memory passes does. + + DistConv's ``DCTensor`` gets the fast kernels too, by unwrapping to the + local shard in front of them rather than by letting the op dispatch through + the wrapper. Both would work -- the Triton kernel is a real dispatcher op, + so ``DCTensor.__torch_dispatch__`` would intercept it, unwrap, run and + rewrap on its own -- but the explicit unwrap is what this module already + does for the compiled kernel, and it is better here for three reasons. + (1) It keeps the subclass policy in one place: ``is_supported`` accepts any + ``torch.Tensor`` *instance*, so relying on dispatch would silently extend + the fast path to every unknown wrapper subclass, which today keeps the + stock kernel. (2) The eligibility predicates then examine the tensor the + kernel will actually touch -- its dtype, device, strides and shape -- rather + than a wrapper's mirrored metadata. (3) The Triton and compiled paths share + one unwrap and one fallback ladder instead of needing two shapes of code, + and a Triton failure can be retried on the compiled kernel without a second + round trip through the wrapper. Semantics are unchanged either way: + DistConv's generic ``__torch_dispatch__`` has no GroupNorm-specific + handling, so statistics are per-shard and no communication happens at any + shard count, exactly as before. """ + def __init__( + self, + num_groups, + num_channels, + eps=1e-5, + affine=True, + device=None, + dtype=None, + activation=None, + ): + if activation not in SUPPORTED_ACTIVATIONS: + raise ValueError( + f"activation must be one of {SUPPORTED_ACTIVATIONS}, got {activation!r}" + ) + super().__init__( + num_groups, num_channels, eps=eps, affine=affine, device=device, dtype=dtype + ) + self.activation = activation + + def extra_repr(self): + base = super().extra_repr() + if self.activation is None: + return base + return f"{base}, activation={self.activation}" + + def _activate(self, out): + """Apply the activation on the two paths that cannot fuse it. + + In place, which is what the ``nn.ReLU(inplace=True)`` this module + absorbed did: ``out`` is a freshly allocated GroupNorm output with no + other consumer, and GroupNorm's backward reads its *input*, never its + output, so overwriting it is safe for autograd as well as for memory. + """ + if self.activation == "relu": + return F.relu(out, inplace=True) + return out + + def _triton_forward(self, local): + """The native channels-last kernel, with the activation fused in.""" + return _get_triton_module().triton_group_norm( + local, self.num_groups, self.weight, self.bias, self.eps, self.activation + ) + + def _compiled_forward(self, local): + return self._activate( + _get_compiled_group_norm()( + local, self.num_groups, self.weight, self.bias, self.eps + ) + ) + + def _eager_forward(self, input): + # super().forward() is the stock kernel; deferring to it keeps the eager + # path identical to nn.GroupNorm's (plus the ReLU) by construction. + return self._activate(super().forward(input)) + def forward(self, input): + global _compile_failed, _triton_failed + distconv = _dctensor_ops(input) # The eligibility checks look at the local shard for a DCTensor (the # peek is a plain attribute read, no autograd involvement) and at the # tensor itself otherwise. local_view = input._tensor if distconv is not None else input + + if _use_triton( + local_view, self.num_groups, self.weight, self.bias, self.activation + ): + try: + return _run_local(input, distconv, self._triton_forward) + except _CONTROL_FLOW_EXCEPTIONS: + raise + except Exception as e: + # A broken or mismatched Triton install, an unwritable JIT cache + # or a shape the kernel mishandles must cost speed, not a + # multi-node run. GroupNorm is pure, so retrying the same call + # on the compiled kernel below is safe -- and the compiled + # kernel, not eager, is the right landing place: it is still + # ~10x the stock one. + _triton_failed = True + logger.warning( + f"Triton GroupNorm failed ({type(e).__name__}: {e}); falling " + "back to the compiled kernel for the rest of this run. " + f"Set {TRITON_ENV_VAR}=0 to skip this attempt entirely." + ) + if not _use_compiled(local_view): - # super().forward() is the stock kernel; deferring to it keeps the - # eager path identical to nn.GroupNorm's by construction. - return super().forward(input) - global _compile_failed + return self._eager_forward(input) try: - if distconv is not None: - # _ToTensor is the autograd-aware unwrap DistConv itself uses; - # DCTensor.from_shard is the public spelling of _FromTensor. - # (There is no public unwrap yet -- upstream ask.) - local = distconv._ToTensor.apply(input) - out = _get_compiled_group_norm()( - local, self.num_groups, self.weight, self.bias, self.eps - ) - return distconv.DCTensor.from_shard(out, input._parallel_strategy) - return _get_compiled_group_norm()( - input, self.num_groups, self.weight, self.bias, self.eps - ) + return _run_local(input, distconv, self._compiled_forward) + except _CONTROL_FLOW_EXCEPTIONS: + raise except Exception as e: # Compilation is an optimization, never a correctness requirement: # a broken Inductor/Triton install, an unwritable cache directory or @@ -241,4 +482,4 @@ def forward(self, input): "falling back to the eager kernel for the rest of this run. " f"Set {COMPILE_ENV_VAR}=0 to skip this attempt entirely." ) - return super().forward(input) + return self._eager_forward(input) diff --git a/ScaFFold/unet/unet_parts.py b/ScaFFold/unet/unet_parts.py index c9e6cb0..9fffe72 100644 --- a/ScaFFold/unet/unet_parts.py +++ b/ScaFFold/unet/unet_parts.py @@ -28,18 +28,35 @@ _outconv_annotate = annotate(fmt="OutConv.{}") -def _group_norm(num_groups, num_channels): +def _group_norm(num_groups, num_channels, activation=None): if num_channels % num_groups != 0: raise ValueError( f"group_norm_groups={num_groups} must evenly divide num_channels={num_channels}" ) - # FastGroupNorm is nn.GroupNorm plus a compiled GPU kernel; it holds the - # same parameters under the same names, so checkpoints are unaffected. - return FastGroupNorm(num_groups, num_channels) + # FastGroupNorm is nn.GroupNorm plus a Triton/compiled GPU kernel; it holds + # the same parameters under the same names, and `activation` is a plain + # attribute rather than a submodule, so checkpoints are unaffected. + return FastGroupNorm(num_groups, num_channels, activation=activation) class DoubleConv(nn.Module): - """(convolution => GroupNorm => ReLU) * 2""" + """(convolution => GroupNorm => ReLU) * 2 + + The ReLU lives *inside* the GroupNorm (``activation="relu"``), because the + Triton GroupNorm kernel folds it into its forward store for free and thereby + removes an entire streaming pass -- 38% of the forward at the shapes that + dominate the step. ``FastGroupNorm`` applies the ReLU on every path, + including eager, so the network's function is unchanged; only the number of + memory passes differs. + + The ``nn.ReLU`` slots are held open by ``nn.Identity`` rather than removed: + ``nn.Sequential`` names its children by position, so deleting them would + renumber the two convolutions and the second GroupNorm and invalidate every + existing checkpoint. Neither ``nn.ReLU`` nor ``nn.Identity`` has parameters + or buffers, so with the placeholders in place the state dict is byte + identical to the pre-fusion model's (pinned by + ``tests/test_groupnorm.py::test_state_dict_matches_plain_groupnorm_model``). + """ def __init__(self, in_channels, out_channels, group_norm_groups, mid_channels=None): super().__init__() @@ -47,11 +64,11 @@ def __init__(self, in_channels, out_channels, group_norm_groups, mid_channels=No mid_channels = out_channels self.double_conv = nn.Sequential( nn.Conv3d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False), - _group_norm(group_norm_groups, mid_channels), - nn.ReLU(inplace=True), + _group_norm(group_norm_groups, mid_channels, activation="relu"), + nn.Identity(), nn.Conv3d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False), - _group_norm(group_norm_groups, out_channels), - nn.ReLU(inplace=True), + _group_norm(group_norm_groups, out_channels, activation="relu"), + nn.Identity(), ) @_doubleconv_annotate diff --git a/tests/test_groupnorm.py b/tests/test_groupnorm.py index dec8b70..fe8ab28 100644 --- a/tests/test_groupnorm.py +++ b/tests/test_groupnorm.py @@ -12,19 +12,28 @@ # # SPDX-License-Identifier: (Apache-2.0) -"""Tests for the compiled GroupNorm fast path (``ScaFFold.unet.group_norm``). +"""Tests for the GroupNorm fast paths (``ScaFFold.unet.group_norm``). The optimization must be invisible everywhere except in the profile: the same state dict as a stock ``nn.GroupNorm`` model (checkpoints stay interchangeable -in both directions), the same numbers within reduction-order noise, and an -eager fallback for every input the compiled kernel cannot or should not take -(CPU, unknown tensor subclasses, a broken compiler). DistConv's ``DCTensor`` -is not in that list: ``forward`` unwraps it to its local shard around the -compiled kernel, so the wrapped production path is served too. That unwrap is -*not* the bare attribute read DistConv's own dispatch does -- dispatch runs -below autograd, where a bare read is safe, while ``forward`` runs above it and -must go through DistConv's ``_ToTensor``/``_FromTensor`` pair to keep the graph -connected. +in both directions), the same numbers within reduction-order noise, and a +fallback for every input the fast kernels cannot or should not take (CPU, +unknown tensor subclasses, a broken Triton or Inductor install). The ladder is +Triton -> compiled -> eager, and a failure at any rung latches that rung off and +drops to the next, never to the bottom. + +DistConv's ``DCTensor`` is not in the rejection list: ``forward`` unwraps it to +its local shard around both fast kernels, so the wrapped production path is +served too. That unwrap is *not* the bare attribute read DistConv's own +dispatch does -- dispatch runs below autograd, where a bare read is safe, while +``forward`` runs above it and must go through DistConv's +``_ToTensor``/``_FromTensor`` pair to keep the graph connected. + +The ReLU that used to follow every GroupNorm now lives inside it +(``activation="relu"``), fused into the Triton store and applied explicitly on +the other two paths. ``DoubleConv`` keeps an ``nn.Identity`` in the vacated +``nn.Sequential`` slot, so the state dict does not move by one key -- which is +what the checkpoint tests here pin. """ from __future__ import annotations @@ -51,39 +60,61 @@ @pytest.fixture(autouse=True) def _restore_compile_state(): - """Keep per-test overrides of the module-level compile state contained.""" - previous = gn_mod.set_compile_enabled(None) - failed = gn_mod._compile_failed + """Keep per-test overrides of the module-level routing state contained.""" + previous_compile = gn_mod.set_compile_enabled(None) + previous_triton = gn_mod.set_triton_enabled(None) + compile_failed = gn_mod._compile_failed + triton_failed = gn_mod._triton_failed yield - gn_mod._compile_override = previous - gn_mod._compile_failed = failed + gn_mod._compile_override = previous_compile + gn_mod._triton_override = previous_triton + gn_mod._compile_failed = compile_failed + gn_mod._triton_failed = triton_failed -def _make_unet(seed: int, group_norm_cls=None): - """Build the worker.py-shaped UNet, optionally with a different norm class.""" +def _make_unet(seed: int): + """Build the worker.py-shaped UNet.""" torch.manual_seed(seed) - if group_norm_cls is None: - return UNet( - n_channels=_N_CHANNELS, - n_classes=_N_CLASSES, - trilinear=False, - layers=2, - group_norm_groups=_GROUPS, - ) - import ScaFFold.unet.unet_parts as parts + return UNet( + n_channels=_N_CHANNELS, + n_classes=_N_CLASSES, + trilinear=False, + layers=2, + group_norm_groups=_GROUPS, + ) - original = parts.FastGroupNorm - parts.FastGroupNorm = group_norm_cls - try: - return UNet( - n_channels=_N_CHANNELS, - n_classes=_N_CLASSES, - trilinear=False, - layers=2, - group_norm_groups=_GROUPS, - ) - finally: - parts.FastGroupNorm = original + +def _make_plain_unet(seed: int): + """The pre-fusion build: stock ``nn.GroupNorm`` followed by ``nn.ReLU``. + + Built by *converting* a normal UNet rather than by patching the norm class + at construction time, because the fusion moved the ReLU into the norm: a + class swap alone would leave ``DoubleConv``'s ``nn.Identity`` placeholders + in place and produce a model with no activations at all, which would make + every numeric comparison below vacuous. Converting reproduces exactly the + module graph this branch replaced -- ``nn.GroupNorm`` where the fast norm + sits, an in-place ``nn.ReLU`` where the placeholder sits -- and consumes no + RNG (``nn.GroupNorm`` initializes to ones/zeros), so the parameters are + bit-identical to what ``_make_unet(seed)`` draws. + """ + model = _make_unet(seed) + for parent in [m for m in model.modules() if isinstance(m, nn.Sequential)]: + for index, child in enumerate(list(parent)): + if isinstance(child, FastGroupNorm): + plain = nn.GroupNorm( + child.num_groups, + child.num_channels, + eps=child.eps, + affine=child.affine, + ) + if child.affine: + with torch.no_grad(): + plain.weight.copy_(child.weight) + plain.bias.copy_(child.bias) + parent[index] = plain + elif isinstance(child, nn.Identity): + parent[index] = nn.ReLU(inplace=True) + return model def _make_input(seed: int = 0, channels: int = _N_CHANNELS, size: int = _N): @@ -103,7 +134,7 @@ def test_state_dict_matches_plain_groupnorm_model(): parameter inventory of the model may not shift by even one key. """ new_model = _make_unet(seed=0) - old_model = _make_unet(seed=0, group_norm_cls=nn.GroupNorm) + old_model = _make_plain_unet(seed=0) new_sd = new_model.state_dict() old_sd = old_model.state_dict() @@ -126,7 +157,7 @@ def test_checkpoint_round_trip_both_directions(tmp_path): script). After each load the two models must agree bit for bit. """ new_model = _make_unet(seed=0) - old_model = _make_unet(seed=1, group_norm_cls=nn.GroupNorm) + old_model = _make_plain_unet(seed=1) x = _make_input(seed=3) old_path = tmp_path / "old.pth" @@ -161,6 +192,119 @@ def test_unet_uses_fast_group_norm(): assert all(isinstance(m, FastGroupNorm) for m in norms) +def test_state_dict_bytes_identical_to_plain_groupnorm_model(): + """Not just the same keys: the serialized checkpoint must be byte identical. + + ``test_state_dict_matches_plain_groupnorm_model`` compares names, shapes and + dtypes; this compares the actual bytes ``torch.save`` writes, which is the + thing that has to stay interchangeable. It is the direct guard on the + ``nn.ReLU`` -> ``nn.Identity`` swap: ``nn.Sequential`` names its children by + position, so *removing* the activation slot rather than holding it open + would renumber ``3.weight`` and ``4.weight``/``4.bias`` and silently + invalidate every checkpoint on disk. + """ + import io + + new_model = _make_unet(seed=0) + old_model = _make_plain_unet(seed=0) + + def blob(model): + buffer = io.BytesIO() + torch.save(model.state_dict(), buffer) + return buffer.getvalue() + + assert blob(new_model) == blob(old_model) + + +def test_double_conv_keeps_the_activation_slots(): + """The fused build keeps six positional slots, with nothing in the spares. + + Pins both halves of the fusion design: the ReLU is *in* the norm + (``activation == "relu"`` at positions 1 and 4) and its old slots (2 and 5) + are parameterless placeholders rather than deletions. + """ + from ScaFFold.unet.unet_parts import DoubleConv + + block = DoubleConv(3, 16, _GROUPS) + children = list(block.double_conv) + assert len(children) == 6 + for norm_index, spare_index in ((1, 2), (4, 5)): + norm = children[norm_index] + assert isinstance(norm, FastGroupNorm) + assert norm.activation == "relu" + spare = children[spare_index] + assert isinstance(spare, nn.Identity) + assert list(spare.parameters()) == [] + assert list(spare.buffers()) == [] + # And the positional key numbering is exactly the pre-fusion one. + assert list(block.state_dict().keys()) == [ + "double_conv.0.weight", + "double_conv.1.weight", + "double_conv.1.bias", + "double_conv.3.weight", + "double_conv.4.weight", + "double_conv.4.bias", + ] + + +def test_double_conv_output_matches_the_explicit_relu_build(): + """Folding the ReLU into the norm may not change a single bit of the output. + + The fused module applies the ReLU itself on every path, so on CPU (eager) + the block must reproduce ``conv -> GroupNorm -> ReLU`` exactly, gradients + included. + """ + from ScaFFold.unet.unet_parts import DoubleConv + + torch.manual_seed(4) + fused = DoubleConv(3, 16, _GROUPS) + reference = DoubleConv(3, 16, _GROUPS) + reference.load_state_dict(fused.state_dict()) + for index in (1, 4): + reference.double_conv[index].activation = None + for index in (2, 5): + reference.double_conv[index] = nn.ReLU(inplace=True) + + x_fused = _make_input(seed=8, channels=3, size=8).requires_grad_(True) + x_reference = x_fused.detach().clone().requires_grad_(True) + + out_fused = fused(x_fused) + out_reference = reference(x_reference) + assert torch.equal(out_fused, out_reference) + # A block whose activation silently vanished would still pass an + # output-equality test against another activation-free block, so assert the + # ReLU is really there. + assert (out_fused < 0).sum() == 0 + assert out_fused.max() > 0 + + out_fused.pow(2).sum().backward() + out_reference.pow(2).sum().backward() + assert torch.equal(x_fused.grad, x_reference.grad) + for (name, a), (_, b) in zip( + fused.named_parameters(), reference.named_parameters() + ): + assert torch.equal(a.grad, b.grad), name + + +def test_activation_argument_is_validated(): + """An unknown activation must fail at construction, not at the first step.""" + with pytest.raises(ValueError, match="activation"): + FastGroupNorm(_GROUPS, 16, activation="gelu") + + +def test_supported_activations_match_the_kernels(): + """The module's activation list may not drift from the kernel's. + + ``group_norm`` spells the tuple out rather than importing it (importing the + kernel module has to stay off the CPU path), so nothing but this test stops + the two copies from diverging into a runtime ``ValueError`` from inside the + custom op. + """ + from ScaFFold.unet import triton_group_norm as triton_mod + + assert set(gn_mod.SUPPORTED_ACTIVATIONS) <= set(triton_mod.SUPPORTED_ACTIVATIONS) + + # --------------------------------------------------------------------------- # CPU behavior: identical numerics, and no compilation at all # --------------------------------------------------------------------------- @@ -245,6 +389,144 @@ def _raises(*args, **kwargs): assert gn_mod._use_compiled(torch.randn(1, 8, 4, 4, 4)) is False +def test_triton_failure_falls_back_to_the_compiled_kernel(monkeypatch, caplog): + """A broken Triton install drops to *compiled*, not all the way to eager. + + The distinction is worth 10x on the shapes that dominate the step, so the + ladder must have three rungs and not two. Simulated by forcing the Triton + predicate on and making its kernel raise; the compiled stand-in must then be + the one that answers, exactly once, with the Triton path latched off. + """ + compiled_calls = [] + + def _raises(*args, **kwargs): + raise RuntimeError("simulated Triton failure") + + def _recording(input, num_groups, weight, bias, eps): + compiled_calls.append(type(input)) + return nn.functional.group_norm(input, num_groups, weight, bias, eps) + + monkeypatch.setattr(gn_mod, "_use_triton", lambda *a, **kw: True) + monkeypatch.setattr(gn_mod, "_get_triton_module", _raises) + monkeypatch.setattr(gn_mod, "_use_compiled", lambda _input: True) + monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _recording) + gn_mod._triton_failed = False + gn_mod._compile_failed = False + + fast = FastGroupNorm(_GROUPS, 64) + x = _make_input(seed=44, channels=64, size=8) + with caplog.at_level(logging.WARNING, logger=gn_mod.__name__): + out = fast(x) + + assert compiled_calls == [torch.Tensor], "compiled kernel was not the fallback" + assert torch.equal( + out, nn.functional.group_norm(x, _GROUPS, fast.weight, fast.bias) + ) + assert any( + "falling back to the compiled kernel" in r.message for r in caplog.records + ) + assert gn_mod._triton_failed is True + assert gn_mod._compile_failed is False + # Latched off: the predicate now refuses even a would-be eligible tensor. + monkeypatch.undo() + assert gn_mod._use_triton(torch.randn(1, 8, 4, 4, 4), 8, None, None, None) is False + + +@pytest.mark.parametrize("rung", ["triton", "compiled"]) +def test_checkpoint_recompute_stop_is_re_raised(monkeypatch, rung): + """``_StopRecomputationError`` is control flow, not a kernel failure. + + ``torch.utils.checkpoint``'s non-reentrant recompute stops itself by raising + it from a saved-tensor *pack hook*, i.e. from inside whichever op is saving + a tensor at that moment -- which, now that the ReLU is fused and nothing + follows GroupNorm in a ``DoubleConv``, is this module. A blanket + ``except Exception`` would swallow it, latch the fast kernel off and drop the + whole model to eager mid-run. (Observed exactly that on + ``test_gpu_activation_checkpointing_matches_eager`` before the re-raise.) + """ + import torch.utils.checkpoint as checkpoint_mod + + stop = checkpoint_mod._StopRecomputationError + + def _raises(*args, **kwargs): + raise stop() + + if rung == "triton": + monkeypatch.setattr(gn_mod, "_use_triton", lambda *a, **kw: True) + monkeypatch.setattr(gn_mod, "_get_triton_module", _raises) + else: + monkeypatch.setattr(gn_mod, "_use_compiled", lambda _input: True) + monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _raises) + gn_mod._triton_failed = False + gn_mod._compile_failed = False + + fast = FastGroupNorm(_GROUPS, 16) + with pytest.raises(stop): + fast(_make_input(seed=45, channels=16, size=4)) + assert gn_mod._triton_failed is False + assert gn_mod._compile_failed is False + + +def test_cpu_activation_checkpointing_keeps_the_fast_path(monkeypatch): + """End-to-end version of the above, on the real model. + + The fast path is forced on for CPU tensors (with the stock kernel standing + in for the compiled one, so only the *routing* is under test) and the model + is run with activation checkpointing. Gradients must match the + non-checkpointed run and the fast path must still be live afterwards -- a + swallowed recompute-stop shows up as a latched-off kernel here. + """ + monkeypatch.setattr(gn_mod, "_use_compiled", lambda t: type(t) is torch.Tensor) + monkeypatch.setattr( + gn_mod, "_get_compiled_group_norm", lambda: nn.functional.group_norm + ) + gn_mod._compile_failed = False + + x = _make_input(seed=46).requires_grad_(True) + + def grads(checkpointing): + model = _make_unet(seed=0) + if checkpointing: + model.use_checkpointing() + model.zero_grad(set_to_none=True) + model(x).pow(2).sum().backward() + return {n: p.grad.detach().clone() for n, p in model.named_parameters()} + + direct = grads(False) + checkpointed = grads(True) + assert gn_mod._compile_failed is False, "recompute-stop was swallowed" + for name in direct: + assert torch.allclose(direct[name], checkpointed[name]), name + + +def test_triton_rejects_unknown_tensor_subclasses(): + """``is_supported`` only asks ``isinstance``, so the type check lives here. + + A ``__torch_dispatch__`` wrapper other than DCTensor has unknown semantics + and must keep the stock kernel, exactly as it does for the compiled path -- + but ``triton_group_norm.is_supported`` would happily accept one, so + ``_use_triton`` has to reject it itself rather than delegating. + """ + + class _Wrapper(torch.Tensor): + pass + + plain = torch.randn(1, 8, 4, 4, 4) + assert gn_mod._use_triton(plain, 8, None, None, None) is False # CPU + assert gn_mod._use_triton(plain.as_subclass(_Wrapper), 8, None, None, None) is False + + +def test_triton_env_opt_out_skips_the_kernel_module_entirely(monkeypatch): + """``SCAFFOLD_GROUPNORM_TRITON=0`` is checked before anything is imported.""" + + def _boom(): + raise AssertionError("the kernel module must not be imported when opted out") + + monkeypatch.setattr(gn_mod, "_get_triton_module", _boom) + gn_mod.set_triton_enabled(False) + assert gn_mod._use_triton(torch.randn(1, 8, 4, 4, 4), 8, None, None, None) is False + + @pytest.mark.parametrize( "value,expected", [ @@ -272,6 +554,116 @@ def test_env_var_unset_means_auto(monkeypatch): assert gn_mod._compile_override is None +@pytest.mark.parametrize( + "value,expected", + [ + ("0", False), + ("false", False), + ("OFF", False), + ("no", False), + ("1", True), + ("true", True), + ("On", True), + ("yes", True), + ("maybe", None), + ], +) +def test_triton_env_var_controls_the_fast_path(monkeypatch, value, expected): + """``SCAFFOLD_GROUPNORM_TRITON`` is parsed exactly like its compile twin.""" + monkeypatch.setenv(gn_mod.TRITON_ENV_VAR, value) + gn_mod.set_triton_enabled(None) + assert gn_mod._triton_override is expected + + +def test_triton_env_var_unset_means_auto(monkeypatch): + """Unset means "on wherever is_supported accepts" -- the production default.""" + monkeypatch.delenv(gn_mod.TRITON_ENV_VAR, raising=False) + gn_mod.set_triton_enabled(None) + assert gn_mod._triton_override is None + + +def test_triton_env_var_garbage_warns(monkeypatch, caplog): + """An unparsable value is ignored *loudly*, like the compile variable.""" + monkeypatch.setenv(gn_mod.TRITON_ENV_VAR, "sometimes") + with caplog.at_level(logging.WARNING, logger=gn_mod.__name__): + gn_mod.set_triton_enabled(None) + assert any(gn_mod.TRITON_ENV_VAR in r.message for r in caplog.records) + assert gn_mod._triton_override is None + + +def test_set_triton_enabled_returns_the_previous_setting(monkeypatch): + """The save/restore contract tests rely on, matching set_compile_enabled.""" + monkeypatch.delenv(gn_mod.TRITON_ENV_VAR, raising=False) + gn_mod.set_triton_enabled(None) + assert gn_mod.set_triton_enabled(False) is None + assert gn_mod.set_triton_enabled(True) is False + assert gn_mod.set_triton_enabled(None) is True + assert gn_mod._triton_override is None + + +def test_cpu_never_imports_triton(fresh_python): + """A CPU-only process must not import triton, nor the kernel module. + + Two separate costs, both of which the CPU unit suite would otherwise pay on + every run: ``import triton`` (seconds, and it is not installed everywhere), + and importing ``ScaFFold.unet.triton_group_norm``, which registers two + dispatcher ops and an autograd formula at import time. ``_use_triton`` + rejects non-CUDA tensors *before* it touches the module, which is what this + pins -- run in a fresh interpreter because the test session itself has long + since imported the kernel module for the kernel's own tests. + """ + out = fresh_python( + "import sys\n" + "import torch\n" + "from ScaFFold.unet.unet_model import UNet\n" + "m = UNet(n_channels=3, n_classes=2, trilinear=False, layers=1, " + "group_norm_groups=8)\n" + "with torch.no_grad():\n" + " m(torch.randn(1, 3, 16, 16, 16))\n" + "print('triton', 'triton' in sys.modules)\n" + "print('kernel', 'ScaFFold.unet.triton_group_norm' in sys.modules)\n" + ) + assert "triton False" in out, out + assert "kernel False" in out, out + + +def test_cpu_activation_is_applied_on_the_eager_path(): + """``activation="relu"`` is a promise of the module, not of the kernel.""" + fast = FastGroupNorm(_GROUPS, 16, activation="relu") + plain = nn.GroupNorm(_GROUPS, 16) + with torch.no_grad(): + plain.weight.copy_(fast.weight) + plain.bias.copy_(fast.bias) + x = _make_input(seed=41, channels=16, size=8) + assert torch.equal(fast(x), torch.relu(plain(x))) + + +def test_cpu_activation_none_leaves_the_output_alone(): + """The default stays a bare GroupNorm -- no accidental global activation.""" + fast = FastGroupNorm(_GROUPS, 16) + x = _make_input(seed=42, channels=16, size=8) + assert torch.equal( + fast(x), nn.functional.group_norm(x, _GROUPS, fast.weight, fast.bias) + ) + + +def test_activation_does_not_allocate_a_second_output(): + """The absorbed ReLU keeps ``nn.ReLU(inplace=True)``'s memory behaviour. + + The old ``nn.Sequential`` spelling mutated the GroupNorm output in place; + an out-of-place ``F.relu`` here would add a full activation-sized allocation + at all 22 sites. Checked by handing the module a stand-in kernel whose + output we still hold: the ReLU must have rewritten *that* tensor. + """ + fast = FastGroupNorm(_GROUPS, 16, activation="relu") + x = _make_input(seed=43, channels=16, size=4) + produced = nn.functional.group_norm(x, _GROUPS, fast.weight, fast.bias, fast.eps) + assert (produced < 0).any(), "test input must have negatives to clamp" + out = fast._activate(produced) + assert out.data_ptr() == produced.data_ptr() + assert (produced < 0).sum() == 0 + + def test_recompile_limit_is_raised_never_lowered(): """Dynamo's stock cap of 8 is below what one UNet needs. @@ -530,8 +922,9 @@ def test_gpu_dctensor_matches_eager_dctensor(dc_cuda, autocast, layout): Both layouts are covered because production requests ``channels_last_3d`` (worker.py) and, with ``PYTORCH_MIOPEN_SUGGEST_NHWC=1`` set as it is there, the convolutions really do hand GroupNorm channels-last activations. Only - parity is asserted, not the output layout: both routes return contiguous - today regardless of the input layout. + parity is asserted, not the output layout: both routes compared here return + contiguous regardless of the input layout (the Triton kernel, which does + not, is pinned off below and covered by its own tests). """ distconv, ps = dc_cuda device = torch.device("cuda") @@ -550,6 +943,9 @@ def plain(t): return t._tensor if isinstance(t, distconv.DCTensor) else t def run(compiled): + # This test is about the compiled rung of the ladder, which the Triton + # one would otherwise pre-empt on the channels-last parametrization. + gn_mod.set_triton_enabled(False) gn_mod.set_compile_enabled(compiled) inp = x.clone().requires_grad_(True) fast.zero_grad(set_to_none=True) @@ -612,6 +1008,9 @@ def test_gpu_compiled_matches_eager(shape, autocast): fast.bias.normal_(0.0, 0.1, generator=generator) def run(compiled): + # Compiled-rung test: the inputs here are contiguous, which the Triton + # kernel declines anyway, but pin it off so the routing cannot drift. + gn_mod.set_triton_enabled(False) gn_mod.set_compile_enabled(compiled) inp = x.clone().requires_grad_(True) fast.zero_grad(set_to_none=True) @@ -649,6 +1048,7 @@ def test_gpu_steady_state_does_not_recompile(): """ from torch._dynamo.utils import counters + gn_mod.set_triton_enabled(False) # this is the compiled rung's guard set gn_mod.set_compile_enabled(True) device = torch.device("cuda") fast = FastGroupNorm(_GROUPS, 64).to(device) @@ -687,6 +1087,7 @@ def test_gpu_activation_checkpointing_matches_eager(): tolerance = 5e-2 def grads(compiled, checkpointing): + gn_mod.set_triton_enabled(False) gn_mod.set_compile_enabled(compiled) model = _make_unet(seed=0).to(device) if checkpointing: @@ -711,3 +1112,232 @@ def assert_agrees(actual, expected, label): assert not gn_mod._compile_failed assert_agrees(compiled, eager, "checkpointed grad") assert_agrees(compiled_nockpt, compiled, "grad") + + +# --------------------------------------------------------------------------- +# GPU behavior: the Triton rung +# --------------------------------------------------------------------------- + + +def _channels_last(t): + return t.is_contiguous(memory_format=torch.channels_last_3d) + + +@pytest.mark.gpu +@pytest.mark.parametrize("activation", [None, "relu"]) +@pytest.mark.parametrize("autocast", [False, True]) +def test_gpu_triton_matches_eager(activation, autocast): + """The Triton kernel is the default for channels-last input and matches eager. + + ``(1, 64, 32^3)`` channels-last is the production shape family at unit-test + size. Three claims at once: the routing really picks Triton when nothing is + forced (the output comes back channels-last, which is the one thing *only* + that path does -- eager and Inductor both return contiguous); the values and + gradients match the eager reference within reduction-order noise; and the + fused activation equals an explicit ReLU on the eager result. + """ + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(11) + x = torch.randn(1, 64, 32, 32, 32, device=device, generator=generator).to( + memory_format=torch.channels_last_3d + ) + grad_out = torch.randn(*x.shape, device=device, generator=generator) + + fast = FastGroupNorm(_GROUPS, 64, activation=activation).to(device) + with torch.no_grad(): + fast.weight.normal_(1.0, 0.1, generator=generator) + fast.bias.normal_(0.0, 0.1, generator=generator) + + def run(triton): + gn_mod.set_triton_enabled(triton) + gn_mod.set_compile_enabled(False) # eager reference, not Inductor + inp = x.clone().requires_grad_(True) + fast.zero_grad(set_to_none=True) + with torch.autocast("cuda", dtype=torch.bfloat16, enabled=autocast): + out = fast(inp) + out.backward(grad_out.to(out.dtype)) + return ( + out.detach(), + inp.grad.detach().clone(), + fast.weight.grad.detach().clone(), + fast.bias.grad.detach().clone(), + ) + + eager = run(False) + triton = run(None) # None = the production default, i.e. no override at all + assert not gn_mod._triton_failed + + assert _channels_last(triton[0]), "Triton path was not taken (output not NDHWC)" + # ... and the control: stock GroupNorm really does return contiguous here, + # so the assertion above is a discriminating signal and not a tautology. + assert not _channels_last(eager[0]) + _assert_close(triton[0], eager[0], 1e-5, "output") + _assert_close(triton[1], eager[1], 1e-4, "d_input") + _assert_close(triton[2], eager[2], 1e-4, "d_weight") + _assert_close(triton[3], eager[3], 1e-4, "d_bias") + # Autocast's fp32 policy for GroupNorm must survive the swap. + assert triton[0].dtype == eager[0].dtype + if activation == "relu": + assert (triton[0] < 0).sum() == 0 + assert triton[0].max() > 0 + + +@pytest.mark.gpu +@pytest.mark.parametrize("activation", [None, "relu"]) +def test_gpu_triton_dctensor_matches_eager_and_stays_wrapped(dc_cuda, activation): + """The production configuration: DCTensor in, DCTensor out, NDHWC preserved. + + worker.py wraps every activation in a DCTensor even at + ``dc_num_shards=[1,1,1]``, so this -- not the plain-tensor case -- is the + path the benchmark actually runs. A producing convolution sits in front so + that the DCTensor handed to GroupNorm is a genuine intermediate: the unwrap + has to be the autograd-aware one or the gradient never reaches the conv. + """ + distconv, ps = dc_cuda + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(53) + x = torch.randn(1, 64, 16, 16, 16, device=device, generator=generator).to( + memory_format=torch.channels_last_3d + ) + + fast = FastGroupNorm(_GROUPS, 64, activation=activation).to(device) + producer = nn.Conv3d(64, 64, 1, bias=False).to( + device, memory_format=torch.channels_last_3d + ) + with torch.no_grad(): + fast.weight.normal_(1.0, 0.1, generator=generator) + fast.bias.normal_(0.0, 0.1, generator=generator) + + def plain(t): + return t._tensor if isinstance(t, distconv.DCTensor) else t + + def run(triton): + gn_mod.set_triton_enabled(triton) + gn_mod.set_compile_enabled(False) + inp = x.clone().requires_grad_(True) + fast.zero_grad(set_to_none=True) + producer.zero_grad(set_to_none=True) + out = fast(producer(distconv.DCTensor.from_shard(inp, ps))) + assert isinstance(out, distconv.DCTensor), "DCTensor did not survive" + local = distconv.distconv._ToTensor.apply(out) + local.float().pow(2).sum().backward() + assert inp.grad is not None, "gradient never reached the input" + assert producer.weight.grad is not None, "gradient never reached the producer" + return ( + local.detach().clone(), + inp.grad.detach().clone(), + plain(producer.weight.grad).detach().clone(), + plain(fast.weight.grad).detach().clone(), + plain(fast.bias.grad).detach().clone(), + ) + + eager = run(False) + triton = run(None) + assert not gn_mod._triton_failed + assert _channels_last(triton[0]), "Triton path was not taken (output not NDHWC)" + + _assert_close(triton[0], eager[0], 1e-5, "output") + for index, what in ( + (1, "d_input"), + (2, "d_producer"), + (3, "d_weight"), + (4, "d_bias"), + ): + _assert_close(triton[index], eager[index], 1e-4, what) + + +@pytest.mark.gpu +def test_gpu_unet_keeps_the_channels_last_chain(monkeypatch): + """The whole point: GroupNorm stops breaking the layout chain in the model. + + Before this kernel, every one of the model's GroupNorms consumed + ``channels_last_3d`` and emitted contiguous, forcing the next convolution to + convert back -- 22 breaks per scale-8 forward. A hook census asserts that + every ``FastGroupNorm`` invocation now takes NDHWC in *and* hands NDHWC out. + + Needs ``PYTORCH_MIOPEN_SUGGEST_NHWC=1`` in the environment for the + convolutions to emit channels-last at all; without it there is nothing to + preserve and the test skips rather than passing vacuously. + """ + device = torch.device("cuda") + model = _make_unet(seed=0).to(device, memory_format=torch.channels_last_3d) + x = _make_input(seed=9).to(device).contiguous(memory_format=torch.channels_last_3d) + + census = [] + + def hook(module, inputs, output): + census.append((_channels_last(inputs[0]), _channels_last(output))) + + for module in model.modules(): + if isinstance(module, FastGroupNorm): + module.register_forward_hook(hook) + + gn_mod.set_triton_enabled(None) + with torch.autocast("cuda", dtype=torch.bfloat16), torch.no_grad(): + model(x) + + assert census, "no GroupNorm ran" + if not any(seen_in for seen_in, _ in census): + pytest.skip( + "convolutions did not emit channels_last_3d; set " + "PYTORCH_MIOPEN_SUGGEST_NHWC=1 (the production setting)" + ) + breaks = [i for i, (seen_in, seen_out) in enumerate(census) if seen_in != seen_out] + assert not breaks, f"GroupNorm broke the layout chain at sites {breaks}" + assert all(seen_out for _, seen_out in census) + + +@pytest.mark.gpu +def test_gpu_unet_triton_matches_the_compiled_build(monkeypatch): + """Whole-model gradients with the Triton kernel vs. without it. + + Compared as relative L2 per parameter against the noise floor + ``test_gpu_activation_checkpointing_matches_eager`` documents: with + ``cudnn.benchmark`` on and bf16 autocast, two *eager* runs of this model + differ by ~4e-3 relative. Anything of that order is the model's own + nondeterminism; a genuinely wrong kernel would be O(1). + + Skips (rather than passing vacuously) when the convolutions are not emitting + channels-last, since the Triton kernel would then never engage. + """ + device = torch.device("cuda") + x = _make_input(seed=9).to(device).contiguous(memory_format=torch.channels_last_3d) + tolerance = 5e-2 + + engaged = [] + original = FastGroupNorm._triton_forward + + def spy(self, local): + engaged.append(tuple(local.shape)) + return original(self, local) + + monkeypatch.setattr(FastGroupNorm, "_triton_forward", spy) + + def grads(triton): + gn_mod.set_triton_enabled(triton) + gn_mod.set_compile_enabled(True) + model = _make_unet(seed=0).to(device, memory_format=torch.channels_last_3d) + model.zero_grad(set_to_none=True) + with torch.autocast("cuda", dtype=torch.bfloat16): + out = model(x) + out.float().pow(2).mean().backward() + return {n: p.grad.detach().clone() for n, p in model.named_parameters()} + + without = grads(False) + assert not engaged + with_triton = grads(None) + if not engaged: + pytest.skip( + "Triton kernel never engaged; set PYTORCH_MIOPEN_SUGGEST_NHWC=1 " + "(the production setting) so the convolutions emit channels-last" + ) + assert not gn_mod._triton_failed + + worst = 0.0 + for name, reference in without.items(): + reference = reference.float() + error = (with_triton[name].float() - reference).norm().item() + relative = error / max(reference.norm().item(), 1e-12) + worst = max(worst, relative) + assert relative < tolerance, f"grad {name}: rel L2 {relative:.3e}" + assert worst < tolerance From c2c41fa5208c78e664b8fd31c9b2a232707f57a5 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Sat, 1 Aug 2026 17:04:26 -0700 Subject: [PATCH 59/62] Fix nine defects found reviewing the GroupNorm wiring The one that matters: the fused ReLU silenced NaN. tl.maximum(y, 0.0) maps NaN and -Inf to zero where F.relu propagates NaN, so a diverged run showed a finite forward and a NaN backward -- and round 2 added the non-finite loss abort precisely so divergence stops the run instead of checkpointing a broken model. Testing the complement (tl.where(y <= 0, 0, y)) keeps NaN on the pass-through side. The backward gate had the same defect and matters as much: threshold_backward passes the gradient where the result is NaN, and pre > 0 was zeroing it. Now bit-identical to F.relu on NaN, both infinities and -0.0, on all three rungs, forward and backward. Cost: +0.18% on the 22-site rollup, against a -0.30% noise floor from the activation=None control. A latch flip between a checkpointed forward and its recompute killed the run. Matching the output memory format across rungs does not fix it -- measured: the rungs save different tensor sets, so it still dies comparing a (64,) to a (1,8). The fix is that a latch no longer demotes a module that has already had a call served by that rung, so a block's forward and its recompute always agree; a broken install still costs exactly one attempt per module. Rungs now preserve the input's memory format anyway, since a fallback that returns contiguous re-breaks the channels-last chain this whole line of work exists to protect. _CONTROL_FLOW_EXCEPTIONS is gone. A denylist of framework mechanisms that legitimately raise through a forward was wrong twice (_StopRecomputationError, then CheckpointError) and is unbounded. Narrowing by exception type does not work either -- a HIP launch failure and a CheckpointError are both RuntimeError. So the narrowing is by scope instead: a decorator tags failures raised inside the kernel call itself, which is a closed region that ends before anything is saved for backward, and the ladder catches only that tag. That also removes the retry's saved-tensor-hook asymmetry structurally rather than defensively. The compiled rung catches TorchDynamoException, whose members are all raised at compile time. Remaining: predicates decline under functorch transforms instead of latching both rungs off; a predicate that cannot answer falls back without latching; a transient OOM neither latches nor falls back, since every fallback allocates an output of the same size; set_*_enabled(True) clears a latch, and the process-local nature of latches is documented with its multi-rank consequence. Whole-module unpickling of a pre-fusion model works again, the fallback no longer graph-breaks under fullgraph, and an unapplicable activation raises instead of silently doing nothing. CPU 386 passed, GPU 267. Mutations: wiring 36/38 (two known-benign), kernel 35/45 with 9 pre-existing not-applicable, opt 48/50. --- ScaFFold/unet/group_norm.py | 411 ++++-- ScaFFold/unet/triton_group_norm.py | 106 +- .../rank_scripts/groupnorm_shards_2rank.py | 4 +- tests/test_groupnorm.py | 79 +- tests/test_groupnorm_wiring_edge.py | 1100 +++++++++++++++++ tests/test_triton_group_norm_edge.py | 179 +++ 6 files changed, 1774 insertions(+), 105 deletions(-) create mode 100644 tests/test_groupnorm_wiring_edge.py diff --git a/ScaFFold/unet/group_norm.py b/ScaFFold/unet/group_norm.py index 9d7e95f..b172b09 100644 --- a/ScaFFold/unet/group_norm.py +++ b/ScaFFold/unet/group_norm.py @@ -41,10 +41,17 @@ one addition is the optional fused ``activation`` (see below), which adds no state either. +All three return the input's memory format, so the rungs are interchangeable in +everything a caller can observe (see :func:`_match_memory_format`). + Routing rejections, in the order they are tested: * an explicit opt-out via ``SCAFFOLD_GROUPNORM_TRITON=0`` / ``SCAFFOLD_GROUPNORM_COMPILE=0``, +* a rung that has failed in this process, for every module that has not already + had a call served by it (see "Latches" below), +* an active ``torch.func`` transform -- a ``vmap``/``grad``/``jvp`` layer is a + routing miss, not a kernel failure, and the stock kernel handles it, * non-CUDA tensors -- the CPU test suite pays neither compile latency nor the Triton import, * tensor subclasses, whose ``__torch_dispatch__`` wrappers have unknown @@ -53,11 +60,54 @@ * for the Triton kernel, anything its ``is_supported`` rejects (a layout, dtype, degenerate shape or affine-parameter dtype it does not serve); for the compiled one, an already-compiled enclosing region (the functional call - inlines instead), -* any failure inside either kernel -- logged once, latched off for the rest of - the process, and retried on the next path down. Both are optimizations, never - correctness requirements: a broken Triton install must degrade a multi-node - run, not kill it. + inlines instead). + +Latches +======= +Both fast rungs are optimizations, never correctness requirements: a broken +Triton install must degrade a multi-node run, not kill it. So a *kernel* +failure is caught, logged once and retried on the next rung down. + +"A kernel failure" is an allowlist, not the absence of one. The Triton rung is +caught on ``triton_group_norm.TritonKernelError``, which that module raises for +anything its launch region produces; the compiled rung on +``torch._dynamo.exc.TorchDynamoException``, the root of every Dynamo and +Inductor compile failure. Everything else propagates -- saved-tensor pack +hooks, ``torch.utils.checkpoint``'s recompute control flow, a user's offloading +hook, ``torch.OutOfMemoryError``, an error from a shape the kernel mishandles +badly enough to corrupt the graph. The previous shape of this code caught +``Exception`` and re-raised a denylist of framework mechanisms, which was wrong +twice (``_StopRecomputationError``, then ``CheckpointError``): the set of things +torch may raise through a forward is open, the set of ways a kernel can be +broken is closed at its own boundary. Both allowlisted exceptions are also +raised strictly *before* their rung saves anything for backward (the Triton op +saves in ``_setup_context``, after its launch region; a Dynamo/Inductor failure +is a compile-time failure, before any execution), so the retry cannot double-fire +saved-tensor hooks. + +A failure latches the rung off **for modules that have never had a call served +by it**. A module that has already run on a rung keeps it. That is not a +performance nicety: ``torch.utils.checkpoint``'s non-reentrant recompute +compares the metadata of every tensor the recomputed forward saves against the +originals, and the three rungs intrinsically save *different tensors* -- Triton +saves ``(input, weight, bias, mean, rstd)``, the other two +``(input, weight, mean, rstd, relu_output)``. A latch that flipped between a +block's forward and its recompute would therefore kill the step with +``CheckpointError: Recomputed values ... have different metadata``, which is the +exact opposite of the contract above (measured; matching the output memory +format is *not* sufficient on its own). Keeping a proven rung pins each +module's choice for the life of the process, so forward and recompute always +agree. + +Note that a latch is process-local: under DDP one rank can end up running a +different kernel from its peers. All three kernels agree to fp32 rounding, not +bitwise, so a rank that latches shifts that rank's gradients and therefore the +all-reduced ones -- a real (measured) change to the job's trajectory, and a +2.1x straggler besides. That is the price of degrading instead of dying, but it +is why the latch is as narrow as it is, and why ``torch.OutOfMemoryError`` -- +transient by nature, and no cheaper on any other rung -- does not latch anything +at all. :func:`set_triton_enabled` / :func:`set_compile_enabled` with ``True`` +clear the latch, which is the supported way to retry after a transient failure. Determinism: all three kernels are bitwise reproducible. Two separate processes running the scale-7 UNet under ``more_determinism`` @@ -88,7 +138,6 @@ import torch import torch.nn as nn import torch.nn.functional as F -import torch.utils.checkpoint # for _CONTROL_FLOW_EXCEPTIONS; torch loads it anyway logger = logging.getLogger(__name__) @@ -132,6 +181,12 @@ # the `triton` import the module itself defers to its first launch. _triton_module = None +# The ladder's two allowlists, resolved on first use of the rung they guard -- +# importing either provider (the kernel module, torch._dynamo) is exactly what +# the lazy _get_* helpers exist to avoid paying for on a CPU-only run. +_TRITON_KERNEL_FAILURES = None +_COMPILED_KERNEL_FAILURES = None + # Set once if the Triton kernel raises; the compiled path is used from then on. _triton_failed = False @@ -167,14 +222,18 @@ def set_compile_enabled(enabled): ``None`` restores the default, which is the environment variable if set and otherwise "compile wherever it is safe". Forcing it on does not override the device and tensor-subclass checks -- those are correctness conditions, - not preferences. Returns the previous setting so callers (tests) can - restore it. + not preferences -- but it *does* clear a failure latch: an explicit "use + this rung" is the supported way to retry after a transient failure, and + leaving the latch set would make this function silently do nothing. + Returns the previous setting so callers (tests) can restore it. """ - global _compile_override + global _compile_override, _compile_failed previous = _compile_override _compile_override = ( _env_override(COMPILE_ENV_VAR) if enabled is None else bool(enabled) ) + if _compile_override is True: + _compile_failed = False return previous @@ -183,15 +242,20 @@ def set_triton_enabled(enabled): The exact counterpart of :func:`set_compile_enabled`: ``None`` restores the default (``SCAFFOLD_GROUPNORM_TRITON`` if set, otherwise "wherever - ``is_supported`` accepts"), forcing it on does not override the device, - subclass or ``is_supported`` checks -- those are correctness conditions -- - and the previous setting is returned so tests can restore it. + ``is_supported`` accepts"), forcing it on clears any failure latch but not + the device, subclass or ``is_supported`` checks -- those are correctness + conditions -- and the previous setting is returned so tests can restore it. + + ``None`` deliberately does *not* clear the latch: it restores a preference, + it does not assert that the kernel works again. """ - global _triton_override + global _triton_override, _triton_failed previous = _triton_override _triton_override = ( _env_override(TRITON_ENV_VAR) if enabled is None else bool(enabled) ) + if _triton_override is True: + _triton_failed = False return previous @@ -250,33 +314,89 @@ def _get_triton_module(): return _triton_module -#: Exceptions torch raises *through* this module as control flow rather than as -#: a kernel failure, and which the fallback ladder must therefore re-raise. -#: -#: ``torch.utils.checkpoint``'s non-reentrant recompute stops itself early by -#: raising ``_StopRecomputationError`` from its saved-tensor *pack hook* -- i.e. -#: from inside whichever op happens to be saving a tensor when the recompute has -#: produced everything the backward needs. In a ``DoubleConv`` that op is this -#: module (GroupNorm saves its input and statistics, and the fused ReLU means -#: nothing follows it in the block), so the exception surfaces inside the -#: ``try``. Swallowing it would latch the fast kernel off, silently drop the -#: model to eager mid-run, and leave the checkpoint machinery waiting for a stop -#: that never came. Private API, so tolerate its absence rather than importing -#: it by name. -_CONTROL_FLOW_EXCEPTIONS = tuple( - exception - for exception in (getattr(torch.utils.checkpoint, "_StopRecomputationError", None),) - if isinstance(exception, type) and issubclass(exception, BaseException) -) - - -def _use_triton(input, num_groups, weight, bias, activation): +def _triton_kernel_failures(): + """The ladder's allowlist for the Triton rung: exactly ``TritonKernelError``. + + The kernel module raises it for every failure of its own launch region -- + a missing or mismatched ``triton``, an unwritable JIT cache, a compile + error, a bad launch -- and for nothing else, so this catches "the kernel is + broken" without also catching the framework mechanisms that legitimately + raise through a forward. See that class's docstring for what is + deliberately left untagged (``OutOfMemoryError``, contract violations). + + Resolved separately from :func:`_get_triton_module` so the except clause is + still available when the thing that failed *is* the module lookup. An + empty tuple (no kernel module at all) means "catch nothing": the ladder + then re-raises, which is right, because with no kernel module there is + nothing that could have failed inside one. + """ + global _TRITON_KERNEL_FAILURES + if _TRITON_KERNEL_FAILURES is None: + try: + from .triton_group_norm import TritonKernelError + + _TRITON_KERNEL_FAILURES = (TritonKernelError,) + except ImportError: # pragma: no cover - the module is in-tree + _TRITON_KERNEL_FAILURES = () + return _TRITON_KERNEL_FAILURES + + +def _compiled_kernel_failures(): + """The compiled rung's allowlist: every Dynamo and Inductor compile failure. + + ``torch._dynamo.exc.TorchDynamoException`` is the root of ``Unsupported`` + (``fullgraph=True`` met something untraceable), ``BackendCompilerFailed`` + and its ``InductorError`` subclass (the backend, and therefore also an + unwritable Inductor cache or a broken C++/Triton toolchain), and + ``InternalTorchDynamoError``. All of them are raised while *compiling*, + i.e. before the compiled callable has executed or saved anything, which is + what makes the fallback safe to retry. + + Resolved on demand and cached: importing ``torch._dynamo`` is precisely the + cost :func:`_get_compiled_group_norm` defers. An empty tuple (a torch + without the module) means "catch nothing", which fails loudly rather than + silently swallowing. + """ + global _COMPILED_KERNEL_FAILURES + if _COMPILED_KERNEL_FAILURES is None: + try: + import torch._dynamo.exc + + _COMPILED_KERNEL_FAILURES = (torch._dynamo.exc.TorchDynamoException,) + except ImportError: # pragma: no cover - torch always ships it + _COMPILED_KERNEL_FAILURES = () + return _COMPILED_KERNEL_FAILURES + + +#: ``True`` while a ``torch.func`` transform (``vmap``/``grad``/``jvp``) is on +#: the stack. Both fast rungs decline then: a functorch layer is a routing +#: question, not a kernel defect, and the stock kernel handles every transform. +#: Not merely a performance choice -- ``is_supported``'s +#: ``is_contiguous(memory_format=...)`` raises outright under ``vmap`` +#: ("NYI: querying is_contiguous inside of vmap"), and the Triton op has no +#: batching rule -- so without this the module is not the drop-in +#: ``nn.GroupNorm`` it claims to be for any caller using ``torch.func``. +_functorch_active = getattr(torch._C, "_are_functorch_transforms_active", lambda: False) + +# Set once if a predicate raised while deciding; see _use_triton. +_predicate_warned = False + + +def _use_triton(input, num_groups, weight, bias, activation, proven=False): """Whether this particular input should take the native Triton kernel. + ``proven`` is the caller's "this module has already had a call served by + this rung", which keeps a proven module on it even after a *global* latch; + see the module docstring's "Latches". + Ordered so that the cheap local tests come first and the module import last: a CPU tensor is rejected before ``_get_triton_module`` is ever called. """ - if _triton_failed or _triton_override is False: + if _triton_override is False: + return False + if _triton_failed and not proven: + return False + if _functorch_active(): return False # Same policy as _use_compiled: an unknown __torch_dispatch__ wrapper has # unknown semantics and keeps the stock kernel. is_supported() would accept @@ -288,15 +408,70 @@ def _use_triton(input, num_groups, weight, bias, activation): if not input.is_cuda: return False # is_supported() is cheap and side-effect free: a handful of attribute reads - # and one stride check, no allocation, no launch, no triton import. - return _get_triton_module().is_supported( - input, num_groups, weight, bias, activation + # and one stride check, no allocation, no launch, no triton import. The + # broad catch is right *here* and nowhere else in this module: a predicate + # that cannot answer has a correct answer available ("no"), it has done no + # work anyone can observe, and the failure is a routing miss rather than a + # broken kernel -- so it must not latch the rung off, which is what letting + # it fall into the ladder's handler used to do. + try: + return _get_triton_module().is_supported( + input, num_groups, weight, bias, activation + ) + except Exception as e: + _warn_once_about_the_predicate(e) + return False + + +def _warn_rung_failure(what, error, fallback, env_var): + """Log a rung failure, without graph-breaking a compiled caller. + + Dynamo cannot trace ``logging.Logger`` methods ("Unsupported: logging.Logger + method not supported for non-export cases"), so a bare ``logger.warning`` + in the ladder's handler turns a *fallback* into a hard Dynamo error for + anyone who wraps this ``forward`` in ``torch.compile(fullgraph=True)`` -- + the one caller for whom the fallback matters most, since the failure it is + reacting to is usually a compile failure. ``is_compiling()`` is a Dynamo + intrinsic that folds to ``True`` at trace time, so the call below becomes + dead code inside a traced region and the fallback traces cleanly. The + latch itself is a global assignment, which Dynamo does replay, so the + fallback is still recorded -- only this message is dropped, and only for a + caller that is compiling this module's forward (nothing in ScaFFold does). + """ + if torch.compiler.is_compiling(): + return + logger.warning( + f"{what} failed ({type(error).__name__}: {error}); falling back to the " + f"{fallback} for modules that have not already used it. " + f"Set {env_var}=0 to skip this attempt entirely." ) -def _use_compiled(input): - """Whether this particular input should take the compiled path.""" - if _compile_failed or _compile_override is False: +def _warn_once_about_the_predicate(error): + """Log the first ``is_supported`` failure; a repeat would log per call.""" + global _predicate_warned + if _predicate_warned: + return + _predicate_warned = True + logger.warning( + f"Triton GroupNorm routing check failed ({type(error).__name__}: " + f"{error}); using the stock kernel for inputs like this one. This is a " + "routing miss, not a kernel failure, so nothing is latched off." + ) + + +def _use_compiled(input, proven=False): + """Whether this particular input should take the compiled path. + + ``proven`` has the same meaning as in :func:`_use_triton`. + """ + if _compile_override is False: + return False + if _compile_failed and not proven: + return False + # Dynamo cannot trace a functorch layer either, and under fullgraph=True + # that is an exception rather than a graph break. + if _functorch_active(): return False # Tensor subclasses route their ops through __torch_dispatch__, which # Dynamo cannot trace. DistConv's DCTensor never reaches this check -- @@ -328,6 +503,38 @@ def _dctensor_ops(input): return None +def _match_memory_format(out, reference): + """Give ``out`` ``reference``'s memory format, copying only if it differs. + + ``F.group_norm`` -- eager or Inductor-compiled -- reads a + ``channels_last_3d`` input through the logical NCDHW order and returns a + *contiguous* tensor, which is the layout break this whole module exists to + avoid: with ``PYTORCH_MIOPEN_SUGGEST_NHWC=1`` every convolution both sides + of it wants channels-last, so one fallback re-breaks the chain for the rest + of the network. One relayout of the GroupNorm output is far cheaper than + the transposes the following convolutions would otherwise insert, and it + makes the three rungs agree on everything a caller can observe rather than + only on the values. + + Free on the Triton rung (already channels-last) and on any contiguous input + (nothing to do); one copy on a fallback from a channels-last input, which is + the only case that reaches the copy at all. + """ + if reference.dim() != 5: + # is_contiguous(memory_format=channels_last_3d) is only defined for 5-D. + return out + if _functorch_active(): + # "NYI: querying is_contiguous inside of vmap for memory_format other + # than torch.contiguous_format" -- and a functorch transform has no + # layout chain to preserve anyway, since both fast rungs decline it. + return out + if not reference.is_contiguous(memory_format=torch.channels_last_3d): + return out + if out.is_contiguous(memory_format=torch.channels_last_3d): + return out + return out.contiguous(memory_format=torch.channels_last_3d) + + def _run_local(input, distconv, kernel): """Run ``kernel`` on a plain tensor, DCTensor in -> DCTensor out. @@ -380,6 +587,21 @@ class FastGroupNorm(nn.GroupNorm): shard count, exactly as before. """ + #: Class-level defaults, so that an instance restored from a *module* + #: pickle written before these attributes existed (``torch.save(model)`` + #: rather than a state dict) still runs. ``nn.Module.__setstate__`` + #: replaces ``__dict__`` wholesale, so anything only ever set in + #: ``__init__`` is simply missing on such an instance. + activation = None + + #: Per-module "a call has been served by this rung". A global latch does + #: not demote a module that has one, which is what keeps a checkpointed + #: block's forward and its recompute on the same rung; see the module + #: docstring's "Latches". Plain attributes, so they are not parameters, + #: buffers or state-dict keys. + _triton_ok = False + _compiled_ok = False + def __init__( self, num_groups, @@ -412,10 +634,23 @@ def _activate(self, out): absorbed did: ``out`` is a freshly allocated GroupNorm output with no other consumer, and GroupNorm's backward reads its *input*, never its output, so overwriting it is safe for autograd as well as for memory. + + Validated *here* rather than only in ``__init__``: ``activation`` is a + plain attribute, so it can be assigned after construction, and the + Triton rung would then fuse an activation this method silently skipped + -- i.e. the network's function would depend on its input's memory + format. This is also the guard that makes adding a third activation to + ``SUPPORTED_ACTIVATIONS`` a loud failure until it is implemented here. """ - if self.activation == "relu": + activation = self.activation + if activation is None: + return out + if activation == "relu": return F.relu(out, inplace=True) - return out + raise ValueError( + f"activation must be one of {SUPPORTED_ACTIVATIONS}, got " + f"{activation!r}; this rung cannot apply it" + ) def _triton_forward(self, local): """The native channels-last kernel, with the activation fused in.""" @@ -425,15 +660,19 @@ def _triton_forward(self, local): def _compiled_forward(self, local): return self._activate( - _get_compiled_group_norm()( - local, self.num_groups, self.weight, self.bias, self.eps + _match_memory_format( + _get_compiled_group_norm()( + local, self.num_groups, self.weight, self.bias, self.eps + ), + local, ) ) def _eager_forward(self, input): # super().forward() is the stock kernel; deferring to it keeps the eager - # path identical to nn.GroupNorm's (plus the ReLU) by construction. - return self._activate(super().forward(input)) + # path identical to nn.GroupNorm's (plus the ReLU and the relayout) by + # construction. + return self._activate(_match_memory_format(super().forward(input), input)) def forward(self, input): global _compile_failed, _triton_failed @@ -445,41 +684,63 @@ def forward(self, input): local_view = input._tensor if distconv is not None else input if _use_triton( - local_view, self.num_groups, self.weight, self.bias, self.activation + local_view, + self.num_groups, + self.weight, + self.bias, + self.activation, + proven=self._triton_ok, ): + triton_failures = _triton_kernel_failures() try: - return _run_local(input, distconv, self._triton_forward) - except _CONTROL_FLOW_EXCEPTIONS: - raise - except Exception as e: + out = _run_local(input, distconv, self._triton_forward) + except triton_failures as e: # A broken or mismatched Triton install, an unwritable JIT cache # or a shape the kernel mishandles must cost speed, not a - # multi-node run. GroupNorm is pure, so retrying the same call - # on the compiled kernel below is safe -- and the compiled - # kernel, not eager, is the right landing place: it is still - # ~10x the stock one. + # multi-node run. GroupNorm is pure and the kernel raises this + # only from its launch region -- before it has saved anything -- + # so retrying the same call on the compiled kernel below is + # safe, and the compiled kernel, not eager, is the right landing + # place: it is still ~10x the stock one. + # + # Logged on the latch's False->True edge only. A module that + # has already used the rung keeps trying it (that is what pins + # a checkpointed block to one rung), so a persistently broken + # kernel would otherwise warn once per call for the rest of the + # run; clearing the latch re-arms the message. + first = not _triton_failed _triton_failed = True - logger.warning( - f"Triton GroupNorm failed ({type(e).__name__}: {e}); falling " - "back to the compiled kernel for the rest of this run. " - f"Set {TRITON_ENV_VAR}=0 to skip this attempt entirely." - ) - - if not _use_compiled(local_view): + if first: + _warn_rung_failure( + "Triton GroupNorm", e, "compiled kernel", TRITON_ENV_VAR + ) + else: + # Only written once: nn.Module.__setattr__ is not free, and + # after the first success this reads a class attribute. + if not self._triton_ok: + self._triton_ok = True + return out + + if not _use_compiled(local_view, proven=self._compiled_ok): return self._eager_forward(input) + compile_failures = _compiled_kernel_failures() try: - return _run_local(input, distconv, self._compiled_forward) - except _CONTROL_FLOW_EXCEPTIONS: - raise - except Exception as e: + out = _run_local(input, distconv, self._compiled_forward) + except compile_failures as e: # Compilation is an optimization, never a correctness requirement: - # a broken Inductor/Triton install, an unwritable cache directory or - # an untraceable input must degrade to the stock kernel, not kill a - # multi-node run. GroupNorm is pure, so retrying eagerly is safe. + # a broken Inductor install, an unwritable cache directory or an + # untraceable input must degrade to the stock kernel, not kill a + # multi-node run. Every exception caught here is a *compile*-time + # one, so nothing ran and retrying eagerly is safe. Same + # once-per-latch-edge logging as the Triton rung above. + first = not _compile_failed _compile_failed = True - logger.warning( - f"torch.compile of GroupNorm failed ({type(e).__name__}: {e}); " - "falling back to the eager kernel for the rest of this run. " - f"Set {COMPILE_ENV_VAR}=0 to skip this attempt entirely." - ) + if first: + _warn_rung_failure( + "torch.compile of GroupNorm", e, "eager kernel", COMPILE_ENV_VAR + ) return self._eager_forward(input) + else: + if not self._compiled_ok: + self._compiled_ok = True + return out diff --git a/ScaFFold/unet/triton_group_norm.py b/ScaFFold/unet/triton_group_norm.py index 2ba2220..61c8105 100644 --- a/ScaFFold/unet/triton_group_norm.py +++ b/ScaFFold/unet/triton_group_norm.py @@ -311,8 +311,8 @@ Fused activation ================ ``activation="relu"`` folds the ReLU into the forward store. In a store-bound -kernel that is free (one ``tl.maximum``) and it removes an entire 2B streaming -pass. Measured against ``F.relu(triton_group_norm(x))``: 39% off the forward +kernel that is free (one compare and one select) and it removes an entire 2B +streaming pass. Measured against ``F.relu(triton_group_norm(x))``: 39% off the forward and 35% off fwd+bwd at ``[1,64,256^3]`` (6.83 -> 4.20 ms and 17.82 -> 11.61 ms), 38%/35% at ``[1,128,128^3]``, 34%/30% at ``[1,256,64^3]``, tapering to 21%/9% at ``[1,512,32^3]`` and below, where the call is host bound and there is @@ -328,6 +328,20 @@ and in bf16/fp16 it would also mis-gate any element whose positive pre-activation rounded to zero on the store. +Both the store and the gate are spelled as the *complement* of the usual test +(``tl.where(y <= 0, 0, y)``, ``tl.where(pre <= 0, 0, dy)``) rather than as +``tl.maximum(y, 0)`` / ``tl.where(pre > 0, dy, 0)``. The two are identical on +every finite value but not on NaN: ``tl.maximum`` returns the *non*-NaN operand +and ``NaN > 0`` is False, so both of the usual spellings silently map a NaN to +0.0, while ``F.relu`` propagates it and ``threshold_backward(grad, result, 0)`` +-- ReLU's real backward -- passes its gradient (``NaN <= 0`` is False too). +Matching ``F.relu`` here is not pedantry: a diverging run whose forward comes +back finite because the fused activation ate the NaN passes straight through +ScaFFold's non-finite-loss abort and checkpoints a broken model. ``+-Inf`` and +``-0.0`` are bit-identical under either spelling (``-0.0`` flushes to ``+0.0``, +as ``F.relu`` does). Cost: nil, measured -- see ``FastGroupNorm``'s tests and +``review/gn-dctensor/wiring-fixes``. + Composition =========== Registered as real dispatcher ops (``scaffold_gn::group_norm`` / @@ -403,8 +417,41 @@ "GNConfig", "default_config", "SUPPORTED_ACTIVATIONS", + "TritonKernelError", ] + +class TritonKernelError(RuntimeError): + """A failure of the Triton kernels themselves, with the original as ``__cause__``. + + Raised in place of whatever ``_forward``/``_backward`` raised -- a missing or + mismatched ``triton``, an unwritable JIT cache, a compile error, a launch + failure, an API change between Triton releases. It exists so that a caller + with a fallback (``ScaFFold.unet.group_norm``'s ladder) can catch *exactly* + "the kernel is broken" and nothing else, instead of catching ``Exception`` + and trying to enumerate every framework mechanism that legitimately raises + through a forward -- saved-tensor pack hooks, ``torch.utils.checkpoint``'s + recompute control flow, functorch, a user's offloading hook. + + Two things are deliberately *not* tagged and therefore propagate unchanged: + + * ``torch.OutOfMemoryError``, which is a resource condition rather than a + defect (every fallback allocates an output of the same size, so retrying + one is a second, differently-shaped OOM at a call site the caller did not + ask about), and + * the ``ValueError``s ``_validate`` raises, which are contract violations by + the caller. ``is_supported`` accepts exactly what ``_validate`` accepts, + so a caller that branches on it can never see one; if one escapes, that + is a bug in this module and must be loud. + + The tagged region contains no autograd-observable work -- allocations and + kernel launches only, with ``save_for_backward`` happening in + ``_setup_context`` strictly *after* ``_forward`` returns -- so an exception + that carries this type is guaranteed to have been raised before the op saved + anything. That is what makes retrying the call on another kernel safe. + """ + + #: The activations that may be fused into the forward store. SUPPORTED_ACTIVATIONS = (None, "relu") @@ -951,7 +998,14 @@ def _normalize_kernel( xhat = (x - mean) * rstd y = xhat * w + b if RELU: - y = tl.maximum(y, 0.0) + # `tl.maximum(y, 0.0)` and `tl.where(y > 0, y, 0.0)` both map NaN + # to 0.0 (the first returns the non-NaN operand, the second + # because `NaN > 0` is False), while `F.relu` propagates it. + # Testing the *complement* keeps NaN on the pass-through side: + # `NaN <= 0` is also False, so NaN falls to `y`. Bit-identical + # to `F.relu` on NaN, +-Inf and -0.0 (which both flush to +0.0), + # for one comparison and one select -- see the module docstring. + y = tl.where(y <= 0.0, 0.0, y) tl.store(Y + base + off, y.to(Y.dtype.element_ty), mask=m) # ------------------------------------------------------------- backward -- @@ -1032,8 +1086,12 @@ def _bwd_partial_kernel( # Identical expression (and therefore identical rounding) to # the forward's pre-activation, so the sign test agrees with # the forward bit for bit. Masked lanes carry dy == 0, so - # gating cannot resurrect them. - dy = tl.where(xhat * w + b > 0.0, dy, 0.0) + # gating cannot resurrect them. Spelled as the *complement* + # (`pre <= 0` zeroes) rather than `pre > 0` passes, so that a + # NaN pre-activation passes the gradient through: that is what + # `threshold_backward(grad, result, 0)` -- ReLU's real backward + # -- does, since `NaN <= 0` is False. See the forward store. + dy = tl.where(xhat * w + b <= 0.0, 0.0, dy) dyw = dy * w acc1 += tl.sum(tl.sum(dyw, 2), 0) acc2 += tl.sum(tl.sum(dyw * xhat, 2), 0) @@ -1178,7 +1236,10 @@ def _dx_kernel( dy = tl.load(DY + base + off, mask=m, other=0.0).to(tl.float32) xhat = (x - mean) * rstd if RELU: - dy = tl.where(xhat * w + b > 0.0, dy, 0.0) + # Same complement spelling as _bwd_partial_kernel: a NaN + # pre-activation must pass the gradient, exactly as + # `threshold_backward(grad, result, 0)` does. + dy = tl.where(xhat * w + b <= 0.0, 0.0, dy) dyw = dy * w dx = rstd * (dyw - c1 - xhat * c2) tl.store(DX + base + off, dx.to(DX.dtype.element_ty), mask=m) @@ -1239,6 +1300,38 @@ def _shape_of(input: torch.Tensor): return n, channels, spatial +def _tag_kernel_failures(fn): + """Re-raise anything ``fn`` raises as :class:`TritonKernelError`. + + Applied to the two functions that do nothing but import Triton, allocate + scratch and launch kernels. The region is *closed*: it runs no + autograd-observable op, so a blanket ``except Exception`` here cannot + swallow framework control flow the way one at the call site would -- there + is no pack hook, no recompute stop and no functorch layer inside it. That + closure is what lets the caller's fallback ladder use a one-element + allowlist instead of an ever-growing denylist. + + ``torch.OutOfMemoryError`` is passed through untagged; see + :class:`TritonKernelError`. + """ + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + try: + return fn(*args, **kwargs) + except torch.OutOfMemoryError: + raise + except TritonKernelError: + raise + except Exception as e: + raise TritonKernelError( + f"{fn.__name__} failed ({type(e).__name__}: {e})" + ) from e + + return wrapper + + +@_tag_kernel_failures def _forward(input, num_groups, weight, bias, eps, activation, out_dtype): _ensure_kernels() n, channels, spatial = _shape_of(input) @@ -1304,6 +1397,7 @@ def _forward(input, num_groups, weight, bias, eps, activation, out_dtype): return out, mean, rstd +@_tag_kernel_failures def _backward(grad_out, input, weight, bias, mean, rstd, num_groups, activation): _ensure_kernels() n, channels, spatial = _shape_of(input) diff --git a/tests/helpers/rank_scripts/groupnorm_shards_2rank.py b/tests/helpers/rank_scripts/groupnorm_shards_2rank.py index a69bcfd..c347365 100644 --- a/tests/helpers/rank_scripts/groupnorm_shards_2rank.py +++ b/tests/helpers/rank_scripts/groupnorm_shards_2rank.py @@ -74,10 +74,10 @@ def forward(compiled): gn_mod._get_compiled_group_norm, ) if compiled: - gn_mod._use_compiled = lambda t: type(t) is torch.Tensor + gn_mod._use_compiled = lambda t, **kw: type(t) is torch.Tensor gn_mod._get_compiled_group_norm = lambda: F.group_norm else: - gn_mod._use_compiled = lambda t: False + gn_mod._use_compiled = lambda t, **kw: False try: norm.zero_grad(set_to_none=True) x = local.clone().requires_grad_(True) diff --git a/tests/test_groupnorm.py b/tests/test_groupnorm.py index fe8ab28..334da73 100644 --- a/tests/test_groupnorm.py +++ b/tests/test_groupnorm.py @@ -364,14 +364,17 @@ class _Wrapper(torch.Tensor): def test_compile_failure_falls_back_to_eager(monkeypatch, caplog): """A broken compiler degrades to the stock kernel instead of killing the run. - Simulated by making the compiled callable raise; the module must return the - eager result, warn once, and stop trying for the rest of the process. + Simulated by making the compiled callable raise the real thing Dynamo and + Inductor raise (``BackendCompilerFailed``/``Unsupported`` share the + ``TorchDynamoException`` root the ladder allowlists); the module must return + the eager result, warn once, and stop trying for the rest of the process. """ + import torch._dynamo.exc def _raises(*args, **kwargs): - raise RuntimeError("simulated Inductor failure") + raise torch._dynamo.exc.Unsupported("simulated Inductor failure") - monkeypatch.setattr(gn_mod, "_use_compiled", lambda _input: True) + monkeypatch.setattr(gn_mod, "_use_compiled", lambda _input, **kw: True) monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _raises) gn_mod._compile_failed = False @@ -397,10 +400,12 @@ def test_triton_failure_falls_back_to_the_compiled_kernel(monkeypatch, caplog): predicate on and making its kernel raise; the compiled stand-in must then be the one that answers, exactly once, with the Triton path latched off. """ + from ScaFFold.unet.triton_group_norm import TritonKernelError + compiled_calls = [] def _raises(*args, **kwargs): - raise RuntimeError("simulated Triton failure") + raise TritonKernelError("simulated Triton failure") def _recording(input, num_groups, weight, bias, eps): compiled_calls.append(type(input)) @@ -408,7 +413,7 @@ def _recording(input, num_groups, weight, bias, eps): monkeypatch.setattr(gn_mod, "_use_triton", lambda *a, **kw: True) monkeypatch.setattr(gn_mod, "_get_triton_module", _raises) - monkeypatch.setattr(gn_mod, "_use_compiled", lambda _input: True) + monkeypatch.setattr(gn_mod, "_use_compiled", lambda _input, **kw: True) monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _recording) gn_mod._triton_failed = False gn_mod._compile_failed = False @@ -455,7 +460,7 @@ def _raises(*args, **kwargs): monkeypatch.setattr(gn_mod, "_use_triton", lambda *a, **kw: True) monkeypatch.setattr(gn_mod, "_get_triton_module", _raises) else: - monkeypatch.setattr(gn_mod, "_use_compiled", lambda _input: True) + monkeypatch.setattr(gn_mod, "_use_compiled", lambda _input, **kw: True) monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _raises) gn_mod._triton_failed = False gn_mod._compile_failed = False @@ -476,7 +481,9 @@ def test_cpu_activation_checkpointing_keeps_the_fast_path(monkeypatch): non-checkpointed run and the fast path must still be live afterwards -- a swallowed recompute-stop shows up as a latched-off kernel here. """ - monkeypatch.setattr(gn_mod, "_use_compiled", lambda t: type(t) is torch.Tensor) + monkeypatch.setattr( + gn_mod, "_use_compiled", lambda t, **kw: type(t) is torch.Tensor + ) monkeypatch.setattr( gn_mod, "_get_compiled_group_norm", lambda: nn.functional.group_norm ) @@ -735,7 +742,9 @@ def _recording(input, num_groups, weight, bias, eps): seen.append(type(input)) return nn.functional.group_norm(input, num_groups, weight, bias, eps) - monkeypatch.setattr(gn_mod, "_use_compiled", lambda t: type(t) is torch.Tensor) + monkeypatch.setattr( + gn_mod, "_use_compiled", lambda t, **kw: type(t) is torch.Tensor + ) monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _recording) fast = _seeded_norm() @@ -760,7 +769,9 @@ def test_dctensor_gradients_reach_the_layer_upstream(monkeypatch, dc_cpu): while GroupNorm's own weight/bias still look healthy. """ distconv, ps = dc_cpu - monkeypatch.setattr(gn_mod, "_use_compiled", lambda t: type(t) is torch.Tensor) + monkeypatch.setattr( + gn_mod, "_use_compiled", lambda t, **kw: type(t) is torch.Tensor + ) monkeypatch.setattr( gn_mod, "_get_compiled_group_norm", lambda: nn.functional.group_norm ) @@ -827,11 +838,14 @@ def _boom(*a, **kw): def test_dctensor_compile_failure_falls_back_to_eager(monkeypatch, caplog, dc_cpu): """A broken compiler degrades the wrapped path to eager, like the plain one.""" distconv, ps = dc_cpu + import torch._dynamo.exc def _raises(*args, **kwargs): - raise RuntimeError("simulated Inductor failure") + raise torch._dynamo.exc.Unsupported("simulated Inductor failure") - monkeypatch.setattr(gn_mod, "_use_compiled", lambda t: type(t) is torch.Tensor) + monkeypatch.setattr( + gn_mod, "_use_compiled", lambda t, **kw: type(t) is torch.Tensor + ) monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _raises) gn_mod._compile_failed = False @@ -1131,10 +1145,15 @@ def test_gpu_triton_matches_eager(activation, autocast): ``(1, 64, 32^3)`` channels-last is the production shape family at unit-test size. Three claims at once: the routing really picks Triton when nothing is - forced (the output comes back channels-last, which is the one thing *only* - that path does -- eager and Inductor both return contiguous); the values and - gradients match the eager reference within reduction-order noise; and the - fused activation equals an explicit ReLU on the eager result. + forced; the values and gradients match the eager reference within + reduction-order noise; and the fused activation equals an explicit ReLU on + the eager result. + + The rung is established by spying on the entry point rather than by + inspecting the output's layout: every rung now returns the *input's* memory + format, deliberately (a fallback that returned contiguous re-broke the + channels-last chain for every convolution after it), so layout no longer + distinguishes them. The layout is asserted separately, of both. """ device = torch.device("cuda") generator = torch.Generator(device=device).manual_seed(11) @@ -1148,11 +1167,15 @@ def test_gpu_triton_matches_eager(activation, autocast): fast.weight.normal_(1.0, 0.1, generator=generator) fast.bias.normal_(0.0, 0.1, generator=generator) + calls = [] + original_triton_forward = FastGroupNorm._triton_forward + def run(triton): gn_mod.set_triton_enabled(triton) gn_mod.set_compile_enabled(False) # eager reference, not Inductor inp = x.clone().requires_grad_(True) fast.zero_grad(set_to_none=True) + before = len(calls) with torch.autocast("cuda", dtype=torch.bfloat16, enabled=autocast): out = fast(inp) out.backward(grad_out.to(out.dtype)) @@ -1161,16 +1184,28 @@ def run(triton): inp.grad.detach().clone(), fast.weight.grad.detach().clone(), fast.bias.grad.detach().clone(), + len(calls) - before, ) - eager = run(False) - triton = run(None) # None = the production default, i.e. no override at all + def spy(self, local): + calls.append(tuple(local.shape)) + return original_triton_forward(self, local) + + FastGroupNorm._triton_forward = spy + try: + eager = run(False) + triton = run(None) # None = the production default, no override at all + finally: + FastGroupNorm._triton_forward = original_triton_forward assert not gn_mod._triton_failed - assert _channels_last(triton[0]), "Triton path was not taken (output not NDHWC)" - # ... and the control: stock GroupNorm really does return contiguous here, - # so the assertion above is a discriminating signal and not a tautology. - assert not _channels_last(eager[0]) + assert triton[4] == 1, "Triton path was not taken" + # ... and the control: the reference really did *not* take it, so the + # comparison below is between two kernels and not one kernel with itself. + assert eager[4] == 0 + # Both preserve the input's channels-last layout; that is the contract now, + # not a rung signature. + assert _channels_last(triton[0]) and _channels_last(eager[0]) _assert_close(triton[0], eager[0], 1e-5, "output") _assert_close(triton[1], eager[1], 1e-4, "d_input") _assert_close(triton[2], eager[2], 1e-4, "d_weight") diff --git a/tests/test_groupnorm_wiring_edge.py b/tests/test_groupnorm_wiring_edge.py new file mode 100644 index 0000000..3d80536 --- /dev/null +++ b/tests/test_groupnorm_wiring_edge.py @@ -0,0 +1,1100 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +"""Edge cases of the GroupNorm *wiring* (``FastGroupNorm``'s three-rung ladder). + +Written as an adversarial review of commit ``ca58664``; ``tests/test_groupnorm.py`` +covers the happy paths and the routing predicates, this file covers the places +where the ladder, the latches and the absorbed ReLU interact with the rest of +torch. See ``review/gn-dctensor/WIRING_REVIEW.md`` for the full write-up. + +The review left ten of these as ``xfail(strict=True)``, one per defect, each +asserting the behaviour the module *should* have. All ten are fixed and the +markers are gone; the tests stay, now as regression guards. The properties +they pin, in the order the defects were found: + +* the fused activation is bit-identical to ``F.relu`` on NaN, the infinities + and ``-0.0``, forward and backward, on all three rungs -- and a NaN produced + under it still reaches the trainer's non-finite-loss abort; +* a latch may not change the rung a checkpointed block is *recomputed* on; +* the ladder catches "the kernel is broken" and nothing else -- not the + checkpoint machinery's control flow, not a user's saved-tensor hook, not an + exception from after the rung already saved something; +* ``torch.func`` is a routing question, not a kernel failure; +* a latch can be cleared; +* ``activation`` is validated where it is used, and a module pickled before it + existed still runs. +""" + +from __future__ import annotations + +import io + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ScaFFold.unet import group_norm as gn_mod +from ScaFFold.unet.group_norm import FastGroupNorm +from ScaFFold.unet.unet_model import UNet + +_GROUPS = 8 + + +@pytest.fixture(autouse=True) +def _restore_routing_state(): + """Keep per-test overrides of the module-level routing state contained. + + Same contract as ``tests/test_groupnorm.py``'s fixture: several tests here + deliberately trip a latch, which is a process global. + """ + previous_compile = gn_mod.set_compile_enabled(None) + previous_triton = gn_mod.set_triton_enabled(None) + compile_failed = gn_mod._compile_failed + triton_failed = gn_mod._triton_failed + yield + gn_mod._compile_override = previous_compile + gn_mod._triton_override = previous_triton + gn_mod._compile_failed = compile_failed + gn_mod._triton_failed = triton_failed + + +def _cl(t): + return t.is_contiguous(memory_format=torch.channels_last_3d) + + +def _cuda_norm(channels=64, activation="relu", size=16, seed=11): + """A seeded ``FastGroupNorm`` plus a channels-last CUDA input for it.""" + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(seed) + x = torch.randn(1, channels, size, size, size, device=device, generator=generator) + x = x.to(memory_format=torch.channels_last_3d) + module = FastGroupNorm(_GROUPS, channels, activation=activation).to(device) + with torch.no_grad(): + module.weight.normal_(1.0, 0.1, generator=generator) + module.bias.normal_(0.0, 0.1, generator=generator) + return module, x + + +def _small_unet(device="cpu", channels_last=False): + torch.manual_seed(0) + model = UNet( + n_channels=3, n_classes=2, trilinear=False, layers=1, group_norm_groups=_GROUPS + ) + if channels_last: + return model.to(device, memory_format=torch.channels_last_3d) + return model.to(device) + + +def _triton_spy(monkeypatch): + """Record every call that actually reached the Triton rung.""" + calls = [] + original = FastGroupNorm._triton_forward + + def spy(self, local): + calls.append(tuple(local.shape)) + return original(self, local) + + monkeypatch.setattr(FastGroupNorm, "_triton_forward", spy) + return calls + + +# --------------------------------------------------------------------------- +# activation semantics: every rung must apply the same function +# --------------------------------------------------------------------------- + + +def test_activate_handles_every_supported_activation(): + """``_activate`` must implement every activation the module advertises. + + ``SUPPORTED_ACTIVATIONS`` is what the *constructor* accepts and what + ``is_supported`` is asked about, but the compiled and eager rungs apply it + through ``_activate``, which tests one literal string. Adding a second + activation to both tuples (the only thing + ``test_supported_activations_match_the_kernels`` checks) would fuse it into + the Triton store and silently drop it everywhere else -- i.e. the network's + function would depend on the memory format of its input. This is the guard + on that: for every non-``None`` activation, ``_activate`` has to *change* + an input that the identity would leave alone. + """ + x = torch.linspace(-2.0, 2.0, 64).reshape(1, 8, 2, 2, 2) + for activation in gn_mod.SUPPORTED_ACTIVATIONS: + module = FastGroupNorm(_GROUPS, 8, activation=activation) + out = module._activate(x.clone()) + if activation is None: + assert torch.equal(out, x) + else: + assert not torch.equal(out, x), ( + f"_activate is a no-op for activation={activation!r}: the " + "compiled and eager rungs would silently skip it while the " + "Triton rung fused it in" + ) + + +def test_affine_false_still_applies_the_activation(): + """``affine=False`` leaves ``weight``/``bias`` ``None`` on every rung.""" + module = FastGroupNorm(_GROUPS, 16, affine=False, activation="relu") + x = torch.randn(1, 16, 4, 4, 4, generator=torch.Generator().manual_seed(5)) + assert torch.equal( + module(x), F.relu(F.group_norm(x, _GROUPS, None, None, module.eps)) + ) + assert list(module.state_dict().keys()) == [] + + +def test_activation_is_not_part_of_the_state(): + """``activation`` may not become a parameter, a buffer or a state-dict key.""" + with_relu = FastGroupNorm(_GROUPS, 16, activation="relu") + without = FastGroupNorm(_GROUPS, 16) + assert list(with_relu.state_dict().keys()) == list(without.state_dict().keys()) + assert list(with_relu.buffers()) == [] + # ... and a checkpoint written by one loads into the other, strict. + result = without.load_state_dict(with_relu.state_dict(), strict=True) + assert not result.missing_keys and not result.unexpected_keys + + +def test_double_conv_bytes_match_a_hand_built_pre_fusion_block(): + """Byte-for-byte state-dict identity against an independently built block. + + ``test_state_dict_bytes_identical_to_plain_groupnorm_model`` compares against + a model produced by *converting* the fused one, which shares its + construction order by definition. This builds the pre-fusion + ``nn.Sequential`` from scratch -- ``Conv3d, GroupNorm, ReLU, Conv3d, + GroupNorm, ReLU`` -- and compares the serialized bytes, which is the + independent version of the same claim. + """ + from ScaFFold.unet.unet_parts import DoubleConv + + torch.manual_seed(17) + fused = DoubleConv(3, 16, _GROUPS) + + torch.manual_seed(17) + reference = nn.Sequential( + nn.Conv3d(3, 16, kernel_size=3, padding=1, bias=False), + nn.GroupNorm(_GROUPS, 16), + nn.ReLU(inplace=True), + nn.Conv3d(16, 16, kernel_size=3, padding=1, bias=False), + nn.GroupNorm(_GROUPS, 16), + nn.ReLU(inplace=True), + ) + + def blob(state_dict): + buffer = io.BytesIO() + torch.save(state_dict, buffer) + return buffer.getvalue() + + assert list(fused.double_conv.state_dict().keys()) == list( + reference.state_dict().keys() + ) + assert blob(fused.double_conv.state_dict()) == blob(reference.state_dict()) + + +def test_module_pickled_before_the_fusion_still_runs(): + """A whole-module pickle predates ``self.activation``; forward must cope. + + ``nn.Module.__setstate__`` replaces ``__dict__`` wholesale, so an instance + restored from a ``torch.save(model)`` written before the fusion has no + ``activation`` at all -- and none of the routing state added since either. + Every attribute this module reads outside ``__init__`` therefore needs a + class-level default. + """ + module = FastGroupNorm(_GROUPS, 16, activation="relu") + state = module.__dict__.copy() + for added_since in ("activation", "_triton_ok", "_compiled_ok"): + state.pop(added_since, None) # exactly what a pre-fusion pickle carries + + restored = FastGroupNorm.__new__(FastGroupNorm) + nn.Module.__setstate__(restored, state) + + x = torch.randn(1, 16, 4, 4, 4, generator=torch.Generator().manual_seed(6)) + out = restored(x) + # A pre-fusion pickle had no activation, so it must behave as one. + assert torch.equal( + out, F.group_norm(x, _GROUPS, restored.weight, restored.bias, restored.eps) + ) + + +def test_unsupported_activation_assigned_after_construction_is_caught(): + """``activation`` is validated where it is *used*, not only at construction. + + It is a plain attribute, so it can be assigned afterwards; ``is_supported`` + would then decline the Triton rung while ``_activate`` silently applied + nothing, i.e. the module would quietly become a bare GroupNorm. The same + hole is the forward-looking risk: adding a third entry to both + ``SUPPORTED_ACTIVATIONS`` tuples without implementing it in ``_activate`` + must not produce a network whose activation depends on its input's memory + format. Failing loudly on the rung that cannot apply it closes both. + """ + module = FastGroupNorm(_GROUPS, 16, activation="relu") + module.activation = "gelu" + x = torch.randn(1, 16, 4, 4, 4, generator=torch.Generator().manual_seed(7)) + with pytest.raises(ValueError, match="activation must be one of"): + module(x) + + +# --------------------------------------------------------------------------- +# the ladder must not swallow torch's own control flow +# --------------------------------------------------------------------------- + + +def test_base_exceptions_are_not_caught(): + """``KeyboardInterrupt``/``SystemExit`` must escape the ladder untouched.""" + for exception in (KeyboardInterrupt, SystemExit): + + def _raises(*args, **kwargs): + raise exception() + + previous = gn_mod._get_triton_module + gn_mod._get_triton_module = _raises + original_use = gn_mod._use_triton + gn_mod._use_triton = lambda *a, **kw: True + gn_mod._triton_failed = False + try: + module = FastGroupNorm(_GROUPS, 16) + with pytest.raises(exception): + module(torch.randn(1, 16, 4, 4, 4)) + assert gn_mod._triton_failed is False + finally: + gn_mod._get_triton_module = previous + gn_mod._use_triton = original_use + + +@pytest.mark.parametrize("rung", ["triton", "compiled"]) +def test_checkpoint_error_is_re_raised(monkeypatch, rung): + """``CheckpointError`` is the checkpoint machinery talking, not a kernel. + + It is raised by the recompute pack hook -- i.e. from inside whichever op is + saving a tensor -- exactly like ``_StopRecomputationError``, and it is a + ``RuntimeError`` subclass, so any handler wide enough to catch "a broken + kernel" by type catches it too. Swallowing it latches the rung off, retries + on the next one and leaves the checkpoint frame in a state the machinery + never expected. The allowlist has to be narrow enough that this propagates + untouched and nothing latches. + """ + import torch.utils.checkpoint as checkpoint_mod + + def _raises(*args, **kwargs): + raise checkpoint_mod.CheckpointError("simulated recompute mismatch") + + if rung == "triton": + monkeypatch.setattr(gn_mod, "_use_triton", lambda *a, **kw: True) + monkeypatch.setattr(gn_mod, "_get_triton_module", _raises) + else: + monkeypatch.setattr(gn_mod, "_use_compiled", lambda _input, **kw: True) + monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _raises) + gn_mod._triton_failed = False + gn_mod._compile_failed = False + + module = FastGroupNorm(_GROUPS, 16) + with pytest.raises(checkpoint_mod.CheckpointError): + module(torch.randn(1, 16, 4, 4, 4)) + assert gn_mod._triton_failed is False + assert gn_mod._compile_failed is False + + +def test_a_rung_failure_does_not_re_fire_saved_tensor_hooks(monkeypatch): + """The retry has to be idempotent with respect to saved-tensor hooks. + + Under non-reentrant activation checkpointing the recompute counts pack-hook + firings and requires the count *and* the metadata to match the forward's, so + a rung that packed some tensors and then failed -- with the fallback packing + its own set on top -- corrupts the frame. A user's offloading hook has the + same problem in a less dramatic way (an offload failure retried by + offloading a second, larger set). + + The property is structural rather than defensive: the ladder catches only + failures that are raised *before* their rung saves anything (the Triton op + saves in ``_setup_context``, after the launch region its ``TritonKernelError`` + comes from; a Dynamo/Inductor error is a compile-time error, before + execution). Both halves are asserted here -- a caught failure packs exactly + what a clean fallback packs, and an exception from *after* the packing is + not the ladder's to swallow. + """ + import torch._dynamo.exc + + class _Boom(Exception): + pass + + def _fails_before_packing(input, num_groups, weight, bias, eps): + raise torch._dynamo.exc.Unsupported("failed while compiling") + + def _fails_after_packing(input, num_groups, weight, bias, eps): + F.group_norm(input, num_groups, weight, bias, eps) + raise _Boom("failed after packing") + + monkeypatch.setattr(gn_mod, "_use_compiled", lambda _input, **kw: True) + module = FastGroupNorm(_GROUPS, 16) + x = torch.randn(1, 16, 4, 4, 4).requires_grad_(True) + + def packs_during(run): + packed = [] + with torch.autograd.graph.saved_tensors_hooks( + lambda t: (packed.append(1), t)[1], lambda t: t + ): + run() + return len(packed) + + baseline = packs_during(lambda: module._eager_forward(x)) + + monkeypatch.setattr( + gn_mod, "_get_compiled_group_norm", lambda: _fails_before_packing + ) + gn_mod._compile_failed = False + retried = packs_during(lambda: module(x)) + assert retried == baseline, ( + f"the failed rung packed {retried - baseline} extra tensors before the " + "fallback ran" + ) + + # ... and a failure that *did* have observable effects is not retried at all. + monkeypatch.setattr( + gn_mod, "_get_compiled_group_norm", lambda: _fails_after_packing + ) + gn_mod._compile_failed = False + with pytest.raises(_Boom): + module(x) + + +# --------------------------------------------------------------------------- +# latches +# --------------------------------------------------------------------------- + + +def test_forcing_a_rung_on_clears_its_failure_latch(): + """``set_*_enabled(True)`` is the documented way to retry after a failure. + + Without this a one-off failure (a transient OOM, a cache-directory hiccup) + costs the rung for the rest of the process with no recovery at all, and the + function's own docstring -- "forcing it on is overridden only by the + correctness checks" -- is false. ``None`` deliberately does *not* clear it: + that restores a preference, it does not assert that the kernel works again. + """ + for setter, latch in ( + (gn_mod.set_triton_enabled, "_triton_failed"), + (gn_mod.set_compile_enabled, "_compile_failed"), + ): + setattr(gn_mod, latch, True) + setter(True) + assert getattr(gn_mod, latch) is False + + setattr(gn_mod, latch, True) + setter(None) + assert getattr(gn_mod, latch) is True + setter(False) + assert getattr(gn_mod, latch) is True + setattr(gn_mod, latch, False) + + +@pytest.mark.parametrize("rung", ["triton", "compiled"]) +def test_out_of_memory_is_not_recorded_as_a_kernel_failure(monkeypatch, rung): + """A transient OOM must propagate, and must not latch a rung off forever. + + ``torch.OutOfMemoryError`` is a resource condition, not a defect: every + fallback allocates an output of the same size, so retrying one is a second, + differently-shaped OOM at a call site the caller never asked about. + Latching on it is worse still -- a per-rank, nondeterministic event that + permanently changes which kernel that rank runs, and therefore (measured) + the all-reduced gradients of the whole job. + """ + from ScaFFold.unet.triton_group_norm import TritonKernelError + + def _oom(*args, **kwargs): + raise torch.OutOfMemoryError("simulated OOM") + + if rung == "triton": + monkeypatch.setattr(gn_mod, "_use_triton", lambda *a, **kw: True) + monkeypatch.setattr(gn_mod, "_get_triton_module", _oom) + else: + monkeypatch.setattr(gn_mod, "_use_compiled", lambda _input, **kw: True) + monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _oom) + gn_mod._triton_failed = False + gn_mod._compile_failed = False + + module = FastGroupNorm(_GROUPS, 16) + with pytest.raises(torch.OutOfMemoryError): + module(torch.randn(1, 16, 4, 4, 4)) + assert gn_mod._triton_failed is False + assert gn_mod._compile_failed is False + # It is a RuntimeError, so a handler that caught the kernel's own error by + # base class would have swallowed it; the allowlist is the tagged type. + assert issubclass(torch.OutOfMemoryError, RuntimeError) + assert not issubclass(torch.OutOfMemoryError, TritonKernelError) + + +def test_a_global_latch_does_not_demote_a_module_that_already_used_the_rung( + monkeypatch, caplog +): + """The unit of the latch is the module, not the process. + + A rung that has already served a module keeps serving it; only modules that + have never used it are steered away. That is what makes a checkpointed + block's forward and its recompute agree (they save different tensors on + different rungs, so a mid-graph change is fatal), and it is why a broken + install still costs exactly one attempt per module rather than one per call. + + The corollary tested here too: because a proven module keeps retrying, the + warning has to be emitted on the latch's edge rather than per call, or a + persistently broken kernel floods the log for the rest of the run. + """ + import logging + + import torch._dynamo.exc + + calls = [] + + def _kernel(input, num_groups, weight, bias, eps): + calls.append(1) + if len(calls) > 1: + raise torch._dynamo.exc.Unsupported("simulated Inductor failure") + return F.group_norm(input, num_groups, weight, bias, eps) + + monkeypatch.setattr(gn_mod, "_use_compiled", gn_mod._use_compiled) + monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _kernel) + monkeypatch.setattr( + gn_mod, + "_use_compiled", + lambda t, proven=False, **kw: (not gn_mod._compile_failed) or proven, + ) + gn_mod._compile_failed = False + + proven = FastGroupNorm(_GROUPS, 16) + fresh = FastGroupNorm(_GROUPS, 16) + x = torch.randn(1, 16, 4, 4, 4) + + proven(x) # succeeds: this module is now proven on the compiled rung + assert proven._compiled_ok is True + + with caplog.at_level(logging.WARNING, logger=gn_mod.__name__): + proven(x) # fails, latches, falls back to eager + assert gn_mod._compile_failed is True + warned = sum("falling back" in r.message for r in caplog.records) + proven(x) # ... and still *tries* the rung, because it is proven + proven(x) + assert sum("falling back" in r.message for r in caplog.records) == warned, ( + "a persistently failing rung warned once per call" + ) + assert len(calls) == 4, "the proven module stopped trying its rung" + + fresh(x) # never used it, so the global latch keeps it away entirely + assert len(calls) == 4 + assert fresh._compiled_ok is False + + +def test_a_latch_flip_mid_forward_is_not_a_numerics_error(monkeypatch): + """Two rungs inside one forward still compose (values, not bits, agree).""" + monkeypatch.setattr( + gn_mod, "_use_compiled", lambda t, **kw: type(t) is torch.Tensor + ) + monkeypatch.setattr( + gn_mod, "_get_compiled_group_norm", lambda: nn.functional.group_norm + ) + gn_mod._compile_failed = False + + module = FastGroupNorm(_GROUPS, 16, activation="relu") + x = torch.randn(1, 16, 4, 4, 4).requires_grad_(True) + first = module(x) + gn_mod._compile_failed = True # latch flips between the two calls + second = module(first) + second.pow(2).sum().backward() + assert torch.isfinite(x.grad).all() + + +# --------------------------------------------------------------------------- +# predicates +# --------------------------------------------------------------------------- + + +def test_predicates_reject_a_parameter_input(): + """``nn.Parameter`` is a subclass, so both fast rungs decline it. + + Not a bug -- the model never feeds a Parameter to a norm -- but it is the + documented consequence of the ``type(input) is torch.Tensor`` policy, and a + regression that loosened it to ``isinstance`` would route real + ``__torch_dispatch__`` wrappers into the kernel. + """ + parameter = nn.Parameter(torch.randn(1, 8, 4, 4, 4)) + assert gn_mod._use_triton(parameter, _GROUPS, None, None, None) is False + assert gn_mod._use_compiled(parameter) is False + + +def test_use_triton_is_side_effect_free_for_rejected_inputs(monkeypatch): + """The predicate may look, but it may not allocate, launch or mutate.""" + module = FastGroupNorm(_GROUPS, 16) + x = torch.randn(1, 16, 4, 4, 4) + before = x.clone() + assert gn_mod._use_triton(x, _GROUPS, module.weight, module.bias, None) is False + assert torch.equal(x, before) + + +# --------------------------------------------------------------------------- +# GPU: the Triton rung in situ +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +def test_gpu_double_backward_fails_loudly(): + """The kernel is first-order only; a second derivative must *raise*. + + ``triton_group_norm``'s backward is itself a custom op with no autograd + formula, so a gradient penalty or an HVP through the wired model has to die + with a clear message rather than silently return a wrong number -- and the + failure must not be mistaken for a broken kernel and latch the rung off. + """ + module, x = _cuda_norm(activation=None) + x = x.clone().requires_grad_(True) + out = module(x) + assert _cl(out), "Triton rung was not taken; the test would be vacuous" + + (first,) = torch.autograd.grad(out.pow(2).sum(), x, create_graph=True) + with pytest.raises(RuntimeError, match="no autograd formula"): + torch.autograd.grad(first.pow(2).sum(), x) + assert gn_mod._triton_failed is False, "a caller error latched the kernel off" + + +@pytest.mark.gpu +def test_gpu_model_double_backward_fails_loudly(): + """Same, through the wired model, which is where a user would hit it.""" + model = _small_unet("cuda", channels_last=True) + x = ( + torch.randn(1, 3, 16, 16, 16, device="cuda") + .contiguous(memory_format=torch.channels_last_3d) + .requires_grad_(True) + ) + gn_mod.set_triton_enabled(None) + out = model(x) + (first,) = torch.autograd.grad(out.pow(2).sum(), x, create_graph=True) + with pytest.raises(RuntimeError, match="no autograd formula"): + first.pow(2).sum().backward() + + +@pytest.mark.gpu +@pytest.mark.parametrize("fullgraph", [True, False]) +def test_gpu_triton_rung_inside_a_compiled_region(monkeypatch, fullgraph): + """The Triton rung is *allowed* inside ``torch.compile``; prove it works. + + ``_use_compiled`` bails out when ``torch.compiler.is_compiling()`` so the + functional GroupNorm inlines, but ``_use_triton`` has no such guard: an + enclosing compiled region traces straight into the custom op. Nothing in + ScaFFold compiles ``FastGroupNorm.forward`` today, so this was untested in + situ. Values, gradients *and* the channels-last output must survive + Dynamo/AOTAutograd unchanged. + """ + import torch._dynamo + + module, x = _cuda_norm(activation="relu") + reference_input = x.clone().requires_grad_(True) + reference = module(reference_input) + reference.pow(2).sum().backward() + reference_grad = reference_input.grad.detach().clone() + module.zero_grad(set_to_none=True) + assert _cl(reference), "Triton rung was not taken; the test would be vacuous" + + torch._dynamo.reset() + calls = _triton_spy(monkeypatch) + compiled = torch.compile(lambda t: module(t), fullgraph=fullgraph, dynamic=False) + + compiled_input = x.clone().requires_grad_(True) + out = compiled(compiled_input) + out.pow(2).sum().backward() + + assert calls, "the Triton rung was not traced inside the compiled region" + assert _cl(out), "the compiled region lost the channels-last output" + assert torch.equal(out, reference) + assert torch.equal(compiled_input.grad, reference_grad) + + +@pytest.mark.gpu +def test_gpu_the_fallback_path_traces_under_fullgraph(monkeypatch): + """A rung failure *while Dynamo is tracing* must still fall back, not die. + + The handler used to call ``logger.warning``, which Dynamo cannot trace + ("Unsupported: logging.Logger method not supported for non-export cases"), + so a caller compiling this forward with ``fullgraph=True`` got a hard error + instead of the fallback -- the one caller for whom the fallback matters + most, since the thing it is reacting to is usually a compile-time failure. + Nothing in ScaFFold compiles ``FastGroupNorm.forward`` today; this pins the + claim that it can. + """ + import torch._dynamo + + from ScaFFold.unet.triton_group_norm import TritonKernelError + + module, x = _cuda_norm(activation="relu") + reference = module(x).detach().clone() + assert _cl(reference), "Triton rung was not taken; the test would be vacuous" + + real = gn_mod._get_triton_module() + + class _BrokenKernelModule: + def __getattr__(self, name): + if name == "triton_group_norm": + + def _raises(*args, **kwargs): + raise TritonKernelError("simulated Triton failure") + + return _raises + return getattr(real, name) + + monkeypatch.setattr(gn_mod, "_get_triton_module", _BrokenKernelModule) + gn_mod._triton_failed = False + module._triton_ok = False + torch._dynamo.reset() + + compiled = torch.compile(lambda t: module(t), fullgraph=True, dynamic=False) + out = compiled(x) + + assert gn_mod._triton_failed is True, "the latch was not recorded" + assert _cl(out), "the fallback rung dropped the channels-last chain" + assert (out - reference).abs().max().item() < 1e-5 + + +@pytest.mark.gpu +def test_gpu_inference_mode_takes_the_triton_rung(monkeypatch): + """``evaluate()`` runs the whole model under ``torch.inference_mode``.""" + module, x = _cuda_norm(activation="relu") + module.eval() + reference = F.relu(F.group_norm(x, _GROUPS, module.weight, module.bias, module.eps)) + calls = _triton_spy(monkeypatch) + with torch.inference_mode(): + out = module(x) + assert calls, "inference_mode fell off the Triton rung" + assert _cl(out) + assert out.is_inference() + assert (out.float() - reference.float()).abs().max().item() < 1e-5 + + +@pytest.mark.gpu +def test_gpu_evaluation_shaped_forward_matches_training_shaped_one(monkeypatch): + """``eval()`` + ``inference_mode`` + autocast is the evaluate() combination.""" + model = _small_unet("cuda", channels_last=True) + x = torch.randn(1, 3, 16, 16, 16, device="cuda").contiguous( + memory_format=torch.channels_last_3d + ) + gn_mod.set_triton_enabled(None) + model.eval() + calls = _triton_spy(monkeypatch) + with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16): + out = model(x) + if not calls: + pytest.skip( + "convolutions did not emit channels_last_3d; set " + "PYTORCH_MIOPEN_SUGGEST_NHWC=1 (the production setting)" + ) + assert torch.isfinite(out.float()).all() + assert gn_mod._triton_failed is False + + +@pytest.mark.gpu +def test_gpu_vmap_over_the_module_still_works(): + """``torch.func`` is a routing question, not a kernel failure. + + ``is_supported``'s ``is_contiguous(memory_format=...)`` raises outright + under a ``vmap`` layer ("NYI"), so a predicate that reaches it -- or a + relayout helper that does -- turns a plain ``nn.GroupNorm`` drop-in into a + hard error for any caller using ``torch.func``. Both fast rungs must + decline while a transform is active and let the stock kernel answer. + """ + module, x = _cuda_norm(activation=None) + batched = torch.stack([x[0], x[0]]) + out = torch.func.vmap(lambda t: module(t.unsqueeze(0)).squeeze(0))(batched) + assert out.shape == batched.shape + assert torch.allclose(out[0], module(x[None, 0]).squeeze(0), atol=1e-5) + + +@pytest.mark.gpu +def test_gpu_a_predicate_that_cannot_answer_falls_back_without_latching( + monkeypatch, caplog +): + """``is_supported`` raising is a routing miss, and "no" is always a valid answer. + + The predicate runs *outside* the ladder's try, so anything it raises escapes + ``forward()`` -- which is how a ``torch.func`` transform used to turn a + drop-in ``nn.GroupNorm`` into a hard error. The functorch check upstream + covers the one caller known to trip it; this covers the shape of the + problem, because ``is_supported`` inspects an *arbitrary* tensor and the set + of wrappers that can make an attribute read raise is not closed. A broad + catch is right here and nowhere else in this module: the predicate has done + no work anyone can observe and a correct answer ("use the stock kernel") is + always available -- so it must fall back, and must not latch, because + nothing about the kernel has been learned. + """ + import logging + + class _Unanswerable: + def __getattr__(self, name): + if name == "is_supported": + + def _raises(*args, **kwargs): + raise RuntimeError("NYI: querying is_contiguous inside of vmap") + + return _raises + return getattr(gn_mod._get_triton_module(), name) + + module, x = _cuda_norm(activation="relu") + reference = module(x).detach().clone() + monkeypatch.setattr(gn_mod, "_get_triton_module", _Unanswerable) + monkeypatch.setattr(gn_mod, "_predicate_warned", False) + gn_mod._triton_failed = False + module._triton_ok = False + + with caplog.at_level(logging.WARNING, logger=gn_mod.__name__): + out = module(x) + assert (out - reference).abs().max().item() < 1e-5 + assert gn_mod._triton_failed is False, "a routing miss latched the kernel off" + assert gn_mod._compile_failed is False + assert any("routing check failed" in r.message for r in caplog.records) + + +@pytest.mark.gpu +def test_gpu_torch_func_grad_does_not_latch_the_rungs_off(): + """A ``torch.func`` call anywhere must not demote the whole process. + + ``torch.func.grad`` used to reach the kernel, fail, and latch *both* fast + rungs off permanently -- i.e. one transform anywhere in a process silently + dropped every GroupNorm in the model to the stock kernel for the rest of + the run. Recording a routing miss as a kernel failure is the general shape + of the bug; this pins the specific instance. + """ + module, x = _cuda_norm(activation=None) + gn_mod._triton_failed = False + gn_mod._compile_failed = False + + grad = torch.func.grad(lambda t: module(t).pow(2).sum())(x) + + assert torch.isfinite(grad).all() + assert gn_mod._triton_failed is False + assert gn_mod._compile_failed is False + # ... and the module is still on the Triton rung afterwards. + assert _cl(module(x)) + + +@pytest.mark.gpu +def test_gpu_a_triton_failure_during_a_checkpointed_step_degrades_not_dies(): + """A latch may not change the rung a checkpointed block is recomputed on. + + Non-reentrant checkpointing compares the metadata of every tensor the + recomputed forward saves against the forward's, and the three rungs save + *different* tensors -- Triton ``(input, weight, bias, mean, rstd)``, the + other two ``(input, weight, mean, rstd, relu_output)``. So a rung change + between a block's forward and its recompute kills the step with a + ``CheckpointError``, which is the exact opposite of the ladder's contract + ("a broken Triton install must degrade a multi-node run, not kill it") and + is reachable whenever the ``activation_checkpointing`` option is on + (``worker.py:230``). Matching the output memory format is *not* enough on + its own -- measured; the saved sets still differ -- so the fix is that a + global latch does not demote a module that has already used the rung. + + The same hazard predates the Triton rung: flipping ``_compile_failed`` + between forward and recompute dies too, which is why both latches are + checked here. + """ + for latch in ("_triton_failed", "_compile_failed"): + model = _small_unet("cuda", channels_last=True) + model.use_checkpointing() + x = ( + torch.randn(1, 3, 16, 16, 16, device="cuda") + .contiguous(memory_format=torch.channels_last_3d) + .requires_grad_(True) + ) + gn_mod.set_triton_enabled(None) + gn_mod._triton_failed = False + gn_mod._compile_failed = False + + out = model(x) + # A one-off failure at any *later* GroupNorm site latches the rung off + # while the blocks already run are waiting to be recomputed. + setattr(gn_mod, latch, True) + out.pow(2).sum().backward() + + assert torch.isfinite(x.grad).all(), latch + + +@pytest.mark.gpu +@pytest.mark.parametrize("channels_last", [True, False]) +def test_gpu_every_rung_returns_the_inputs_memory_format(channels_last): + """All three rungs must agree on the output layout, not just the values. + + ``F.group_norm`` -- eager or Inductor-compiled -- returns a *contiguous* + tensor whatever it was given, so a single fallback used to re-break the + channels-last chain for every convolution after it, which is the exact + thing this module exists to prevent. It also made the rungs distinguishable + to anything that inspects metadata (``torch.utils.checkpoint``, a compiled + caller's guards), which is a correctness problem rather than a speed one. + """ + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(9) + x = torch.randn(1, 64, 8, 8, 8, device=device, generator=generator) + x = x.to(memory_format=torch.channels_last_3d) if channels_last else x.contiguous() + module = FastGroupNorm(_GROUPS, 64, activation="relu").to(device) + + outputs = _run_on_every_rung(module, x) + for label, out in outputs.items(): + assert _cl(out) is channels_last, ( + f"the {label} rung returned " + f"{'channels_last_3d' if _cl(out) else 'contiguous'} for a " + f"{'channels_last_3d' if channels_last else 'contiguous'} input" + ) + reference = outputs["eager"].detach().float() + for label, out in outputs.items(): + assert (out.detach().float() - reference).abs().max().item() < 1e-5, label + + +@pytest.mark.gpu +def test_gpu_triton_rejects_a_cuda_tensor_subclass(): + """The subclass check has to be tested on a tensor that would otherwise pass. + + ``test_triton_rejects_unknown_tensor_subclasses`` hands ``_use_triton`` a + *CPU* subclass, which the ``is_cuda`` check rejects one line later -- so it + cannot tell whether the ``type(input) is torch.Tensor`` test exists at all + (a mutation deleting that line survives the whole suite). The check is + load-bearing: ``is_supported`` only asks ``isinstance``, so without it every + unknown ``__torch_dispatch__`` wrapper would be routed into the kernel. + """ + + class _Wrapper(torch.Tensor): + pass + + x = torch.randn(1, 64, 8, 8, 8, device="cuda").to( + memory_format=torch.channels_last_3d + ) + wrapped = x.as_subclass(_Wrapper) + # The control: everything *except* the subclass test accepts this tensor. + from ScaFFold.unet import triton_group_norm as triton_mod + + assert triton_mod.is_supported(wrapped, _GROUPS, None, None, None) is True + assert gn_mod._use_triton(wrapped, _GROUPS, None, None, None) is False + assert gn_mod._use_compiled(wrapped) is False + + +#: NaN, +Inf, -Inf, -0.0 and four ordinary values -- everything the fused +#: activation has to agree with ``F.relu`` on. ``tl.maximum(y, 0)`` and +#: ``tl.where(y > 0, y, 0)`` both map NaN to 0.0 (the first returns the non-NaN +#: operand, the second because ``NaN > 0`` is False); ``F.relu`` propagates it. +_SPECIAL_VALUES = [float("nan"), float("inf"), float("-inf"), -0.0, 0.0, -1.0, 1.0, 2.0] + + +def _run_on_every_rung(module, x): + """``{rung: output}`` for the same module and input on all three rungs.""" + results = {} + for label, triton, compiled in ( + ("triton", True, False), + ("compiled", False, True), + ("eager", False, False), + ): + gn_mod.set_triton_enabled(triton) + gn_mod.set_compile_enabled(compiled) + results[label] = module(x) + return results + + +def _bits(t): + return t.detach().float().cpu().contiguous().view(torch.int32) + + +@pytest.mark.gpu +@pytest.mark.parametrize("activation", ["relu", None]) +@pytest.mark.parametrize("poison", [float("nan"), float("inf"), float("-inf")]) +def test_gpu_all_rungs_agree_on_nan_and_inf(poison, activation): + """One non-finite input value must poison the same elements on every rung. + + The Triton store used ``tl.maximum(y, 0.0)``, which returns the *non*-NaN + operand, so a diverging activation came back from the Triton rung as a + finite 0.0 while ``F.relu`` on the other two kept it NaN. That is worse + than a numerics discrepancy: the forward looks finite while the backward is + still NaN, so the run sails past ScaFFold's non-finite-loss abort and + checkpoints a broken model -- and the model's output becomes a function of + its input's memory format. ``activation=None`` is the control: all three + agreed there even before the fix, which is what localizes the divergence to + the fused activation. + """ + device = torch.device("cuda") + x = torch.randn(1, 64, 4, 4, 4, device=device) + x.view(-1)[0] = poison + x = x.to(memory_format=torch.channels_last_3d) + module = FastGroupNorm(_GROUPS, 64, activation=activation).to(device) + + results = { + label: out.detach().float().cpu().contiguous().isnan() + for label, out in _run_on_every_rung(module, x).items() + } + assert int(results["eager"].sum()) > 0, "the poison did not reach the output" + assert torch.equal(results["compiled"], results["eager"]) + assert torch.equal(results["triton"], results["eager"]), ( + f"the fused activation turned {int(results['eager'].sum())} NaNs into " + f"{int(results['triton'].sum())}" + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize("activation", ["relu", None]) +def test_gpu_fused_activation_is_bit_identical_to_relu(activation): + """Every special value of the *pre-activation*, on every rung, bit for bit. + + Poisoning the input can only produce NaN pre-activations (one NaN or Inf + makes the whole group's statistics NaN), so the four values that actually + distinguish the spellings of ReLU are reached the other way round: a zero + ``weight`` makes the pre-activation exactly ``bias``, elementwise, so the + bias vector chooses what the activation sees. Expected, per + ``F.relu``: NaN stays NaN, ``+Inf`` stays ``+Inf``, ``-Inf`` and both zeros + become ``+0.0`` (never ``-0.0``). + + With ``activation=None`` the zeros are normalized (``+ 0.0`` maps ``-0.0`` + to ``+0.0`` and leaves NaN, the infinities and every normal value alone) + before the same bitwise comparison: a ``-0.0`` bias survives to the output + there, and whether ``xhat * 0 + (-0.0)`` keeps the sign depends on whether + the kernel contracted the multiply-add into an FMA -- true of the Triton + *and* the Inductor rung, false of eager, and nothing to do with the + activation. ``torch.equal`` is no use for either case: it reports NaN as + unequal to itself. + """ + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(3) + x = torch.randn(1, 64, 4, 4, 4, device=device, generator=generator).to( + memory_format=torch.channels_last_3d + ) + module = FastGroupNorm(_GROUPS, 64, activation=activation).to(device) + with torch.no_grad(): + module.weight.zero_() + module.bias.copy_(torch.tensor(_SPECIAL_VALUES * 8, device=device)) + + reference = F.group_norm(x, _GROUPS, module.weight, module.bias, module.eps) + if activation == "relu": + reference = F.relu(reference) + + for label, out in _run_on_every_rung(module, x).items(): + if activation == "relu": + assert torch.equal(_bits(out), _bits(reference)), ( + f"{label} rung differs from F.relu(F.group_norm(...)) in the " + "bit pattern of at least one special value" + ) + else: + assert torch.equal(_bits(out + 0.0), _bits(reference + 0.0)), label + + +@pytest.mark.gpu +def test_gpu_fused_relu_backward_gates_like_threshold_backward(): + """ReLU's backward passes the gradient where the output is NaN, too. + + ``threshold_backward(grad, result, 0)`` zeroes where ``result <= 0``, and + ``NaN <= 0`` is False -- so a NaN pre-activation passes its gradient. The + kernel recomputes the pre-activation and must gate with the same + complement; ``pre > 0 ? dy : 0`` would silently zero it. + """ + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(4) + base = torch.randn(1, 64, 4, 4, 4, device=device, generator=generator).to( + memory_format=torch.channels_last_3d + ) + module = FastGroupNorm(_GROUPS, 64, activation="relu").to(device) + with torch.no_grad(): + module.weight.zero_() + module.bias.copy_(torch.tensor(_SPECIAL_VALUES * 8, device=device)) + + grads = {} + for label, triton in (("triton", True), ("eager", False)): + gn_mod.set_triton_enabled(triton) + gn_mod.set_compile_enabled(False) + x = base.clone().requires_grad_(True) + module.zero_grad(set_to_none=True) + module(x).sum().backward() + grads[label] = (x.grad.detach().clone(), module.bias.grad.detach().clone()) + + # d_bias is exactly the gate: one per element that passed. + assert torch.equal(_bits(grads["triton"][1]), _bits(grads["eager"][1])) + assert grads["eager"][1][0].item() > 0, "the NaN lane's gradient was gated off" + assert torch.equal( + _bits(grads["triton"][0].float()), _bits(grads["eager"][0].float()) + ) + + +@pytest.mark.gpu +def test_gpu_fused_relu_nan_still_trips_the_trainers_non_finite_guard( + monkeypatch, tiny_trainer +): + """End to end: a NaN produced under the fused path reaches the abort. + + ScaFFold aborts a run whose reduced epoch losses are non-finite, precisely + so a diverged run stops instead of overwriting ``checkpoint_last.pth`` with + NaN weights. A fused activation that ate the NaN would hand that guard a + finite loss and let the run continue on a model whose *gradients* are still + NaN. The loss below is the real one: a real UNet, on the GPU, with the + Triton rung verified to have served every GroupNorm in it. + """ + from ScaFFold.utils import trainer as trainer_mod + + model = _small_unet("cuda", channels_last=True) + poisoned = torch.randn(1, 3, 16, 16, 16, device="cuda") + poisoned.view(-1)[0] = float("nan") + x = poisoned.contiguous(memory_format=torch.channels_last_3d).requires_grad_(True) + gn_mod.set_triton_enabled(None) + calls = _triton_spy(monkeypatch) + loss = model(x).float().mean() + if not calls: + pytest.skip( + "convolutions did not emit channels_last_3d; set " + "PYTORCH_MIOPEN_SUGGEST_NHWC=1 (the production setting)" + ) + assert not torch.isfinite(loss).item(), ( + "the fused activation swallowed the NaN: the forward is finite while " + "the backward is not, which is exactly what hides divergence" + ) + + trainer = tiny_trainer(config_overrides={"checkpoint_interval": 1, "epochs": 3}) + monkeypatch.setattr( + trainer, + "_run_training_batch", + lambda batch, **kw: (1, loss.detach().cpu(), torch.tensor(0.0)), + ) + # A *finite* validation loss, so the abort can only come from the model's. + monkeypatch.setattr( + trainer_mod, "evaluate", lambda *a, **kw: (7.4e-10, 0.5, 0.5, 2, 2) + ) + trainer.cleanup_or_resume() + with pytest.raises(ValueError, match="[Nn]on-finite"): + trainer.train() + assert not trainer.checkpoint_manager.last_ckpt_path.exists() + + +@pytest.mark.gpu +def test_gpu_ddp_wrapped_model_takes_the_triton_rung(monkeypatch): + """DDP's module-tree walk must not be confused by the nn.Identity slots.""" + import torch.distributed as dist + from torch.nn.parallel import DistributedDataParallel + + created = False + if not dist.is_initialized(): + import os + + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29623") + dist.init_process_group(backend="nccl", rank=0, world_size=1) + created = True + try: + model = _small_unet("cuda", channels_last=True) + wrapped = DistributedDataParallel(model, device_ids=[0]) + x = torch.randn(1, 3, 16, 16, 16, device="cuda").contiguous( + memory_format=torch.channels_last_3d + ) + gn_mod.set_triton_enabled(None) + calls = _triton_spy(monkeypatch) + wrapped(x).pow(2).sum().backward() + if not calls: + pytest.skip( + "convolutions did not emit channels_last_3d; set " + "PYTORCH_MIOPEN_SUGGEST_NHWC=1 (the production setting)" + ) + assert all(torch.isfinite(p.grad).all() for p in model.parameters()) + finally: + if created and dist.is_initialized(): + dist.destroy_process_group() diff --git a/tests/test_triton_group_norm_edge.py b/tests/test_triton_group_norm_edge.py index be88cd4..5998c5f 100644 --- a/tests/test_triton_group_norm_edge.py +++ b/tests/test_triton_group_norm_edge.py @@ -1520,3 +1520,182 @@ def test_dweight_blocks_are_covered_when_there_are_more_of_them_than_tiles( ) assert plan.grid_dx == plan.dwdb_progs _parity(shape, groups, seed=13) + + +# --------------------------------------------------------------------------- +# the kernel-failure boundary +# --------------------------------------------------------------------------- + + +def test_kernel_failures_are_tagged_and_carry_their_cause(): + """Everything the launch region raises comes out as ``TritonKernelError``. + + The tag is what lets a caller with a fallback (``FastGroupNorm``'s ladder) + catch *exactly* "the kernel is broken" instead of catching ``Exception`` and + then trying to enumerate every framework mechanism -- saved-tensor pack + hooks, ``torch.utils.checkpoint``'s recompute control flow, functorch -- + that legitimately raises through a forward. The region it wraps is closed + (allocations and launches, no autograd-observable op), so a blanket catch + inside it is sound where one at the call site is not. + + The tag must survive the *type* of the original error, whatever it was: a + mismatched Triton release raises ``TypeError``/``AttributeError`` from a + changed signature, an unwritable JIT cache ``OSError``, a bad launch + ``RuntimeError``. + """ + for original in ( + RuntimeError("launch failed"), + TypeError("triton API changed"), + AttributeError("no such attribute"), + OSError("unwritable cache dir"), + ImportError("no module named triton"), + ): + + @tgn._tag_kernel_failures + def _boom(): + raise original + + with pytest.raises(tgn.TritonKernelError) as caught: + _boom() + assert caught.value.__cause__ is original + assert type(original).__name__ in str(caught.value) + + +def test_out_of_memory_is_not_tagged_as_a_kernel_failure(): + """An OOM is a resource condition, and every fallback allocates as much. + + Tagging it would make the ladder retry on a rung that is about to OOM in + the same place, and would latch a rung off for the rest of the process on a + transient, per-rank event. It has to come out unchanged. + """ + + @tgn._tag_kernel_failures + def _oom(): + raise torch.OutOfMemoryError("simulated OOM") + + with pytest.raises(torch.OutOfMemoryError): + _oom() + assert not issubclass(torch.OutOfMemoryError, tgn.TritonKernelError) + + +def test_contract_violations_are_not_tagged(): + """``_validate``'s ``ValueError``s are caller errors, and stay loud. + + ``is_supported`` accepts exactly what ``_validate`` accepts, so a caller + that branches on the predicate can never see one; if the two ever disagree, + the failure must not be laundered into "the kernel is broken" and silently + fall back. + """ + with pytest.raises(ValueError, match="activation must be one of"): + tgn._validate(torch.zeros(1, 8, 2, 2, 2), 8, None, None, "gelu") + with pytest.raises(ValueError, match="expected a 5-D"): + tgn._validate(torch.zeros(1, 8, 2, 2), 8, None, None, None) + + +@pytest.mark.gpu +def test_a_real_launch_failure_is_tagged(monkeypatch): + """End to end: break the launch and the public op raises the tagged type.""" + x = torch.randn(1, 64, 4, 4, 4, device="cuda").to(memory_format=CL) + + def _broken(*args, **kwargs): + raise RuntimeError("simulated HIP launch failure") + + tgn._ensure_kernels() + monkeypatch.setattr(tgn, "_stats_partial_kernel", _Unlaunchable(_broken)) + with pytest.raises(tgn.TritonKernelError): + torch.ops.scaffold_gn.group_norm(x, 8, None, None, EPS, None, None) + + +class _Unlaunchable: + """A stand-in for a ``triton.jit`` kernel whose launch raises.""" + + def __init__(self, fn): + self._fn = fn + + def __getitem__(self, grid): + return self._fn + + +# --------------------------------------------------------------------------- +# the fused activation on non-finite values +# --------------------------------------------------------------------------- + +#: NaN, +Inf, -Inf, -0.0 and four ordinary values. ``tl.maximum(y, 0)`` returns +#: the non-NaN operand and ``tl.where(y > 0, y, 0)`` fails ``NaN > 0``, so both +#: of the obvious spellings map NaN to 0.0 where ``F.relu`` propagates it. +_SPECIALS = [float("nan"), float("inf"), float("-inf"), -0.0, 0.0, -1.0, 1.0, 2.0] + + +def _zero_weight_case(activation, seed=3): + """A case whose pre-activation is exactly ``bias``, elementwise. + + Poisoning the *input* can only produce NaN pre-activations -- one non-finite + value makes the whole group's statistics NaN -- so the values that actually + distinguish the spellings of ReLU have to be placed directly. A zero + ``weight`` does that: ``xhat * 0 + bias == bias``. + """ + x = torch.randn( + 1, + 64, + 4, + 4, + 4, + device="cuda", + generator=torch.Generator("cuda").manual_seed(seed), + ).to(memory_format=CL) + weight = torch.zeros(64, device="cuda") + bias = torch.tensor(_SPECIALS * 8, device="cuda") + reference = F.group_norm(x, 8, weight, bias, EPS) + if activation == "relu": + reference = F.relu(reference) + return x, weight, bias, reference + + +@pytest.mark.gpu +def test_fused_relu_matches_f_relu_on_nan_inf_and_negative_zero(): + """The fused store must be ``F.relu``, bit for bit, on every special value. + + NaN in, NaN out -- and that matters beyond numerics: ScaFFold aborts a run + whose loss goes non-finite, so an activation that turns a diverging NaN into + a finite 0.0 makes the forward look healthy while the backward is still NaN, + and the run checkpoints a broken model. ``-Inf`` and both signed zeros must + come out as ``+0.0``, never ``-0.0``. + """ + x, weight, bias, reference = _zero_weight_case("relu") + out, _mean, _rstd = torch.ops.scaffold_gn.group_norm( + x, 8, weight, bias, EPS, "relu", None + ) + assert torch.equal(out.cpu().view(torch.int32), reference.cpu().view(torch.int32)) + # ... and the control: without the fusion the same values pass through. + plain, _m, _r = torch.ops.scaffold_gn.group_norm( + x, 8, weight, bias, EPS, None, None + ) + assert plain[0, 0, 0, 0, 0].isnan() and plain[0, 1, 0, 0, 0].isinf() + + +@pytest.mark.gpu +def test_fused_relu_backward_gates_like_threshold_backward(): + """ReLU's backward is ``result <= 0 ? 0 : grad``, so a NaN passes. + + The kernel recomputes the pre-activation and must gate with the same + complement: ``pre > 0 ? dy : 0`` reads identically on every finite value and + silently zeroes the NaN lane, which is the backward half of the same defect. + """ + x, weight, bias, _reference = _zero_weight_case("relu") + grad_out = torch.ones_like(x) + + weight = weight.requires_grad_(True) + bias = bias.requires_grad_(True) + xg = x.clone().requires_grad_(True) + reference = F.relu(F.group_norm(xg, 8, weight, bias, EPS)) + reference.backward(grad_out) + ref_dbias = bias.grad.detach().clone() + ref_dx = xg.grad.detach().clone() + + weight.grad = bias.grad = xg.grad = None + triton_group_norm(xg, 8, weight, bias, EPS, "relu").backward(grad_out) + + # d_bias counts exactly the elements whose gradient the gate let through. + assert ref_dbias[0].item() == 64, "the reference gated the NaN lane off" + assert torch.equal(bias.grad.cpu(), ref_dbias.cpu()) + assert torch.equal(xg.grad.cpu(), ref_dx.cpu()) From 6b8670f080f4266929aca11dadb59672fc974b6f Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Sat, 1 Aug 2026 20:23:49 -0700 Subject: [PATCH 60/62] Make the compiled GroupNorm's recompile limit reach the thread that compiles activation_checkpointing with the compiled rung and a DCTensor input crashed: Dynamo hit a recompile limit reported as 8 while _raise_recompile_limit() had set 64, and setting 256 by hand did not help either. torch._dynamo.config's user overrides are ContextVar-backed, so they are per-thread. Non-reentrant checkpointing recomputes inside the backward pass, which runs on the autograd engine's device worker thread, and that thread reads the stock default. Instrumenting the limit check shows it directly: MainThread limit_seen=64 exceeded=0, worker thread limit_seen=8 exceeded=1. The compile that overflows is the one our setting cannot reach. Passing recompile_limit to torch.compile fixes it, because Dynamo applies that with config.patch() around the compile itself, on whichever thread compiles. What made it overflow is DistConv. TORCH_LOGS=recompiles shows the recompute's first miss is GLOBAL_STATE changed: torch_function -- DistConv's backward runs below __torch_function__, so the recompute misses all five forward entries and compiles a second parallel set. Five shapes times two torch-function states is ten, over the limit of eight. The count is bounded, so the fix is to make the limit real rather than to chase a guard: it now stabilizes at ten. FailOnRecompileLimitHit also escaped the ladder entirely, being the one Dynamo compile-time failure that derives from Exception rather than TorchDynamoException. It is named explicitly now. Catching it is not sufficient on its own, and this was measured rather than reasoned: with the exception merely caught, the real UNet died with CheckpointError and the synthetic case with a memory fault, because falling back mid-replay makes a checkpointed block's recompute disagree with its forward. So a proven rung now re-raises instead of degrading while a backward is in flight, detected with the graph-task id that checkpoint's own unpack hook uses. A first failure still degrades, which is where the contract about not killing a multi-node run actually lives. Two upstream bugs worth filing: the global recompile_limit does not apply to compiles triggered from the autograd worker thread, and its warning prints the limit it enforced while the config reads the value we set; and FailOnRecompileLimitHit sits outside the TorchDynamoException root that handlers are written against. CPU 391 passed, GPU 270. Wiring mutations 37/38, the survivor equivalent. Interleaved A/B on the Triton hot path: 0.1651 vs 0.1651 ms. --- ScaFFold/unet/group_norm.py | 145 ++++++++++++++++-- tests/test_groupnorm.py | 227 ++++++++++++++++++++++++++++ tests/test_groupnorm_wiring_edge.py | 77 +++++++++- 3 files changed, 436 insertions(+), 13 deletions(-) diff --git a/ScaFFold/unet/group_norm.py b/ScaFFold/unet/group_norm.py index b172b09..f6b134f 100644 --- a/ScaFFold/unet/group_norm.py +++ b/ScaFFold/unet/group_norm.py @@ -72,7 +72,9 @@ caught on ``triton_group_norm.TritonKernelError``, which that module raises for anything its launch region produces; the compiled rung on ``torch._dynamo.exc.TorchDynamoException``, the root of every Dynamo and -Inductor compile failure. Everything else propagates -- saved-tensor pack +Inductor compile failure, *plus* ``FailOnRecompileLimitHit``, which despite the +name derives from ``Exception`` and not from that root (see +:func:`_compiled_kernel_failures`). Everything else propagates -- saved-tensor pack hooks, ``torch.utils.checkpoint``'s recompute control flow, a user's offloading hook, ``torch.OutOfMemoryError``, an error from a shape the kernel mishandles badly enough to corrupt the graph. The previous shape of this code caught @@ -99,6 +101,17 @@ module's choice for the life of the process, so forward and recompute always agree. +The same reasoning bounds the *fallback* itself, which the latch alone does not: +a proven module still has to answer the call its rung just failed, and answering +it eagerly is exactly the flip the paragraph above forbids -- if that call is a +checkpoint recompute. So the fallback is declined in the one case where it +would corrupt rather than degrade: a module proven on the rung, failing while an +autograd graph task is in flight (:func:`_replaying_a_forward`), re-raises. +Every other failure -- and in particular every *first* failure, which is what a +broken Triton install, an unwritable Inductor cache or a missing compiler +produce -- still degrades, which is where the "must not kill a multi-node run" +contract actually lives. + Note that a latch is process-local: under DDP one rank can end up running a different kernel from its peers. All three kernels agree to fp32 rounding, not bitwise, so a rank that latches shifts that rank's gradients and therefore the @@ -161,8 +174,13 @@ #: UNet presents one entry per distinct activation shape (5 at scale 7) times #: grad-enabled/no-grad (training vs. evaluation), i.e. 10 -- above the stock #: limit of 8, which would silently drop the whole model back to eager mid-run. -#: The traced function is a single ``F.group_norm`` call, so the extra entries -#: cost only their one-time compilation. +#: ``activation_checkpointing`` on a ``DCTensor`` doubles that again: the +#: recompute reaches this module with ``__torch_function__`` subclass handling +#: *disabled* (DistConv's backward runs below it), which is part of Dynamo's +#: ``GLOBAL_STATE`` guard, so the recomputed forward misses every entry the +#: original forward built and compiles a second set beside it -- 20 for the same +#: 5 shapes (measured). The traced function is a single ``F.group_norm`` call, +#: so the extra entries cost only their one-time compilation. _MIN_RECOMPILE_LIMIT = 64 # Lazily built on the first eligible forward: importing ScaFFold must not drag @@ -265,12 +283,25 @@ def _group_norm(input, num_groups, weight, bias, eps): def _raise_recompile_limit(): - """Lift Dynamo's per-function recompile cap to cover every UNet GN shape. + """Lift Dynamo's *global* recompile cap to cover every UNet GN shape. Only ever raises it, so a caller that deliberately set a larger limit keeps theirs -- but note the converse: a limit deliberately set *smaller* than ours is clobbered up to ``_MIN_RECOMPILE_LIMIT``. ``cache_size_limit`` is the older spelling of ``recompile_limit``; set whichever exists. + + This is the *portable* half of the mitigation and, on its own, not a + sufficient one: ``torch._dynamo.config`` stores user overrides in a + ``ContextVar`` (``torch/utils/_config_module.py``: "User overrides are + thread-local"), so an assignment made here is invisible to every other + thread, which keeps reading the stock default of 8. That matters because + ``torch.utils.checkpoint``'s non-reentrant recompute runs inside the + backward pass, i.e. on the autograd engine's device worker thread, and a + recompute that has to compile -- which it does on a ``DCTensor``, see + ``_MIN_RECOMPILE_LIMIT`` -- would hit 8 there no matter what this function + wrote on the main thread. :func:`_compile_group_norm` therefore also asks + ``torch.compile`` for a per-region limit, which Dynamo applies on whichever + thread is compiling. """ config = torch._dynamo.config for name in ("recompile_limit", "cache_size_limit"): @@ -279,6 +310,29 @@ def _raise_recompile_limit(): setattr(config, name, _MIN_RECOMPILE_LIMIT) +def _compile_group_norm(): + """``torch.compile`` :func:`_group_norm` with a thread-proof recompile cap. + + ``recompile_limit=`` is the per-region spelling of the cap: Dynamo applies + it with ``config.patch()`` around the compile itself, on whatever thread + that compile happens on, which is the only spelling that survives the + autograd worker thread (see :func:`_raise_recompile_limit`). Older torches + have no such keyword -- there the global assignment is all there is, and the + checkpoint-recompute case is simply out of reach. + """ + try: + return torch.compile( + _group_norm, + dynamic=False, + fullgraph=True, + recompile_limit=_MIN_RECOMPILE_LIMIT, + ) + except TypeError: + # A torch too old for the keyword: still compile, because the rung is + # worth far more than the one configuration the keyword rescues. + return torch.compile(_group_norm, dynamic=False, fullgraph=True) + + def _get_compiled_group_norm(): """Build (once) the compiled functional GroupNorm shared by every module. @@ -292,7 +346,7 @@ def _get_compiled_group_norm(): global _compiled_group_norm if _compiled_group_norm is None: _raise_recompile_limit() - _compiled_group_norm = torch.compile(_group_norm, dynamic=False, fullgraph=True) + _compiled_group_norm = _compile_group_norm() return _compiled_group_norm @@ -348,9 +402,19 @@ def _compiled_kernel_failures(): (``fullgraph=True`` met something untraceable), ``BackendCompilerFailed`` and its ``InductorError`` subclass (the backend, and therefore also an unwritable Inductor cache or a broken C++/Triton toolchain), and - ``InternalTorchDynamoError``. All of them are raised while *compiling*, - i.e. before the compiled callable has executed or saved anything, which is - what makes the fallback safe to retry. + ``InternalTorchDynamoError``. + + ``FailOnRecompileLimitHit`` -- raised when a frame needs more cache entries + than the recompile limit allows, which under ``fullgraph=True`` is a hard + error rather than a drop to eager -- is *not* under that root: it derives + straight from ``Exception`` (``torch/_dynamo/exc.py``), so catching only + ``TorchDynamoException`` lets it kill the run. It is named separately + rather than assumed, and only added when it really is outside the root, so + a torch that later reparents it does not produce a duplicate entry. + + All of these are raised while *compiling*, i.e. before the compiled callable + has executed or saved anything, which is what makes the fallback safe to + retry. Resolved on demand and cached: importing ``torch._dynamo`` is precisely the cost :func:`_get_compiled_group_norm` defers. An empty tuple (a torch @@ -360,14 +424,56 @@ def _compiled_kernel_failures(): global _COMPILED_KERNEL_FAILURES if _COMPILED_KERNEL_FAILURES is None: try: - import torch._dynamo.exc - - _COMPILED_KERNEL_FAILURES = (torch._dynamo.exc.TorchDynamoException,) + import torch._dynamo.exc as dynamo_exc except ImportError: # pragma: no cover - torch always ships it _COMPILED_KERNEL_FAILURES = () + else: + failures = [dynamo_exc.TorchDynamoException] + limit_hit = getattr(dynamo_exc, "FailOnRecompileLimitHit", None) + if isinstance(limit_hit, type) and not issubclass( + limit_hit, dynamo_exc.TorchDynamoException + ): + failures.append(limit_hit) + _COMPILED_KERNEL_FAILURES = tuple(failures) return _COMPILED_KERNEL_FAILURES +def _replaying_a_forward(): + """``True`` while this thread is executing inside an autograd graph task. + + ``torch._C._current_graph_task_id()`` is ``-1`` outside a backward pass and + the running task's id inside one; it is the same signal + ``torch.utils.checkpoint`` keys its own recompute bookkeeping on + (``torch/utils/checkpoint.py``'s ``unpack_hook``). + + A GroupNorm *forward* that runs while a backward is in flight is not a new + call: it is a checkpoint recompute (or a double backward) replaying a + forward that has already happened and whose saved tensors are already held. + That is the one place where quietly answering on a different rung than the + original forward used is not a fallback but a corruption -- the rungs save + different tensors for backward, so the recompute's saved set no longer + matches (measured: ``CheckpointError: Recomputed values ... have different + metadata``, and on one shape a GPU memory fault instead). See + :meth:`FastGroupNorm.forward`. + + ``is_compiling()`` first, for the same reason :func:`_warn_rung_failure` + checks it: the probe below is a ``torch._C`` builtin returning an ``int``, + which Dynamo cannot trace ("Unsupported torch.* op returned non-Tensor"), + so a caller who wraps this ``forward`` in ``torch.compile(fullgraph=True)`` + would get a hard error where the fallback belongs. Dynamo folds it to + ``True`` at trace time, leaving ``False`` here as a constant -- which is + also the right answer: tracing is not replaying, and the recompute this + guards against runs with Dynamo disabled anyway + (``torch.utils.checkpoint``'s ``_run_fn_with_dynamo_disabled``). + """ + if torch.compiler.is_compiling(): + return False + task_id = getattr(torch._C, "_current_graph_task_id", None) + if task_id is None: # pragma: no cover - every supported torch has it + return False + return task_id() != -1 + + #: ``True`` while a ``torch.func`` transform (``vmap``/``grad``/``jvp``) is on #: the stack. Both fast rungs decline then: a functorch layer is a routing #: question, not a kernel defect, and the stock kernel handles every transform. @@ -714,6 +820,12 @@ def forward(self, input): _warn_rung_failure( "Triton GroupNorm", e, "compiled kernel", TRITON_ENV_VAR ) + # ... with one exception, shared with the compiled rung below + # and explained there: a module already proven on this rung must + # not be answered from a different one while a backward is + # replaying its forward. + if self._triton_ok and _replaying_a_forward(): + raise else: # Only written once: nn.Module.__setattr__ is not free, and # after the first success this reads a class attribute. @@ -739,6 +851,17 @@ def forward(self, input): _warn_rung_failure( "torch.compile of GroupNorm", e, "eager kernel", COMPILE_ENV_VAR ) + # The one call this rung must not answer eagerly: a module already + # proven on it, failing while a backward is in flight, is a + # checkpoint recompute of a forward that *did* run compiled. The + # rungs save different tensors, so handing back the eager result + # makes the recomputed saved set disagree with the saved one and + # torch rejects the step -- a `CheckpointError`, or worse (both + # measured). Degrading is for modules with nothing to contradict; + # here the honest answer is the original exception, which at least + # names the rung and the shape that could not be served. + if self._compiled_ok and _replaying_a_forward(): + raise return self._eager_forward(input) else: if not self._compiled_ok: diff --git a/tests/test_groupnorm.py b/tests/test_groupnorm.py index 334da73..eff199f 100644 --- a/tests/test_groupnorm.py +++ b/tests/test_groupnorm.py @@ -697,6 +697,95 @@ def test_recompile_limit_is_raised_never_lowered(): setattr(config, name, original) +def test_the_compiled_region_carries_its_own_recompile_limit(monkeypatch): + """The global limit is thread-local, so the region must carry one too. + + ``torch._dynamo.config`` keeps user overrides in a ``ContextVar`` + ("User overrides are thread-local", ``torch/utils/_config_module.py``), so + what :func:`_raise_recompile_limit` writes is invisible from every *other* + thread -- and one of those threads matters: ``torch.utils.checkpoint``'s + non-reentrant recompute runs inside the backward pass, i.e. on the autograd + engine's device worker thread. On a ``DCTensor`` that recompute has to + compile (it reaches this module with ``__torch_function__`` subclass + handling disabled, which is part of Dynamo's ``GLOBAL_STATE`` guard, so it + misses every entry the forward built), and there the limit read the stock 8 + however large the global had been set -- ``FailOnRecompileLimitHit``, run + over. ``torch.compile``'s ``recompile_limit=`` is applied by Dynamo around + the compile itself, on whichever thread that compile happens on, which is + the only spelling that reaches the worker; this pins that we ask for it. + """ + seen = {} + + def _fake_compile(fn, **kwargs): + seen.update(kwargs) + return fn + + monkeypatch.setattr(torch, "compile", _fake_compile) + assert gn_mod._compile_group_norm() is gn_mod._group_norm + assert seen.get("recompile_limit") == gn_mod._MIN_RECOMPILE_LIMIT + assert seen.get("fullgraph") is True and seen.get("dynamic") is False + + +def test_compiling_still_works_without_a_per_region_limit(monkeypatch): + """A torch too old for ``recompile_limit=`` must still get a callable. + + The keyword is the fix for the worker thread, not a requirement for + compiling at all; dropping the whole rung on a ``TypeError`` would be a far + bigger regression than the case it addresses. + """ + calls = [] + + def _fake_compile(fn, **kwargs): + calls.append(kwargs) + if "recompile_limit" in kwargs: + raise TypeError("compile() got an unexpected keyword 'recompile_limit'") + return fn + + monkeypatch.setattr(torch, "compile", _fake_compile) + assert gn_mod._compile_group_norm() is gn_mod._group_norm + assert len(calls) == 2 and "recompile_limit" not in calls[1] + + +def test_a_recompile_limit_hit_is_a_kernel_failure_not_a_crash(monkeypatch, caplog): + """``FailOnRecompileLimitHit`` has to land in the ladder, not in the run. + + It is what ``fullgraph=True`` raises when a frame needs more cache entries + than the recompile limit allows, and -- unlike every other Dynamo failure -- + it derives straight from ``Exception`` rather than from + ``TorchDynamoException``, so an allowlist that names only the latter lets it + escape and kill the step (observed at ``5943389``). It is raised while + compiling, before the callable has run or saved anything, so the eager + retry underneath it is safe. + """ + import torch._dynamo.exc + + limit_hit = torch._dynamo.exc.FailOnRecompileLimitHit + assert not issubclass(limit_hit, torch._dynamo.exc.TorchDynamoException), ( + "naming it separately is only needed while it sits outside that root" + ) + + def _kernel(*args, **kwargs): + raise limit_hit("simulated recompile limit hit") + + monkeypatch.setattr( + gn_mod, "_use_compiled", lambda t, **kw: type(t) is torch.Tensor + ) + monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _kernel) + gn_mod._compile_failed = False + + fast = _seeded_norm() + x = _make_input(seed=51, channels=16, size=4) + expected = nn.functional.group_norm(x, _GROUPS, fast.weight, fast.bias, fast.eps) + + with caplog.at_level(logging.WARNING, logger=gn_mod.__name__): + out = fast(x) + + assert torch.allclose(out, expected), "the eager fallback did not run" + assert gn_mod._compile_failed is True, "the failure did not latch the rung off" + assert fast._compiled_ok is False + assert any("falling back" in record.message for record in caplog.records) + + # --------------------------------------------------------------------------- # DCTensor routing: unwrap -> compiled kernel -> rewrap # --------------------------------------------------------------------------- @@ -1128,6 +1217,144 @@ def assert_agrees(actual, expected, label): assert_agrees(compiled_nockpt, compiled, "grad") +@pytest.mark.gpu +def test_gpu_the_recompile_limit_holds_on_a_worker_thread(): + """Past Dynamo's stock 8 entries, compiling from a thread that never set it. + + ``torch._dynamo.config``'s user overrides live in a ``ContextVar``, so the + limit :func:`_raise_recompile_limit` writes on the main thread is not the + limit another thread reads -- and the compiles that matter happen on + another thread, because ``torch.utils.checkpoint``'s recompute runs inside + the backward pass, on the autograd engine's device worker. This drives the + same shape of traffic directly: entries live on ``_group_norm``'s code + object and are shared between threads, so a worker that pushes the count + past 8 is exactly the situation the recompute creates. Before the + per-region ``recompile_limit=``, the ninth compile raised + ``FailOnRecompileLimitHit`` here. + + ``torch._dynamo.reset()`` first because those entries also accumulate + across the whole test session, which would otherwise decide the outcome. + """ + import threading + + torch._dynamo.reset() + gn_mod.set_triton_enabled(False) # this is the compiled rung's limit + gn_mod.set_compile_enabled(True) + gn_mod._compile_failed = False + + device = torch.device("cuda") + # Nine distinct channel counts: nine cache entries, one more than the stock + # limit allows, and the one that overflows must land on the worker thread. + norms = [FastGroupNorm(_GROUPS, 8 * n).to(device) for n in range(1, 10)] + failures = [] + + def run(subset): + try: + for norm in subset: + norm(torch.randn(1, norm.num_channels, 2, 2, 2, device=device)) + except BaseException as error: # noqa: BLE001 - re-raised below + failures.append(error) + + run(norms[:2]) + worker = threading.Thread(target=run, args=(norms[2:],)) + worker.start() + worker.join() + + if failures: + raise AssertionError(f"compiling off the main thread failed: {failures[0]}") + assert gn_mod._compile_failed is False, "the rung latched itself off" + assert all(norm._compiled_ok for norm in norms), ( + "some module never had a call served by the compiled rung" + ) + + +@pytest.mark.gpu +def test_gpu_checkpointed_dctensor_recompute_keeps_the_compiled_rung(dc_cuda): + """The three-way combination that used to die: ckpt + compiled rung + DCTensor. + + ``activation_checkpointing: true`` with ``SCAFFOLD_GROUPNORM_TRITON=0`` on + DistConv activations is a supported configuration and it crashed: the + recompute reaches this module with ``__torch_function__`` subclass handling + *disabled* (DistConv's backward runs below it), which is part of Dynamo's + ``GLOBAL_STATE`` guard, so it misses every cache entry the forward built and + compiles a second set beside them -- twice the shapes, past 8 -- on the + autograd worker thread, where the module's raised limit was invisible. Each + pair of the three is fine on its own; all three together raised + ``FailOnRecompileLimitHit`` (at ``5943389``) or, once the ladder caught it + and dropped a *proven* module to eager mid-recompute, ``CheckpointError``. + + Five norms is the smallest count that reproduces it: 5 forward entries plus + 5 recompute entries is 10, and the ninth compile is the one that overflows. + The convolutions are what make the block's backward run below torch-function + (a bare unwrap does not), and the loss is taken on the ``DCTensor`` for the + same reason the trainer's is. + """ + import threading + + import torch.utils.checkpoint + + distconv, ps = dc_cuda + device = torch.device("cuda") + channels = (8, 16, 24, 32, 40) + + torch._dynamo.reset() + gn_mod.set_triton_enabled(False) + gn_mod.set_compile_enabled(True) + gn_mod._compile_failed = False + + torch.manual_seed(5) + norms = [FastGroupNorm(_GROUPS, c).to(device) for c in channels] + convs = [ + nn.Conv3d(previous, c, 1, bias=False).to(device) + for previous, c in zip((1,) + channels[:-1], channels) + ] + tail = nn.Conv3d(channels[-1], 1, 1, bias=False).to(device) + + # Where each GroupNorm call happens, as Dynamo's GLOBAL_STATE guard sees it. + states = set() + + def block(t): + for conv, norm in zip(convs, norms): + states.add( + ( + threading.current_thread() is threading.main_thread(), + torch._C._is_torch_function_enabled(), + ) + ) + t = norm(conv(t)) + return tail(t) + + x = torch.randn(1, 1, 4, 4, 4, device=device) + + def step(checkpointing): + for parameter in [x] + [ + p for m in convs + norms + [tail] for p in m.parameters() + ]: + parameter.grad = None + x.requires_grad_(True) + wrapped = distconv.DCTensor.from_shard(x, ps) + if checkpointing: + out = torch.utils.checkpoint.checkpoint(block, wrapped, use_reentrant=False) + else: + out = block(wrapped) + out.float().square().mean().backward() + return [norm.weight.grad.detach().clone() for norm in norms] + + checkpointed = step(True) + step(True) # a second step must not compile anything new either + direct = step(False) + + assert (False, False) in states, ( + "the recompute did not run below torch-function off the main thread; " + "this configuration no longer reproduces the guard split it targets" + ) + assert gn_mod._compile_failed is False, "the compiled rung latched itself off" + assert all(norm._compiled_ok for norm in norms), "a norm never ran compiled" + for index, (recomputed, plain) in enumerate(zip(checkpointed, direct)): + assert torch.isfinite(recomputed).all(), index + _assert_close(recomputed, plain, 1e-4, f"norm {index} weight grad") + + # --------------------------------------------------------------------------- # GPU behavior: the Triton rung # --------------------------------------------------------------------------- diff --git a/tests/test_groupnorm_wiring_edge.py b/tests/test_groupnorm_wiring_edge.py index 3d80536..66d3c06 100644 --- a/tests/test_groupnorm_wiring_edge.py +++ b/tests/test_groupnorm_wiring_edge.py @@ -491,6 +491,72 @@ def _kernel(input, num_groups, weight, bias, eps): assert fresh._compiled_ok is False +@pytest.mark.parametrize("rung", ["triton", "compiled"]) +def test_a_proven_rung_does_not_degrade_while_a_backward_replays_it(monkeypatch, rung): + """A fallback *during a recompute* corrupts rather than degrades. + + The latch already refuses to demote a module that has used a rung, but the + fallback itself sidestepped that: the failing call still got answered from + the next rung down, and if that call is ``torch.utils.checkpoint``'s + recompute of a forward that ran on the failing rung, the recomputed forward + saves a different set of tensors than the original did and torch rejects + the whole step (``CheckpointError``; on one measured shape a GPU memory + fault instead). Neither is a degradation, so this one case re-raises. + + It is narrow on purpose -- ``_replaying_a_forward()`` is false in an + ordinary forward, where ``test_a_global_latch_does_not_demote_a_module_...`` + still requires the fallback -- and the second half here pins the other side + of the narrowness: a module that has *not* used the rung degrades even + inside the backward, because the forward it is replaying went down the + ladder too and the two agree. + """ + import torch._dynamo.exc + import torch.utils.checkpoint as checkpoint_mod + + from ScaFFold.unet.triton_group_norm import TritonKernelError + + failure = TritonKernelError if rung == "triton" else torch._dynamo.exc.Unsupported + latch = "_triton_failed" if rung == "triton" else "_compile_failed" + proven_flag = "_triton_ok" if rung == "triton" else "_compiled_ok" + + def make_kernel(fail_always): + def _kernel(input, num_groups, weight, bias, eps, *activation): + if fail_always or gn_mod._replaying_a_forward(): + raise failure("simulated kernel failure") + return F.group_norm(input, num_groups, weight, bias, eps) + + return _kernel + + def run(fail_always): + gn_mod._triton_failed = False + gn_mod._compile_failed = False + kernel = make_kernel(fail_always) + if rung == "triton": + monkeypatch.setattr(gn_mod, "_use_triton", lambda *a, **kw: True) + module_stub = type("_Stub", (), {"triton_group_norm": staticmethod(kernel)}) + monkeypatch.setattr(gn_mod, "_get_triton_module", lambda: module_stub) + else: + monkeypatch.setattr(gn_mod, "_use_compiled", lambda t, **kw: True) + monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: kernel) + module = FastGroupNorm(_GROUPS, 16) + x = torch.randn(1, 16, 4, 4, 4, requires_grad=True) + out = checkpoint_mod.checkpoint(module, x, use_reentrant=False) + out.pow(2).sum().backward() + return module, x + + # Proven in the forward, failing in the recompute: answering from another + # rung would be the metadata mismatch, so the failure has to come back out. + with pytest.raises(failure): + run(fail_always=False) + + # Never served by the rung: the forward already went down the ladder, so + # the recompute doing the same agrees with it and the step survives. + module, x = run(fail_always=True) + assert getattr(module, proven_flag) is False + assert getattr(gn_mod, latch) is True + assert torch.isfinite(x.grad).all() + + def test_a_latch_flip_mid_forward_is_not_a_numerics_error(monkeypatch): """Two rungs inside one forward still compose (values, not bits, agree).""" monkeypatch.setattr( @@ -615,7 +681,8 @@ def test_gpu_triton_rung_inside_a_compiled_region(monkeypatch, fullgraph): @pytest.mark.gpu -def test_gpu_the_fallback_path_traces_under_fullgraph(monkeypatch): +@pytest.mark.parametrize("proven", [False, True]) +def test_gpu_the_fallback_path_traces_under_fullgraph(monkeypatch, proven): """A rung failure *while Dynamo is tracing* must still fall back, not die. The handler used to call ``logger.warning``, which Dynamo cannot trace @@ -625,6 +692,12 @@ def test_gpu_the_fallback_path_traces_under_fullgraph(monkeypatch): most, since the thing it is reacting to is usually a compile-time failure. Nothing in ScaFFold compiles ``FastGroupNorm.forward`` today; this pins the claim that it can. + + Both halves of the handler's guard have to trace, which is why ``proven`` + is parametrized: with ``_triton_ok`` false Dynamo folds the ``and`` away + without ever looking at ``_replaying_a_forward()``, so only the ``True`` + arm reaches it -- and a probe Dynamo cannot trace there would be the same + defect as the logging call, reintroduced. """ import torch._dynamo @@ -648,7 +721,7 @@ def _raises(*args, **kwargs): monkeypatch.setattr(gn_mod, "_get_triton_module", _BrokenKernelModule) gn_mod._triton_failed = False - module._triton_ok = False + module._triton_ok = proven torch._dynamo.reset() compiled = torch.compile(lambda t: module(t), fullgraph=True, dynamic=False) From 964afd673c994dead1b779ae97f949ec763df0a0 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Sun, 2 Aug 2026 00:30:59 -0700 Subject: [PATCH 61/62] Concatenate the decoder skip in one channels-last pass The premise this started from was wrong in two ways, both worth recording. Most of the profile's 4.77 ms aten::cat is not the skip concatenation at all. forward_halo_exchange ends in a cat and distconv_forward calls it once per spatial dim, so every 3x3x3 convolution pays three full-tensor concatenations to materialize a (D+2, H+2, W+2) copy -- including at dc_num_shards=1, where it is an open-coded padding=1. At the largest decoder block that is 1.90 ms of halo cat against 0.62 ms of skip cat. And ATen's cat is not making the mistake Inductor made with GroupNorm: on the forward it runs at 100.5 / 104.6 / 107.5% of the 3.35 TB/s streaming roofline at the three large decoder shapes. A Triton kernel with the same dtype behaviour measures -0.03 ms, i.e. nothing. What is actually wrong is the dtype and the backward. Under autocast the skip arrives fp32 (GroupNorm's fp32 policy) and the upsampled half bf16; cat carries the promote policy, so it widens bf16 to fp32, concatenates at fp32, and the convolution narrows it back -- three full-resolution passes to deliver one. And cat's backward hands out views that its consumers then force contiguous, at 51-63% of roofline. skip_concat does it in one channels-last pass at the consumer's dtype, with a split backward that addresses its cotangent by stride instead of calling contiguous() on it. That last part is the whole difference: under DCTensor the cotangent is a narrowed view of the halo-padded tensor, and the obvious kernel lost to plain torch (-1.06 vs -1.85 ms) until it stopped materializing it. Whole scale-7 step 92.783 -> 91.605 ms (1.3%), peak 7.215 -> 6.715 GiB. The four decoder blocks in isolation 53.053 -> 50.732 ms (4.4%); the largest block's skip path 2.43 -> 0.91 ms, with convolution_backward flat at 17.43/17.40 so nothing moved into the convolutions. Every conv output, block output and input gradient is bitwise identical and so is the whole-model loss. Measured and rejected: splitting the first convolution by input channel, which is algebraically equivalent and 5.5% slower at every block -- two half-channel convolutions plus an add cost more than one convolution plus a concatenation. CPU 408 passed, GPU 339. --- ScaFFold/unet/triton_cat.py | 824 ++++++++++++++++++++++++++++++++++++ ScaFFold/unet/unet_parts.py | 31 +- tests/test_triton_cat.py | 478 +++++++++++++++++++++ tests/test_unet.py | 142 +++++++ 4 files changed, 1472 insertions(+), 3 deletions(-) create mode 100644 ScaFFold/unet/triton_cat.py create mode 100644 tests/test_triton_cat.py diff --git a/ScaFFold/unet/triton_cat.py b/ScaFFold/unet/triton_cat.py new file mode 100644 index 0000000..bb52b97 --- /dev/null +++ b/ScaFFold/unet/triton_cat.py @@ -0,0 +1,824 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +"""Channels-last-native Triton channel concatenation (NDHWC in, NDHWC out). + +Why this exists +=============== +Every ``Up`` block of the UNet joins the decoder's upsampled activation to the +encoder's skip activation with ``torch.cat([skip, up], dim=1)`` and feeds the +result to a convolution. Profiled at scale 7 (channels-last, bf16 autocast, +DCTensor) that concatenation and the copies autocast wraps around it cost +**~10.4 ms of a 92 ms step, 11.6%** -- the largest remaining item in ScaFFold's +own code once the GroupNorm kernel landed. It is not one problem but two, and +neither is inherent: + +**1. The dtypes are wrong on the wire.** Under ``torch.autocast`` the two +inputs do not have the same dtype. GroupNorm carries autocast's ``fp32`` cast +policy, so the skip tensor -- which is a ``DoubleConv`` output, i.e. a +GroupNorm output -- arrives as **fp32**, while the upsampled tensor comes +straight out of a ``ConvTranspose3d`` and is **bf16**. ``aten::cat`` carries +the ``promote`` policy, so autocast *widens the bf16 input to fp32*, cats at +fp32, and then the following convolution -- ``lower_precision_fp`` policy -- +casts the whole double-width result straight back down to bf16. At the largest +decoder block ([1,64,128^3] skip + [1,64,128^3] up) that is + + up->fp32 read 268 MB write 537 MB + cat read 1074 MB write 1074 MB + conv cast read 1074 MB write 537 MB = 4.56 GB of traffic + +to deliver 537 MB of bf16 to the convolution. Emitting bf16 *directly from the +concatenation* is **bitwise identical** to what the convolution receives today +-- the fp32 intermediate holds exact copies of an fp32 tensor and of a widened +bf16 tensor, so rounding it to bf16 recovers exactly ``(bf16(skip), up)`` -- +and costs one pass: read 537 + 268, write 537 = **1.34 GB, 3.4x less**. + +**2. The kernel iterates the wrong order.** This is the same defect the +GroupNorm kernel exists to fix. A ``channels_last_3d`` tensor ``(N, C, D, H, +W)`` is *physically* a dense ``(N*D*H*W, C)`` array, so a channel +concatenation is, in memory, ``out[m, :Ca] = a[m, :]`` and ``out[m, Ca:] = +b[m, :]`` -- a pure streaming join of two dense arrays into one, 2 reads and 1 +write, perfectly coalescable. ATen reaches it through the *logical* NCDHW +order and lands on TensorIterator's generic offset-calculator path, where the +output tile is never contiguous. + +The kernels below own a ``(BLOCK_M, C)`` tile of the *physical* array: they +issue one fully contiguous store per tile and gather the two halves of it from +the two sources, whose valid lanes are themselves contiguous runs. The dtype +conversion rides along in registers, so the retyping in point 1 is free. + +Public API +========== +``cat_channels(a, b, out_dtype=None)`` + ``torch.cat([a, b], dim=1)`` with an optional output dtype override. + ``out_dtype=None`` reproduces ``torch.cat``'s own promotion exactly. + Anything :func:`is_supported` declines is served by ``torch.cat`` itself. + +``is_supported(a, b, out_dtype=None)`` + Cheap, side-effect-free predicate: ``True`` exactly when the Triton kernel + will run. + +``skip_concat(skip, upsampled)`` + The UNet decoder's call: chooses the output dtype (see + :func:`consumer_dtype`), unwraps DistConv's ``DCTensor`` to its local shard + and rewraps the result. This is the only function ``unet_parts`` calls. + +Contract +======== +For every input :func:`is_supported` accepts, ``cat_channels(a, b, dt)`` equals +``torch.cat([a, b], dim=1).to(dt)`` **bitwise**, with: + +* **memory format** -- the output is ``channels_last_3d``-contiguous, which is + also what ``torch.cat`` returns for channels-last inputs. +* **dtype** -- exactly ``out_dtype``, or ``torch.promote_types(a.dtype, + b.dtype)`` when that is ``None``. Narrowing is a *single* rounding of each + source value: every supported dtype widens exactly into fp32, which is the + kernel's compute type, so rounding fp32->bf16 once in the kernel is the same + bits as ATen's widen-then-narrow. +* **autograd** -- first order, ``d_a`` and ``d_b`` in the *inputs'* dtypes, + which is what autograd requires and what the current chain produces after + autocast's cast nodes run their own backward. Second-order raises, exactly + as :mod:`ScaFFold.unet.triton_group_norm` does and for the same reason: the + backward is itself a custom op with no autograd formula. +* **determinism** -- trivially bitwise reproducible. There is no reduction, + no atomic and no autotuning; every output element is a copy of exactly one + input element and the tile shape is a pure function of the channel count. +* **device** -- the kernels run on the inputs' device whatever device is + current; see ``_device_guard``, which exists because a Triton launch follows + the *current* device and not its arguments'. +* **rejections** -- :func:`is_supported` declines anything the kernel cannot + serve physically (non-channels-last, non-5-D, unsupported dtype, CPU, + mismatched spatial extent, empty), and :func:`cat_channels` then routes it to + ``torch.cat``, so the function is total. + +Measured cost +============= +See ``review/skip-path/RESULTS.md`` for the interleaved A/B, the op-level +before/after and the gradient check. +""" + +from __future__ import annotations + +import contextlib +import functools +import importlib.util +import sys +from typing import Optional, Tuple + +import torch + +__all__ = [ + "cat_channels", + "is_supported", + "skip_concat", + "consumer_dtype", + "CatKernelError", +] + + +class CatKernelError(RuntimeError): + """A failure of the Triton kernels themselves, with the original as ``__cause__``. + + Mirrors :class:`ScaFFold.unet.triton_group_norm.TritonKernelError`: it is + raised only from a *closed* region that allocates and launches and runs no + autograd-observable op, so a caller may retry the call on ``torch.cat`` + without worrying that half a graph was already recorded. + ``torch.OutOfMemoryError`` is passed through untagged -- it is a resource + condition, not a defect, and the fallback would allocate the same bytes. + """ + + +#: Dtypes the kernels read and write directly. All three widen exactly into +#: fp32, which is what makes the single-rounding claim in the docstring hold. +SUPPORTED_DTYPES = (torch.float32, torch.bfloat16, torch.float16) + +_CL_FORMAT = torch.channels_last_3d +_INT32_MAX = 2**31 - 1 + +#: Elements per tile and the cap on voxels per program. Both are pure +#: functions of the channel count, so a run cannot change the tiling underneath +#: a comparison. Chosen by a sweep of BLOCK_M in {1..64} x num_warps in +#: {1,2,4,8} at the four scale-7 decoder shapes +#: (``review/skip-path/logs/cat_bench_tune.log``, +#: ``split_bench_tune.log``): with ``BLOCK_M = clamp(4096 // next_pow2(C), 1, +#: 32)`` and four warps the kernels are within 0.8% of the per-shape optimum at +#: the two shapes that dominate and within 4% at the two launch-bound ones, +#: which is not worth a frozen table. +_TILE_ELEMS = 4096 +_MAX_BLOCK_M = 32 +_NUM_WARPS = 4 + + +# --------------------------------------------------------------------------- # +# Triton kernels +# --------------------------------------------------------------------------- # +triton = None +tl = None +_cat_kernel = None +_split_kernel = None + +_TRITON_AVAILABLE: Optional[bool] = None + + +def triton_available() -> bool: + """Whether ``triton`` can be imported, cached, without importing it.""" + global _TRITON_AVAILABLE + if _TRITON_AVAILABLE is None: + try: + _TRITON_AVAILABLE = importlib.util.find_spec("triton") is not None + except (ImportError, ValueError): + _TRITON_AVAILABLE = False + return _TRITON_AVAILABLE + + +def _build_kernels(): + """Import Triton and install the JIT kernels into this module's globals. + + Defined inside a function purely so ``import triton`` is deferred to the + first GPU call, and written into ``globals()`` because Triton resolves + names through ``fn.__globals__``. + """ + global triton, tl + import triton as _triton + import triton.language as _tl + + triton = _triton + tl = _tl + + @_triton.jit + def _cat_kernel( + A, + B, + OUT, + M, + CA: tl.constexpr, + CB: tl.constexpr, + C: tl.constexpr, + CAP: tl.constexpr, + CBP: tl.constexpr, + BLOCK_M: tl.constexpr, + INT64: tl.constexpr, + ): + """One program per ``BLOCK_M`` voxels: join two dense rows into one. + + Each source keeps its **own** tile width (``next_pow2`` of its channel + count) and is copied by its own load/store pair, rather than both being + gathered into one ``(BLOCK_M, C)`` tile and stored once. The single + fully contiguous store is the more obvious design and was measured + first; it is **20% slower** end to end over the four decoder shapes + (0.718 vs 0.571 ms, ``review/skip-path/logs/cat_bench_tune.log``). The + reason is that the one-store form has to mask *both* loads down to + complementary halves of a double-width lane space, which halves the + useful work per instruction and defeats vectorization, and it buys only + a contiguous store -- whereas a store of ``CA`` contiguous elements at a + stride of ``C`` already covers whole cache lines whenever + ``CA * itemsize`` is a multiple of the line, which it is for every + channel count this network uses. The strided store is not the problem; + the mask was. + """ + pid = tl.program_id(0) + m0 = pid * BLOCK_M + if INT64: + wide = m0.to(tl.int64) + base_a = wide * CA + base_b = wide * CB + base_o = wide * C + else: + base_a = m0 * CA + base_b = m0 * CB + base_o = m0 * C + + rows = tl.arange(0, BLOCK_M) + rmask = rows < M - m0 + ca = tl.arange(0, CAP) + cb = tl.arange(0, CBP) + # Clamp the padding lanes' column index: their loads and stores are + # masked and never touch memory, but keeping the arithmetic inside the + # allocation avoids forming a pointer the compiler may treat as poison. + cam = ca < CA + cbm = cb < CB + ca = tl.where(cam, ca, 0) + cb = tl.where(cbm, cb, 0) + am = rmask[:, None] & cam[None, :] + bm = rmask[:, None] & cbm[None, :] + + av = tl.load(A + base_a + rows[:, None] * CA + ca[None, :], mask=am, other=0.0) + tl.store( + OUT + base_o + rows[:, None] * C + ca[None, :], + av.to(OUT.dtype.element_ty), + mask=am, + ) + bv = tl.load(B + base_b + rows[:, None] * CB + cb[None, :], mask=bm, other=0.0) + tl.store( + OUT + base_o + rows[:, None] * C + CA + cb[None, :], + bv.to(OUT.dtype.element_ty), + mask=bm, + ) + + @_triton.jit + def _split_kernel( + G, + DA, + DB, + SN, + SD, + SH, + D, + H, + W, + CA: tl.constexpr, + CB: tl.constexpr, + C: tl.constexpr, + CAP: tl.constexpr, + CBP: tl.constexpr, + BLOCK_W: tl.constexpr, + WANT_A: tl.constexpr, + WANT_B: tl.constexpr, + INT64: tl.constexpr, + ): + """The transpose of :func:`_cat_kernel`: two strided loads, two stores. + + This is the pass ATen is genuinely bad at. ``cat``'s backward hands the + consumer a *narrowed view*, and every consumer then forces it + contiguous, so the work happens as a generic strided ``copy_`` that + reaches only **51-63%** of this device's streaming roofline at the four + decoder shapes; this kernel reaches **90-103%** + (``review/skip-path/logs/split_bench_tune.log``). + + **The incoming gradient is addressed by its strides, not assumed dense**, + and that is not a nicety. Under DistConv -- which production uses even + at ``dc_num_shards=1`` -- the convolution that consumes the + concatenation is reached through a halo exchange that materialises a + *padded* tensor, so the cotangent that comes back here is a narrowed + view of a ``(D+2, H+2, W+2)`` one: channels-last within each voxel and + contiguous along W, but with a gap at every H and D boundary. An + earlier version simply called ``.contiguous(memory_format=channels_last_3d)`` + on it, which costs a **whole extra full-resolution pass** that ATen's + view-based backward never pays -- 0.46 ms per step at the largest + decoder shape, enough on its own to turn this kernel from a win into a + loss against plain ``torch.cat`` at the right dtype (measured: the + isolated four-block sum went from 53.40 ms with the relayout to 52.10 + without). + + The requirement is therefore only that channels are innermost + (``stride(1) == 1``) and that a voxel's neighbours along W are one + channel-run apart (``stride(4) == C``); everything above W is addressed + through ``SN``/``SD``/``SH``. That admits a dense channels-last tensor + and any narrowing of one on D, H or W, which is every case this op + sees. The driver falls back to a relayout for anything else. + """ + pid = tl.program_id(0) + h = tl.program_id(1) + nd = tl.program_id(2) + n = nd // D + d = nd % D + + w0 = pid * BLOCK_W + ws = w0 + tl.arange(0, BLOCK_W) + wmask = ws < W + + if INT64: + gbase = n.to(tl.int64) * SN + d.to(tl.int64) * SD + h.to(tl.int64) * SH + row0 = ((n.to(tl.int64) * D + d) * H + h) * W + w0 + else: + gbase = n * SN + d * SD + h * SH + row0 = ((n * D + d) * H + h) * W + w0 + gbase = gbase + w0 * C + + wl = tl.arange(0, BLOCK_W) + ca = tl.arange(0, CAP) + cb = tl.arange(0, CBP) + cam = ca < CA + cbm = cb < CB + ca = tl.where(cam, ca, 0) + cb = tl.where(cbm, cb, 0) + am = wmask[:, None] & cam[None, :] + bm = wmask[:, None] & cbm[None, :] + + if WANT_A: + ga = tl.load(G + gbase + wl[:, None] * C + ca[None, :], mask=am, other=0.0) + tl.store( + DA + row0 * CA + wl[:, None] * CA + ca[None, :], + ga.to(DA.dtype.element_ty), + mask=am, + ) + if WANT_B: + gb = tl.load( + G + gbase + wl[:, None] * C + CA + cb[None, :], mask=bm, other=0.0 + ) + tl.store( + DB + row0 * CB + wl[:, None] * CB + cb[None, :], + gb.to(DB.dtype.element_ty), + mask=bm, + ) + + globals().update(_cat_kernel=_cat_kernel, _split_kernel=_split_kernel) + + +def _ensure_kernels(): + if _cat_kernel is None: + _build_kernels() + + +# --------------------------------------------------------------------------- # +# python drivers +# --------------------------------------------------------------------------- # +_NO_GUARD = contextlib.nullcontext() + + +def _device_guard(device: torch.device): + """Make ``device`` current for the kernel launches inside the ``with``. + + A Triton launch goes to whatever device is *current*, not to the device its + arguments live on; without this a tensor on ``cuda:1`` while ``cuda:0`` is + current makes the kernel dereference another device's pointers and the + process dies with ``Memory access fault by GPU node-N``. ATen ops carry a + ``DeviceGuard`` and handle the same call, so this is required for the + drop-in contract. The ``current_device()`` test keeps the common + (already-current) path free; see the same helper in + :mod:`ScaFFold.unet.triton_group_norm` for the measurement. + """ + if device.index == torch.cuda.current_device(): + return _NO_GUARD + return torch.cuda.device(device) + + +def _next_pow2(x: int) -> int: + return 1 << (x - 1).bit_length() if x > 1 else 1 + + +def _cdiv(a: int, b: int) -> int: + return -(-a // b) + + +@functools.lru_cache(maxsize=64) +def _block_m(channels_pow2: int) -> int: + """Voxels per program. A pure function of the padded channel count. + + Bitwise determinism does not actually depend on this -- a copy has no + reduction order to perturb -- but keeping it a pure function of the shape + means a run cannot change tiling underneath a comparison, which is the + property the GroupNorm kernel's frozen table exists to give and is worth + having for free here too. + """ + return max(1, min(_MAX_BLOCK_M, _TILE_ELEMS // channels_pow2)) + + +@functools.lru_cache(maxsize=64) +def _block_w(channels_pow2: int) -> int: + """Voxels per program for the backward, which tiles along W within a line. + + Same rule and same sweep as :func:`_block_m`; the backward cannot tile over + a flat voxel index because its source is only guaranteed contiguous *within* + a ``(n, d, h)`` line -- see ``_split_kernel``. + """ + return max(1, min(_MAX_BLOCK_M, _TILE_ELEMS // channels_pow2)) + + +def _rows_of(t: torch.Tensor) -> int: + """``N * D * H * W`` -- the number of voxels in the physical (rows, C) view.""" + rows = t.shape[0] + for d in t.shape[2:]: + rows *= d + return rows + + +def _tag_kernel_failures(fn): + """Re-raise anything ``fn`` raises as :class:`CatKernelError`. + + Applied to the two functions that do nothing but allocate and launch, so + the tagged region cannot swallow framework control flow (there is no pack + hook, no recompute stop and no functorch layer inside it). + """ + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + try: + return fn(*args, **kwargs) + except torch.OutOfMemoryError: + raise + except CatKernelError: + raise + except Exception as e: # noqa: BLE001 -- re-raised, see docstring + raise CatKernelError( + f"{fn.__name__} failed ({type(e).__name__}: {e})" + ) from e + + return wrapper + + +@_tag_kernel_failures +def _forward(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype): + _ensure_kernels() + ca, cb = a.shape[1], b.shape[1] + c = ca + cb + rows = _rows_of(a) + cp = _next_pow2(c) + block_m = _block_m(cp) + + with _device_guard(a.device): + out = torch.empty( + (a.shape[0], c, *a.shape[2:]), + dtype=out_dtype, + device=a.device, + memory_format=_CL_FORMAT, + ) + _cat_kernel[(_cdiv(rows, block_m),)]( + a, + b, + out, + rows, + CA=ca, + CB=cb, + C=c, + CAP=_next_pow2(ca), + CBP=_next_pow2(cb), + BLOCK_M=block_m, + INT64=rows * c > _INT32_MAX, + num_warps=_NUM_WARPS, + ) + return out + + +@_tag_kernel_failures +def _backward(grad, ca, cb, a_dtype, b_dtype, want_a, want_b): + _ensure_kernels() + c = ca + cb + n, _, d, h, w = grad.shape + sn, _, sd, sh, _ = grad.stride() + block_w = _block_w(_next_pow2(c)) + + with _device_guard(grad.device): + da = torch.empty( + (n, ca, d, h, w) if want_a else (0,), + dtype=a_dtype, + device=grad.device, + **({"memory_format": _CL_FORMAT} if want_a else {}), + ) + db = torch.empty( + (n, cb, d, h, w) if want_b else (0,), + dtype=b_dtype, + device=grad.device, + **({"memory_format": _CL_FORMAT} if want_b else {}), + ) + if want_a or want_b: + _split_kernel[(_cdiv(w, block_w), h, n * d)]( + grad, + da, + db, + sn, + sd, + sh, + d, + h, + w, + CA=ca, + CB=cb, + C=c, + CAP=_next_pow2(ca), + CBP=_next_pow2(cb), + BLOCK_W=block_w, + WANT_A=want_a, + WANT_B=want_b, + # The widest index the kernel forms is the *source* base, which + # spans the (possibly padded) parent tensor, so it is bounded by + # the largest stride and not by this tensor's own element count. + INT64=max(sn * n, n * d * h * w * c) > _INT32_MAX, + num_warps=_NUM_WARPS, + ) + return da, db + + +# --------------------------------------------------------------------------- # +# torch.library registration +# --------------------------------------------------------------------------- # +def _line_addressable(grad: torch.Tensor) -> bool: + """Whether ``_split_kernel`` can read ``grad`` in place. + + It needs channels innermost and voxels one channel-run apart along W; the + D and H axes are addressed through their own strides, so any narrowing of a + channels-last tensor qualifies -- which is what DistConv's halo-padded + cotangent is. A dense channels-last tensor is the special case where the + strides happen to be exact multiples. + """ + if grad.dim() != 5: + return False + stride = grad.stride() + return stride[1] == 1 and stride[4] == grad.shape[1] + + +def _validate(a, b, out_dtype): + if a.dim() != 5 or b.dim() != 5: + raise ValueError( + f"expected two 5-D NCDHW tensors, got {tuple(a.shape)} and {tuple(b.shape)}" + ) + if a.shape[0] != b.shape[0] or tuple(a.shape[2:]) != tuple(b.shape[2:]): + raise ValueError( + f"shapes must agree except on dim 1, got {tuple(a.shape)} and " + f"{tuple(b.shape)}" + ) + if a.shape[1] < 1 or b.shape[1] < 1 or a.numel() == 0 or b.numel() == 0: + raise ValueError("both inputs must have at least one channel and be non-empty") + for name, t in (("a", a), ("b", b)): + if t.dtype not in SUPPORTED_DTYPES: + raise ValueError(f"unsupported dtype {t.dtype} for {name}") + if not t.is_contiguous(memory_format=_CL_FORMAT): + raise ValueError( + f"{name} must be channels_last_3d-contiguous; use cat_channels() " + "which falls back to torch.cat for other layouts" + ) + if out_dtype is not None and out_dtype not in SUPPORTED_DTYPES: + raise ValueError(f"unsupported out_dtype {out_dtype}") + + +@torch.library.custom_op( + "scaffold_cat::cat_channels", mutates_args=(), device_types="cuda" +) +def _cat_op( + a: torch.Tensor, b: torch.Tensor, out_dtype: Optional[torch.dtype] +) -> torch.Tensor: + """``torch.cat([a, b], dim=1)`` for channels-last-3d inputs.""" + _validate(a, b, out_dtype) + return _forward(a, b, out_dtype or torch.promote_types(a.dtype, b.dtype)) + + +@_cat_op.register_fake +def _(a, b, out_dtype): + return torch.empty( + (a.shape[0], a.shape[1] + b.shape[1], *a.shape[2:]), + dtype=out_dtype or torch.promote_types(a.dtype, b.dtype), + device=a.device, + memory_format=_CL_FORMAT, + ) + + +@torch.library.custom_op( + "scaffold_cat::split_channels", mutates_args=(), device_types="cuda" +) +def _split_op( + grad: torch.Tensor, + ca: int, + cb: int, + a_dtype: Optional[torch.dtype], + b_dtype: Optional[torch.dtype], + want_a: bool, + want_b: bool, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Split ``grad`` back into its ``ca``- and ``cb``-channel halves. + + Returns a zero-element placeholder for a half whose ``want_*`` is False. + + ``grad`` is relaid out only when it is not *line addressable* -- see + :func:`_line_addressable`, and see ``_split_kernel`` for why paying that + copy unconditionally is what made an earlier version of this op lose to + plain ``torch.cat``. + """ + if not _line_addressable(grad): + grad = grad.contiguous(memory_format=_CL_FORMAT) + return _backward( + grad, ca, cb, a_dtype or grad.dtype, b_dtype or grad.dtype, want_a, want_b + ) + + +@_split_op.register_fake +def _(grad, ca, cb, a_dtype, b_dtype, want_a, want_b): + spatial = tuple(grad.shape[2:]) + da = ( + torch.empty( + (grad.shape[0], ca, *spatial), + dtype=a_dtype or grad.dtype, + device=grad.device, + memory_format=_CL_FORMAT, + ) + if want_a + else grad.new_empty(0, dtype=a_dtype or grad.dtype) + ) + db = ( + torch.empty( + (grad.shape[0], cb, *spatial), + dtype=b_dtype or grad.dtype, + device=grad.device, + memory_format=_CL_FORMAT, + ) + if want_b + else grad.new_empty(0, dtype=b_dtype or grad.dtype) + ) + return da, db + + +def _setup_context(ctx, inputs, output): + a, b, _out_dtype = inputs + # Only metadata is saved: the backward of a concatenation does not read + # either input, so saving them would pin two full-resolution activations + # for the whole backward pass for nothing. + ctx.ca = a.shape[1] + ctx.cb = b.shape[1] + ctx.a_dtype = a.dtype + ctx.b_dtype = b.dtype + ctx.needs = (ctx.needs_input_grad[0], ctx.needs_input_grad[1]) + + +def _autograd_backward(ctx, grad_out): + need_a, need_b = ctx.needs + if not (need_a or need_b): + return None, None, None + da, db = torch.ops.scaffold_cat.split_channels( + grad_out, ctx.ca, ctx.cb, ctx.a_dtype, ctx.b_dtype, need_a, need_b + ) + return (da if need_a else None), (db if need_b else None), None + + +torch.library.register_autograd( + "scaffold_cat::cat_channels", _autograd_backward, setup_context=_setup_context +) + + +# --------------------------------------------------------------------------- # +# public API +# --------------------------------------------------------------------------- # +def is_supported(a, b, out_dtype: Optional[torch.dtype] = None) -> bool: + """Whether the Triton kernel can serve ``cat_channels(a, b, out_dtype)``. + + Cheap (attribute reads and two stride checks) and side-effect free: it does + not import Triton, allocate or launch. ``False`` means "use + ``torch.cat``". + + Note that this accepts any ``torch.Tensor`` *instance*, so it must be + called on the tensor the kernel will actually touch. ``skip_concat`` + unwraps DistConv's ``DCTensor`` before asking, for exactly the reason + :class:`ScaFFold.unet.group_norm.FastGroupNorm` documents: a wrapper + subclass's mirrored metadata is not the shard's. + """ + if not (isinstance(a, torch.Tensor) and isinstance(b, torch.Tensor)): + return False + if a.device.type != "cuda" or a.device != b.device: + return False + if not triton_available(): + return False + if a.dim() != 5 or b.dim() != 5: + return False + if a.dtype not in SUPPORTED_DTYPES or b.dtype not in SUPPORTED_DTYPES: + return False + if out_dtype is not None and out_dtype not in SUPPORTED_DTYPES: + return False + if a.shape[0] != b.shape[0] or tuple(a.shape[2:]) != tuple(b.shape[2:]): + return False + if a.shape[1] < 1 or b.shape[1] < 1: + return False + if a.numel() == 0 or b.numel() == 0: + return False + if not a.is_contiguous(memory_format=_CL_FORMAT): + return False + if not b.is_contiguous(memory_format=_CL_FORMAT): + return False + return True + + +def cat_channels(a, b, out_dtype: Optional[torch.dtype] = None): + """``torch.cat([a, b], dim=1)``, optionally emitting ``out_dtype`` directly. + + Total: anything :func:`is_supported` declines is served by ``torch.cat`` + (followed by a cast when ``out_dtype`` asks for one), so this is a drop-in + on CPU, on non-channels-last input and without Triton. + + The fallback casts the *inputs* rather than the concatenated result when + ``out_dtype`` is narrower. That is bitwise the same answer -- every + supported dtype widens exactly into the promoted type, so narrowing before + or after the copy rounds the same values once -- and it does not + materialize a double-width intermediate. + """ + if not is_supported(a, b, out_dtype): + if out_dtype is None: + return torch.cat([a, b], dim=1) + return torch.cat([a.to(out_dtype), b.to(out_dtype)], dim=1) + return torch.ops.scaffold_cat.cat_channels(a, b, out_dtype) + + +def consumer_dtype(*tensors) -> torch.dtype: + """The dtype the convolution that consumes the concatenation will see. + + Under an enabled autocast region the answer is autocast's ``dtype``, + because ``aten::convolution`` carries the ``lower_precision_fp`` cast + policy and casts whatever it is handed. Producing that dtype from the + concatenation is therefore *bitwise identical* to producing ATen's promoted + dtype and letting the convolution narrow it -- the promoted tensor holds + exact widenings of both sources -- while writing (and reading back) half + the bytes, and it removes autocast's own widening of the narrower input on + the way in. + + Outside autocast the answer is ``torch.cat``'s ordinary promotion, so the + result is unchanged for eval, ``inference_mode`` and pure-fp32 runs. + """ + dtype = tensors[0].dtype + for t in tensors[1:]: + dtype = torch.promote_types(dtype, t.dtype) + device_type = tensors[0].device.type + try: + if torch.is_autocast_enabled(device_type): + autocast_dtype = torch.get_autocast_dtype(device_type) + else: + return dtype + except (RuntimeError, TypeError): # a device type autocast does not know + return dtype + # Only ever narrow: if autocast's dtype is not one the kernel can hold, or + # is wider than the inputs, keep the promotion ATen would have done. + if autocast_dtype not in SUPPORTED_DTYPES: + return dtype + if torch.promote_types(autocast_dtype, dtype) is autocast_dtype: + return dtype + return autocast_dtype + + +def _dctensor_ops(input): + """The ``distconv.distconv`` module when ``input`` is a DCTensor, else None. + + Resolved through ``sys.modules`` rather than an import, so this module + stays importable (and the CPU suite runnable) without DistConv installed. + """ + distconv = sys.modules.get("distconv.distconv") + if distconv is not None and isinstance(input, distconv.DCTensor): + return distconv + return None + + +def skip_concat(skip, upsampled): + """Join a decoder skip activation to an upsampled one, as ``Up.forward`` needs. + + Equivalent to ``torch.cat([skip, upsampled], dim=1)`` as the following + convolution observes it -- see :func:`consumer_dtype` for why the dtype may + legitimately differ from ``torch.cat``'s own. + + DistConv's ``DCTensor`` is unwrapped to its local shard in front of the + kernel and rewrapped after, rather than being left to + ``DCTensor.__torch_dispatch__``. Both work -- these are real dispatcher + ops -- but the explicit unwrap is what ``FastGroupNorm`` already does, and + it keeps the subclass policy in one place: ``is_supported`` accepts any + ``torch.Tensor`` instance, so relying on dispatch would silently extend the + fast path to every unknown wrapper subclass. The unwrap goes through + DistConv's ``_ToTensor``/``from_shard`` autograd pair and not a bare + ``._tensor`` read, which would sever the graph back to the producing + convolution. DistConv has no concatenation-specific handling, so the + semantics are identical at every shard count. + """ + distconv = _dctensor_ops(skip) + if distconv is None or _dctensor_ops(upsampled) is None: + return cat_channels(skip, upsampled, consumer_dtype(skip, upsampled)) + if skip._parallel_strategy != upsampled._parallel_strategy: + raise ValueError( + "skip and upsampled tensors have different parallel strategies" + ) + local_skip = distconv._ToTensor.apply(skip) + local_up = distconv._ToTensor.apply(upsampled) + out = cat_channels(local_skip, local_up, consumer_dtype(local_skip, local_up)) + return distconv.DCTensor.from_shard(out, skip._parallel_strategy) diff --git a/ScaFFold/unet/unet_parts.py b/ScaFFold/unet/unet_parts.py index 9fffe72..6f5eac9 100644 --- a/ScaFFold/unet/unet_parts.py +++ b/ScaFFold/unet/unet_parts.py @@ -14,13 +14,13 @@ """Parts of the U-Net model""" -import torch import torch.nn as nn import torch.nn.functional as F from ScaFFold.utils.perf_measure import annotate from .group_norm import FastGroupNorm +from .triton_cat import skip_concat _doubleconv_annotate = annotate(fmt="DoubleConv.{}") _down_annotate = annotate(fmt="Down.{}") @@ -92,7 +92,30 @@ def forward(self, x): class Up(nn.Module): - """Upscaling then double conv""" + """Upscaling then double conv + + The skip concatenation goes through :func:`ScaFFold.unet.triton_cat.skip_concat` + rather than ``torch.cat``. That does two things, both of which leave the + tensor the following convolution reads *bitwise* unchanged (verified in + ``tests/test_triton_cat.py``): + + * It emits the dtype the convolution will use instead of ``torch.cat``'s + promoted one. Under ``torch.autocast`` the two halves do not have the + same dtype -- the skip comes from a GroupNorm, an fp32-policy op, and the + upsampled half comes from a ``ConvTranspose3d`` and is bf16 -- so ``cat`` + widens the bf16 half to fp32, concatenates at fp32, and the convolution + then narrows the whole double-width result straight back down. Three + full-resolution passes to deliver one. + * It runs a channels-last-native kernel, which matters most for the + *backward*: ``cat``'s backward is a narrowed view that every consumer + then forces contiguous, and that strided copy reaches only 51-63% of this + device's streaming roofline against the kernel's 90-103%. + + Both rest on ``self.conv`` beginning with a convolution, which the + constructor below guarantees on either branch. ``skip_concat`` is total: + anything its ``is_supported`` declines -- CPU, non-channels-last, no Triton + -- falls back to ``torch.cat``. + """ def __init__(self, in_channels, out_channels, group_norm_groups, trilinear=True): super().__init__() @@ -135,7 +158,9 @@ def forward(self, x1, x2): # if you have padding issues, see # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd - x = torch.cat([x2, x1], dim=1) + # torch.cat([x2, x1], dim=1) with the dtype and the layout the + # convolution below actually wants; see the class docstring. + x = skip_concat(x2, x1) return self.conv(x) diff --git a/tests/test_triton_cat.py b/tests/test_triton_cat.py new file mode 100644 index 0000000..84a3864 --- /dev/null +++ b/tests/test_triton_cat.py @@ -0,0 +1,478 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +"""Tests for the channels-last Triton channel concatenation (``triton_cat``). + +The op replaces ``torch.cat([a, b], dim=1)``, whose result is a *copy* rather +than an arithmetic expression, so unlike the GroupNorm kernel these are not +tolerance tests: every parity assertion here is **bitwise**, forward and +backward, including the mixed-dtype and narrowing cases. A tolerance would +hide exactly the bugs this kernel can have -- an off-by-one in the channel +split, a lost tail row, a double rounding. + +The other things being pinned down: + +* the **dtype rule**. ``consumer_dtype`` is what licenses emitting bf16 from a + concatenation whose ``torch.cat`` result would be fp32; the tests check both + that it is what the following convolution receives and that it never *widens* + anything. +* the **fallback**. ``is_supported`` must decline everything the kernel cannot + serve physically, and ``cat_channels`` must then still answer -- so the CPU + suite exercises the whole public API with no GPU and no Triton. +* **composition**: a ``DCTensor`` round trip with the autograd graph intact, + ``torch.utils.checkpoint`` recompute, ``inference_mode``, and the + ``trilinear`` branch of ``Up``. + +CPU runs never touch Triton: the module defers ``import triton`` to the first +call that reaches a kernel, which ``test_import_does_not_pull_in_triton`` +checks in a fresh interpreter. +""" + +from __future__ import annotations + +import itertools +import subprocess +import sys + +import pytest +import torch + +from ScaFFold.unet import triton_cat +from ScaFFold.unet.triton_cat import ( + cat_channels, + consumer_dtype, + is_supported, + skip_concat, +) + +CL = torch.channels_last_3d +DTYPES = (torch.float32, torch.bfloat16, torch.float16) + +gpu = pytest.mark.gpu +requires_cuda = pytest.mark.skipif( + not torch.cuda.is_available(), reason="needs a CUDA device" +) + + +def cl_tensor(*shape, dtype=torch.float32, device="cpu", seed=0): + """A deterministic channels-last-3d tensor.""" + generator = torch.Generator().manual_seed(seed) + return ( + torch.randn(*shape, generator=generator) + .to(device=device, dtype=dtype) + .contiguous(memory_format=CL) + ) + + +# --------------------------------------------------------------------------- # +# The public API is total: it must work with no GPU and no Triton at all. +# --------------------------------------------------------------------------- # +def test_cpu_falls_back_and_matches_torch_cat_bitwise(): + """On CPU ``is_supported`` declines and ``cat_channels`` still answers.""" + a = cl_tensor(2, 5, 3, 4, 5, seed=1) + b = cl_tensor(2, 3, 3, 4, 5, seed=2) + assert not is_supported(a, b) + got = cat_channels(a, b) + assert torch.equal(got, torch.cat([a, b], dim=1)) + assert got.is_contiguous(memory_format=CL) + + +def test_cpu_fallback_honours_out_dtype_bitwise(): + """A narrowing ``out_dtype`` must round each value exactly once. + + ``torch.cat([a, b]).to(bf16)`` widens then narrows; the fallback narrows + the inputs first. Those agree bitwise because the promoted dtype is an + exact widening of both, and that is the property the kernel relies on too. + """ + a = cl_tensor(1, 4, 2, 3, 4, seed=3) + b = cl_tensor(1, 4, 2, 3, 4, dtype=torch.bfloat16, seed=4) + for out_dtype in DTYPES: + got = cat_channels(a, b, out_dtype) + assert got.dtype is out_dtype + assert torch.equal(got, torch.cat([a, b], dim=1).to(out_dtype)) + + +def test_cpu_fallback_backward_matches_torch_cat_bitwise(): + a = cl_tensor(1, 6, 2, 2, 2, seed=5).requires_grad_(True) + b = cl_tensor(1, 2, 2, 2, 2, dtype=torch.bfloat16, seed=6).requires_grad_(True) + grad = cl_tensor(1, 8, 2, 2, 2, seed=7) + + torch.cat([a, b], dim=1).backward(grad) + ref = (a.grad.clone(), b.grad.clone()) + a.grad = b.grad = None + + cat_channels(a, b).backward(grad) + assert torch.equal(a.grad, ref[0]) + assert torch.equal(b.grad, ref[1]) + assert a.grad.dtype is torch.float32 + assert b.grad.dtype is torch.bfloat16 + + +@pytest.mark.parametrize( + "make", + [ + pytest.param(lambda: (torch.randn(1, 2, 2, 2, 2), None), id="not-a-tensor"), + pytest.param(lambda: (torch.randn(2, 2), torch.randn(2, 2)), id="rank-2"), + pytest.param( + lambda: (torch.randn(1, 2, 2, 2, 2).double(), torch.randn(1, 2, 2, 2, 2)), + id="float64", + ), + pytest.param( + lambda: (torch.randn(1, 2, 2, 2, 2), torch.randn(1, 2, 2, 2, 3)), + id="spatial-mismatch", + ), + pytest.param( + lambda: (torch.randn(0, 2, 2, 2, 2), torch.randn(0, 2, 2, 2, 2)), + id="empty", + ), + ], +) +def test_is_supported_declines_what_the_kernel_cannot_serve(make): + a, b = make() + assert not is_supported(a, b) + + +def test_import_does_not_pull_in_triton(): + """Importing the module must not import Triton (it is deferred to first use).""" + code = ( + "import sys; import ScaFFold.unet.triton_cat as m; " + "assert 'triton' not in sys.modules, sorted(k for k in sys.modules " + "if k.startswith('triton')); print('ok')" + ) + out = subprocess.run( + [sys.executable, "-c", code], capture_output=True, text=True, check=True + ) + assert out.stdout.strip().endswith("ok") + + +# --------------------------------------------------------------------------- # +# The dtype rule. +# --------------------------------------------------------------------------- # +def test_consumer_dtype_outside_autocast_is_torch_cats_promotion(): + a = cl_tensor(1, 2, 2, 2, 2) + b = cl_tensor(1, 2, 2, 2, 2, dtype=torch.bfloat16) + assert consumer_dtype(a, b) is torch.promote_types(torch.float32, torch.bfloat16) + assert consumer_dtype(a, a) is torch.float32 + assert consumer_dtype(b, b) is torch.bfloat16 + + +@requires_cuda +@gpu +def test_consumer_dtype_under_autocast_is_the_autocast_dtype(): + """The mixed-dtype case the decoder actually presents. + + fp32 skip (a GroupNorm output, fp32 cast policy) + bf16 upsampled (a + ConvTranspose3d output). ``torch.cat`` would answer fp32; the convolution + that consumes it narrows to bf16, so bf16 is the honest answer. + """ + a = cl_tensor(1, 4, 2, 2, 2, device="cuda") + b = cl_tensor(1, 4, 2, 2, 2, dtype=torch.bfloat16, device="cuda") + with torch.autocast("cuda", dtype=torch.bfloat16): + assert consumer_dtype(a, b) is torch.bfloat16 + # never widens: two bf16 inputs stay bf16, and an fp16 autocast over + # bf16 inputs must not silently change their width either + assert consumer_dtype(b, b) is torch.bfloat16 + assert consumer_dtype(a, b) is torch.float32 + + +@requires_cuda +@gpu +def test_skip_concat_is_bitwise_what_the_convolution_receives(): + """The load-bearing claim: the dtype shortcut changes no bits downstream.""" + a = cl_tensor(1, 8, 3, 4, 5, device="cuda") + b = cl_tensor(1, 8, 3, 4, 5, dtype=torch.bfloat16, device="cuda", seed=11) + with torch.autocast("cuda", dtype=torch.bfloat16): + got = skip_concat(a, b) + # what the convolution would have been handed by the old chain: + # torch.cat promotes to fp32, then autocast narrows for the conv. + reference = torch.cat([a, b], dim=1).to(torch.bfloat16) + assert got.dtype is torch.bfloat16 + assert torch.equal(got, reference) + + +# --------------------------------------------------------------------------- # +# GPU parity: bitwise, forward and backward. +# --------------------------------------------------------------------------- # +_SHAPES = [ + (1, 8, 8, 4, 4, 4), # power-of-two channels, the UNet's case + (1, 3, 5, 2, 3, 4), # odd channel counts, C = 8 + (1, 5, 6, 2, 2, 2), # C = 11, not a power of two + (2, 4, 4, 2, 2, 2), # batch > 1 + (1, 1, 1, 1, 1, 1), # degenerate +] + + +@requires_cuda +@gpu +@pytest.mark.parametrize("shape", _SHAPES, ids=lambda s: "x".join(str(v) for v in s)) +@pytest.mark.parametrize("dtypes", list(itertools.product(DTYPES, DTYPES))) +def test_forward_is_bitwise_torch_cat(shape, dtypes): + n, ca, cb, d, h, w = shape + da, db = dtypes + a = cl_tensor(n, ca, d, h, w, dtype=da, device="cuda", seed=21) + b = cl_tensor(n, cb, d, h, w, dtype=db, device="cuda", seed=22) + assert is_supported(a, b) + got = cat_channels(a, b) + ref = torch.cat([a, b], dim=1) + assert got.dtype is ref.dtype + assert got.shape == ref.shape + assert got.is_contiguous(memory_format=CL) + assert torch.equal(got, ref) + + +@requires_cuda +@gpu +@pytest.mark.parametrize("out_dtype", DTYPES) +def test_forward_out_dtype_is_bitwise_cat_then_cast(out_dtype): + a = cl_tensor(1, 6, 3, 4, 5, device="cuda", seed=23) + b = cl_tensor(1, 10, 3, 4, 5, dtype=torch.bfloat16, device="cuda", seed=24) + got = cat_channels(a, b, out_dtype) + assert got.dtype is out_dtype + assert torch.equal(got, torch.cat([a, b], dim=1).to(out_dtype)) + + +@requires_cuda +@gpu +@pytest.mark.parametrize("dtypes", list(itertools.product(DTYPES, DTYPES))) +def test_backward_is_bitwise_torch_cat(dtypes): + da, db = dtypes + a = cl_tensor(1, 6, 2, 3, 4, dtype=da, device="cuda", seed=25).requires_grad_(True) + b = cl_tensor(1, 10, 2, 3, 4, dtype=db, device="cuda", seed=26).requires_grad_(True) + grad = cl_tensor( + 1, 16, 2, 3, 4, dtype=torch.promote_types(da, db), device="cuda", seed=27 + ) + + torch.cat([a, b], dim=1).backward(grad) + ref = (a.grad.clone(), b.grad.clone()) + a.grad = b.grad = None + + cat_channels(a, b).backward(grad) + assert a.grad.dtype is da and b.grad.dtype is db + assert a.grad.is_contiguous(memory_format=CL) + assert b.grad.is_contiguous(memory_format=CL) + assert torch.equal(a.grad, ref[0]) + assert torch.equal(b.grad, ref[1]) + + +@requires_cuda +@gpu +def test_backward_with_only_one_side_requiring_grad(): + """The ``WANT_A``/``WANT_B`` kernel flags must not corrupt the other half.""" + a = cl_tensor(1, 6, 2, 2, 2, device="cuda", seed=28).requires_grad_(True) + b = cl_tensor(1, 2, 2, 2, 2, device="cuda", seed=29) + grad = cl_tensor(1, 8, 2, 2, 2, device="cuda", seed=30) + cat_channels(a, b).backward(grad) + assert b.grad is None + assert torch.equal(a.grad, grad[:, :6]) + + a2 = cl_tensor(1, 6, 2, 2, 2, device="cuda", seed=28) + b2 = cl_tensor(1, 2, 2, 2, 2, device="cuda", seed=29).requires_grad_(True) + cat_channels(a2, b2).backward(grad) + assert a2.grad is None + assert torch.equal(b2.grad, grad[:, 6:]) + + +@requires_cuda +@gpu +def test_backward_accepts_a_non_channels_last_gradient(): + """Nothing guarantees the consumer hands back a channels-last cotangent.""" + a = cl_tensor(1, 6, 2, 2, 2, device="cuda", seed=31).requires_grad_(True) + b = cl_tensor(1, 2, 2, 2, 2, device="cuda", seed=32).requires_grad_(True) + grad = torch.randn(1, 8, 2, 2, 2, device="cuda") # plain contiguous + cat_channels(a, b).backward(grad) + assert torch.equal(a.grad, grad[:, :6]) + assert torch.equal(b.grad, grad[:, 6:]) + + +@requires_cuda +@gpu +@pytest.mark.parametrize("pad", [1, 2]) +def test_backward_reads_a_halo_padded_gradient_in_place(pad): + """The production cotangent: a narrowed view of a halo-padded tensor. + + DistConv reaches the convolution that consumes the concatenation through a + halo exchange that materialises a ``(D+2, H+2, W+2)`` tensor, so what comes + back here is a *narrow* of one -- channels-last per voxel and contiguous + along W, but with a gap at every H and D boundary. Relaying that out costs + a whole extra full-resolution pass, which is exactly what made an earlier + version of this op lose to ``torch.cat``; the kernel must read it in place. + """ + a = cl_tensor(1, 6, 3, 4, 5, device="cuda", seed=51).requires_grad_(True) + b = cl_tensor(1, 10, 3, 4, 5, dtype=torch.bfloat16, device="cuda", seed=52) + b.requires_grad_(True) + parent = cl_tensor( + 1, 16, 3 + 2 * pad, 4 + 2 * pad, 5 + 2 * pad, device="cuda", seed=53 + ) + grad = parent[:, :, pad:-pad, pad:-pad, pad:-pad] + assert grad.shape == (1, 16, 3, 4, 5) + assert not grad.is_contiguous(memory_format=CL) + assert triton_cat._line_addressable(grad), "must take the in-place path" + + torch.cat([a, b], dim=1).backward(grad) + ref = (a.grad.clone(), b.grad.clone()) + a.grad = b.grad = None + + cat_channels(a, b).backward(grad) + assert torch.equal(a.grad, ref[0]) + assert torch.equal(b.grad, ref[1]) + assert a.grad.is_contiguous(memory_format=CL) + assert b.grad.is_contiguous(memory_format=CL) + + +def test_line_addressable_accepts_narrows_and_declines_permutations(): + """The predicate that decides whether the backward can skip its relayout.""" + dense = cl_tensor(1, 8, 4, 5, 6) + assert triton_cat._line_addressable(dense) + parent = cl_tensor(1, 8, 6, 7, 8) + assert triton_cat._line_addressable(parent[:, :, 1:-1, 1:-1, 1:-1]) + assert triton_cat._line_addressable(parent[:, :, :, :, 1:-1]) + # a plain contiguous (NCDHW) tensor has channels *outermost* + assert not triton_cat._line_addressable(torch.randn(1, 8, 4, 5, 6)) + # a spatial permutation breaks the "W neighbours are one run apart" rule + assert not triton_cat._line_addressable(dense.transpose(3, 4)) + # narrowing the channel axis breaks stride(4) == C + assert not triton_cat._line_addressable(dense[:, :4]) + + +@requires_cuda +@gpu +def test_non_channels_last_input_falls_back_to_torch_cat(): + a = torch.randn(1, 6, 2, 2, 2, device="cuda") # contiguous, not CL + b = cl_tensor(1, 2, 2, 2, 2, device="cuda", seed=33) + assert not is_supported(a, b) + assert torch.equal(cat_channels(a, b), torch.cat([a, b], dim=1)) + + +@requires_cuda +@gpu +def test_repeated_calls_are_bitwise_identical(): + """A copy has no reduction order, and the tiling is a pure function of C.""" + a = cl_tensor(1, 6, 3, 4, 5, device="cuda", seed=34).requires_grad_(True) + b = cl_tensor(1, 10, 3, 4, 5, device="cuda", seed=35).requires_grad_(True) + grad = cl_tensor(1, 16, 3, 4, 5, device="cuda", seed=36) + first = cat_channels(a, b).clone() + cat_channels(a, b).backward(grad) + ga, gb = a.grad.clone(), b.grad.clone() + a.grad = b.grad = None + second = cat_channels(a, b).clone() + cat_channels(a, b).backward(grad) + assert torch.equal(first, second) + assert torch.equal(a.grad, ga) and torch.equal(b.grad, gb) + + +@requires_cuda +@gpu +def test_runs_on_a_non_current_device(): + """A Triton launch follows the *current* device; the guard must override it.""" + if torch.cuda.device_count() < 2: + pytest.skip("needs two CUDA devices") + a = cl_tensor(1, 6, 2, 2, 2, device="cuda:1", seed=37) + b = cl_tensor(1, 2, 2, 2, 2, device="cuda:1", seed=38) + with torch.cuda.device(0): + got = cat_channels(a, b) + assert got.device == a.device + assert torch.equal(got, torch.cat([a, b], dim=1)) + + +@requires_cuda +@gpu +def test_second_order_raises_rather_than_returning_garbage(): + """First order only, exactly like the Triton GroupNorm; it must fail loudly. + + The loss has to be *nonlinear* for this to bite. A concatenation is a + copy, so its own second derivative is identically zero: differentiating + ``sum(cat(a, b))`` twice gives a cotangent that does not depend on ``a`` at + all, and autograd correctly reports "does not require grad" without ever + reaching this op. With a square in the way the cotangent *does* depend on + ``a``, the split has to be differentiated, and there is no formula for it. + """ + a = cl_tensor(1, 4, 2, 2, 2, device="cuda", seed=39).requires_grad_(True) + b = cl_tensor(1, 4, 2, 2, 2, device="cuda", seed=40).requires_grad_(True) + out = cat_channels(a, b) + (grad_a,) = torch.autograd.grad((out * out).sum(), a, create_graph=True) + with pytest.raises(RuntimeError, match="no autograd formula was registered"): + torch.autograd.grad(grad_a.sum(), a) + + +@requires_cuda +@gpu +def test_dctensor_round_trip_keeps_the_graph(): + """Production wraps activations in a DCTensor even at ``dc_num_shards=1``.""" + distconv = pytest.importorskip("distconv") + if not torch.distributed.is_initialized(): + pytest.skip("needs an initialized process group") + ps = distconv.ParallelStrategy( + num_shards=(1, 1, 1), shard_dim=(2, 3, 4), device_type="cuda" + ) + a = cl_tensor(1, 4, 2, 2, 2, device="cuda", seed=41).requires_grad_(True) + b = cl_tensor(1, 4, 2, 2, 2, device="cuda", seed=42).requires_grad_(True) + out = skip_concat( + distconv.DCTensor.from_shard(a, ps), distconv.DCTensor.from_shard(b, ps) + ) + assert isinstance(out, distconv.DCTensor) + distconv.distconv._ToTensor.apply(out).pow(2).sum().backward() + assert a.grad is not None and b.grad is not None + + +@requires_cuda +@gpu +def test_survives_activation_checkpoint_recompute(): + """The block's forward is replayed inside backward under checkpointing.""" + a = cl_tensor(1, 4, 2, 2, 2, device="cuda", seed=43).requires_grad_(True) + b = cl_tensor(1, 4, 2, 2, 2, device="cuda", seed=44).requires_grad_(True) + + def block(x, y): + return cat_channels(x, y) * 2.0 + + ref = block(a, b) + ref.pow(2).sum().backward() + ga, gb = a.grad.clone(), b.grad.clone() + a.grad = b.grad = None + + out = torch.utils.checkpoint.checkpoint(block, a, b, use_reentrant=False) + out.pow(2).sum().backward() + assert torch.equal(a.grad, ga) + assert torch.equal(b.grad, gb) + + +@requires_cuda +@gpu +def test_inference_mode_and_no_grad(): + a = cl_tensor(1, 4, 2, 2, 2, device="cuda", seed=45) + b = cl_tensor(1, 4, 2, 2, 2, device="cuda", seed=46) + ref = torch.cat([a, b], dim=1) + with torch.no_grad(): + assert torch.equal(cat_channels(a, b), ref) + with torch.inference_mode(): + assert torch.equal(cat_channels(a, b), ref) + + +@requires_cuda +@gpu +def test_kernel_failure_is_tagged_so_a_caller_can_fall_back(): + """``CatKernelError`` must be what escapes when the launch itself breaks.""" + a = cl_tensor(1, 4, 2, 2, 2, device="cuda", seed=47) + b = cl_tensor(1, 4, 2, 2, 2, device="cuda", seed=48) + original = triton_cat._forward.__wrapped__ + + def boom(*args, **kwargs): + raise ValueError("simulated launch failure") + + triton_cat._forward.__wrapped__ = boom + try: + wrapped = triton_cat._tag_kernel_failures(boom) + with pytest.raises(triton_cat.CatKernelError): + wrapped(a, b, torch.float32) + finally: + triton_cat._forward.__wrapped__ = original diff --git a/tests/test_unet.py b/tests/test_unet.py index b61ae21..a17b65c 100644 --- a/tests/test_unet.py +++ b/tests/test_unet.py @@ -198,3 +198,145 @@ def counting_pad(tensor, pad, *args, **kwargs): f"Guard not yet in place: {exact_match_pad_calls} pad calls with diffs=0 " f"(expected 0 when fixed). This is the RED baseline." ) + + +# --------------------------------------------------------------------------- # +# The decoder skip concatenation. +# +# ``Up.forward`` no longer calls ``torch.cat`` directly; it goes through +# ``ScaFFold.unet.triton_cat.skip_concat``, which may legitimately emit a +# narrower dtype than ``torch.cat`` would when autocast is on (see the ``Up`` +# docstring). Everything below pins the part that must NOT change: outside +# autocast the block is bitwise what it was, the ``F.pad`` path still works, +# and both the ``trilinear`` and ``ConvTranspose3d`` branches agree with an +# explicit ``torch.cat`` reference. The kernel's own parity tests live in +# ``tests/test_triton_cat.py``. +# --------------------------------------------------------------------------- # +@pytest.mark.parametrize("trilinear", [False, True]) +def test_up_matches_an_explicit_torch_cat_reference(trilinear): + """``Up.forward`` must equal ``conv(cat([x2, up(x1)]))``, bitwise, on CPU.""" + from ScaFFold.unet.unet_parts import Up + + up = Up(in_channels=32, out_channels=16, group_norm_groups=8, trilinear=trilinear) + up.eval() + generator = torch.Generator().manual_seed(11) + # Either branch must hand ``self.conv`` ``in_channels`` channels: the + # transposed convolution halves 32 -> 16, while ``nn.Upsample`` changes no + # channels, so its input already carries 16. + x1 = torch.randn(1, 16 if trilinear else 32, 8, 8, 8, generator=generator) + x2 = torch.randn(1, 16, 16, 16, 16, generator=generator) + + with torch.no_grad(): + got = up(x1, x2) + reference = up.conv(torch.cat([x2, up.up(x1)], dim=1)) + + assert got.shape == reference.shape + assert torch.equal(got, reference), ( + "the skip concatenation must be bitwise torch.cat outside autocast" + ) + + +def test_up_still_pads_and_concatenates_when_shapes_disagree(): + """The non-power-of-two path: ``F.pad`` fires and the result still matches.""" + import torch.nn.functional as F + + from ScaFFold.unet.unet_parts import Up + + up = Up(in_channels=32, out_channels=16, group_norm_groups=8, trilinear=False) + up.eval() + generator = torch.Generator().manual_seed(12) + x1 = torch.randn(1, 32, 7, 7, 7, generator=generator) # -> 14^3 after up + x2 = torch.randn(1, 16, 16, 16, 16, generator=generator) # 16^3: diff = 2 + + with torch.no_grad(): + got = up(x1, x2) + padded = F.pad(up.up(x1), [1, 1, 1, 1, 1, 1]) + reference = up.conv(torch.cat([x2, padded], dim=1)) + + assert got.shape == (1, 16, 16, 16, 16) + assert torch.equal(got, reference) + + +def test_up_gradients_match_an_explicit_torch_cat_reference(): + """Backward through the skip concatenation, bitwise, on CPU.""" + from ScaFFold.unet.unet_parts import Up + + up = Up(in_channels=32, out_channels=16, group_norm_groups=8, trilinear=False) + generator = torch.Generator().manual_seed(13) + x1 = torch.randn(1, 32, 8, 8, 8, generator=generator) + x2 = torch.randn(1, 16, 16, 16, 16, generator=generator) + + a, b = x1.clone().requires_grad_(True), x2.clone().requires_grad_(True) + up.zero_grad(set_to_none=True) + up(a, b).pow(2).sum().backward() + got = (a.grad.clone(), b.grad.clone()) + got_params = {n: p.grad.clone() for n, p in up.named_parameters()} + + c, d = x1.clone().requires_grad_(True), x2.clone().requires_grad_(True) + up.zero_grad(set_to_none=True) + up.conv(torch.cat([d, up.up(c)], dim=1)).pow(2).sum().backward() + + assert torch.equal(got[0], c.grad) + assert torch.equal(got[1], d.grad) + for name, param in up.named_parameters(): + assert torch.equal(got_params[name], param.grad), name + + +def test_up_concatenation_keeps_channels_last(): + """The concatenation must not break the layout chain it exists to preserve. + + Asserted on ``skip_concat`` with two channels-last halves rather than on a + whole ``Up`` block: on CPU ``nn.ConvTranspose3d`` returns a *contiguous* + tensor whatever it is handed, so the block's own inputs to the + concatenation are not both channels-last there and the block-level + assertion would be measuring the convolution's layout policy, not this + one's. On GPU with ``PYTORCH_MIOPEN_SUGGEST_NHWC=1`` -- the configuration + the kernel exists for -- both halves are channels-last and this is the + property ``Up`` relies on. + """ + from ScaFFold.unet.triton_cat import skip_concat + + generator = torch.Generator().manual_seed(14) + x1 = torch.randn(1, 16, 16, 16, 16, generator=generator).contiguous( + memory_format=torch.channels_last_3d + ) + x2 = torch.randn(1, 16, 16, 16, 16, generator=generator).contiguous( + memory_format=torch.channels_last_3d + ) + out = skip_concat(x2, x1) + assert out.shape == (1, 32, 16, 16, 16) + assert out.is_contiguous(memory_format=torch.channels_last_3d) + assert torch.equal(out, torch.cat([x2, x1], dim=1)) + + +def test_whole_model_forward_and_backward_still_agree_with_a_cat_based_up(): + """End to end: swapping the concatenation back must change nothing on CPU.""" + import torch as _torch + + from ScaFFold.unet import unet_parts + + def cat_forward(self, x1, x2): + x1 = self.up(x1) + return self.conv(_torch.cat([x2, x1], dim=1)) + + model = UNet( + n_channels=_N_CHANNELS, n_classes=_N_CLASSES, trilinear=False, layers=2 + ) + x = _make_input(seed=15).requires_grad_(True) + + model.zero_grad(set_to_none=True) + model(x).pow(2).sum().backward() + grads = {n: p.grad.clone() for n, p in model.named_parameters()} + x_grad = x.grad.clone() + + original = unet_parts.Up.forward + try: + unet_parts.Up.forward = cat_forward + x2 = _make_input(seed=15).requires_grad_(True) + model.zero_grad(set_to_none=True) + model(x2).pow(2).sum().backward() + for name, param in model.named_parameters(): + assert torch.equal(grads[name], param.grad), name + assert torch.equal(x_grad, x2.grad) + finally: + unet_parts.Up.forward = original From 86087b4ae68ce7ece4072d3a2a448a5fd598e3e6 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Sun, 2 Aug 2026 00:42:48 -0700 Subject: [PATCH 62/62] Drop the concatenation kernel, keep the dtype fix Of the 1.18 ms ff2813c took off the scale-7 step, 1.09 ms and all 0.50 GiB of the memory came from concatenating at the dtype the following convolution reads instead of at torch.cat's promoted one. The Triton kernel was worth the remaining 0.08 ms -- under 0.1% of the step. That does not pay for 824 lines of hand-written kernel plus its 69 tests in a benchmark whose value depends on other people trusting it, particularly so soon after a review round found nine defects in the first one. The dtype computation moves into unet_parts as _consumer_dtype/_skip_concat, which is what the kernel's own fallback path already was. Behaviour outside autocast is unchanged, so eval, inference_mode and pure-fp32 runs still get torch.cat's ordinary promotion. What the kernel measured is kept in the Up docstring rather than the code: cat's backward is a narrowed view that consumers force contiguous, at 51-63% of this device's streaming roofline against the kernel's 90-103%. If the skip path is ever worth revisiting, that is where the remaining headroom is. CPU 397 passed, GPU 270. --- ScaFFold/unet/triton_cat.py | 824 ------------------------------------ ScaFFold/unet/unet_parts.py | 86 +++- tests/test_triton_cat.py | 478 --------------------- tests/test_unet.py | 23 +- 4 files changed, 74 insertions(+), 1337 deletions(-) delete mode 100644 ScaFFold/unet/triton_cat.py delete mode 100644 tests/test_triton_cat.py diff --git a/ScaFFold/unet/triton_cat.py b/ScaFFold/unet/triton_cat.py deleted file mode 100644 index bb52b97..0000000 --- a/ScaFFold/unet/triton_cat.py +++ /dev/null @@ -1,824 +0,0 @@ -# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. -# Produced at the Lawrence Livermore National Laboratory. -# Written by the LBANN Research Team (B. Van Essen, et al.) listed in -# the CONTRIBUTORS file. See the top-level LICENSE file for details. -# -# LLNL-CODE-697807. -# All rights reserved. -# -# This file is part of LBANN: Livermore Big Artificial Neural Network -# Toolkit. For details, see http://software.llnl.gov/LBANN or -# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. -# -# SPDX-License-Identifier: (Apache-2.0) - -"""Channels-last-native Triton channel concatenation (NDHWC in, NDHWC out). - -Why this exists -=============== -Every ``Up`` block of the UNet joins the decoder's upsampled activation to the -encoder's skip activation with ``torch.cat([skip, up], dim=1)`` and feeds the -result to a convolution. Profiled at scale 7 (channels-last, bf16 autocast, -DCTensor) that concatenation and the copies autocast wraps around it cost -**~10.4 ms of a 92 ms step, 11.6%** -- the largest remaining item in ScaFFold's -own code once the GroupNorm kernel landed. It is not one problem but two, and -neither is inherent: - -**1. The dtypes are wrong on the wire.** Under ``torch.autocast`` the two -inputs do not have the same dtype. GroupNorm carries autocast's ``fp32`` cast -policy, so the skip tensor -- which is a ``DoubleConv`` output, i.e. a -GroupNorm output -- arrives as **fp32**, while the upsampled tensor comes -straight out of a ``ConvTranspose3d`` and is **bf16**. ``aten::cat`` carries -the ``promote`` policy, so autocast *widens the bf16 input to fp32*, cats at -fp32, and then the following convolution -- ``lower_precision_fp`` policy -- -casts the whole double-width result straight back down to bf16. At the largest -decoder block ([1,64,128^3] skip + [1,64,128^3] up) that is - - up->fp32 read 268 MB write 537 MB - cat read 1074 MB write 1074 MB - conv cast read 1074 MB write 537 MB = 4.56 GB of traffic - -to deliver 537 MB of bf16 to the convolution. Emitting bf16 *directly from the -concatenation* is **bitwise identical** to what the convolution receives today --- the fp32 intermediate holds exact copies of an fp32 tensor and of a widened -bf16 tensor, so rounding it to bf16 recovers exactly ``(bf16(skip), up)`` -- -and costs one pass: read 537 + 268, write 537 = **1.34 GB, 3.4x less**. - -**2. The kernel iterates the wrong order.** This is the same defect the -GroupNorm kernel exists to fix. A ``channels_last_3d`` tensor ``(N, C, D, H, -W)`` is *physically* a dense ``(N*D*H*W, C)`` array, so a channel -concatenation is, in memory, ``out[m, :Ca] = a[m, :]`` and ``out[m, Ca:] = -b[m, :]`` -- a pure streaming join of two dense arrays into one, 2 reads and 1 -write, perfectly coalescable. ATen reaches it through the *logical* NCDHW -order and lands on TensorIterator's generic offset-calculator path, where the -output tile is never contiguous. - -The kernels below own a ``(BLOCK_M, C)`` tile of the *physical* array: they -issue one fully contiguous store per tile and gather the two halves of it from -the two sources, whose valid lanes are themselves contiguous runs. The dtype -conversion rides along in registers, so the retyping in point 1 is free. - -Public API -========== -``cat_channels(a, b, out_dtype=None)`` - ``torch.cat([a, b], dim=1)`` with an optional output dtype override. - ``out_dtype=None`` reproduces ``torch.cat``'s own promotion exactly. - Anything :func:`is_supported` declines is served by ``torch.cat`` itself. - -``is_supported(a, b, out_dtype=None)`` - Cheap, side-effect-free predicate: ``True`` exactly when the Triton kernel - will run. - -``skip_concat(skip, upsampled)`` - The UNet decoder's call: chooses the output dtype (see - :func:`consumer_dtype`), unwraps DistConv's ``DCTensor`` to its local shard - and rewraps the result. This is the only function ``unet_parts`` calls. - -Contract -======== -For every input :func:`is_supported` accepts, ``cat_channels(a, b, dt)`` equals -``torch.cat([a, b], dim=1).to(dt)`` **bitwise**, with: - -* **memory format** -- the output is ``channels_last_3d``-contiguous, which is - also what ``torch.cat`` returns for channels-last inputs. -* **dtype** -- exactly ``out_dtype``, or ``torch.promote_types(a.dtype, - b.dtype)`` when that is ``None``. Narrowing is a *single* rounding of each - source value: every supported dtype widens exactly into fp32, which is the - kernel's compute type, so rounding fp32->bf16 once in the kernel is the same - bits as ATen's widen-then-narrow. -* **autograd** -- first order, ``d_a`` and ``d_b`` in the *inputs'* dtypes, - which is what autograd requires and what the current chain produces after - autocast's cast nodes run their own backward. Second-order raises, exactly - as :mod:`ScaFFold.unet.triton_group_norm` does and for the same reason: the - backward is itself a custom op with no autograd formula. -* **determinism** -- trivially bitwise reproducible. There is no reduction, - no atomic and no autotuning; every output element is a copy of exactly one - input element and the tile shape is a pure function of the channel count. -* **device** -- the kernels run on the inputs' device whatever device is - current; see ``_device_guard``, which exists because a Triton launch follows - the *current* device and not its arguments'. -* **rejections** -- :func:`is_supported` declines anything the kernel cannot - serve physically (non-channels-last, non-5-D, unsupported dtype, CPU, - mismatched spatial extent, empty), and :func:`cat_channels` then routes it to - ``torch.cat``, so the function is total. - -Measured cost -============= -See ``review/skip-path/RESULTS.md`` for the interleaved A/B, the op-level -before/after and the gradient check. -""" - -from __future__ import annotations - -import contextlib -import functools -import importlib.util -import sys -from typing import Optional, Tuple - -import torch - -__all__ = [ - "cat_channels", - "is_supported", - "skip_concat", - "consumer_dtype", - "CatKernelError", -] - - -class CatKernelError(RuntimeError): - """A failure of the Triton kernels themselves, with the original as ``__cause__``. - - Mirrors :class:`ScaFFold.unet.triton_group_norm.TritonKernelError`: it is - raised only from a *closed* region that allocates and launches and runs no - autograd-observable op, so a caller may retry the call on ``torch.cat`` - without worrying that half a graph was already recorded. - ``torch.OutOfMemoryError`` is passed through untagged -- it is a resource - condition, not a defect, and the fallback would allocate the same bytes. - """ - - -#: Dtypes the kernels read and write directly. All three widen exactly into -#: fp32, which is what makes the single-rounding claim in the docstring hold. -SUPPORTED_DTYPES = (torch.float32, torch.bfloat16, torch.float16) - -_CL_FORMAT = torch.channels_last_3d -_INT32_MAX = 2**31 - 1 - -#: Elements per tile and the cap on voxels per program. Both are pure -#: functions of the channel count, so a run cannot change the tiling underneath -#: a comparison. Chosen by a sweep of BLOCK_M in {1..64} x num_warps in -#: {1,2,4,8} at the four scale-7 decoder shapes -#: (``review/skip-path/logs/cat_bench_tune.log``, -#: ``split_bench_tune.log``): with ``BLOCK_M = clamp(4096 // next_pow2(C), 1, -#: 32)`` and four warps the kernels are within 0.8% of the per-shape optimum at -#: the two shapes that dominate and within 4% at the two launch-bound ones, -#: which is not worth a frozen table. -_TILE_ELEMS = 4096 -_MAX_BLOCK_M = 32 -_NUM_WARPS = 4 - - -# --------------------------------------------------------------------------- # -# Triton kernels -# --------------------------------------------------------------------------- # -triton = None -tl = None -_cat_kernel = None -_split_kernel = None - -_TRITON_AVAILABLE: Optional[bool] = None - - -def triton_available() -> bool: - """Whether ``triton`` can be imported, cached, without importing it.""" - global _TRITON_AVAILABLE - if _TRITON_AVAILABLE is None: - try: - _TRITON_AVAILABLE = importlib.util.find_spec("triton") is not None - except (ImportError, ValueError): - _TRITON_AVAILABLE = False - return _TRITON_AVAILABLE - - -def _build_kernels(): - """Import Triton and install the JIT kernels into this module's globals. - - Defined inside a function purely so ``import triton`` is deferred to the - first GPU call, and written into ``globals()`` because Triton resolves - names through ``fn.__globals__``. - """ - global triton, tl - import triton as _triton - import triton.language as _tl - - triton = _triton - tl = _tl - - @_triton.jit - def _cat_kernel( - A, - B, - OUT, - M, - CA: tl.constexpr, - CB: tl.constexpr, - C: tl.constexpr, - CAP: tl.constexpr, - CBP: tl.constexpr, - BLOCK_M: tl.constexpr, - INT64: tl.constexpr, - ): - """One program per ``BLOCK_M`` voxels: join two dense rows into one. - - Each source keeps its **own** tile width (``next_pow2`` of its channel - count) and is copied by its own load/store pair, rather than both being - gathered into one ``(BLOCK_M, C)`` tile and stored once. The single - fully contiguous store is the more obvious design and was measured - first; it is **20% slower** end to end over the four decoder shapes - (0.718 vs 0.571 ms, ``review/skip-path/logs/cat_bench_tune.log``). The - reason is that the one-store form has to mask *both* loads down to - complementary halves of a double-width lane space, which halves the - useful work per instruction and defeats vectorization, and it buys only - a contiguous store -- whereas a store of ``CA`` contiguous elements at a - stride of ``C`` already covers whole cache lines whenever - ``CA * itemsize`` is a multiple of the line, which it is for every - channel count this network uses. The strided store is not the problem; - the mask was. - """ - pid = tl.program_id(0) - m0 = pid * BLOCK_M - if INT64: - wide = m0.to(tl.int64) - base_a = wide * CA - base_b = wide * CB - base_o = wide * C - else: - base_a = m0 * CA - base_b = m0 * CB - base_o = m0 * C - - rows = tl.arange(0, BLOCK_M) - rmask = rows < M - m0 - ca = tl.arange(0, CAP) - cb = tl.arange(0, CBP) - # Clamp the padding lanes' column index: their loads and stores are - # masked and never touch memory, but keeping the arithmetic inside the - # allocation avoids forming a pointer the compiler may treat as poison. - cam = ca < CA - cbm = cb < CB - ca = tl.where(cam, ca, 0) - cb = tl.where(cbm, cb, 0) - am = rmask[:, None] & cam[None, :] - bm = rmask[:, None] & cbm[None, :] - - av = tl.load(A + base_a + rows[:, None] * CA + ca[None, :], mask=am, other=0.0) - tl.store( - OUT + base_o + rows[:, None] * C + ca[None, :], - av.to(OUT.dtype.element_ty), - mask=am, - ) - bv = tl.load(B + base_b + rows[:, None] * CB + cb[None, :], mask=bm, other=0.0) - tl.store( - OUT + base_o + rows[:, None] * C + CA + cb[None, :], - bv.to(OUT.dtype.element_ty), - mask=bm, - ) - - @_triton.jit - def _split_kernel( - G, - DA, - DB, - SN, - SD, - SH, - D, - H, - W, - CA: tl.constexpr, - CB: tl.constexpr, - C: tl.constexpr, - CAP: tl.constexpr, - CBP: tl.constexpr, - BLOCK_W: tl.constexpr, - WANT_A: tl.constexpr, - WANT_B: tl.constexpr, - INT64: tl.constexpr, - ): - """The transpose of :func:`_cat_kernel`: two strided loads, two stores. - - This is the pass ATen is genuinely bad at. ``cat``'s backward hands the - consumer a *narrowed view*, and every consumer then forces it - contiguous, so the work happens as a generic strided ``copy_`` that - reaches only **51-63%** of this device's streaming roofline at the four - decoder shapes; this kernel reaches **90-103%** - (``review/skip-path/logs/split_bench_tune.log``). - - **The incoming gradient is addressed by its strides, not assumed dense**, - and that is not a nicety. Under DistConv -- which production uses even - at ``dc_num_shards=1`` -- the convolution that consumes the - concatenation is reached through a halo exchange that materialises a - *padded* tensor, so the cotangent that comes back here is a narrowed - view of a ``(D+2, H+2, W+2)`` one: channels-last within each voxel and - contiguous along W, but with a gap at every H and D boundary. An - earlier version simply called ``.contiguous(memory_format=channels_last_3d)`` - on it, which costs a **whole extra full-resolution pass** that ATen's - view-based backward never pays -- 0.46 ms per step at the largest - decoder shape, enough on its own to turn this kernel from a win into a - loss against plain ``torch.cat`` at the right dtype (measured: the - isolated four-block sum went from 53.40 ms with the relayout to 52.10 - without). - - The requirement is therefore only that channels are innermost - (``stride(1) == 1``) and that a voxel's neighbours along W are one - channel-run apart (``stride(4) == C``); everything above W is addressed - through ``SN``/``SD``/``SH``. That admits a dense channels-last tensor - and any narrowing of one on D, H or W, which is every case this op - sees. The driver falls back to a relayout for anything else. - """ - pid = tl.program_id(0) - h = tl.program_id(1) - nd = tl.program_id(2) - n = nd // D - d = nd % D - - w0 = pid * BLOCK_W - ws = w0 + tl.arange(0, BLOCK_W) - wmask = ws < W - - if INT64: - gbase = n.to(tl.int64) * SN + d.to(tl.int64) * SD + h.to(tl.int64) * SH - row0 = ((n.to(tl.int64) * D + d) * H + h) * W + w0 - else: - gbase = n * SN + d * SD + h * SH - row0 = ((n * D + d) * H + h) * W + w0 - gbase = gbase + w0 * C - - wl = tl.arange(0, BLOCK_W) - ca = tl.arange(0, CAP) - cb = tl.arange(0, CBP) - cam = ca < CA - cbm = cb < CB - ca = tl.where(cam, ca, 0) - cb = tl.where(cbm, cb, 0) - am = wmask[:, None] & cam[None, :] - bm = wmask[:, None] & cbm[None, :] - - if WANT_A: - ga = tl.load(G + gbase + wl[:, None] * C + ca[None, :], mask=am, other=0.0) - tl.store( - DA + row0 * CA + wl[:, None] * CA + ca[None, :], - ga.to(DA.dtype.element_ty), - mask=am, - ) - if WANT_B: - gb = tl.load( - G + gbase + wl[:, None] * C + CA + cb[None, :], mask=bm, other=0.0 - ) - tl.store( - DB + row0 * CB + wl[:, None] * CB + cb[None, :], - gb.to(DB.dtype.element_ty), - mask=bm, - ) - - globals().update(_cat_kernel=_cat_kernel, _split_kernel=_split_kernel) - - -def _ensure_kernels(): - if _cat_kernel is None: - _build_kernels() - - -# --------------------------------------------------------------------------- # -# python drivers -# --------------------------------------------------------------------------- # -_NO_GUARD = contextlib.nullcontext() - - -def _device_guard(device: torch.device): - """Make ``device`` current for the kernel launches inside the ``with``. - - A Triton launch goes to whatever device is *current*, not to the device its - arguments live on; without this a tensor on ``cuda:1`` while ``cuda:0`` is - current makes the kernel dereference another device's pointers and the - process dies with ``Memory access fault by GPU node-N``. ATen ops carry a - ``DeviceGuard`` and handle the same call, so this is required for the - drop-in contract. The ``current_device()`` test keeps the common - (already-current) path free; see the same helper in - :mod:`ScaFFold.unet.triton_group_norm` for the measurement. - """ - if device.index == torch.cuda.current_device(): - return _NO_GUARD - return torch.cuda.device(device) - - -def _next_pow2(x: int) -> int: - return 1 << (x - 1).bit_length() if x > 1 else 1 - - -def _cdiv(a: int, b: int) -> int: - return -(-a // b) - - -@functools.lru_cache(maxsize=64) -def _block_m(channels_pow2: int) -> int: - """Voxels per program. A pure function of the padded channel count. - - Bitwise determinism does not actually depend on this -- a copy has no - reduction order to perturb -- but keeping it a pure function of the shape - means a run cannot change tiling underneath a comparison, which is the - property the GroupNorm kernel's frozen table exists to give and is worth - having for free here too. - """ - return max(1, min(_MAX_BLOCK_M, _TILE_ELEMS // channels_pow2)) - - -@functools.lru_cache(maxsize=64) -def _block_w(channels_pow2: int) -> int: - """Voxels per program for the backward, which tiles along W within a line. - - Same rule and same sweep as :func:`_block_m`; the backward cannot tile over - a flat voxel index because its source is only guaranteed contiguous *within* - a ``(n, d, h)`` line -- see ``_split_kernel``. - """ - return max(1, min(_MAX_BLOCK_M, _TILE_ELEMS // channels_pow2)) - - -def _rows_of(t: torch.Tensor) -> int: - """``N * D * H * W`` -- the number of voxels in the physical (rows, C) view.""" - rows = t.shape[0] - for d in t.shape[2:]: - rows *= d - return rows - - -def _tag_kernel_failures(fn): - """Re-raise anything ``fn`` raises as :class:`CatKernelError`. - - Applied to the two functions that do nothing but allocate and launch, so - the tagged region cannot swallow framework control flow (there is no pack - hook, no recompute stop and no functorch layer inside it). - """ - - @functools.wraps(fn) - def wrapper(*args, **kwargs): - try: - return fn(*args, **kwargs) - except torch.OutOfMemoryError: - raise - except CatKernelError: - raise - except Exception as e: # noqa: BLE001 -- re-raised, see docstring - raise CatKernelError( - f"{fn.__name__} failed ({type(e).__name__}: {e})" - ) from e - - return wrapper - - -@_tag_kernel_failures -def _forward(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype): - _ensure_kernels() - ca, cb = a.shape[1], b.shape[1] - c = ca + cb - rows = _rows_of(a) - cp = _next_pow2(c) - block_m = _block_m(cp) - - with _device_guard(a.device): - out = torch.empty( - (a.shape[0], c, *a.shape[2:]), - dtype=out_dtype, - device=a.device, - memory_format=_CL_FORMAT, - ) - _cat_kernel[(_cdiv(rows, block_m),)]( - a, - b, - out, - rows, - CA=ca, - CB=cb, - C=c, - CAP=_next_pow2(ca), - CBP=_next_pow2(cb), - BLOCK_M=block_m, - INT64=rows * c > _INT32_MAX, - num_warps=_NUM_WARPS, - ) - return out - - -@_tag_kernel_failures -def _backward(grad, ca, cb, a_dtype, b_dtype, want_a, want_b): - _ensure_kernels() - c = ca + cb - n, _, d, h, w = grad.shape - sn, _, sd, sh, _ = grad.stride() - block_w = _block_w(_next_pow2(c)) - - with _device_guard(grad.device): - da = torch.empty( - (n, ca, d, h, w) if want_a else (0,), - dtype=a_dtype, - device=grad.device, - **({"memory_format": _CL_FORMAT} if want_a else {}), - ) - db = torch.empty( - (n, cb, d, h, w) if want_b else (0,), - dtype=b_dtype, - device=grad.device, - **({"memory_format": _CL_FORMAT} if want_b else {}), - ) - if want_a or want_b: - _split_kernel[(_cdiv(w, block_w), h, n * d)]( - grad, - da, - db, - sn, - sd, - sh, - d, - h, - w, - CA=ca, - CB=cb, - C=c, - CAP=_next_pow2(ca), - CBP=_next_pow2(cb), - BLOCK_W=block_w, - WANT_A=want_a, - WANT_B=want_b, - # The widest index the kernel forms is the *source* base, which - # spans the (possibly padded) parent tensor, so it is bounded by - # the largest stride and not by this tensor's own element count. - INT64=max(sn * n, n * d * h * w * c) > _INT32_MAX, - num_warps=_NUM_WARPS, - ) - return da, db - - -# --------------------------------------------------------------------------- # -# torch.library registration -# --------------------------------------------------------------------------- # -def _line_addressable(grad: torch.Tensor) -> bool: - """Whether ``_split_kernel`` can read ``grad`` in place. - - It needs channels innermost and voxels one channel-run apart along W; the - D and H axes are addressed through their own strides, so any narrowing of a - channels-last tensor qualifies -- which is what DistConv's halo-padded - cotangent is. A dense channels-last tensor is the special case where the - strides happen to be exact multiples. - """ - if grad.dim() != 5: - return False - stride = grad.stride() - return stride[1] == 1 and stride[4] == grad.shape[1] - - -def _validate(a, b, out_dtype): - if a.dim() != 5 or b.dim() != 5: - raise ValueError( - f"expected two 5-D NCDHW tensors, got {tuple(a.shape)} and {tuple(b.shape)}" - ) - if a.shape[0] != b.shape[0] or tuple(a.shape[2:]) != tuple(b.shape[2:]): - raise ValueError( - f"shapes must agree except on dim 1, got {tuple(a.shape)} and " - f"{tuple(b.shape)}" - ) - if a.shape[1] < 1 or b.shape[1] < 1 or a.numel() == 0 or b.numel() == 0: - raise ValueError("both inputs must have at least one channel and be non-empty") - for name, t in (("a", a), ("b", b)): - if t.dtype not in SUPPORTED_DTYPES: - raise ValueError(f"unsupported dtype {t.dtype} for {name}") - if not t.is_contiguous(memory_format=_CL_FORMAT): - raise ValueError( - f"{name} must be channels_last_3d-contiguous; use cat_channels() " - "which falls back to torch.cat for other layouts" - ) - if out_dtype is not None and out_dtype not in SUPPORTED_DTYPES: - raise ValueError(f"unsupported out_dtype {out_dtype}") - - -@torch.library.custom_op( - "scaffold_cat::cat_channels", mutates_args=(), device_types="cuda" -) -def _cat_op( - a: torch.Tensor, b: torch.Tensor, out_dtype: Optional[torch.dtype] -) -> torch.Tensor: - """``torch.cat([a, b], dim=1)`` for channels-last-3d inputs.""" - _validate(a, b, out_dtype) - return _forward(a, b, out_dtype or torch.promote_types(a.dtype, b.dtype)) - - -@_cat_op.register_fake -def _(a, b, out_dtype): - return torch.empty( - (a.shape[0], a.shape[1] + b.shape[1], *a.shape[2:]), - dtype=out_dtype or torch.promote_types(a.dtype, b.dtype), - device=a.device, - memory_format=_CL_FORMAT, - ) - - -@torch.library.custom_op( - "scaffold_cat::split_channels", mutates_args=(), device_types="cuda" -) -def _split_op( - grad: torch.Tensor, - ca: int, - cb: int, - a_dtype: Optional[torch.dtype], - b_dtype: Optional[torch.dtype], - want_a: bool, - want_b: bool, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Split ``grad`` back into its ``ca``- and ``cb``-channel halves. - - Returns a zero-element placeholder for a half whose ``want_*`` is False. - - ``grad`` is relaid out only when it is not *line addressable* -- see - :func:`_line_addressable`, and see ``_split_kernel`` for why paying that - copy unconditionally is what made an earlier version of this op lose to - plain ``torch.cat``. - """ - if not _line_addressable(grad): - grad = grad.contiguous(memory_format=_CL_FORMAT) - return _backward( - grad, ca, cb, a_dtype or grad.dtype, b_dtype or grad.dtype, want_a, want_b - ) - - -@_split_op.register_fake -def _(grad, ca, cb, a_dtype, b_dtype, want_a, want_b): - spatial = tuple(grad.shape[2:]) - da = ( - torch.empty( - (grad.shape[0], ca, *spatial), - dtype=a_dtype or grad.dtype, - device=grad.device, - memory_format=_CL_FORMAT, - ) - if want_a - else grad.new_empty(0, dtype=a_dtype or grad.dtype) - ) - db = ( - torch.empty( - (grad.shape[0], cb, *spatial), - dtype=b_dtype or grad.dtype, - device=grad.device, - memory_format=_CL_FORMAT, - ) - if want_b - else grad.new_empty(0, dtype=b_dtype or grad.dtype) - ) - return da, db - - -def _setup_context(ctx, inputs, output): - a, b, _out_dtype = inputs - # Only metadata is saved: the backward of a concatenation does not read - # either input, so saving them would pin two full-resolution activations - # for the whole backward pass for nothing. - ctx.ca = a.shape[1] - ctx.cb = b.shape[1] - ctx.a_dtype = a.dtype - ctx.b_dtype = b.dtype - ctx.needs = (ctx.needs_input_grad[0], ctx.needs_input_grad[1]) - - -def _autograd_backward(ctx, grad_out): - need_a, need_b = ctx.needs - if not (need_a or need_b): - return None, None, None - da, db = torch.ops.scaffold_cat.split_channels( - grad_out, ctx.ca, ctx.cb, ctx.a_dtype, ctx.b_dtype, need_a, need_b - ) - return (da if need_a else None), (db if need_b else None), None - - -torch.library.register_autograd( - "scaffold_cat::cat_channels", _autograd_backward, setup_context=_setup_context -) - - -# --------------------------------------------------------------------------- # -# public API -# --------------------------------------------------------------------------- # -def is_supported(a, b, out_dtype: Optional[torch.dtype] = None) -> bool: - """Whether the Triton kernel can serve ``cat_channels(a, b, out_dtype)``. - - Cheap (attribute reads and two stride checks) and side-effect free: it does - not import Triton, allocate or launch. ``False`` means "use - ``torch.cat``". - - Note that this accepts any ``torch.Tensor`` *instance*, so it must be - called on the tensor the kernel will actually touch. ``skip_concat`` - unwraps DistConv's ``DCTensor`` before asking, for exactly the reason - :class:`ScaFFold.unet.group_norm.FastGroupNorm` documents: a wrapper - subclass's mirrored metadata is not the shard's. - """ - if not (isinstance(a, torch.Tensor) and isinstance(b, torch.Tensor)): - return False - if a.device.type != "cuda" or a.device != b.device: - return False - if not triton_available(): - return False - if a.dim() != 5 or b.dim() != 5: - return False - if a.dtype not in SUPPORTED_DTYPES or b.dtype not in SUPPORTED_DTYPES: - return False - if out_dtype is not None and out_dtype not in SUPPORTED_DTYPES: - return False - if a.shape[0] != b.shape[0] or tuple(a.shape[2:]) != tuple(b.shape[2:]): - return False - if a.shape[1] < 1 or b.shape[1] < 1: - return False - if a.numel() == 0 or b.numel() == 0: - return False - if not a.is_contiguous(memory_format=_CL_FORMAT): - return False - if not b.is_contiguous(memory_format=_CL_FORMAT): - return False - return True - - -def cat_channels(a, b, out_dtype: Optional[torch.dtype] = None): - """``torch.cat([a, b], dim=1)``, optionally emitting ``out_dtype`` directly. - - Total: anything :func:`is_supported` declines is served by ``torch.cat`` - (followed by a cast when ``out_dtype`` asks for one), so this is a drop-in - on CPU, on non-channels-last input and without Triton. - - The fallback casts the *inputs* rather than the concatenated result when - ``out_dtype`` is narrower. That is bitwise the same answer -- every - supported dtype widens exactly into the promoted type, so narrowing before - or after the copy rounds the same values once -- and it does not - materialize a double-width intermediate. - """ - if not is_supported(a, b, out_dtype): - if out_dtype is None: - return torch.cat([a, b], dim=1) - return torch.cat([a.to(out_dtype), b.to(out_dtype)], dim=1) - return torch.ops.scaffold_cat.cat_channels(a, b, out_dtype) - - -def consumer_dtype(*tensors) -> torch.dtype: - """The dtype the convolution that consumes the concatenation will see. - - Under an enabled autocast region the answer is autocast's ``dtype``, - because ``aten::convolution`` carries the ``lower_precision_fp`` cast - policy and casts whatever it is handed. Producing that dtype from the - concatenation is therefore *bitwise identical* to producing ATen's promoted - dtype and letting the convolution narrow it -- the promoted tensor holds - exact widenings of both sources -- while writing (and reading back) half - the bytes, and it removes autocast's own widening of the narrower input on - the way in. - - Outside autocast the answer is ``torch.cat``'s ordinary promotion, so the - result is unchanged for eval, ``inference_mode`` and pure-fp32 runs. - """ - dtype = tensors[0].dtype - for t in tensors[1:]: - dtype = torch.promote_types(dtype, t.dtype) - device_type = tensors[0].device.type - try: - if torch.is_autocast_enabled(device_type): - autocast_dtype = torch.get_autocast_dtype(device_type) - else: - return dtype - except (RuntimeError, TypeError): # a device type autocast does not know - return dtype - # Only ever narrow: if autocast's dtype is not one the kernel can hold, or - # is wider than the inputs, keep the promotion ATen would have done. - if autocast_dtype not in SUPPORTED_DTYPES: - return dtype - if torch.promote_types(autocast_dtype, dtype) is autocast_dtype: - return dtype - return autocast_dtype - - -def _dctensor_ops(input): - """The ``distconv.distconv`` module when ``input`` is a DCTensor, else None. - - Resolved through ``sys.modules`` rather than an import, so this module - stays importable (and the CPU suite runnable) without DistConv installed. - """ - distconv = sys.modules.get("distconv.distconv") - if distconv is not None and isinstance(input, distconv.DCTensor): - return distconv - return None - - -def skip_concat(skip, upsampled): - """Join a decoder skip activation to an upsampled one, as ``Up.forward`` needs. - - Equivalent to ``torch.cat([skip, upsampled], dim=1)`` as the following - convolution observes it -- see :func:`consumer_dtype` for why the dtype may - legitimately differ from ``torch.cat``'s own. - - DistConv's ``DCTensor`` is unwrapped to its local shard in front of the - kernel and rewrapped after, rather than being left to - ``DCTensor.__torch_dispatch__``. Both work -- these are real dispatcher - ops -- but the explicit unwrap is what ``FastGroupNorm`` already does, and - it keeps the subclass policy in one place: ``is_supported`` accepts any - ``torch.Tensor`` instance, so relying on dispatch would silently extend the - fast path to every unknown wrapper subclass. The unwrap goes through - DistConv's ``_ToTensor``/``from_shard`` autograd pair and not a bare - ``._tensor`` read, which would sever the graph back to the producing - convolution. DistConv has no concatenation-specific handling, so the - semantics are identical at every shard count. - """ - distconv = _dctensor_ops(skip) - if distconv is None or _dctensor_ops(upsampled) is None: - return cat_channels(skip, upsampled, consumer_dtype(skip, upsampled)) - if skip._parallel_strategy != upsampled._parallel_strategy: - raise ValueError( - "skip and upsampled tensors have different parallel strategies" - ) - local_skip = distconv._ToTensor.apply(skip) - local_up = distconv._ToTensor.apply(upsampled) - out = cat_channels(local_skip, local_up, consumer_dtype(local_skip, local_up)) - return distconv.DCTensor.from_shard(out, skip._parallel_strategy) diff --git a/ScaFFold/unet/unet_parts.py b/ScaFFold/unet/unet_parts.py index 6f5eac9..5df1cee 100644 --- a/ScaFFold/unet/unet_parts.py +++ b/ScaFFold/unet/unet_parts.py @@ -14,13 +14,13 @@ """Parts of the U-Net model""" +import torch import torch.nn as nn import torch.nn.functional as F from ScaFFold.utils.perf_measure import annotate from .group_norm import FastGroupNorm -from .triton_cat import skip_concat _doubleconv_annotate = annotate(fmt="DoubleConv.{}") _down_annotate = annotate(fmt="Down.{}") @@ -39,6 +39,53 @@ def _group_norm(num_groups, num_channels, activation=None): return FastGroupNorm(num_groups, num_channels, activation=activation) +def _consumer_dtype(*tensors): + """The dtype the convolution consuming a concatenation will actually see. + + Inside an enabled autocast region the answer is autocast's dtype, because + ``aten::convolution`` carries the ``lower_precision_fp`` cast policy and + casts whatever it is handed. Producing that dtype from the concatenation + is *bitwise identical* to producing ATen's promoted dtype and letting the + convolution narrow it -- the promoted tensor holds exact widenings of both + sources, so narrowing before or after the copy rounds the same values once + -- while writing and reading back half the bytes. + + Outside autocast the answer is ``torch.cat``'s ordinary promotion, so eval, + ``inference_mode`` and pure-fp32 runs are unchanged. + """ + dtype = tensors[0].dtype + for tensor in tensors[1:]: + dtype = torch.promote_types(dtype, tensor.dtype) + device_type = tensors[0].device.type + try: + if not torch.is_autocast_enabled(device_type): + return dtype + autocast_dtype = torch.get_autocast_dtype(device_type) + except (RuntimeError, TypeError): # a device type autocast does not know + return dtype + # Only ever narrow: if autocast's dtype is the wider of the two, keep the + # promotion ATen would have done. + if torch.promote_types(autocast_dtype, dtype) is autocast_dtype: + return dtype + return autocast_dtype + + +def _skip_concat(skip, upsampled): + """``torch.cat([skip, upsampled], dim=1)`` at the consumer's dtype. + + Under ``torch.autocast`` the two halves do not share a dtype: the skip + comes from a GroupNorm, an fp32-policy op, while the upsampled half comes + from a ``ConvTranspose3d`` and is bf16. ``torch.cat`` carries the + ``promote`` policy, so it widens the bf16 half to fp32, concatenates at + fp32, and the following convolution narrows the whole double-width result + straight back down -- three full-resolution passes to deliver one. Casting + the inputs first collapses that to one, and the convolution reads the same + bits either way (see :func:`_consumer_dtype`). + """ + dtype = _consumer_dtype(skip, upsampled) + return torch.cat([skip.to(dtype), upsampled.to(dtype)], dim=1) + + class DoubleConv(nn.Module): """(convolution => GroupNorm => ReLU) * 2 @@ -94,27 +141,20 @@ def forward(self, x): class Up(nn.Module): """Upscaling then double conv - The skip concatenation goes through :func:`ScaFFold.unet.triton_cat.skip_concat` - rather than ``torch.cat``. That does two things, both of which leave the - tensor the following convolution reads *bitwise* unchanged (verified in - ``tests/test_triton_cat.py``): - - * It emits the dtype the convolution will use instead of ``torch.cat``'s - promoted one. Under ``torch.autocast`` the two halves do not have the - same dtype -- the skip comes from a GroupNorm, an fp32-policy op, and the - upsampled half comes from a ``ConvTranspose3d`` and is bf16 -- so ``cat`` - widens the bf16 half to fp32, concatenates at fp32, and the convolution - then narrows the whole double-width result straight back down. Three - full-resolution passes to deliver one. - * It runs a channels-last-native kernel, which matters most for the - *backward*: ``cat``'s backward is a narrowed view that every consumer - then forces contiguous, and that strided copy reaches only 51-63% of this - device's streaming roofline against the kernel's 90-103%. - - Both rest on ``self.conv`` beginning with a convolution, which the - constructor below guarantees on either branch. ``skip_concat`` is total: - anything its ``is_supported`` declines -- CPU, non-channels-last, no Triton - -- falls back to ``torch.cat``. + The skip concatenation goes through :func:`_skip_concat` rather than + ``torch.cat`` directly, so that it emits the dtype the following + convolution will use instead of ``torch.cat``'s promoted one. The tensor + that convolution reads is bitwise unchanged either way; it is written and + read back at half the width. This rests on ``self.conv`` beginning with a + convolution, which the constructor below guarantees on either branch. + + Measured at scale 7: 1.09 ms of a 92.8 ms step, and 0.50 GiB of peak + memory. A channels-last-native Triton concatenation kernel was built and + measured too -- ``cat``'s *backward* is a narrowed view that consumers force + contiguous, at 51-63% of this device's streaming roofline against the + kernel's 90-103% -- but it was worth a further 0.08 ms of the step, which + did not justify a second hand-written kernel in a benchmark other people + have to trust. ``review/skip-path/RESULTS.md`` has the numbers. """ def __init__(self, in_channels, out_channels, group_norm_groups, trilinear=True): @@ -160,7 +200,7 @@ def forward(self, x1, x2): # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd # torch.cat([x2, x1], dim=1) with the dtype and the layout the # convolution below actually wants; see the class docstring. - x = skip_concat(x2, x1) + x = _skip_concat(x2, x1) return self.conv(x) diff --git a/tests/test_triton_cat.py b/tests/test_triton_cat.py deleted file mode 100644 index 84a3864..0000000 --- a/tests/test_triton_cat.py +++ /dev/null @@ -1,478 +0,0 @@ -# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. -# Produced at the Lawrence Livermore National Laboratory. -# Written by the LBANN Research Team (B. Van Essen, et al.) listed in -# the CONTRIBUTORS file. See the top-level LICENSE file for details. -# -# LLNL-CODE-697807. -# All rights reserved. -# -# This file is part of LBANN: Livermore Big Artificial Neural Network -# Toolkit. For details, see http://software.llnl.gov/LBANN or -# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. -# -# SPDX-License-Identifier: (Apache-2.0) - -"""Tests for the channels-last Triton channel concatenation (``triton_cat``). - -The op replaces ``torch.cat([a, b], dim=1)``, whose result is a *copy* rather -than an arithmetic expression, so unlike the GroupNorm kernel these are not -tolerance tests: every parity assertion here is **bitwise**, forward and -backward, including the mixed-dtype and narrowing cases. A tolerance would -hide exactly the bugs this kernel can have -- an off-by-one in the channel -split, a lost tail row, a double rounding. - -The other things being pinned down: - -* the **dtype rule**. ``consumer_dtype`` is what licenses emitting bf16 from a - concatenation whose ``torch.cat`` result would be fp32; the tests check both - that it is what the following convolution receives and that it never *widens* - anything. -* the **fallback**. ``is_supported`` must decline everything the kernel cannot - serve physically, and ``cat_channels`` must then still answer -- so the CPU - suite exercises the whole public API with no GPU and no Triton. -* **composition**: a ``DCTensor`` round trip with the autograd graph intact, - ``torch.utils.checkpoint`` recompute, ``inference_mode``, and the - ``trilinear`` branch of ``Up``. - -CPU runs never touch Triton: the module defers ``import triton`` to the first -call that reaches a kernel, which ``test_import_does_not_pull_in_triton`` -checks in a fresh interpreter. -""" - -from __future__ import annotations - -import itertools -import subprocess -import sys - -import pytest -import torch - -from ScaFFold.unet import triton_cat -from ScaFFold.unet.triton_cat import ( - cat_channels, - consumer_dtype, - is_supported, - skip_concat, -) - -CL = torch.channels_last_3d -DTYPES = (torch.float32, torch.bfloat16, torch.float16) - -gpu = pytest.mark.gpu -requires_cuda = pytest.mark.skipif( - not torch.cuda.is_available(), reason="needs a CUDA device" -) - - -def cl_tensor(*shape, dtype=torch.float32, device="cpu", seed=0): - """A deterministic channels-last-3d tensor.""" - generator = torch.Generator().manual_seed(seed) - return ( - torch.randn(*shape, generator=generator) - .to(device=device, dtype=dtype) - .contiguous(memory_format=CL) - ) - - -# --------------------------------------------------------------------------- # -# The public API is total: it must work with no GPU and no Triton at all. -# --------------------------------------------------------------------------- # -def test_cpu_falls_back_and_matches_torch_cat_bitwise(): - """On CPU ``is_supported`` declines and ``cat_channels`` still answers.""" - a = cl_tensor(2, 5, 3, 4, 5, seed=1) - b = cl_tensor(2, 3, 3, 4, 5, seed=2) - assert not is_supported(a, b) - got = cat_channels(a, b) - assert torch.equal(got, torch.cat([a, b], dim=1)) - assert got.is_contiguous(memory_format=CL) - - -def test_cpu_fallback_honours_out_dtype_bitwise(): - """A narrowing ``out_dtype`` must round each value exactly once. - - ``torch.cat([a, b]).to(bf16)`` widens then narrows; the fallback narrows - the inputs first. Those agree bitwise because the promoted dtype is an - exact widening of both, and that is the property the kernel relies on too. - """ - a = cl_tensor(1, 4, 2, 3, 4, seed=3) - b = cl_tensor(1, 4, 2, 3, 4, dtype=torch.bfloat16, seed=4) - for out_dtype in DTYPES: - got = cat_channels(a, b, out_dtype) - assert got.dtype is out_dtype - assert torch.equal(got, torch.cat([a, b], dim=1).to(out_dtype)) - - -def test_cpu_fallback_backward_matches_torch_cat_bitwise(): - a = cl_tensor(1, 6, 2, 2, 2, seed=5).requires_grad_(True) - b = cl_tensor(1, 2, 2, 2, 2, dtype=torch.bfloat16, seed=6).requires_grad_(True) - grad = cl_tensor(1, 8, 2, 2, 2, seed=7) - - torch.cat([a, b], dim=1).backward(grad) - ref = (a.grad.clone(), b.grad.clone()) - a.grad = b.grad = None - - cat_channels(a, b).backward(grad) - assert torch.equal(a.grad, ref[0]) - assert torch.equal(b.grad, ref[1]) - assert a.grad.dtype is torch.float32 - assert b.grad.dtype is torch.bfloat16 - - -@pytest.mark.parametrize( - "make", - [ - pytest.param(lambda: (torch.randn(1, 2, 2, 2, 2), None), id="not-a-tensor"), - pytest.param(lambda: (torch.randn(2, 2), torch.randn(2, 2)), id="rank-2"), - pytest.param( - lambda: (torch.randn(1, 2, 2, 2, 2).double(), torch.randn(1, 2, 2, 2, 2)), - id="float64", - ), - pytest.param( - lambda: (torch.randn(1, 2, 2, 2, 2), torch.randn(1, 2, 2, 2, 3)), - id="spatial-mismatch", - ), - pytest.param( - lambda: (torch.randn(0, 2, 2, 2, 2), torch.randn(0, 2, 2, 2, 2)), - id="empty", - ), - ], -) -def test_is_supported_declines_what_the_kernel_cannot_serve(make): - a, b = make() - assert not is_supported(a, b) - - -def test_import_does_not_pull_in_triton(): - """Importing the module must not import Triton (it is deferred to first use).""" - code = ( - "import sys; import ScaFFold.unet.triton_cat as m; " - "assert 'triton' not in sys.modules, sorted(k for k in sys.modules " - "if k.startswith('triton')); print('ok')" - ) - out = subprocess.run( - [sys.executable, "-c", code], capture_output=True, text=True, check=True - ) - assert out.stdout.strip().endswith("ok") - - -# --------------------------------------------------------------------------- # -# The dtype rule. -# --------------------------------------------------------------------------- # -def test_consumer_dtype_outside_autocast_is_torch_cats_promotion(): - a = cl_tensor(1, 2, 2, 2, 2) - b = cl_tensor(1, 2, 2, 2, 2, dtype=torch.bfloat16) - assert consumer_dtype(a, b) is torch.promote_types(torch.float32, torch.bfloat16) - assert consumer_dtype(a, a) is torch.float32 - assert consumer_dtype(b, b) is torch.bfloat16 - - -@requires_cuda -@gpu -def test_consumer_dtype_under_autocast_is_the_autocast_dtype(): - """The mixed-dtype case the decoder actually presents. - - fp32 skip (a GroupNorm output, fp32 cast policy) + bf16 upsampled (a - ConvTranspose3d output). ``torch.cat`` would answer fp32; the convolution - that consumes it narrows to bf16, so bf16 is the honest answer. - """ - a = cl_tensor(1, 4, 2, 2, 2, device="cuda") - b = cl_tensor(1, 4, 2, 2, 2, dtype=torch.bfloat16, device="cuda") - with torch.autocast("cuda", dtype=torch.bfloat16): - assert consumer_dtype(a, b) is torch.bfloat16 - # never widens: two bf16 inputs stay bf16, and an fp16 autocast over - # bf16 inputs must not silently change their width either - assert consumer_dtype(b, b) is torch.bfloat16 - assert consumer_dtype(a, b) is torch.float32 - - -@requires_cuda -@gpu -def test_skip_concat_is_bitwise_what_the_convolution_receives(): - """The load-bearing claim: the dtype shortcut changes no bits downstream.""" - a = cl_tensor(1, 8, 3, 4, 5, device="cuda") - b = cl_tensor(1, 8, 3, 4, 5, dtype=torch.bfloat16, device="cuda", seed=11) - with torch.autocast("cuda", dtype=torch.bfloat16): - got = skip_concat(a, b) - # what the convolution would have been handed by the old chain: - # torch.cat promotes to fp32, then autocast narrows for the conv. - reference = torch.cat([a, b], dim=1).to(torch.bfloat16) - assert got.dtype is torch.bfloat16 - assert torch.equal(got, reference) - - -# --------------------------------------------------------------------------- # -# GPU parity: bitwise, forward and backward. -# --------------------------------------------------------------------------- # -_SHAPES = [ - (1, 8, 8, 4, 4, 4), # power-of-two channels, the UNet's case - (1, 3, 5, 2, 3, 4), # odd channel counts, C = 8 - (1, 5, 6, 2, 2, 2), # C = 11, not a power of two - (2, 4, 4, 2, 2, 2), # batch > 1 - (1, 1, 1, 1, 1, 1), # degenerate -] - - -@requires_cuda -@gpu -@pytest.mark.parametrize("shape", _SHAPES, ids=lambda s: "x".join(str(v) for v in s)) -@pytest.mark.parametrize("dtypes", list(itertools.product(DTYPES, DTYPES))) -def test_forward_is_bitwise_torch_cat(shape, dtypes): - n, ca, cb, d, h, w = shape - da, db = dtypes - a = cl_tensor(n, ca, d, h, w, dtype=da, device="cuda", seed=21) - b = cl_tensor(n, cb, d, h, w, dtype=db, device="cuda", seed=22) - assert is_supported(a, b) - got = cat_channels(a, b) - ref = torch.cat([a, b], dim=1) - assert got.dtype is ref.dtype - assert got.shape == ref.shape - assert got.is_contiguous(memory_format=CL) - assert torch.equal(got, ref) - - -@requires_cuda -@gpu -@pytest.mark.parametrize("out_dtype", DTYPES) -def test_forward_out_dtype_is_bitwise_cat_then_cast(out_dtype): - a = cl_tensor(1, 6, 3, 4, 5, device="cuda", seed=23) - b = cl_tensor(1, 10, 3, 4, 5, dtype=torch.bfloat16, device="cuda", seed=24) - got = cat_channels(a, b, out_dtype) - assert got.dtype is out_dtype - assert torch.equal(got, torch.cat([a, b], dim=1).to(out_dtype)) - - -@requires_cuda -@gpu -@pytest.mark.parametrize("dtypes", list(itertools.product(DTYPES, DTYPES))) -def test_backward_is_bitwise_torch_cat(dtypes): - da, db = dtypes - a = cl_tensor(1, 6, 2, 3, 4, dtype=da, device="cuda", seed=25).requires_grad_(True) - b = cl_tensor(1, 10, 2, 3, 4, dtype=db, device="cuda", seed=26).requires_grad_(True) - grad = cl_tensor( - 1, 16, 2, 3, 4, dtype=torch.promote_types(da, db), device="cuda", seed=27 - ) - - torch.cat([a, b], dim=1).backward(grad) - ref = (a.grad.clone(), b.grad.clone()) - a.grad = b.grad = None - - cat_channels(a, b).backward(grad) - assert a.grad.dtype is da and b.grad.dtype is db - assert a.grad.is_contiguous(memory_format=CL) - assert b.grad.is_contiguous(memory_format=CL) - assert torch.equal(a.grad, ref[0]) - assert torch.equal(b.grad, ref[1]) - - -@requires_cuda -@gpu -def test_backward_with_only_one_side_requiring_grad(): - """The ``WANT_A``/``WANT_B`` kernel flags must not corrupt the other half.""" - a = cl_tensor(1, 6, 2, 2, 2, device="cuda", seed=28).requires_grad_(True) - b = cl_tensor(1, 2, 2, 2, 2, device="cuda", seed=29) - grad = cl_tensor(1, 8, 2, 2, 2, device="cuda", seed=30) - cat_channels(a, b).backward(grad) - assert b.grad is None - assert torch.equal(a.grad, grad[:, :6]) - - a2 = cl_tensor(1, 6, 2, 2, 2, device="cuda", seed=28) - b2 = cl_tensor(1, 2, 2, 2, 2, device="cuda", seed=29).requires_grad_(True) - cat_channels(a2, b2).backward(grad) - assert a2.grad is None - assert torch.equal(b2.grad, grad[:, 6:]) - - -@requires_cuda -@gpu -def test_backward_accepts_a_non_channels_last_gradient(): - """Nothing guarantees the consumer hands back a channels-last cotangent.""" - a = cl_tensor(1, 6, 2, 2, 2, device="cuda", seed=31).requires_grad_(True) - b = cl_tensor(1, 2, 2, 2, 2, device="cuda", seed=32).requires_grad_(True) - grad = torch.randn(1, 8, 2, 2, 2, device="cuda") # plain contiguous - cat_channels(a, b).backward(grad) - assert torch.equal(a.grad, grad[:, :6]) - assert torch.equal(b.grad, grad[:, 6:]) - - -@requires_cuda -@gpu -@pytest.mark.parametrize("pad", [1, 2]) -def test_backward_reads_a_halo_padded_gradient_in_place(pad): - """The production cotangent: a narrowed view of a halo-padded tensor. - - DistConv reaches the convolution that consumes the concatenation through a - halo exchange that materialises a ``(D+2, H+2, W+2)`` tensor, so what comes - back here is a *narrow* of one -- channels-last per voxel and contiguous - along W, but with a gap at every H and D boundary. Relaying that out costs - a whole extra full-resolution pass, which is exactly what made an earlier - version of this op lose to ``torch.cat``; the kernel must read it in place. - """ - a = cl_tensor(1, 6, 3, 4, 5, device="cuda", seed=51).requires_grad_(True) - b = cl_tensor(1, 10, 3, 4, 5, dtype=torch.bfloat16, device="cuda", seed=52) - b.requires_grad_(True) - parent = cl_tensor( - 1, 16, 3 + 2 * pad, 4 + 2 * pad, 5 + 2 * pad, device="cuda", seed=53 - ) - grad = parent[:, :, pad:-pad, pad:-pad, pad:-pad] - assert grad.shape == (1, 16, 3, 4, 5) - assert not grad.is_contiguous(memory_format=CL) - assert triton_cat._line_addressable(grad), "must take the in-place path" - - torch.cat([a, b], dim=1).backward(grad) - ref = (a.grad.clone(), b.grad.clone()) - a.grad = b.grad = None - - cat_channels(a, b).backward(grad) - assert torch.equal(a.grad, ref[0]) - assert torch.equal(b.grad, ref[1]) - assert a.grad.is_contiguous(memory_format=CL) - assert b.grad.is_contiguous(memory_format=CL) - - -def test_line_addressable_accepts_narrows_and_declines_permutations(): - """The predicate that decides whether the backward can skip its relayout.""" - dense = cl_tensor(1, 8, 4, 5, 6) - assert triton_cat._line_addressable(dense) - parent = cl_tensor(1, 8, 6, 7, 8) - assert triton_cat._line_addressable(parent[:, :, 1:-1, 1:-1, 1:-1]) - assert triton_cat._line_addressable(parent[:, :, :, :, 1:-1]) - # a plain contiguous (NCDHW) tensor has channels *outermost* - assert not triton_cat._line_addressable(torch.randn(1, 8, 4, 5, 6)) - # a spatial permutation breaks the "W neighbours are one run apart" rule - assert not triton_cat._line_addressable(dense.transpose(3, 4)) - # narrowing the channel axis breaks stride(4) == C - assert not triton_cat._line_addressable(dense[:, :4]) - - -@requires_cuda -@gpu -def test_non_channels_last_input_falls_back_to_torch_cat(): - a = torch.randn(1, 6, 2, 2, 2, device="cuda") # contiguous, not CL - b = cl_tensor(1, 2, 2, 2, 2, device="cuda", seed=33) - assert not is_supported(a, b) - assert torch.equal(cat_channels(a, b), torch.cat([a, b], dim=1)) - - -@requires_cuda -@gpu -def test_repeated_calls_are_bitwise_identical(): - """A copy has no reduction order, and the tiling is a pure function of C.""" - a = cl_tensor(1, 6, 3, 4, 5, device="cuda", seed=34).requires_grad_(True) - b = cl_tensor(1, 10, 3, 4, 5, device="cuda", seed=35).requires_grad_(True) - grad = cl_tensor(1, 16, 3, 4, 5, device="cuda", seed=36) - first = cat_channels(a, b).clone() - cat_channels(a, b).backward(grad) - ga, gb = a.grad.clone(), b.grad.clone() - a.grad = b.grad = None - second = cat_channels(a, b).clone() - cat_channels(a, b).backward(grad) - assert torch.equal(first, second) - assert torch.equal(a.grad, ga) and torch.equal(b.grad, gb) - - -@requires_cuda -@gpu -def test_runs_on_a_non_current_device(): - """A Triton launch follows the *current* device; the guard must override it.""" - if torch.cuda.device_count() < 2: - pytest.skip("needs two CUDA devices") - a = cl_tensor(1, 6, 2, 2, 2, device="cuda:1", seed=37) - b = cl_tensor(1, 2, 2, 2, 2, device="cuda:1", seed=38) - with torch.cuda.device(0): - got = cat_channels(a, b) - assert got.device == a.device - assert torch.equal(got, torch.cat([a, b], dim=1)) - - -@requires_cuda -@gpu -def test_second_order_raises_rather_than_returning_garbage(): - """First order only, exactly like the Triton GroupNorm; it must fail loudly. - - The loss has to be *nonlinear* for this to bite. A concatenation is a - copy, so its own second derivative is identically zero: differentiating - ``sum(cat(a, b))`` twice gives a cotangent that does not depend on ``a`` at - all, and autograd correctly reports "does not require grad" without ever - reaching this op. With a square in the way the cotangent *does* depend on - ``a``, the split has to be differentiated, and there is no formula for it. - """ - a = cl_tensor(1, 4, 2, 2, 2, device="cuda", seed=39).requires_grad_(True) - b = cl_tensor(1, 4, 2, 2, 2, device="cuda", seed=40).requires_grad_(True) - out = cat_channels(a, b) - (grad_a,) = torch.autograd.grad((out * out).sum(), a, create_graph=True) - with pytest.raises(RuntimeError, match="no autograd formula was registered"): - torch.autograd.grad(grad_a.sum(), a) - - -@requires_cuda -@gpu -def test_dctensor_round_trip_keeps_the_graph(): - """Production wraps activations in a DCTensor even at ``dc_num_shards=1``.""" - distconv = pytest.importorskip("distconv") - if not torch.distributed.is_initialized(): - pytest.skip("needs an initialized process group") - ps = distconv.ParallelStrategy( - num_shards=(1, 1, 1), shard_dim=(2, 3, 4), device_type="cuda" - ) - a = cl_tensor(1, 4, 2, 2, 2, device="cuda", seed=41).requires_grad_(True) - b = cl_tensor(1, 4, 2, 2, 2, device="cuda", seed=42).requires_grad_(True) - out = skip_concat( - distconv.DCTensor.from_shard(a, ps), distconv.DCTensor.from_shard(b, ps) - ) - assert isinstance(out, distconv.DCTensor) - distconv.distconv._ToTensor.apply(out).pow(2).sum().backward() - assert a.grad is not None and b.grad is not None - - -@requires_cuda -@gpu -def test_survives_activation_checkpoint_recompute(): - """The block's forward is replayed inside backward under checkpointing.""" - a = cl_tensor(1, 4, 2, 2, 2, device="cuda", seed=43).requires_grad_(True) - b = cl_tensor(1, 4, 2, 2, 2, device="cuda", seed=44).requires_grad_(True) - - def block(x, y): - return cat_channels(x, y) * 2.0 - - ref = block(a, b) - ref.pow(2).sum().backward() - ga, gb = a.grad.clone(), b.grad.clone() - a.grad = b.grad = None - - out = torch.utils.checkpoint.checkpoint(block, a, b, use_reentrant=False) - out.pow(2).sum().backward() - assert torch.equal(a.grad, ga) - assert torch.equal(b.grad, gb) - - -@requires_cuda -@gpu -def test_inference_mode_and_no_grad(): - a = cl_tensor(1, 4, 2, 2, 2, device="cuda", seed=45) - b = cl_tensor(1, 4, 2, 2, 2, device="cuda", seed=46) - ref = torch.cat([a, b], dim=1) - with torch.no_grad(): - assert torch.equal(cat_channels(a, b), ref) - with torch.inference_mode(): - assert torch.equal(cat_channels(a, b), ref) - - -@requires_cuda -@gpu -def test_kernel_failure_is_tagged_so_a_caller_can_fall_back(): - """``CatKernelError`` must be what escapes when the launch itself breaks.""" - a = cl_tensor(1, 4, 2, 2, 2, device="cuda", seed=47) - b = cl_tensor(1, 4, 2, 2, 2, device="cuda", seed=48) - original = triton_cat._forward.__wrapped__ - - def boom(*args, **kwargs): - raise ValueError("simulated launch failure") - - triton_cat._forward.__wrapped__ = boom - try: - wrapped = triton_cat._tag_kernel_failures(boom) - with pytest.raises(triton_cat.CatKernelError): - wrapped(a, b, torch.float32) - finally: - triton_cat._forward.__wrapped__ = original diff --git a/tests/test_unet.py b/tests/test_unet.py index a17b65c..fa5f47b 100644 --- a/tests/test_unet.py +++ b/tests/test_unet.py @@ -204,13 +204,12 @@ def counting_pad(tensor, pad, *args, **kwargs): # The decoder skip concatenation. # # ``Up.forward`` no longer calls ``torch.cat`` directly; it goes through -# ``ScaFFold.unet.triton_cat.skip_concat``, which may legitimately emit a -# narrower dtype than ``torch.cat`` would when autocast is on (see the ``Up`` -# docstring). Everything below pins the part that must NOT change: outside -# autocast the block is bitwise what it was, the ``F.pad`` path still works, -# and both the ``trilinear`` and ``ConvTranspose3d`` branches agree with an -# explicit ``torch.cat`` reference. The kernel's own parity tests live in -# ``tests/test_triton_cat.py``. +# ``unet_parts._skip_concat``, which may legitimately emit a narrower dtype +# than ``torch.cat`` would when autocast is on (see the ``Up`` docstring). +# Everything below pins the part that must NOT change: outside autocast the +# block is bitwise what it was, the ``F.pad`` path still works, and both the +# ``trilinear`` and ``ConvTranspose3d`` branches agree with an explicit +# ``torch.cat`` reference. # --------------------------------------------------------------------------- # @pytest.mark.parametrize("trilinear", [False, True]) def test_up_matches_an_explicit_torch_cat_reference(trilinear): @@ -285,16 +284,16 @@ def test_up_gradients_match_an_explicit_torch_cat_reference(): def test_up_concatenation_keeps_channels_last(): """The concatenation must not break the layout chain it exists to preserve. - Asserted on ``skip_concat`` with two channels-last halves rather than on a + Asserted on ``_skip_concat`` with two channels-last halves rather than on a whole ``Up`` block: on CPU ``nn.ConvTranspose3d`` returns a *contiguous* tensor whatever it is handed, so the block's own inputs to the concatenation are not both channels-last there and the block-level assertion would be measuring the convolution's layout policy, not this - one's. On GPU with ``PYTORCH_MIOPEN_SUGGEST_NHWC=1`` -- the configuration - the kernel exists for -- both halves are channels-last and this is the - property ``Up`` relies on. + one's. On GPU with ``PYTORCH_MIOPEN_SUGGEST_NHWC=1`` -- the production + configuration -- both halves are channels-last and this is the property + ``Up`` relies on. """ - from ScaFFold.unet.triton_cat import skip_concat + from ScaFFold.unet.unet_parts import _skip_concat as skip_concat generator = torch.Generator().manual_seed(14) x1 = torch.randn(1, 16, 16, 16, 16, generator=generator).contiguous(